Update app.py

app.py

@@ -2,7 +2,7 @@ import os
 import torch
 import numpy as np
 from fastapi import FastAPI, UploadFile, File, Form, HTTPException
-from fastapi.responses import StreamingResponse, HTMLResponse, JSONResponse
+from fastapi.responses import StreamingResponse, HTMLResponse
 from PIL import Image
 from io import BytesIO
 import requests
@@ -10,27 +10,18 @@ from transformers import AutoModelForImageSegmentation
 import uvicorn
 
 # ---------------------------------------------------------
-# Optional HEIC support
-# ---------------------------------------------------------
-try:
-    import pillow_heif
-    pillow_heif.register_heif_opener()
-except ImportError:
-    pass
-
-# ---------------------------------------------------------
-# Performance tuning (CPU)
+# CPU optimization (important for HF Spaces)
 # ---------------------------------------------------------
 os.environ["OMP_NUM_THREADS"] = "1"
 os.environ["MKL_NUM_THREADS"] = "1"
 torch.set_num_threads(1)
 
 # ---------------------------------------------------------
-# Constants (optimized)
+# Config (speed focused)
 # ---------------------------------------------------------
-TARGET_SIZE = (384, 384)          # 🔥 faster than 512
-MAX_SIDE = 3000
+TARGET_SIZE = (320, 320)          # 🔥 faster inference
 MAX_FILE_SIZE = 5 * 1024 * 1024   # 5MB
+MAX_COMPRESS_DIM = 1400           # aggressive resize
 
 # ---------------------------------------------------------
 # Load model
@@ -41,76 +32,53 @@ 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 model...")
 
-birefnet = AutoModelForImageSegmentation.from_pretrained(
+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()
-print("Model ready.")
+model.to(device, dtype=dtype).eval()
+
+print("Model ready")
 
 # ---------------------------------------------------------
-# Helpers
+# Image 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 Exception:
-        raise HTTPException(status_code=400, detail="Invalid image URL")
-
+def load_image_from_url(url: str):
+    r = requests.get(url, timeout=10)
+    r.raise_for_status()
+    return Image.open(BytesIO(r.content)).convert("RGB")
 
-def compress_if_needed(img: Image.Image, raw_bytes: bytes) -> Image.Image:
-    size = len(raw_bytes)
 
-    if size <= MAX_FILE_SIZE:
+# 🔥 FAST compression (key part)
+def compress_if_needed(img: Image.Image, raw_bytes: bytes):
+    if len(raw_bytes) <= MAX_FILE_SIZE:
         return img
 
-    print(f"[INFO] Compressing {size/1024/1024:.2f}MB image...")
+    print("[INFO] Compressing image >5MB")
 
     img = img.convert("RGB")
 
+    # Resize aggressively
     w, h = img.size
-    scale = min(1.0, 1600 / max(w, h))  # 🔥 more aggressive
+    scale = min(1.0, MAX_COMPRESS_DIM / max(w, h))
     img = img.resize((int(w * scale), int(h * scale)), Image.BILINEAR)
 
-    quality = 85
-
-    while True:
-        buffer = BytesIO()
-        img.save(buffer, format="JPEG", quality=quality, optimize=True)
-        compressed_size = buffer.tell()
-
-        if compressed_size <= MAX_FILE_SIZE or quality <= 40:
-            print(f"[INFO] Final size: {compressed_size/1024/1024:.2f}MB")
-            buffer.seek(0)
-            return Image.open(buffer).convert("RGB")
-
-        quality -= 10
-
-
-def auto_downscale(img: Image.Image) -> Image.Image:
-    w, h = img.size
-
-    if max(w, h) <= MAX_SIDE:
-        return img
+    # Reduce quality quickly (no loop → faster)
+    buffer = BytesIO()
+    img.save(buffer, format="JPEG", quality=70, optimize=True)
+    buffer.seek(0)
 
-    scale = MAX_SIDE / max(w, h)
-    new_size = (int(w * scale), int(h * scale))
+    return Image.open(buffer).convert("RGB")
 
-    print(f"[INFO] Downscaling {w} → {new_size}")
-    return img.resize(new_size, Image.BILINEAR)
 
-
-def transform(img: Image.Image) -> torch.Tensor:
+def transform(img):
     img = img.resize(TARGET_SIZE, Image.BILINEAR)
 
-    arr = np.array(img).astype(np.float32) / 255.0
+    arr = np.asarray(img, dtype=np.float32) / 255.0
 
     mean = np.array([0.485, 0.456, 0.406])
     std = np.array([0.229, 0.224, 0.225])
@@ -121,12 +89,13 @@ def transform(img: Image.Image) -> torch.Tensor:
     return torch.from_numpy(arr).unsqueeze(0).to(device=device, dtype=dtype)
 
 
-def run_inference(img: Image.Image) -> Image.Image:
+# 🔥 FAST inference
+def remove_background(img: Image.Image):
     orig_size = img.size
     tensor = transform(img)
 
-    with torch.inference_mode():  # 🔥 faster
-        pred = birefnet(tensor)
+    with torch.inference_mode():
+        pred = model(tensor)
         pred = pred[-1] if isinstance(pred, (list, tuple)) else pred
         pred = pred.sigmoid()[0, 0].cpu()
 
@@ -139,39 +108,28 @@ def run_inference(img: Image.Image) -> Image.Image:
 
 
 # ---------------------------------------------------------
-# FastAPI app
+# FastAPI
 # ---------------------------------------------------------
-app = FastAPI(title="Fast Background Remover API")
+app = FastAPI()
 
-# ---------------------------------------------------------
-# Redirect GET
-# ---------------------------------------------------------
-@app.get("/remove-background")
-async def redirect_to_post():
-    return JSONResponse(
-        {"detail": "Use POST /remove-background"},
-        status_code=405
-    )
-
-# ---------------------------------------------------------
-# Main endpoint
-# ---------------------------------------------------------
 @app.post("/remove-background")
 async def remove_bg(file: UploadFile = File(None), image_url: str = Form(None)):
     try:
         if file:
             raw = await file.read()
             img = Image.open(BytesIO(raw)).convert("RGB")
+
+            # ✅ Step 1: compress if >5MB
             img = compress_if_needed(img, raw)
 
         elif image_url:
             img = load_image_from_url(image_url)
 
         else:
-            raise HTTPException(status_code=400, detail="Provide file or image_url")
+            raise HTTPException(400, "Provide file or URL")
 
-        img = auto_downscale(img)
-        result = run_inference(img)
+        # ✅ Step 2: remove background
+        result = remove_background(img)
 
         buf = BytesIO()
         result.save(buf, format="PNG")
@@ -180,16 +138,16 @@ async def remove_bg(file: UploadFile = File(None), image_url: str = Form(None)):
         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
+# Simple UI
 # ---------------------------------------------------------
 @app.get("/", response_class=HTMLResponse)
-async def ui():
+async def home():
     return """
-    <html>
+<html>
     <head>
         <title>Fast Background Remover</title>
         <link rel='stylesheet'
@@ -265,6 +223,7 @@ async def ui():
     """
 
 
+
 # ---------------------------------------------------------
 # Run
 # ---------------------------------------------------------