Update app.py

app.py

@@ -1,93 +1,84 @@
-from fastapi import FastAPI, File, UploadFile, Form
-from fastapi.responses import JSONResponse
-from PIL import Image
-import torch
+import os
 import io
+import torch
+from PIL import Image
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import JSONResponse
 from transformers import AutoProcessor, AutoModelForCausalLM
-import subprocess
 
-# Attempt to install flash-attn (if needed)
+# Auto-install flash-attn if needed
+import subprocess
 try:
-    subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, check=True, shell=True)
+    subprocess.run(
+        'pip install flash-attn --no-build-isolation',
+        env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"},
+        check=True,
+        shell=True
+    )
 except subprocess.CalledProcessError as e:
-    print(f"Error installing flash-attn: {e}")
-    print("Continuing without flash-attn.")
+    print(f"Flash-attn install failed: {e}")
+    print("Continuing without flash-attn...")
 
-# Determine device
+# Device setup
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load Florence-2 Base
-try:
-    vision_language_model_base = AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True, attn_implementation="eager").to(device).eval()
-    vision_language_processor_base = AutoProcessor.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True)
-except Exception as e:
-    print(f"Error loading base model: {e}")
-    vision_language_model_base = None
-    vision_language_processor_base = None
-
-# Load Florence-2 Large
+# Load Florence-2-base 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)
+    model = AutoModelForCausalLM.from_pretrained(
+        'microsoft/Florence-2-base',
+        trust_remote_code=True,
+        attn_implementation="eager"
+    ).to(device).eval()
+    processor = AutoProcessor.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True)
 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-base: {e}")
+    model = None
+    processor = None
 
-# Initialize FastAPI
-app = FastAPI()
+# FastAPI setup
+app = FastAPI(title="Florence-2 Image Captioning API")
 
 @app.post("/describe-image")
-async def describe_image(
-    file: UploadFile = File(...),
-    model_choice: str = Form("Base")
-):
+async def describe_image(file: UploadFile = File(...)):
+    if model is None or processor is None:
+        return JSONResponse(status_code=500, content={"error": "Model not loaded"})
+
     if not file.filename.lower().endswith((".jpg", ".jpeg", ".png")):
-        return JSONResponse(status_code=400, content={"error": "Invalid image file type."})
+        return JSONResponse(status_code=400, content={"error": "Invalid file type. Please upload an image."})
 
     try:
-        image_bytes = await file.read()
-        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-    except Exception as e:
-        return JSONResponse(status_code=400, content={"error": f"Failed to process image: {str(e)}"})
+        # Load image from upload
+        image_data = await file.read()
+        image = Image.open(io.BytesIO(image_data)).convert("RGB")
 
-    if model_choice == "Base":
-        if vision_language_model_base is None:
-            return JSONResponse(status_code=500, content={"error": "Base model not loaded."})
-        model = vision_language_model_base
-        processor = vision_language_processor_base
-    elif model_choice == "Large":
-        if vision_language_model_large is None:
-            return JSONResponse(status_code=500, content={"error": "Large model not loaded."})
-        model = vision_language_model_large
-        processor = vision_language_processor_large
-    else:
-        return JSONResponse(status_code=400, content={"error": "Invalid model choice."})
+        # Prepare inputs
+        inputs = processor(
+            text="<MORE_DETAILED_CAPTION>",
+            images=image,
+            return_tensors="pt"
+        ).to(device)
 
-    try:
-        inputs = processor(text="<MORE_DETAILED_CAPTION>", images=image, return_tensors="pt").to(device)
+        # Generate caption
         with torch.no_grad():
             generated_ids = model.generate(
                 input_ids=inputs["input_ids"],
                 pixel_values=inputs["pixel_values"],
-                max_new_tokens=1024,
-                early_stopping=False,
-                do_sample=False,
-                num_beams=3,
+                max_new_tokens=512,
+                num_beams=3
             )
 
         generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
-        processed_description = processor.post_process_generation(
+        processed = processor.post_process_generation(
             generated_text,
             task="<MORE_DETAILED_CAPTION>",
             image_size=(image.width, image.height)
         )
-        image_description = processed_description["<MORE_DETAILED_CAPTION>"]
-        return JSONResponse(content={"description": image_description})
+        description = processed["<MORE_DETAILED_CAPTION>"]
+        return {"description": description}
 
     except Exception as e:
-        return JSONResponse(status_code=500, content={"error": f"Image processing failed: {str(e)}"})
+        return JSONResponse(status_code=500, content={"error": str(e)})
 
 @app.get("/health")
 def health():
-    return {"status": "ok", "device": device}
+    return {"status": "ok", "model": "florence-2-base"}