Update app.py

app.py

@@ -1,12 +1,11 @@
-# app.py - FastAPI BiRefNet Background Remover API
+# app.py - Fixed FastAPI for HF Spaces
 import torch
 import numpy as np
 from PIL import Image
 from safetensors.torch import load_file
 import io
-from typing import Optional
 
-from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi import FastAPI, File, UploadFile, Form, HTTPException, Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import StreamingResponse
 import uvicorn
@@ -14,7 +13,7 @@ import uvicorn
 from birefnet import BiRefNet
 from BiRefNet_config import BiRefNetConfig
 
-# Global model (loaded once)
+# Global model
 device = "cpu"
 config = BiRefNetConfig()
 model = BiRefNet(config)
@@ -44,10 +43,10 @@ def remove_bg(image: Image.Image) -> Image.Image:
 
 app = FastAPI(title="BiRefNet Background Remover API")
 
-# CORS for NextJS/Vercel
+# CORS - allow all for testing (restrict in production)
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Update with your domain in production
+    allow_origins=["*"],
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
@@ -55,32 +54,47 @@ app.add_middleware(
 
 @app.get("/")
 async def root():
-    return {"message": "BiRefNet Background Remover API", "status": "ready"}
+    return {"message": "BiRefNet Background Remover API ✅", "endpoints": ["/remove-bg"]}
 
+# FIXED: Accept both UploadFile and raw bytes
 @app.post("/remove-bg")
-async def remove_background(
-    file: UploadFile = File(..., description="Image file (PNG/JPG)")
-):
-    if not file.content_type.startswith("image/"):
-        raise HTTPException(400, detail="File must be an image")
-    
+async def remove_background(request: Request, file: UploadFile = File(None)):
     try:
-        # Read and process image
-        contents = await file.read()
-        image = Image.open(io.BytesIO(contents))
+        # Handle different content types
+        if file is None:
+            # Fallback: read raw body if no file param
+            body = await request.body()
+            if not body:
+                raise HTTPException(400, detail="No image data received")
+            image = Image.open(io.BytesIO(body))
+        else:
+            contents = await file.read()
+            if not contents:
+                raise HTTPException(400, detail="Empty file")
+            image = Image.open(io.BytesIO(contents))
+        
+        # Validate image
+        if not image.format or image.format not in ["JPEG", "PNG", "JPG"]:
+            raise HTTPException(400, detail="Invalid image format")
+        
+        # Process
         result = remove_bg(image)
         
-        # Save to bytes
+        # Return PNG
         img_byte_arr = io.BytesIO()
         result.save(img_byte_arr, format="PNG")
         img_byte_arr.seek(0)
         
         return StreamingResponse(
-            img_byte_arr, 
+            img_byte_arr,
             media_type="image/png",
-            headers={"Content-Disposition": "inline; filename=removed-bg.png"}
+            headers={
+                "Content-Disposition": "inline; filename=removed-bg.png",
+                "Cache-Control": "public, max-age=3600"
+            }
         )
     except Exception as e:
+        print(f"❌ Error: {str(e)}")
         raise HTTPException(500, detail=f"Processing failed: {str(e)}")
 
 if __name__ == "__main__":