Update app.py

app.py

@@ -1,27 +1,34 @@
 import os
-from io import BytesIO
-
-import numpy as np
-import requests
+import threading
 import torch
-from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Query
+import numpy as np
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException
 from fastapi.responses import StreamingResponse, HTMLResponse
-from fastapi.middleware.cors import CORSMiddleware
 from PIL import Image
+from io import BytesIO
+import requests
 from transformers import AutoModelForImageSegmentation
+import uvicorn
 
 # ---------------------------------------------------------
-# Optional HEIC Support
+# Optional HEIC/HEIF support
 # ---------------------------------------------------------
 try:
     import pillow_heif
     pillow_heif.register_heif_opener()
-    print("HEIC/HEIF supported")
-except:
-    print("Install pillow-heif for HEIC support")
+    print("HEIC/HEIF supported.")
+except ImportError:
+    print("Install pillow-heif for HEIC support.")
 
 # ---------------------------------------------------------
-# Load Model
+# Performance settings (especially for CPU on HF Spaces)
+# ---------------------------------------------------------
+os.environ["OMP_NUM_THREADS"] = "1"
+os.environ["MKL_NUM_THREADS"] = "1"
+torch.set_num_threads(1)
+
+# ---------------------------------------------------------
+# Load model
 # ---------------------------------------------------------
 MODEL_DIR = "models/BiRefNet"
 os.makedirs(MODEL_DIR, exist_ok=True)
@@ -29,7 +36,7 @@ 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...")
+print("Loading BiRefNet model...")
 birefnet = AutoModelForImageSegmentation.from_pretrained(
     "ZhengPeng7/BiRefNet",
     cache_dir=MODEL_DIR,
@@ -37,284 +44,183 @@ birefnet = AutoModelForImageSegmentation.from_pretrained(
     revision="main"
 )
 birefnet.to(device, dtype=dtype).eval()
-print("Model Ready")
+print("Model loaded.")
+
+# Global lock to protect the model during inference
+inference_lock = threading.Lock()
 
 # ---------------------------------------------------------
-# FastAPI App
+# FastAPI app
 # ---------------------------------------------------------
 app = FastAPI(title="Background Remover API")
 
-# Allow requests from any app
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
 # ---------------------------------------------------------
-# Utility Functions
+# Helper functions
 # ---------------------------------------------------------
-def load_image_from_url(url: str):
+def load_image_from_url(url: str) -> Image.Image:
     try:
-        resp = requests.get(url, timeout=10)
-        resp.raise_for_status()
-        return Image.open(BytesIO(resp.content)).convert("RGB")
+        r = requests.get(url, timeout=10)
+        r.raise_for_status()
+        return Image.open(BytesIO(r.content)).convert("RGB")
     except Exception as e:
-        raise HTTPException(status_code=400, detail=f"Error loading image URL: {str(e)}")
+        raise HTTPException(status_code=400, detail=f"Cannot load image from URL: {str(e)}")
+
 
-def transform_image(image: Image.Image, resolution: int):
+def transform_image(image: Image.Image, resolution: int = 512) -> torch.Tensor:
     image = image.resize((resolution, resolution))
-    img = np.array(image).astype("float32") / 255.0
-    mean = np.array([0.485, 0.456, 0.406], dtype="float32")
-    std = np.array([0.229, 0.224, 0.225], dtype="float32")
-    img = (img - mean) / std
-    img = img.transpose((2, 0, 1))
-    return torch.from_numpy(img).unsqueeze(0).to(device=device, dtype=dtype)
+    arr = np.array(image).astype(np.float32) / 255.0
+    mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
+    std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
+
+    arr = (arr - mean) / std
+    arr = np.transpose(arr, (2, 0, 1))
 
-def process_image(image: Image.Image, resolution: int):
-    original_size = image.size
-    tensor = transform_image(image, resolution)
+    tensor = torch.from_numpy(arr).unsqueeze(0).to(device=device, dtype=dtype)
+    return tensor
 
-    with torch.no_grad():
-        mask_pred = birefnet(tensor)[-1].sigmoid().cpu()[0, 0]
 
-    mask = Image.fromarray((mask_pred.numpy() * 255).astype("uint8"))
-    mask = mask.resize(original_size)
+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 torch.no_grad():
+            try:
+                preds = birefnet(input_tensor)[-1].sigmoid().cpu()
+            except Exception as e:
+                raise RuntimeError(f"Inference error: {str(e)}")
+
+    pred = preds[0, 0]
+    mask = Image.fromarray((pred.numpy() * 255).astype(np.uint8)).resize(orig_size)
 
     image = image.convert("RGBA")
     image.putalpha(mask)
     return image
 
 # ---------------------------------------------------------
-# GET + POST endpoint
+# /remove-background endpoint
 # ---------------------------------------------------------
-@app.api_route("/remove-background", methods=["GET", "POST"])
+@app.post("/remove-background")
 async def remove_background(
     file: UploadFile = File(None),
     image_url: str = Form(None),
-    resolution: int = Form(512),
-    get_url: str = Query(None),
-    get_res: int = Query(512),
+    resolution: int = Form(512)
 ):
     try:
-        # GET mode: /remove-background?get_url=...&get_res=512
-        if get_url:
-            img = load_image_from_url(get_url)
-            resolution = get_res
-
-        # POST mode - file upload
-        elif file:
+        if file:
             data = await file.read()
-            if not data:
-                raise HTTPException(status_code=400, detail="Empty file")
-            img = Image.open(BytesIO(data)).convert("RGB")
-
-        # POST mode - URL
+            image = Image.open(BytesIO(data)).convert("RGB")
         elif image_url:
-            img = load_image_from_url(image_url)
-
+            image = load_image_from_url(image_url)
         else:
-            raise HTTPException(status_code=400, detail="No image provided")
+            raise HTTPException(status_code=400, detail="Provide either file or image_url.")
 
-        result = process_image(img, resolution)
+        result = run_inference(image, resolution)
 
-        buffer = BytesIO()
-        result.save(buffer, format="PNG")
-        buffer.seek(0)
+        buf = BytesIO()
+        result.save(buf, format="PNG")
+        buf.seek(0)
 
-        return StreamingResponse(
-            buffer,
-            media_type="image/png",
-            headers={"Content-Disposition": "inline; filename=result.png"},
-        )
+        return StreamingResponse(buf, media_type="image/png")
 
     except HTTPException:
         raise
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Internal error: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Processing error: {str(e)}")
 
-# ---------------------------------------------------------
-# Favicon (stop 404 logs)
-# ---------------------------------------------------------
-@app.get("/favicon.ico")
-async def favicon():
-    return HTMLResponse("")
 
 # ---------------------------------------------------------
-# UI for testing (POST method)
+# Test UI
 # ---------------------------------------------------------
 @app.get("/", response_class=HTMLResponse)
 async def index():
-    return """
-<!DOCTYPE html>
-<html>
-<head>
-    <title>Background Remover API</title>
-    <meta name="viewport" content="width=device-width, initial-scale=1">
-
-    <style>
-        body {
-            font-family: Arial, sans-serif;
-            padding: 20px;
-            margin: 0;
-            background: #f5f5f5;
-        }
-        h2 {
-            text-align: center;
-            margin-bottom: 25px;
-        }
-
-        .container {
-            max-width: 700px;
-            margin: auto;
-            background: white;
-            padding: 25px;
-            border-radius: 12px;
-            box-shadow: 0 4px 10px rgba(0,0,0,0.1);
-        }
-
-        .section {
-            margin-bottom: 30px;
-        }
-
-        label {
-            font-weight: bold;
-        }
-
-        input[type="file"],
-        input[type="text"],
-        input[type="number"] {
-            width: 100%;
-            padding: 10px;
-            margin-top: 6px;
-            border-radius: 6px;
-            border: 1px solid #ccc;
-        }
-
-        button {
-            padding: 12px 18px;
-            margin-top: 12px;
-            width: 100%;
-            border: none;
-            background: #007bff;
-            color: white;
-            border-radius: 6px;
-            cursor: pointer;
-            font-size: 16px;
-        }
-
-        button:hover {
-            background: #005dc4;
-        }
-
-        #resultWrapper {
-            text-align: center;
-            margin-top: 20px;
-        }
-
-        img {
-            max-width: 100%;
-            border-radius: 10px;
-            margin-top: 15px;
-        }
-
-        /* Responsive fixes */
-        @media (max-width: 600px) {
-            .container {
-                padding: 15px;
-            }
-            button {
-                font-size: 15px;
-            }
-        }
-    </style>
-
-</head>
-<body>
-
-    <h2>Background Remover API Tester</h2>
-
-    <div class="container">
-
-        <!-- Upload Section -->
-        <div class="section">
-            <label>Upload Image</label>
-            <form id="uploadForm" enctype="multipart/form-data">
-                <input type="file" id="file" name="file" accept="image/*">
-
-                <label>Resolution</label>
-                <input type="number" id="resFile" value="512" min="64" max="2048">
-
-                <button type="submit">Remove Background</button>
+    html = """
+    <!DOCTYPE html>
+    <html>
+    <head>
+        <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>
+    </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>
             </form>
-        </div>
 
-        <hr>
+            <div class="mb-4">OR</div>
 
-        <!-- URL Section -->
-        <div class="section">
-            <label>Image URL</label>
             <form id="urlForm">
-                <input type="text" id="imgUrl" placeholder="https://example.com/image.jpg">
-
-                <label>Resolution</label>
-                <input type="number" id="resUrl" value="512" min="64" max="2048">
-
-                <button type="submit">Remove Background</button>
+                <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>
-        </div>
 
-        <div id="resultWrapper">
-            <h3>Result</h3>
-            <img id="result" />
+            <h5 class="mt-4">Result:</h5>
+            <img id="resultImg" />
         </div>
 
-    </div>
-
-    <script>
-        const resultImg = document.getElementById("result");
-
-        // Upload File
-        document.getElementById("uploadForm").onsubmit = async (e) => {
-            e.preventDefault();
-            const file = document.getElementById("file").files[0];
-            const res = document.getElementById("resFile").value;
-
-            if (!file) return alert("Please select a file");
+        <script>
+            const uploadForm = document.getElementById("uploadForm");
+            const urlForm = document.getElementById("urlForm");
+            const resultImg = document.getElementById("resultImg");
 
-            const fd = new FormData();
-            fd.append("file", file);
-            fd.append("resolution", res);
+            uploadForm.addEventListener("submit", async e => {
+                e.preventDefault();
+                const file = document.getElementById("fileInput").files[0];
+                if (!file) return alert("Choose an image");
 
-            const r = await fetch("/remove-background", { method: "POST", body: fd });
-            resultImg.src = URL.createObjectURL(await r.blob());
-        };
+                const res = document.getElementById("resInput").value;
+                const form = new FormData();
+                form.append("file", file);
+                form.append("resolution", res);
 
-        // URL Mode
-        document.getElementById("urlForm").onsubmit = async (e) => {
-            e.preventDefault();
-            const url = document.getElementById("imgUrl").value;
-            const res = document.getElementById("resUrl").value;
+                const r = await fetch("/remove-background", { method: "POST", body: form });
+                const blob = await r.blob();
+                resultImg.src = URL.createObjectURL(blob);
+            });
 
-            if (!url.trim()) return alert("Please enter an image URL");
+            urlForm.addEventListener("submit", async e => {
+                e.preventDefault();
+                const url = document.getElementById("urlInput").value.trim();
+                if (!url) return alert("Enter URL");
 
-            const fd = new FormData();
-            fd.append("image_url", url);
-            fd.append("resolution", res);
+                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: fd });
-            resultImg.src = URL.createObjectURL(await r.blob());
-        };
-    </script>
+                const r = await fetch("/remove-background", { method: "POST", body: form });
+                const blob = await r.blob();
+                resultImg.src = URL.createObjectURL(blob);
+            });
+        </script>
+    </body>
+    </html>
+    """
+    return HTMLResponse(html)
 
-</body>
-</html>
-"""
 
 # ---------------------------------------------------------
-# Run App
+# Run server
 # ---------------------------------------------------------
 if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run("app:app", host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)