Update app.py

app.py

@@ -1,84 +1,78 @@
-import os
 import io
+import os
 import torch
 from PIL import Image
-from fastapi import FastAPI, File, UploadFile
+from fastapi import FastAPI, UploadFile, File
 from fastapi.responses import JSONResponse
 from transformers import AutoProcessor, AutoModelForCausalLM
 
-# 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
-    )
-except subprocess.CalledProcessError as e:
-    print(f"Flash-attn install failed: {e}")
-    print("Continuing without flash-attn...")
-
-# Device setup
+# Setup
 device = "cuda" if torch.cuda.is_available() else "cpu"
+app = FastAPI(title="Florence-2 Base Image Captioning API")
 
-# Load Florence-2-base model and processor
+# Load Florence-2 base model
 try:
-    model = AutoModelForCausalLM.from_pretrained(
+    vision_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 Florence-2-base: {e}")
-    model = None
-    processor = None
 
-# FastAPI setup
-app = FastAPI(title="Florence-2 Image Captioning API")
+    vision_processor = AutoProcessor.from_pretrained(
+        'microsoft/Florence-2-base',
+        trust_remote_code=True
+    )
+except Exception as e:
+    vision_model = None
+    vision_processor = None
+    print(f"Model loading error: {e}")
 
 @app.post("/describe-image")
 async def describe_image(file: UploadFile = File(...)):
-    if model is None or processor is None:
+    if vision_model is None or vision_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 file type. Please upload an image."})
-
     try:
-        # Load image from upload
-        image_data = await file.read()
-        image = Image.open(io.BytesIO(image_data)).convert("RGB")
+        contents = await file.read()
+        image = Image.open(io.BytesIO(contents)).convert("RGB")
 
-        # Prepare inputs
-        inputs = processor(
+        # Preprocess
+        inputs = vision_processor(
             text="<MORE_DETAILED_CAPTION>",
             images=image,
             return_tensors="pt"
         ).to(device)
 
-        # Generate caption
         with torch.no_grad():
-            generated_ids = model.generate(
+            generated_ids = vision_model.generate(
                 input_ids=inputs["input_ids"],
                 pixel_values=inputs["pixel_values"],
-                max_new_tokens=512,
-                num_beams=3
+                max_new_tokens=1024,
+                num_beams=3,
             )
 
-        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
-        processed = processor.post_process_generation(
+        generated_text = vision_processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+        processed = vision_processor.post_process_generation(
             generated_text,
             task="<MORE_DETAILED_CAPTION>",
-            image_size=(image.width, image.height)
+            image_size=image.size
         )
-        description = processed["<MORE_DETAILED_CAPTION>"]
-        return {"description": description}
+        caption = processed["<MORE_DETAILED_CAPTION>"]
+
+        return JSONResponse(content={
+            "filename": file.filename,
+            "description": caption
+        })
 
     except Exception as e:
         return JSONResponse(status_code=500, content={"error": str(e)})
 
-@app.get("/health")
-def health():
-    return {"status": "ok", "model": "florence-2-base"}
+@app.get("/")
+def root():
+    return {"message": "Florence-2 Base Image Captioning API is running"}
+
+# Run the app when executed directly
+if __name__ == "__main__":
+    import uvicorn
+    port = int(os.getenv("PORT", 7860))  # Spaces set PORT env var
+    uvicorn.run("app:app", host="0.0.0.0", port=port)