Update app.py

app.py

@@ -5,8 +5,8 @@ from fastapi.responses import JSONResponse
 from io import BytesIO
 from PIL import Image
 
-# --- Load model and processor ---
 model_id = "HPAI-BSC/Aloe-Vision-7B-AR"
+
 processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
 model = AutoModelForVision2Seq.from_pretrained(
     model_id,
@@ -17,25 +17,39 @@ model = AutoModelForVision2Seq.from_pretrained(
 
 app = FastAPI(title="Aloe Vision 7B AR API")
 
-# --- Inference endpoint ---
 @app.post("/predict")
 async def predict(
-    file: UploadFile = File(...),
-    question: str = Form("What do you see?")
+    file: UploadFile = File(None),
+    question: str = Form(None)
 ):
     try:
-        image = Image.open(BytesIO(await file.read())).convert("RGB")
+        # --- Case 1: both image and text ---
+        if file and question:
+            image = Image.open(BytesIO(await file.read())).convert("RGB")
+            messages = [
+                {"role": "user", "content": [
+                    {"type": "image", "image": image},
+                    {"type": "text", "text": question}
+                ]}
+            ]
+
+        # --- Case 2: text only ---
+        elif question and not file:
+            messages = [{"role": "user", "content": [{"type": "text", "text": question}]}]
 
-        messages = [
-            {
-                "role": "user",
-                "content": [
+        # --- Case 3: image only ---
+        elif file and not question:
+            image = Image.open(BytesIO(await file.read())).convert("RGB")
+            messages = [
+                {"role": "user", "content": [
                     {"type": "image", "image": image},
-                    {"type": "text", "text": question},
-                ],
-            }
-        ]
+                    {"type": "text", "text": "Describe this image briefly."}
+                ]}
+            ]
+        else:
+            return JSONResponse({"error": "You must provide an image, text, or both."}, status_code=400)
 
+        # --- Process ---
         text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
         image_inputs = processor.process_vision_info(messages)
         inputs = processor(text=[text], **image_inputs, return_tensors="pt").to(model.device)