Update app.py

app.py

@@ -1,9 +1,8 @@
 import os
-import threading
 import torch
 import numpy as np
-from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Request
-from fastapi.responses import StreamingResponse, HTMLResponse, RedirectResponse, JSONResponse
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException
+from fastapi.responses import StreamingResponse, HTMLResponse, JSONResponse
 from PIL import Image
 from io import BytesIO
 import requests
@@ -11,7 +10,7 @@ from transformers import AutoModelForImageSegmentation
 import uvicorn
 
 # ---------------------------------------------------------
-# Optional HEIC/HEIF
+# HEIC/HEIF SUPPORT
 # ---------------------------------------------------------
 try:
     import pillow_heif
@@ -20,41 +19,47 @@ except ImportError:
     pass
 
 # ---------------------------------------------------------
-# Performance settings for HF CPU
+# CPU PERFORMANCE (FIXED)
 # ---------------------------------------------------------
-os.environ["OMP_NUM_THREADS"] = "1"
-os.environ["MKL_NUM_THREADS"] = "1"
-torch.set_num_threads(1)
+CPU_THREADS = min(4, os.cpu_count() or 2)
+os.environ["OMP_NUM_THREADS"] = str(CPU_THREADS)
+os.environ["MKL_NUM_THREADS"] = str(CPU_THREADS)
+torch.set_num_threads(CPU_THREADS)
 
 # ---------------------------------------------------------
-# Constants
+# SETTINGS
 # ---------------------------------------------------------
-TARGET_SIZE = (512, 512)      # Faster inference
-MAX_SIDE = 3000               # Auto-downscale for huge uploads
+TARGET_SIZE = (512, 512)
+MAX_SIDE = 2000
 
 # ---------------------------------------------------------
-# Load model
+# LOAD MODEL
 # ---------------------------------------------------------
 MODEL_DIR = "models/BiRefNet"
 os.makedirs(MODEL_DIR, exist_ok=True)
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-
-print("Loading BiRefNet…")
-birefnet = AutoModelForImageSegmentation.from_pretrained(
+print("Loading model...")
+model = AutoModelForImageSegmentation.from_pretrained(
     "ZhengPeng7/BiRefNet",
     cache_dir=MODEL_DIR,
-    trust_remote_code=True,
-    revision="main",
+    trust_remote_code=True
 )
-birefnet.to(device, dtype=dtype).eval()
+model.eval()
 print("Model ready.")
 
-lock = threading.Lock()
+# ---------------------------------------------------------
+# WARMUP
+# ---------------------------------------------------------
+def warmup():
+    dummy = torch.randn(1, 3, 512, 512)
+    with torch.no_grad():
+        _ = model(dummy)
+
+warmup()
+print("Warmup done.")
 
 # ---------------------------------------------------------
-# Helpers
+# HELPERS
 # ---------------------------------------------------------
 def load_image_from_url(url: str) -> Image.Image:
     try:
@@ -62,7 +67,7 @@ def load_image_from_url(url: str) -> Image.Image:
         r.raise_for_status()
         return Image.open(BytesIO(r.content)).convert("RGB")
     except Exception:
-        raise HTTPException(status_code=400, detail="Invalid image URL")
+        raise HTTPException(400, "Invalid image URL")
 
 
 def auto_downscale(img: Image.Image) -> Image.Image:
@@ -71,35 +76,27 @@ def auto_downscale(img: Image.Image) -> Image.Image:
         return img
 
     scale = MAX_SIDE / max(w, h)
-    new_w = int(w * scale)
-    new_h = int(h * scale)
-
-    print(f"[INFO] Downscaling {w}×{h} → {new_w}×{new_h}")
-    return img.resize((new_w, new_h), Image.LANCZOS)
+    return img.resize((int(w * scale), int(h * scale)), Image.LANCZOS)
 
 
-def transform(img: Image.Image) -> torch.Tensor:
+def transform(img: Image.Image):
     img = img.resize(TARGET_SIZE)
 
-    arr = np.array(img).astype(np.float32) / 255.0
-    mean = np.array([0.485, 0.456, 0.406])
-    std = np.array([0.229, 0.224, 0.225])
-    arr = (arr - mean) / std
-    arr = np.transpose(arr, (2, 0, 1))
+    arr = np.asarray(img, dtype=np.float32) / 255.0
+    arr = (arr - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]
+    arr = arr.transpose(2, 0, 1)
 
-    t = torch.from_numpy(arr).unsqueeze(0).to(device=device, dtype=dtype)
-    return t
+    return torch.from_numpy(arr).unsqueeze(0)
 
 
 def run_inference(img: Image.Image) -> Image.Image:
     orig_size = img.size
     tensor = transform(img)
 
-    with lock:
-        with torch.no_grad():
-            pred = birefnet(tensor)[-1].sigmoid().cpu()[0, 0]
+    with torch.no_grad():
+        pred = model(tensor)[-1].sigmoid()[0, 0].cpu().numpy()
 
-    mask = Image.fromarray((pred.numpy() * 255).astype(np.uint8)).resize(orig_size)
+    mask = Image.fromarray((pred * 255).astype(np.uint8)).resize(orig_size)
 
     img = img.convert("RGBA")
     img.putalpha(mask)
@@ -107,22 +104,22 @@ def run_inference(img: Image.Image) -> Image.Image:
 
 
 # ---------------------------------------------------------
-# FastAPI app
+# FASTAPI
 # ---------------------------------------------------------
 app = FastAPI(title="Background Remover API")
 
 # ---------------------------------------------------------
-# Redirect GET → POST logic
+# GET redirect
 # ---------------------------------------------------------
 @app.get("/remove-background")
-async def redirect_to_post():
+def redirect():
     return JSONResponse(
-        {"detail": "This endpoint only supports POST. Use POST /remove-background"},
+        {"detail": "Use POST /remove-background"},
         status_code=405
     )
 
 # ---------------------------------------------------------
-# Main POST endpoint
+# MAIN ENDPOINT (UNCHANGED)
 # ---------------------------------------------------------
 @app.post("/remove-background")
 async def remove_bg(file: UploadFile = File(None), image_url: str = Form(None)):
@@ -130,101 +127,115 @@ async def remove_bg(file: UploadFile = File(None), image_url: str = Form(None)):
         if file:
             raw = await file.read()
             img = Image.open(BytesIO(raw)).convert("RGB")
+
         elif image_url:
             img = load_image_from_url(image_url)
+
         else:
-            raise HTTPException(status_code=400, detail="Upload file or image_url required")
+            raise HTTPException(400, "Provide file or image_url")
 
         img = auto_downscale(img)
         result = run_inference(img)
 
         buf = BytesIO()
-        result.save(buf, format="PNG")
+        result.save(buf, format="PNG", optimize=True)
         buf.seek(0)
 
         return StreamingResponse(buf, media_type="image/png")
 
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        raise HTTPException(500, str(e))
 
 
 # ---------------------------------------------------------
-# UI: Show INPUT + OUTPUT (big preview)
+# UI (IMPROVED BUT SAME LOGIC)
 # ---------------------------------------------------------
 @app.get("/", response_class=HTMLResponse)
-async def ui():
+def ui():
     return """
     <html>
     <head>
-        <title>Background Remover – Test UI</title>
+        <title>Background Remover</title>
         <link rel='stylesheet'
               href='https://cdn.jsdelivr.net/npm/bootstrap@5.3.2/dist/css/bootstrap.min.css'>
     </head>
     <body class='bg-light'>
         <div class='container py-4 text-center'>
 
-            <h2 class='mb-4'>API Test Panel (POST Only)</h2>
+            <h2 class='mb-4'>Background Remover</h2>
 
             <div class='row'>
                 <div class='col-md-6'>
-                    <h5>Input Image</h5>
+                    <h5>Input</h5>
                     <img id='inputImg' style='max-width:100%; border-radius:10px;'>
                 </div>
                 <div class='col-md-6'>
-                    <h5>Output Image</h5>
+                    <h5>Output</h5>
                     <img id='outputImg' style='max-width:100%; border-radius:10px;'>
                 </div>
             </div>
 
             <hr>
 
-            <h4>Upload Test</h4>
-            <form id="uploadForm" enctype='multipart/form-data'>
+            <h4>Upload Image</h4>
+            <form id="uploadForm">
                 <input type='file' id='fileInput' class='form-control mb-3'>
-                <button class='btn btn-primary'>Send POST</button>
+                <button class='btn btn-primary'>Remove Background</button>
             </form>
 
             <hr>
 
-            <h4>URL Test</h4>
+            <h4>Image URL</h4>
             <form id='urlForm'>
-                <input id='urlInput' class='form-control mb-3' placeholder='https://example.com/image.jpg'>
-                <button class='btn btn-success'>Send POST</button>
+                <input id='urlInput' class='form-control mb-3' placeholder='https://image.jpg'>
+                <button class='btn btn-success'>Remove Background</button>
             </form>
+
         </div>
 
         <script>
         const inputImg = document.getElementById("inputImg");
         const outputImg = document.getElementById("outputImg");
 
-        // FILE TEST
+        async function sendRequest(formData) {
+            const res = await fetch("/remove-background", {
+                method: "POST",
+                body: formData
+            });
+
+            if (!res.ok) {
+                alert("Error processing image");
+                return;
+            }
+
+            const blob = await res.blob();
+            outputImg.src = URL.createObjectURL(blob);
+        }
+
         document.getElementById("uploadForm").addEventListener("submit", async e => {
             e.preventDefault();
             const file = document.getElementById("fileInput").files[0];
-            if (!file) return alert("Select a file first.");
+            if (!file) return alert("Select file");
 
             inputImg.src = URL.createObjectURL(file);
 
             const fd = new FormData();
             fd.append("file", file);
 
-            const r = await fetch("/remove-background", { method:"POST", body:fd });
-            outputImg.src = URL.createObjectURL(await r.blob());
+            sendRequest(fd);
         });
 
-        // URL TEST
         document.getElementById("urlForm").addEventListener("submit", async e => {
             e.preventDefault();
             const url = document.getElementById("urlInput").value.trim();
-            if (!url) return alert("Enter an image URL first.");
+            if (!url) return alert("Enter URL");
 
             inputImg.src = url;
 
             const fd = new FormData();
             fd.append("image_url", url);
 
-            const r = await fetch("/remove-background", { method:"POST", body:fd });
-            outputImg.src = URL.createObjectURL(await r.blob());
+            sendRequest(fd);
         });
         </script>
 
@@ -232,8 +243,9 @@ async def ui():
     </html>
     """
 
+
 # ---------------------------------------------------------
-# Run app
+# RUN
 # ---------------------------------------------------------
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)