Update app.py

app.py

@@ -5,7 +5,7 @@ from typing import Dict
 from PIL import Image
 from io import BytesIO
 import torch
-from transformers import AutoModelForCausalLM, AutoProcessor
+from transformers import AutoModelForVision2Seq, AutoProcessor
 from fastapi import FastAPI, File, UploadFile
 from fastapi.responses import JSONResponse
 import uvicorn
@@ -39,15 +39,21 @@ def load_florence_model():
         try:
             log_message("[*] Loading Florence-2 model and processor...")
             
-            # Load base model
-            base_model = AutoModelForCausalLM.from_pretrained(MODEL_ID, trust_remote_code=True)
+            # Load model and processor with specific configuration
+            processor = AutoProcessor.from_pretrained(
+                MODEL_ID,
+                trust_remote_code=True,
+                revision="9a515b7",  # Pin to a specific version
+            )
+            
+            model = AutoModelForVision2Seq.from_pretrained(
+                MODEL_ID,
+                trust_remote_code=True,
+                revision="9a515b7",  # Pin to a specific version
+                torch_dtype=torch.float32,
+            ).to(DEVICE)
             
-            # Load directly to CPU
-            model = base_model.to(DEVICE)
-            log_message("[ ] Model loaded on CPU")
-                
             model.eval()
-            processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
             log_message("[ ] Florence-2 loaded and ready.")
         except Exception as e:
             log_message(f"[ERROR] Failed to load Florence-2 model: {e}")