Update app.py

app.py

@@ -12,18 +12,18 @@ import uvicorn
 DEVICE = "cpu"                      # Use CPU for compatibility
 RESIZE_DIM = (512, 512)             # Resize images to this resolution
 MAX_IMAGE_SIZE = 10 * 1024 * 1024   # 10MB max image size
+TASK = "<MORE_DETAILED_CAPTION>"    # Hardcoded task
 
 # ===== FastAPI App =====
 app = FastAPI(
     title="Florence-2 Image Analysis API",
-    description="Analyze images using Microsoft's Florence-2 model",
+    description="Analyze images using Microsoft's Florence-2 model with detailed captions",
     version="1.0.0"
 )
 
 # ===== Request/Response Models =====
 class ImageAnalysisRequest(BaseModel):
     image_url: HttpUrl
-    task: str = "<MORE_DETAILED_CAPTION>"  # Default task
 
 class ImageAnalysisResponse(BaseModel):
     caption: str
@@ -38,8 +38,9 @@ try:
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_ID,
         trust_remote_code=True,
-        attn_implementation="eager"
-    ).to(DEVICE).eval()
+        torch_dtype=torch.float32,
+        device_map="auto"
+    ).eval()
     print("[INFO] Model loaded successfully!")
 except Exception as e:
     print(f"[ERROR] Failed to load model: {e}")
@@ -75,8 +76,8 @@ def download_image(url: str) -> Image.Image:
     except Exception as e:
         raise ValueError(f"Failed to process image: {e}")
 
-def analyze_image(image: Image.Image, task: str = "<MORE_DETAILED_CAPTION>") -> str:
-    """Analyze image using Florence-2 model"""
+def analyze_image(image: Image.Image) -> str:
+    """Analyze image using Florence-2 model with hardcoded task"""
     if not processor or not model:
         raise ValueError("Model not loaded properly")
     
@@ -84,9 +85,9 @@ def analyze_image(image: Image.Image, task: str = "<MORE_DETAILED_CAPTION>") ->
         # Resize image for faster processing
         image = image.resize(RESIZE_DIM, Image.BILINEAR)
 
-        # Prepare inputs
+        # Prepare inputs with hardcoded task
         inputs = processor(
-            text=task,
+            text=TASK,
             images=image,
             return_tensors="pt"
         ).to(DEVICE)
@@ -101,10 +102,14 @@ def analyze_image(image: Image.Image, task: str = "<MORE_DETAILED_CAPTION>") ->
                 do_sample=False
             )
 
-        # Decode
+        # Decode and clean output
         generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-        print(f"[DEBUG] Generated text: {generated_text}")
-        return generated_text.strip()
+        
+        # Remove the task prompt from the beginning if present
+        if generated_text.startswith(TASK):
+            generated_text = generated_text[len(TASK):].strip()
+            
+        return generated_text
 
     except Exception as e:
         print(f"[ERROR] Exception in analyze_image: {e}")
@@ -117,7 +122,8 @@ async def root():
     return {
         "message": "Florence-2 Image Analysis API",
         "status": "running",
-        "model_loaded": processor is not None and model is not None
+        "model_loaded": processor is not None and model is not None,
+        "task": TASK
     }
 
 @app.get("/health")
@@ -127,39 +133,21 @@ async def health_check():
         "status": "healthy" if (processor and model) else "unhealthy",
         "model_loaded": processor is not None and model is not None,
         "device": DEVICE,
-        "available_tasks": [
-            "<MORE_DETAILED_CAPTION>",
-            "<DETAILED_CAPTION>",
-            "<CAPTION>",
-            "<OD>",  # Object Detection
-            "<DENSE_REGION_CAPTION>",
-            "<REGION_PROPOSAL>"
-        ]
+        "task": TASK
     }
 
 @app.post("/analyze", response_model=ImageAnalysisResponse)
 async def analyze_image_endpoint(request: ImageAnalysisRequest):
     """
     Analyze an image from a URL using Florence-2 model
-    
-    Available tasks:
-    - <MORE_DETAILED_CAPTION>: Generate detailed image captions
-    - <DETAILED_CAPTION>: Generate detailed captions
-    - <CAPTION>: Generate basic captions
-    - <OD>: Object detection
-    - <DENSE_REGION_CAPTION>: Dense region captioning
-    - <REGION_PROPOSAL>: Region proposal
+    Always uses <MORE_DETAILED_CAPTION> task for detailed image descriptions
     """
     try:
-        # Validate task
-        valid_tasks = [
-            "<MORE_DETAILED_CAPTION>", "<DETAILED_CAPTION>", "<CAPTION>",
-            "<OD>", "<DENSE_REGION_CAPTION>", "<REGION_PROPOSAL>"
-        ]
-        if request.task not in valid_tasks:
+        # Validate model is loaded
+        if not processor or not model:
             raise HTTPException(
-                status_code=400,
-                detail=f"Invalid task. Available tasks: {valid_tasks}"
+                status_code=503,
+                detail="Model not loaded. Please check server logs."
             )
         
         # Download and process image
@@ -167,15 +155,17 @@ async def analyze_image_endpoint(request: ImageAnalysisRequest):
         image = download_image(request.image_url)
         print(f"[INFO] Image downloaded successfully: {image.size}")
         
-        # Analyze image
-        caption = analyze_image(image, request.task)
-        print(f"[INFO] Analysis complete: {caption}")
+        # Analyze image with hardcoded task
+        caption = analyze_image(image)
+        print(f"[INFO] Analysis complete")
         
         return ImageAnalysisResponse(
             caption=caption,
             success=True
         )
         
+    except HTTPException:
+        raise
     except ValueError as e:
         print(f"[ERROR] ValueError: {e}")
         return ImageAnalysisResponse(
@@ -185,27 +175,35 @@ async def analyze_image_endpoint(request: ImageAnalysisRequest):
         )
     except Exception as e:
         print(f"[ERROR] Unexpected error: {e}")
-        raise HTTPException(status_code=500, detail=f"Internal server error: {e}")
+        return ImageAnalysisResponse(
+            caption="",
+            success=False,
+            error_message=f"Internal server error: {str(e)}"
+        )
 
 @app.get("/analyze")
-async def analyze_image_get(image_url: str, task: str = "<MORE_DETAILED_CAPTION>"):
+async def analyze_image_get(image_url: str):
     """
     GET endpoint for quick image analysis
-    Usage: /analyze?image_url=https://example.com/image.jpg&task=<MORE_DETAILED_CAPTION>
+    Usage: /analyze?image_url=https://example.com/image.jpg
     """
-    request = ImageAnalysisRequest(image_url=image_url, task=task)
-    return await analyze_image_endpoint(request)
+    try:
+        request = ImageAnalysisRequest(image_url=image_url)
+        return await analyze_image_endpoint(request)
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=str(e))
 
 # ===== Main Execution =====
 if __name__ == "__main__":
     port = int(os.getenv("PORT", 7860))
     print(f"[INFO] Starting server on port {port}")
     print(f"[INFO] Model status: {'Loaded' if (processor and model) else 'Failed to load'}")
+    print(f"[INFO] Task: {TASK}")
     print(f"[INFO] API Documentation: http://localhost:{port}/docs")
     
     uvicorn.run(
-        "app:app",
+        app,
         host="0.0.0.0",
         port=port,
-        reload=False  # Set to True for development
+        reload=False
     )