Implementa suporte a imagens no endpoint de previsão

Atualiza o processamento de mensagens para lidar com diferentes tipos de entrada, como imagens via URL, base64 e caminhos locais. Essa mudança aprimora a flexibilidade do modelo, permitindo que ele receba entradas visuais diretamente, o que pode melhorar a qualidade das previsões em cenários que exigem contexto visual. Além disso, ajusta o processamento de texto para suportar múltiplos conteúdos, garantindo uma integração mais eficaz com as funcionalidades do modelo.

app.py

@@ -3,19 +3,45 @@ from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
 from qwen_vl_utils import process_vision_info
 import torch
 from typing import List, Dict
+import base64
+import requests
+from PIL import Image
+from io import BytesIO
 
 app = FastAPI()
 
 model = Qwen2VLForConditionalGeneration.from_pretrained("Qwen/Qwen2-VL-2B-Instruct", torch_dtype="auto", device_map="auto")
-processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct")
+min_pixels = 256 * 28 * 28
+max_pixels = 1280 * 28 * 28
+processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct", min_pixels=min_pixels, max_pixels=max_pixels)
+
+def process_image(image_data):
+    if image_data.startswith("http://") or image_data.startswith("https://"):
+        response = requests.get(image_data)
+        img = Image.open(BytesIO(response.content))
+    elif image_data.startswith("data:image"):
+        img_data = base64.b64decode(image_data.split(",")[1])
+        img = Image.open(BytesIO(img_data))
+    else:  # Assume it's a local file path
+        img = Image.open(image_data)
+    return img
 
 @app.post("/predict")
 async def predict(messages: List[Dict] = Body(...)):
     # Processamento e inferência
-    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    image_inputs, video_inputs = process_vision_info(messages)
+    texts = []
+    image_inputs = []
+    video_inputs = []
+
+    for message in messages:
+        for content in message["content"]:
+            if content["type"] == "text":
+                texts.append(processor.apply_chat_template(content["text"], tokenize=False, add_generation_prompt=True))
+            elif content["type"] == "image":
+                image_inputs.append(process_image(content["image"]))
+    
     inputs = processor(
-        text=[text],
+        text=texts,
         images=image_inputs,
         videos=video_inputs,
         padding=True,
@@ -24,6 +50,10 @@ async def predict(messages: List[Dict] = Body(...)):
     inputs = inputs.to("cpu")  # Altere para "cuda" se tiver GPU disponível
     
     generated_ids = model.generate(**inputs, max_new_tokens=128)
-    generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
-    output_text = processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)
-    return {"response": output_text}
+    generated_ids_trimmed = [
+        out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+    ]
+    output_texts = processor.batch_decode(
+        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+    )
+    return {"response": output_texts}