Update app.py

app.py

@@ -1,49 +1,10 @@
-# 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 import FastAPI, File, UploadFile, Form, Request, HTTPException
 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")
-
-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
+import io
+from PIL import Image
 
-app = FastAPI(title="BiRefNet Background Remover API")
+app = FastAPI()
 
-# CORS - allow all for testing (restrict in production)
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -52,50 +13,43 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-@app.get("/")
-async def root():
-    return {"message": "BiRefNet Background Remover API ✅", "endpoints": ["/remove-bg"]}
+# Your model loading code here...
+print("✅ Model loaded")
 
-# FIXED: Accept both UploadFile and raw bytes
 @app.post("/remove-bg")
-async def remove_background(
-    file: UploadFile = File(None),
-    request: Request = None
-):
+async def remove_bg_endpoint(request: Request):
+    """Accepts ANY form field name - "file", "image", etc."""
     try:
-        contents = None
+        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
         
-        # Try UploadFile first
-        if file:
-            contents = await file.read()
-            print(f"✅ File received: {file.filename}, {len(contents)} bytes")
-        else:
-            # Fallback: check raw body
-            body = await request.body()
-            contents = body
-            print(f"✅ Raw body received: {len(contents)} bytes")
+        if not file_data:
+            raise HTTPException(400, "No image file found. Send as 'file' or 'image'")
         
-        if not contents or len(contents) == 0:
-            raise HTTPException(400, detail="No file uploaded")
+        image = Image.open(io.BytesIO(file_data))
+        print(f"✅ Image: {image.size}")
         
-        image = Image.open(io.BytesIO(contents))
-        print(f"✅ Image loaded: {image.size}, format: {image.format}")
+        # Your remove_bg(image) function here
+        result = remove_bg(image)  # Your existing function
         
-        result = remove_bg(image)
+        img_bytes = io.BytesIO()
+        result.save(img_bytes, "PNG")
+        img_bytes.seek(0)
         
-        img_byte_arr = io.BytesIO()
-        result.save(img_byte_arr, format="PNG")
-        img_byte_arr.seek(0)
+        return StreamingResponse(img_bytes, media_type="image/png")
         
-        print("✅ Processing complete")
-        return StreamingResponse(
-            img_byte_arr,
-            media_type="image/png",
-            headers={"Content-Disposition": "inline; filename=removed-bg.png"}
-        )
     except Exception as e:
-        print(f"❌ Detailed error: {str(e)}")
-        raise HTTPException(400, detail=f"Processing failed: {str(e)}")
+        print(f"❌ Error: {e}")
+        raise HTTPException(400, f"Failed: {str(e)}")
 
-if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+@app.get("/")
+def root():
+    return {"status": "OK", "endpoint": "/remove-bg POST with file=image or file=file"}