Update app.py

app.py

@@ -1,57 +1,87 @@
+# app.py - FastAPI BiRefNet Background Remover API
 import torch
-import gradio as gr
 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.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse
+import uvicorn
 
 from birefnet import BiRefNet
 from BiRefNet_config import BiRefNetConfig
 
-
+# Global model (loaded once)
 device = "cpu"
-
-# Load config + model
 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.fromarray(image)
+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")
 
-demo = gr.Interface(
-    fn=remove_bg,
-    inputs=gr.Image(type="numpy"),
-    outputs=gr.Image(type="pil"),
-    title="BiRefNet Lite – Background Remover",
+# CORS for NextJS/Vercel
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Update with your domain in production
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
 )
 
-demo.launch()
+@app.get("/")
+async def root():
+    return {"message": "BiRefNet Background Remover API", "status": "ready"}
+
+@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")
+    
+    try:
+        # Read and process image
+        contents = await file.read()
+        image = Image.open(io.BytesIO(contents))
+        result = remove_bg(image)
+        
+        # Save to bytes
+        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"}
+        )
+    except Exception as e:
+        raise HTTPException(500, detail=f"Processing failed: {str(e)}")
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860)