Refine image processing and input handling

Melhora a função de processamento de imagens para incluir verificação de erros durante as requisições HTTP e clarifica o tratamento de entradas na função de previsão. AGora suporta tanto strings quanto listas para o conteúdo, disponibilizando feedback para formatação inválida. Essas alterações aumentam a robustez do sistema e melhoram a experiência do usuário ao prevenir falhas inesperadas.

app.py

@@ -2,7 +2,7 @@ from fastapi import FastAPI, Body
 from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
 from qwen_vl_utils import process_vision_info
 import torch
-from typing import List, Dict
+from typing import List, Dict, Union
 import base64
 import requests
 from PIL import Image
@@ -15,31 +15,41 @@ 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):
+def process_image(image_data: str) -> Image.Image:
     if image_data.startswith("http://") or image_data.startswith("https://"):
         response = requests.get(image_data)
+        response.raise_for_status()  # Adiciona verificação de erro na requisição
         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
+    else:  # Assume que é um caminho de arquivo local
         img = Image.open(image_data)
     return img
 
 @app.post("/predict")
-async def predict(messages: List[Dict] = Body(...)):
+async def predict(messages: List[Dict[str, Union[str, List[Dict[str, str]]]]] = Body(...)):
     # Processamento e inferência
     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"]))
-    
+        content = message.get("content")
+        if isinstance(content, str):
+            texts.append(processor.apply_chat_template(content, tokenize=False, add_generation_prompt=True))
+        elif isinstance(content, list):
+            for item in content:
+                if isinstance(item, dict) and "type" in item:
+                    if item["type"] == "text":
+                        texts.append(processor.apply_chat_template(item["text"], tokenize=False, add_generation_prompt=True))
+                    elif item["type"] == "image":
+                        image_inputs.append(process_image(item["image"]))
+                else:
+                    raise ValueError(f"Formato inválido para o item: {item}")
+        else:
+            raise ValueError(f"Formato inválido para o conteúdo: {content}")
+
     inputs = processor(
         text=texts,
         images=image_inputs,