Update app.py

app.py

@@ -1,28 +1,57 @@
-# app.py - Fixed FastAPI for HF Spaces
+import os
+import io
 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, HTTPException, Request, Depends
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import StreamingResponse
+from fastapi.security import APIKeyHeader
 import uvicorn
 
 from birefnet import BiRefNet
 from BiRefNet_config import BiRefNetConfig
 
-# Global model
-device = "cpu"
+# =========================
+# HUGGING FACE SECRET
+# =========================
+API_KEY = os.getenv("BIREFNET_API_KEY")
+
+if not API_KEY:
+    raise RuntimeError("❌ BIREFNET_API_KEY not found in HF Space Secrets")
+
+DEVICE = "cpu"
+
+# =========================
+# LOAD MODEL
+# =========================
 config = BiRefNetConfig()
 model = BiRefNet(config)
+
 state_dict = load_file("model.safetensors")
 model.load_state_dict(state_dict, strict=False)
-model.to(device)
+model.to(DEVICE)
 model.eval()
+
 print("✅ BiRefNet Lite loaded")
 
+# =========================
+# API KEY AUTH
+# =========================
+api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
+
+def verify_api_key(api_key: str = Depends(api_key_header)):
+    if api_key != API_KEY:
+        raise HTTPException(
+            status_code=401,
+            detail="Invalid or missing API key"
+        )
+
+# =========================
+# IMAGE PIPELINE
+# =========================
 def preprocess(img: Image.Image):
     img = img.convert("RGB").resize((1024, 1024))
     arr = np.array(img).astype(np.float32) / 255.0
@@ -31,71 +60,86 @@ def preprocess(img: Image.Image):
 
 @torch.no_grad()
 def remove_bg(image: Image.Image) -> Image.Image:
-    x = preprocess(image).to(device)
+    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
 
+# =========================
+# FASTAPI APP
+# =========================
 app = FastAPI(title="BiRefNet Background Remover API")
 
-# CORS - allow all for testing (restrict in production)
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],
+    allow_origins=["*"],  # Restrict later
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
+# =========================
+# ROUTES
+# =========================
 @app.get("/")
 async def root():
-    return {"message": "BiRefNet Background Remover API ✅", "endpoints": ["/remove-bg"]}
+    return {
+        "status": "ok",
+        "secured": True,
+        "endpoint": "/remove-bg"
+    }
 
-# FIXED: Accept both UploadFile and raw bytes
 @app.post("/remove-bg")
-async def remove_background(request: Request, file: UploadFile = File(None)):
+async def remove_background(
+    request: Request,
+    file: UploadFile = File(None),
+    _: str = Depends(verify_api_key)
+):
     try:
-        # 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")
+                raise HTTPException(400, "No image data received")
             image = Image.open(io.BytesIO(body))
         else:
             contents = await file.read()
             if not contents:
-                raise HTTPException(400, detail="Empty file")
+                raise HTTPException(400, "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
+
+        if image.format not in ["JPEG", "JPG", "PNG"]:
+            raise HTTPException(400, "Invalid image format")
+
         result = remove_bg(image)
-        
-        # Return PNG
-        img_byte_arr = io.BytesIO()
-        result.save(img_byte_arr, format="PNG")
-        img_byte_arr.seek(0)
-        
+
+        img_bytes = io.BytesIO()
+        result.save(img_bytes, format="PNG")
+        img_bytes.seek(0)
+
         return StreamingResponse(
-            img_byte_arr,
+            img_bytes,
             media_type="image/png",
             headers={
-                "Content-Disposition": "inline; filename=removed-bg.png",
-                "Cache-Control": "public, max-age=3600"
+                "Content-Disposition": "inline; filename=removed-bg.png"
             }
         )
+
     except Exception as e:
-        print(f"❌ Error: {str(e)}")
-        raise HTTPException(500, detail=f"Processing failed: {str(e)}")
+        print("❌ Error:", e)
+        raise HTTPException(500, "Processing failed")
 
+# =========================
+# HF DOCKER ENTRYPOINT
+# =========================
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(
+        app,
+        host="0.0.0.0",
+        port=int(os.environ.get("PORT", 7860))
+    )