Update app.py

app.py

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