Create app.py

app.py

@@ -0,0 +1,127 @@
+import os
+import json
+import time
+from typing import Dict
+from PIL import Image
+from io import BytesIO
+import torch
+from transformers import AutoModelForCausalLM, AutoProcessor
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import JSONResponse
+import uvicorn
+
+# Disable SDPA if not supported 
+
+# ==== CONFIGURATION ====
+# Florence-2 Configuration
+MODEL_ID = "microsoft/Florence-2-large"
+DEVICE = "cpu"  # Using CPU instead of GPU
+
+# Create FastAPI app
+app = FastAPI(title="Florence-2 Image Captioning API")
+
+# Florence-2 Model (will be loaded once)
+model = None
+processor = None
+
+def log_message(message: str):
+    """Simple logging function"""
+    timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
+    print(f"[{timestamp}] {message}")
+
+def load_florence_model():
+    """Load Florence-2 model and processor"""
+    global model, processor
+    if model is None or processor is None:
+        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)
+            model.eval()
+            
+            processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
+            log_message("[ ] Florence-2 loaded and ready on CPU")
+        except Exception as e:
+            log_message(f"[ERROR] Failed to load Florence-2 model: {e}")
+            raise
+
+def caption_image(image: Image.Image) -> str:
+    """Generate detailed caption for an image using Florence-2"""
+    if model is None or processor is None:
+        return "Model not loaded."
+    
+    task_prompt = "<MORE_DETAILED_CAPTION>"
+    prompt = task_prompt
+
+    try:
+        # Process image
+        inputs = processor(
+            text=prompt,
+            images=image,
+            return_tensors="pt",
+            padding=True,
+            truncation=True
+        ).to(DEVICE)
+        
+        with torch.no_grad():
+            generated_ids = model.generate(
+                input_ids=inputs["input_ids"],
+                pixel_values=inputs["pixel_values"],
+                max_new_tokens=1350,
+                do_sample=True,
+                temperature=0.7,
+                top_p=0.9,
+                num_beams=3,
+            )
+
+        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
+        return generated_text
+
+    except Exception as e:
+        log_message(f"[!] Caption generation failed: {e}")
+        return "Captioning error."
+
+@app.on_event("startup")
+async def startup_event():
+    """Load model on startup"""
+    load_florence_model()
+
+@app.post("/caption")
+async def create_caption(file: UploadFile = File(...)) -> Dict:
+    """
+    API endpoint to receive an image and return its caption
+    """
+    try:
+        log_message(f"[API] Received image: {file.filename}")
+        
+        # Read and validate image
+        contents = await file.read()
+        image = Image.open(BytesIO(contents)).convert("RGB")
+        
+        # Generate caption
+        log_message(f"[API] Generating caption for {file.filename}")
+        caption = caption_image(image)
+        
+        log_message(f"[API] Caption generated for {file.filename}: {caption[:100]}...")
+        
+        return {
+            "status": "success",
+            "filename": file.filename,
+            "caption": caption
+        }
+        
+    except Exception as e:
+        error_msg = f"Error processing image: {str(e)}"
+        log_message(f"[ERROR] {error_msg}")
+        return JSONResponse(
+            status_code=500,
+            content={
+                "status": "error",
+                "message": error_msg
+            }
+        )
+
+if __name__ == "__main__":
+    log_message("Starting Florence-2 Vision Analysis API Server")
+    uvicorn.run(app, host="0.0.0.0", port=8000)