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, AutoModelForVision2Seq
+from transformers import AutoProcessor, AutoModelForCausalLM
 import uvicorn
 
 # ===== CONFIG =====
@@ -35,10 +35,9 @@ print("[INFO] Loading Florence-2 model on CPU...")
 try:
     MODEL_ID = "microsoft/Florence-2-large"
     processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
-    model = AutoModelForVision2Seq.from_pretrained(
+    model = AutoModelForCausalLM.from_pretrained(
         MODEL_ID,
-        trust_remote_code=True,
-        torch_dtype=torch.float32
+        trust_remote_code=True
     ).to(DEVICE).eval()
     print("[INFO] Model loaded successfully!")
 except Exception as e:
@@ -88,9 +87,7 @@ def analyze_image(image: Image.Image, task: str = "<MORE_DETAILED_CAPTION>") ->
         inputs = processor(
             text=task,
             images=image,
-            return_tensors="pt",
-            padding=True,
-            truncation=True
+            return_tensors="pt"
         ).to(DEVICE)
 
         # Generate caption
@@ -98,27 +95,15 @@ def analyze_image(image: Image.Image, task: str = "<MORE_DETAILED_CAPTION>") ->
             generated_ids = model.generate(
                 input_ids=inputs["input_ids"],
                 pixel_values=inputs["pixel_values"],
-                max_new_tokens=1024,       # Maximum token length
+                max_new_tokens=1024,
                 num_beams=3,
-                do_sample=False,
-                repetition_penalty=1.2     # Helps avoid repetitive outputs
+                do_sample=False
             )
 
-        # Decode and post-process
+        # Decode
         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.")
+        return generated_text.strip()
 
     except Exception as e:
         print(f"[ERROR] Exception in analyze_image: {e}")