Update app.py

app.py

@@ -29,25 +29,15 @@ except subprocess.CalledProcessError as e:
 # Determine the device to use
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load the base model and processor
+# Load Florence-2-large model and processor
 try:
-    vision_language_model_base = AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True).to(device).eval()
-    vision_language_processor_base = AutoProcessor.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True)
-    print("✓ Base model loaded successfully")
+    vision_language_model = AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-large', trust_remote_code=True).to(device).eval()
+    vision_language_processor = AutoProcessor.from_pretrained('microsoft/Florence-2-large', trust_remote_code=True)
+    print("✓ Florence-2-large model loaded successfully")
 except Exception as e:
-    print(f"Error loading base model: {e}")
-    vision_language_model_base = None
-    vision_language_processor_base = None
-
-# Load the large model and processor
-try:
-    vision_language_model_large = AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-large', trust_remote_code=True).to(device).eval()
-    vision_language_processor_large = AutoProcessor.from_pretrained('microsoft/Florence-2-large', trust_remote_code=True)
-    print("✓ Large model loaded successfully")
-except Exception as e:
-    print(f"Error loading large model: {e}")
-    vision_language_model_large = None
-    vision_language_processor_large = None
+    print(f"Error loading Florence-2-large model: {e}")
+    vision_language_model = None
+    vision_language_processor = None
 
 def load_image_from_url(image_url):
     """Load an image from a URL."""
@@ -88,47 +78,32 @@ def describe_image(uploaded_image, model_choice):
     if uploaded_image is None:
         return "Please upload an image."
 
-    if model_choice == "Florence-2-base":
-        if vision_language_model_base is None:
-            return "Base model failed to load."
-        model = vision_language_model_base
-        processor = vision_language_processor_base
-    elif model_choice == "Florence-2-large":
-        if vision_language_model_large is None:
-            return "Large model failed to load."
-        model = vision_language_model_large
-        processor = vision_language_processor_large
-    else:
-        return "Invalid model choice."
+    if vision_language_model is None:
+        return "Florence-2-large model failed to load."
+    
+    model = vision_language_model
+    processor = vision_language_processor
 
     try:
         return process_image_description(model, processor, uploaded_image)
     except Exception as e:
         return f"Error generating caption: {str(e)}"
 
-def describe_image_from_url(image_url, model_choice):
+def describe_image_from_url(image_url, model_choice=None):
     """Generate description from image URL."""
     try:
         if not image_url:
             return {"error": "image_url is required"}
         
-        if model_choice not in ["Florence-2-base", "Florence-2-large"]:
-            return {"error": "Invalid model choice. Use 'Florence-2-base' or 'Florence-2-large'"}
-        
+        if vision_language_model is None:
+            return {"error": "Florence-2-large model not available"}
+            
         # Load image from URL
         image = load_image_from_url(image_url)
         
-        # Select model and processor
-        if model_choice == "Florence-2-base":
-            if vision_language_model_base is None:
-                return {"error": "Base model not available"}
-            model = vision_language_model_base
-            processor = vision_language_processor_base
-        else:
-            if vision_language_model_large is None:
-                return {"error": "Large model not available"}
-            model = vision_language_model_large
-            processor = vision_language_processor_large
+        # Use the loaded large model
+        model = vision_language_model
+        processor = vision_language_processor
         
         # Generate caption
         caption = process_image_description(model, processor, image)
@@ -147,7 +122,7 @@ def describe_image_from_url(image_url, model_choice):
 IMAGE_SERVER_BASE = os.getenv("IMAGE_SERVER_BASE", " ")
 DATA_COLLECTION_BASE = os.getenv("DATA_COLLECTION_BASE", "https://fred808-flow.hf.space")
 REQUESTER_ID = os.getenv("FLO_REQUESTER_ID", f"florence-2-{os.getpid()}")
-MODEL_CHOICE = os.getenv("FLO_MODEL_CHOICE", "Florence-2-base")
+MODEL_CHOICE = "Florence-2-large"  # Always use large model
 
 
 def sanitize_name(name: str, max_len: int = 200) -> str:
@@ -259,9 +234,7 @@ def background_worker():
     # Wait for model to be ready
     waited = 0
     while waited < 120:
-        if MODEL_CHOICE == "Florence-2-base" and vision_language_model_base:
-            break
-        elif MODEL_CHOICE == "Florence-2-large" and vision_language_model_large:
+        if vision_language_model is not None:
             break
         time.sleep(1)
         waited += 1
@@ -344,12 +317,8 @@ def background_worker():
                         try:
                             pil_img = Image.open(BytesIO(img_bytes)).convert('RGB')
                             
-                            if MODEL_CHOICE == "Florence-2-base":
-                                model = vision_language_model_base
-                                processor = vision_language_processor_base
-                            else:
-                                model = vision_language_model_large
-                                processor = vision_language_processor_large
+                            model = vision_language_model
+                            processor = vision_language_processor
                             
                             print(f"[BACKGROUND] Generating caption for {filename}")
                             caption = process_image_description(model, processor, pil_img)
@@ -410,8 +379,8 @@ async def root():
     return {
         "name": "Florence-2 Image Captioning Server",
         "status": "running",
-        "model_base": vision_language_model_base is not None,
-        "model_large": vision_language_model_large is not None,
+        "model": "Florence-2-large",
+        "model_loaded": vision_language_model is not None,
         "device": device
     }
 
@@ -419,21 +388,12 @@ async def root():
 async def health():
     return {
         "status": "healthy",
-        "model_base": vision_language_model_base is not None,
-        "model_large": vision_language_model_large is not None,
+        "model": "Florence-2-large",
+        "model_loaded": vision_language_model is not None,
         "device": device,
         "model_choice": MODEL_CHOICE
     }
 
-# Start background worker thread (daemon) so it doesn't block shutdown
-def _start_worker_thread():
-    t = threading.Thread(target=background_worker, daemon=True)
-    t.start()
-
-# Start background worker when FastAPI starts
-@app.on_event("startup")
-async def startup_event():
-    _start_worker_thread()
 
 
 @app.get("/analyze")
@@ -459,7 +419,6 @@ async def analyze_get(image_url: str = None, model_choice: str = None):
 async def analyze_post(file: UploadFile = File(None), model_choice: str = Form(None)):
     """Analyze an uploaded image (multipart/form-data). Returns caption JSON."""
     try:
-        mc = model_choice or MODEL_CHOICE
         if file is None:
             raise HTTPException(status_code=400, detail="file is required")
 
@@ -469,17 +428,11 @@ async def analyze_post(file: UploadFile = File(None), model_choice: str = Form(N
         except Exception as e:
             raise HTTPException(status_code=400, detail=f"Failed to read uploaded image: {e}")
 
-        # Choose model
-        if mc == "Florence-2-large":
-            if vision_language_model_large is None:
-                raise HTTPException(status_code=503, detail="Base model not loaded")
-            model = vision_language_model_large
-            processor = vision_language_processor_large
-        else:
-            if vision_language_model_large is None:
-                raise HTTPException(status_code=503, detail="Large model not loaded")
-            model = vision_language_model_large
-            processor = vision_language_processor_large
+        if vision_language_model is None:
+            raise HTTPException(status_code=503, detail="Florence-2-large model not loaded")
+        
+        model = vision_language_model
+        processor = vision_language_processor
 
         caption = process_image_description(model, processor, pil_img)
         return JSONResponse(content={"success": True, "caption": caption})