Views
No views yet
1import os
2import requests
3import torch
4from PIL import Image
5from transformers import MllamaForConditionalGeneration, AutoProcessor
6
7# Path to your locally saved merged multimodal model
8model_id = "miike-ai/r1-11b-vision"
9
10# Load the model and processor
11model = MllamaForConditionalGeneration.from_pretrained(
12 model_id,
13 torch_dtype=torch.bfloat16,
14 device_map="auto",
15)
16processor = AutoProcessor.from_pretrained(model_id)
17
18def multimodal_inference(text_prompt, image_path=None):
19 """
20 Runs a single inference on the multimodal model.
21 Args:
22 text_prompt (str): The user prompt for text-based input.
23 image_path (str, optional): Path or URL to an image (if any).
24 Returns:
25 str: Model-generated response.
26 """
27
28 # Prepare user message
29 user_message = {"role": "user", "content": [{"type": "text", "text": text_prompt}]}
30
31 # Load image if provided
32 image = None
33 if image_path:
34 try:
35 if image_path.startswith("http"): # Handle URLs
36 image = Image.open(requests.get(image_path, stream=True).raw)
37 else: # Handle local file
38 image = Image.open(image_path)
39
40 print(f"📷 Loaded image: {image.size}") # Debugging
41 user_message["content"].insert(0, {"type": "image"}) # Add image token to message
42 except Exception as e:
43 print(f"⚠️ Error loading image: {e}")
44 return None
45
46 # Format input for the model
47 input_text = processor.apply_chat_template([user_message], add_generation_prompt=True)
48
49 # Convert input to model tensors
50 if image is not None:
51 inputs = processor(images=[image], text=[input_text], add_special_tokens=True, return_tensors="pt").to(model.device)
52 else:
53 inputs = processor(text=[input_text], add_special_tokens=True, return_tensors="pt").to(model.device)
54
55 # Generate response
56 with torch.no_grad():
57 output = model.generate(**inputs, max_new_tokens=256)
58
59 # Decode model output
60 response_text = processor.decode(output[0], skip_special_tokens=True)
61
62 return response_text
63
64# Example usage
65text_input = "What is in this image?"
66image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/0052a70beed5bf71b92610a43a52df6d286cd5f3/diffusers/rabbit.jpg" # Can be a URL or local file path
67
68response = multimodal_inference(text_input, image_path)
69print("\n🧠 Assistant:", response)