Update app.py

app.py

@@ -1,10 +1,49 @@
-from fastapi import FastAPI, File, UploadFile, Form, Request, HTTPException
-from fastapi.middleware.cors import CORSMiddleware
-import io
+# 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 fastapi import FastAPI, File, UploadFile, Form, HTTPException, Request
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse
+import uvicorn
+
+from birefnet import BiRefNet
+from BiRefNet_config import BiRefNetConfig
+
+# Global model
+device = "cpu"
+config = BiRefNetConfig()
+model = BiRefNet(config)
+state_dict = load_file("model.safetensors")
+model.load_state_dict(state_dict, strict=False)
+model.to(device)
+model.eval()
+print("✅ BiRefNet Lite loaded")
 
-app = FastAPI()
+def preprocess(img: Image.Image):
+    img = img.convert("RGB").resize((1024, 1024))
+    arr = np.array(img).astype(np.float32) / 255.0
+    arr = arr.transpose(2, 0, 1)
+    return torch.from_numpy(arr).unsqueeze(0)
 
+@torch.no_grad()
+def remove_bg(image: Image.Image) -> Image.Image:
+    x = preprocess(image).to(device)
+    pred = model(x)[0]
+    pred = torch.sigmoid(pred)
+    mask = pred.squeeze().cpu().numpy()
+    mask = (mask * 255).astype(np.uint8)
+    mask = Image.fromarray(mask).resize(image.size)
+    out = image.convert("RGBA")
+    out.putalpha(mask)
+    return out
+
+app = FastAPI(title="BiRefNet Background Remover API")
+
+# CORS - allow all for testing (restrict in production)
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -13,43 +52,50 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Your model loading code here...
-print("✅ Model loaded")
+@app.get("/")
+async def root():
+    return {"message": "BiRefNet Background Remover API ✅", "endpoints": ["/remove-bg"]}
 
+# FIXED: Accept both UploadFile and raw bytes
 @app.post("/remove-bg")
-async def remove_bg_endpoint(request: Request):
-    """Accepts ANY form field name - "file", "image", etc."""
+async def remove_background(request: Request, file: UploadFile = File(None)):
     try:
-        form = await request.form()
-        print("📋 All form fields:", list(form.keys()))
-        
-        # Try common field names
-        file_data = None
-        for field_name in ['file', 'image', 'data', 'upload']:
-            if field_name in form:
-                file_data = await form[field_name].read()
-                print(f"✅ Found {field_name}: {len(file_data)} bytes")
-                break
+        # 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))
         
-        if not file_data:
-            raise HTTPException(400, "No image file found. Send as 'file' or 'image'")
+        # Validate image
+        if not image.format or image.format not in ["JPEG", "PNG", "JPG"]:
+            raise HTTPException(400, detail="Invalid image format")
         
-        image = Image.open(io.BytesIO(file_data))
-        print(f"✅ Image: {image.size}")
+        # Process
+        result = remove_bg(image)
         
-        # Your remove_bg(image) function here
-        result = remove_bg(image)  # Your existing function
-        
-        img_bytes = io.BytesIO()
-        result.save(img_bytes, "PNG")
-        img_bytes.seek(0)
-        
-        return StreamingResponse(img_bytes, media_type="image/png")
+        # Return PNG
+        img_byte_arr = io.BytesIO()
+        result.save(img_byte_arr, format="PNG")
+        img_byte_arr.seek(0)
         
+        return StreamingResponse(
+            img_byte_arr,
+            media_type="image/png",
+            headers={
+                "Content-Disposition": "inline; filename=removed-bg.png",
+                "Cache-Control": "public, max-age=3600"
+            }
+        )
     except Exception as e:
-        print(f"❌ Error: {e}")
-        raise HTTPException(400, f"Failed: {str(e)}")
+        print(f"❌ Error: {str(e)}")
+        raise HTTPException(500, detail=f"Processing failed: {str(e)}")
 
-@app.get("/")
-def root():
-    return {"status": "OK", "endpoint": "/remove-bg POST with file=image or file=file"}
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860)