Update app.py

app.py

@@ -5,16 +5,25 @@ from typing import Dict
 from PIL import Image
 from io import BytesIO
 import torch
-from transformers import AutoModelForCausalLM, AutoProcessor
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoProcessor
+from transformers import PretrainedConfig
 from fastapi import FastAPI, File, UploadFile
 from fastapi.responses import JSONResponse
 import uvicorn
 
-# Disable SDPA if not supported 
+# Configure PyTorch settings
+torch.backends.cuda.enable_flash_sdp(False)
+torch.backends.cuda.enable_math_sdp(False)
+torch.backends.cuda.enable_mem_efficient_sdp(False)
 
 # ==== CONFIGURATION ====
+class SimpleFlorenceConfig(PretrainedConfig):
+    model_type = "florence"
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
 # Florence-2 Configuration
-MODEL_ID = "microsoft/Florence-2-large"
+MODEL_ID = "microsoft/Florence-2-base"  # Using base model for stability
 DEVICE = "cpu"  # Using CPU instead of GPU
 
 # Create FastAPI app
@@ -36,11 +45,27 @@ def load_florence_model():
         try:
             log_message("[*] Loading Florence-2 model and processor...")
             
-            # Load model on CPU
-            model = AutoModelForCausalLM.from_pretrained(MODEL_ID, trust_remote_code=True).to(DEVICE)
+            # Create simple configuration
+            config = SimpleFlorenceConfig()
+            config.use_flash_attention = False
+            config.architectures = ["Florence2ForConditionalGeneration"]
+            
+            # Load model with simplified config
+            model = AutoModelForCausalLM.from_pretrained(
+                MODEL_ID,
+                config=config,
+                trust_remote_code=True,
+                torch_dtype=torch.float32,
+                use_flash_attention_2=False,
+                revision="main"  # Explicitly use main branch
+            ).to(DEVICE)
             model.eval()
             
-            processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
+            processor = AutoProcessor.from_pretrained(
+                MODEL_ID,
+                trust_remote_code=True,
+                revision="main"  # Explicitly use main branch
+            )
             log_message("[ ] Florence-2 loaded and ready on CPU")
         except Exception as e:
             log_message(f"[ERROR] Failed to load Florence-2 model: {e}")