Update app.py

app.py

@@ -34,18 +34,43 @@ class ImageAnalysisResponse(BaseModel):
 print("[INFO] Loading Florence-2 model on CPU...")
 try:
     MODEL_ID = "microsoft/Florence-2-large"
+    
+    # Load processor
     processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
+    
+    # Load model with specific parameters to avoid SDPA issues
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_ID,
         trust_remote_code=True,
         torch_dtype=torch.float32,
-        device_map="auto"
-    ).eval()
+        attn_implementation="eager",  # Force eager attention to avoid SDPA issues
+        device_map=None  # Explicitly set to None for CPU
+    )
+    
+    # Move to device manually
+    model = model.to(DEVICE)
+    model.eval()
+    
     print("[INFO] Model loaded successfully!")
 except Exception as e:
     print(f"[ERROR] Failed to load model: {e}")
-    processor = None
-    model = None
+    # Try fallback to base model if large fails
+    try:
+        print("[INFO] Trying Florence-2-base as fallback...")
+        MODEL_ID = "microsoft/Florence-2-base"
+        processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_ID,
+            trust_remote_code=True,
+            torch_dtype=torch.float32,
+            attn_implementation="eager",
+            device_map=None
+        ).to(DEVICE).eval()
+        print("[INFO] Fallback model loaded successfully!")
+    except Exception as fallback_error:
+        print(f"[ERROR] Fallback also failed: {fallback_error}")
+        processor = None
+        model = None
 
 # ===== Helper Functions =====
 def download_image(url: str) -> Image.Image:
@@ -92,14 +117,16 @@ def analyze_image(image: Image.Image) -> str:
             return_tensors="pt"
         ).to(DEVICE)
 
-        # Generate caption
+        # Generate caption with error handling
         with torch.no_grad():
             generated_ids = model.generate(
                 input_ids=inputs["input_ids"],
                 pixel_values=inputs["pixel_values"],
-                max_new_tokens=1024,
+                max_new_tokens=256,  # Reduced for stability
                 num_beams=3,
-                do_sample=False
+                do_sample=False,
+                early_stopping=True,
+                pad_token_id=processor.tokenizer.eos_token_id
             )
 
         # Decode and clean output