Update app.py

app.py

@@ -5,7 +5,7 @@ from PIL import Image
 from io import BytesIO
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel, HttpUrl
-from transformers import AutoProcessor, AutoModelForCausalLM
+from transformers import AutoProcessor, AutoModelForVision2Seq
 import uvicorn
 
 # ===== CONFIG =====
@@ -31,13 +31,14 @@ class ImageAnalysisResponse(BaseModel):
     error_message: str = None
 
 # ===== Load Florence-2 Base Model =====
-print("[INFO] Loading Florence-2-base model on CPU...")
+print("[INFO] Loading Florence-2 model on CPU...")
 try:
-    processor = AutoProcessor.from_pretrained("microsoft/Florence-2-large", trust_remote_code=True)
-    model = AutoModelForCausalLM.from_pretrained(
-        "microsoft/Florence-2-Large",
+    MODEL_ID = "microsoft/Florence-2-large"
+    processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
+    model = AutoModelForVision2Seq.from_pretrained(
+        MODEL_ID,
         trust_remote_code=True,
-        attn_implementation="eager"
+        torch_dtype=torch.float32
     ).to(DEVICE).eval()
     print("[INFO] Model loaded successfully!")
 except Exception as e:
@@ -82,14 +83,16 @@ def analyze_image(image: Image.Image, task: str = "<MORE_DETAILED_CAPTION>") ->
     try:
         # Resize image for faster processing
         image = image.resize(RESIZE_DIM, Image.BILINEAR)
-        
+
         # Prepare inputs
         inputs = processor(
             text=task,
             images=image,
-            return_tensors="pt"
+            return_tensors="pt",
+            padding=True,
+            truncation=True
         ).to(DEVICE)
-        
+
         # Generate caption
         with torch.no_grad():
             generated_ids = model.generate(
@@ -100,21 +103,25 @@ def analyze_image(image: Image.Image, task: str = "<MORE_DETAILED_CAPTION>") ->
                 do_sample=False,
                 repetition_penalty=1.2     # Helps avoid repetitive outputs
             )
-        
+
         # Decode and post-process
         generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+        print(f"[DEBUG] Generated text: {generated_text}")
         result = processor.post_process_generation(
             generated_text,
             task=task,
             image_size=RESIZE_DIM
         )
-        
+
         if result is None:
+            print("[ERROR] Post-processing returned None. The model may not have generated a valid output for the given task.")
             raise ValueError("Post-processing returned None. The model may not have generated a valid output for the given task.")
 
+        print(f"[DEBUG] Post-processed result: {result}")
         return result.get(task, "No caption generated.")
-        
+
     except Exception as e:
+        print(f"[ERROR] Exception in analyze_image: {e}")
         raise ValueError(f"Failed to analyze image: {e}")
         
 # ===== API Endpoints =====