Update app.py

app.py

@@ -16,19 +16,18 @@ import uvicorn
 try:
     import pillow_heif
     pillow_heif.register_heif_opener()
-    print("HEIC/HEIF supported.")
 except ImportError:
-    print("Install pillow-heif for HEIC support.")
+    pass
 
 # ---------------------------------------------------------
-# Performance settings (especially for CPU on HF Spaces)
+# Performance settings for CPU (HF Spaces)
 # ---------------------------------------------------------
 os.environ["OMP_NUM_THREADS"] = "1"
 os.environ["MKL_NUM_THREADS"] = "1"
 torch.set_num_threads(1)
 
 # ---------------------------------------------------------
-# Load model
+# Model load
 # ---------------------------------------------------------
 MODEL_DIR = "models/BiRefNet"
 os.makedirs(MODEL_DIR, exist_ok=True)
@@ -36,7 +35,6 @@ 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 model...")
 birefnet = AutoModelForImageSegmentation.from_pretrained(
     "ZhengPeng7/BiRefNet",
     cache_dir=MODEL_DIR,
@@ -44,16 +42,10 @@ birefnet = AutoModelForImageSegmentation.from_pretrained(
     revision="main"
 )
 birefnet.to(device, dtype=dtype).eval()
-print("Model loaded.")
 
-# Global lock to protect the model during inference
+# Thread lock to protect inference on CPU
 inference_lock = threading.Lock()
 
-# ---------------------------------------------------------
-# FastAPI app
-# ---------------------------------------------------------
-app = FastAPI(title="Background Remover API")
-
 # ---------------------------------------------------------
 # Helper functions
 # ---------------------------------------------------------
@@ -66,6 +58,23 @@ def load_image_from_url(url: str) -> Image.Image:
         raise HTTPException(status_code=400, detail=f"Cannot load image from URL: {str(e)}")
 
 
+def auto_downscale(image: Image.Image, max_side: int = 3000) -> Image.Image:
+    """Downscale very large images to speed up CPU inference."""
+    w, h = image.size
+
+    if max(w, h) <= max_side:
+        return image
+
+    scale = max_side / max(w, h)
+    new_w = int(w * scale)
+    new_h = int(h * scale)
+
+    image = image.resize((new_w, new_h), Image.LANCZOS)
+    print(f"[INFO] Downscaled large image from {w}x{h} to {new_w}x{new_h}")
+
+    return image
+
+
 def transform_image(image: Image.Image, resolution: int = 512) -> torch.Tensor:
     image = image.resize((resolution, resolution))
     arr = np.array(image).astype(np.float32) / 255.0
@@ -81,14 +90,12 @@ def transform_image(image: Image.Image, resolution: int = 512) -> torch.Tensor:
 
 def run_inference(image: Image.Image, resolution: int = 512) -> Image.Image:
     orig_size = image.size
+
     input_tensor = transform_image(image, resolution)
 
-    with inference_lock:               # prevents CPU thread crashes
+    with inference_lock:
         with torch.no_grad():
-            try:
-                preds = birefnet(input_tensor)[-1].sigmoid().cpu()
-            except Exception as e:
-                raise RuntimeError(f"Inference error: {str(e)}")
+            preds = birefnet(input_tensor)[-1].sigmoid().cpu()
 
     pred = preds[0, 0]
     mask = Image.fromarray((pred.numpy() * 255).astype(np.uint8)).resize(orig_size)
@@ -97,9 +104,12 @@ def run_inference(image: Image.Image, resolution: int = 512) -> Image.Image:
     image.putalpha(mask)
     return image
 
+
 # ---------------------------------------------------------
-# /remove-background endpoint
+# API endpoint
 # ---------------------------------------------------------
+@app = FastAPI(title="Background Remover API")
+
 @app.post("/remove-background")
 async def remove_background(
     file: UploadFile = File(None),
@@ -113,12 +123,15 @@ async def remove_background(
         elif image_url:
             image = load_image_from_url(image_url)
         else:
-            raise HTTPException(status_code=400, detail="Provide either file or image_url.")
+            raise HTTPException(status_code=400, detail="Provide file or image_url.")
+
+        # Auto-downscale for large images → much faster
+        image = auto_downscale(image)
 
         result = run_inference(image, resolution)
 
         buf = BytesIO()
-        result.save(buf, format="PNG")
+        result.save(buf, format="PNG", optimize=True)
         buf.seek(0)
 
         return StreamingResponse(buf, media_type="image/png")
@@ -130,93 +143,69 @@ async def remove_background(
 
 
 # ---------------------------------------------------------
-# Test UI
+# Developer test UI
 # ---------------------------------------------------------
 @app.get("/", response_class=HTMLResponse)
 async def index():
-    html = """
-    <!DOCTYPE html>
+    return """
     <html>
     <head>
-        <meta charset="UTF-8">
+        <meta charset='utf-8' />
         <title>Background Remover</title>
-        <link href="https://cdn.jsdelivr.net/npm/bootstrap@5.3.2/dist/css/bootstrap.min.css" rel="stylesheet">
-        <style>
-            body { background: #f8f9fa; padding-top: 40px; }
-            .container { max-width: 700px; }
-            img { max-width: 100%; margin-top: 20px; border-radius: 10px; }
-        </style>
+        <link rel='stylesheet' href='https://cdn.jsdelivr.net/npm/bootstrap@5.3.2/dist/css/bootstrap.min.css'>
     </head>
-
-    <body>
-        <div class="container text-center">
-            <h2 class="mb-4">Background Remover API</h2>
-
-            <form id="uploadForm" class="mb-4" enctype="multipart/form-data">
-                <div class="mb-3">
-                    <input class="form-control" type="file" id="fileInput" name="file" accept="image/*">
-                </div>
-                <div class="mb-3">
-                    <input class="form-control" type="number" id="resInput" name="resolution" value="512" min="64" max="2048">
-                </div>
-                <button class="btn btn-primary">Upload</button>
+    <body style='background:#f8f9fa;padding-top:40px;'>
+        <div class='container text-center'>
+            <h2>Background Remover API</h2>
+
+            <form id='uploadForm' class='mb-4' enctype='multipart/form-data'>
+                <input class='form-control mb-2' type='file' id='fileInput' name='file'>
+                <input class='form-control mb-2' type='number' id='resInput' name='resolution' value='512'>
+                <button class='btn btn-primary'>Upload</button>
             </form>
 
-            <div class="mb-4">OR</div>
+            <div class='mb-3'>OR</div>
 
-            <form id="urlForm">
-                <div class="mb-3">
-                    <input class="form-control" type="text" id="urlInput" placeholder="https://example.com/img.jpg">
-                </div>
-                <div class="mb-3">
-                    <input class="form-control" type="number" id="urlResInput" value="512" min="64" max="2048">
-                </div>
-                <button class="btn btn-success">Use URL</button>
+            <form id='urlForm'>
+                <input class='form-control mb-2' id='urlInput' placeholder='Image URL'>
+                <input class='form-control mb-2' id='urlResInput' type='number' value='512'>
+                <button class='btn btn-success'>Use URL</button>
             </form>
 
-            <h5 class="mt-4">Result:</h5>
-            <img id="resultImg" />
+            <h5 class='mt-4'>Result:</h5>
+            <img id='resultImg' style='max-width:100%;border-radius:10px;'/>
         </div>
 
         <script>
-            const uploadForm = document.getElementById("uploadForm");
-            const urlForm = document.getElementById("urlForm");
             const resultImg = document.getElementById("resultImg");
 
-            uploadForm.addEventListener("submit", async e => {
+            document.getElementById("uploadForm").addEventListener("submit", async e =>{
                 e.preventDefault();
                 const file = document.getElementById("fileInput").files[0];
                 if (!file) return alert("Choose an image");
-
                 const res = document.getElementById("resInput").value;
-                const form = new FormData();
-                form.append("file", file);
-                form.append("resolution", res);
-
-                const r = await fetch("/remove-background", { method: "POST", body: form });
-                const blob = await r.blob();
-                resultImg.src = URL.createObjectURL(blob);
+                const f = new FormData();
+                f.append("file", file);
+                f.append("resolution", res);
+                const r = await fetch("/remove-background", { method:"POST", body:f });
+                resultImg.src = URL.createObjectURL(await r.blob());
             });
 
-            urlForm.addEventListener("submit", async e => {
+            document.getElementById("urlForm").addEventListener("submit", async e =>{
                 e.preventDefault();
                 const url = document.getElementById("urlInput").value.trim();
                 if (!url) return alert("Enter URL");
-
                 const res = document.getElementById("urlResInput").value;
-                const form = new FormData();
-                form.append("image_url", url);
-                form.append("resolution", res);
-
-                const r = await fetch("/remove-background", { method: "POST", body: form });
-                const blob = await r.blob();
-                resultImg.src = URL.createObjectURL(blob);
+                const f = new FormData();
+                f.append("image_url", url);
+                f.append("resolution", res);
+                const r = await fetch("/remove-background", { method:"POST", body:f });
+                resultImg.src = URL.createObjectURL(await r.blob());
             });
         </script>
     </body>
     </html>
     """
-    return HTMLResponse(html)
 
 
 # ---------------------------------------------------------