Update app.py

app.py

@@ -1,17 +1,56 @@
-from fastapi import FastAPI, UploadFile
-from transformers import pipeline
+import torch
+from transformers import AutoProcessor, AutoModelForVision2Seq
+from fastapi import FastAPI, File, UploadFile, Form
+from fastapi.responses import JSONResponse
+from io import BytesIO
 from PIL import Image
-import io
-import uvicorn
 
-app = FastAPI()
-pipe = pipeline("image-classification", model="google/vit-base-patch16-224")
+# --- 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,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    trust_remote_code=True,
+)
 
+app = FastAPI(title="Aloe Vision 7B AR API")
+
+# --- Inference endpoint ---
 @app.post("/predict")
-async def predict(file: UploadFile):
-    img = Image.open(io.BytesIO(await file.read()))
-    preds = pipe(img)
-    return preds
+async def predict(
+    file: UploadFile = File(...),
+    question: str = Form("What do you see?")
+):
+    try:
+        image = Image.open(BytesIO(await file.read())).convert("RGB")
+
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image", "image": image},
+                    {"type": "text", "text": question},
+                ],
+            }
+        ]
+
+        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)
+
+        generated = model.generate(
+            **inputs,
+            max_new_tokens=256,
+            do_sample=False,
+            eos_token_id=processor.tokenizer.eos_token_id,
+        )
+
+        output_text = processor.batch_decode(generated, skip_special_tokens=True)[0]
+        answer = output_text.split(text)[-1].strip()
+
+        return JSONResponse({"answer": answer})
 
-if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    except Exception as e:
+        return JSONResponse({"error": str(e)}, status_code=500)