Update app.py

app.py

@@ -10,27 +10,29 @@ from transformers import AutoModelForImageSegmentation
 import uvicorn
 
 # ---------------------------------------------------------
-# HEIC/HEIF SUPPORT
+# HEIC SUPPORT
 # ---------------------------------------------------------
 try:
     import pillow_heif
     pillow_heif.register_heif_opener()
-except ImportError:
+except:
     pass
 
 # ---------------------------------------------------------
-# CPU PERFORMANCE (FIXED)
+# CPU OPTIMIZATION
 # ---------------------------------------------------------
 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)
 
 # ---------------------------------------------------------
 # SETTINGS
 # ---------------------------------------------------------
 TARGET_SIZE = (512, 512)
-MAX_SIDE = 2000
+MAX_SIDE = 1800
 
 # ---------------------------------------------------------
 # LOAD MODEL
@@ -39,24 +41,42 @@ MODEL_DIR = "models/BiRefNet"
 os.makedirs(MODEL_DIR, exist_ok=True)
 
 print("Loading model...")
+
 model = AutoModelForImageSegmentation.from_pretrained(
     "ZhengPeng7/BiRefNet",
     cache_dir=MODEL_DIR,
     trust_remote_code=True
 )
+
+# ✅ 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)
+
 print("Model ready.")
 
 # ---------------------------------------------------------
 # WARMUP
 # ---------------------------------------------------------
 def warmup():
-    dummy = torch.randn(1, 3, 512, 512)
+    dummy = torch.randn(1, 3, 512, 512).to(memory_format=torch.channels_last)
     with torch.no_grad():
         _ = model(dummy)
 
 warmup()
-print("Warmup done.")
 
 # ---------------------------------------------------------
 # HELPERS
@@ -66,7 +86,7 @@ def load_image_from_url(url: str) -> Image.Image:
         r = requests.get(url, timeout=10)
         r.raise_for_status()
         return Image.open(BytesIO(r.content)).convert("RGB")
-    except Exception:
+    except:
         raise HTTPException(400, "Invalid image URL")
 
 
@@ -83,14 +103,20 @@ def transform(img: Image.Image):
     img = img.resize(TARGET_SIZE)
 
     arr = np.asarray(img, dtype=np.float32) / 255.0
-    arr = (arr - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]
+    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)
 
-    return torch.from_numpy(arr).unsqueeze(0)
+    tensor = torch.from_numpy(arr).unsqueeze(0).float()
+
+    # ✅ channels last
+    return tensor.to(memory_format=torch.channels_last)
 
 
 def run_inference(img: Image.Image) -> Image.Image:
     orig_size = img.size
+
     tensor = transform(img)
 
     with torch.no_grad():
@@ -100,13 +126,14 @@ def run_inference(img: Image.Image) -> Image.Image:
 
     img = img.convert("RGBA")
     img.putalpha(mask)
+
     return img
 
 
 # ---------------------------------------------------------
 # FASTAPI
 # ---------------------------------------------------------
-app = FastAPI(title="Background Remover API")
+app = FastAPI(title="Fast Background Remover")
 
 # ---------------------------------------------------------
 # GET redirect
@@ -135,6 +162,7 @@ async def remove_bg(file: UploadFile = File(None), image_url: str = Form(None)):
             raise HTTPException(400, "Provide file or image_url")
 
         img = auto_downscale(img)
+
         result = run_inference(img)
 
         buf = BytesIO()
@@ -148,7 +176,7 @@ async def remove_bg(file: UploadFile = File(None), image_url: str = Form(None)):
 
 
 # ---------------------------------------------------------
-# UI (IMPROVED BUT SAME LOGIC)
+# UI (UNCHANGED BUT CLEAN)
 # ---------------------------------------------------------
 @app.get("/", response_class=HTMLResponse)
 def ui():
@@ -157,38 +185,36 @@ def ui():
     <head>
         <title>Background Remover</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 class='mb-4'>Background Remover</h2>
+            <h2>Background Remover</h2>
 
             <div class='row'>
                 <div class='col-md-6'>
                     <h5>Input</h5>
-                    <img id='inputImg' style='max-width:100%; border-radius:10px;'>
+                    <img id='inputImg' style='max-width:100%'>
                 </div>
                 <div class='col-md-6'>
                     <h5>Output</h5>
-                    <img id='outputImg' style='max-width:100%; border-radius:10px;'>
+                    <img id='outputImg' style='max-width:100%'>
                 </div>
             </div>
 
             <hr>
 
-            <h4>Upload Image</h4>
             <form id="uploadForm">
                 <input type='file' id='fileInput' class='form-control mb-3'>
-                <button class='btn btn-primary'>Remove Background</button>
+                <button class='btn btn-primary'>Upload</button>
             </form>
 
             <hr>
 
-            <h4>Image URL</h4>
             <form id='urlForm'>
-                <input id='urlInput' class='form-control mb-3' placeholder='https://image.jpg'>
-                <button class='btn btn-success'>Remove Background</button>
+                <input id='urlInput' class='form-control mb-3' placeholder='Image URL'>
+                <button class='btn btn-success'>Send URL</button>
             </form>
 
         </div>
@@ -197,53 +223,41 @@ def ui():
         const inputImg = document.getElementById("inputImg");
         const outputImg = document.getElementById("outputImg");
 
-        async function sendRequest(formData) {
-            const res = await fetch("/remove-background", {
-                method: "POST",
-                body: formData
+        async function send(fd){
+            const r = await fetch("/remove-background", {
+                method:"POST",
+                body:fd
             });
 
-            if (!res.ok) {
-                alert("Error processing image");
-                return;
-            }
-
-            const blob = await res.blob();
+            const blob = await r.blob();
             outputImg.src = URL.createObjectURL(blob);
         }
 
-        document.getElementById("uploadForm").addEventListener("submit", async e => {
+        document.getElementById("uploadForm").onsubmit = async e=>{
             e.preventDefault();
-            const file = document.getElementById("fileInput").files[0];
-            if (!file) return alert("Select file");
-
+            const file = fileInput.files[0];
             inputImg.src = URL.createObjectURL(file);
 
             const fd = new FormData();
             fd.append("file", file);
+            send(fd);
+        };
 
-            sendRequest(fd);
-        });
-
-        document.getElementById("urlForm").addEventListener("submit", async e => {
+        document.getElementById("urlForm").onsubmit = async e=>{
             e.preventDefault();
-            const url = document.getElementById("urlInput").value.trim();
-            if (!url) return alert("Enter URL");
-
+            const url = urlInput.value;
             inputImg.src = url;
 
             const fd = new FormData();
             fd.append("image_url", url);
-
-            sendRequest(fd);
-        });
+            send(fd);
+        };
         </script>
 
     </body>
     </html>
     """
 
-
 # ---------------------------------------------------------
 # RUN
 # ---------------------------------------------------------