Update app.py

app.py

@@ -6,15 +6,11 @@ import torch
 import torch.nn.functional as F
 import numpy as np
 import io
+import sys
 import os
 import asyncio
 from torchvision.transforms.functional import normalize
-
-# ── BiRefNet / RMBG-2.0 Model Import ──────────────────────────────
-# RMBG-2.0 is Bria's production release built on the BiRefNet architecture.
-# It scores ~93 on DIS5K benchmarks vs ~86 for RMBG-1.4 — the jump that
-# takes hair/fur edge quality from "Good" to "Excellent".
-from transformers import AutoModelForImageSegmentation
+from safetensors.torch import load_file
 
 # ── CPU Thread Tuning ──────────────────────────────────────────────
 # Must be set before any model loading
@@ -32,7 +28,7 @@ import torch._dynamo
 torch._dynamo.config.suppress_errors = True
 
 # ── App Setup ──────────────────────────────────────────────────────
-app = FastAPI(title="Lampdra API", version="2.0")
+app = FastAPI(title="Lampdra API", version="1.0")
 
 ALLOWED_ORIGINS = [
     "https://lampdra.com",
@@ -74,36 +70,31 @@ MAX_PIXELS    = 10000 * 10000       # 100MP — generous for real photos
 POOL_SIZE = 1
 DEVICE    = "cuda" if torch.cuda.is_available() else "cpu"
 
-# ── BiRefNet Preprocessing Constants ──────────────────────────────
-# IMPORTANT CHANGE from RMBG-1.4:
-# RMBG-1.4 used mean=[0.5,0.5,0.5] / std=[1.0,1.0,1.0] (simple centre-scale).
-# RMBG-2.0 / BiRefNet uses standard ImageNet normalisation.
-# Using the wrong constants produces grey/washed-out masks — hair especially.
-BIREFNET_MEAN = [0.485, 0.456, 0.406]
-BIREFNET_STD  = [0.229, 0.224, 0.225]
-INPUT_SIZE    = (1024, 1024)   # BiRefNet native resolution — unchanged from RMBG-1.4
+sys.path.insert(0, "/app")
+from rmbg.briarmbg import BriaRMBG
+
+# Load weights from disk once — all pool copies share this read
+print(f"Device: {DEVICE}")
+print(f"Loading weights from disk...")
+state_dict = load_file("/app/rmbg/model.safetensors")
 
 # Pool queue — asyncio.Queue is thread-safe and blocks callers
 # when all copies are busy until one becomes free
 model_pool = asyncio.Queue()
 
-def build_model_copy() -> torch.nn.Module:
-    """
-    Load one copy of RMBG-2.0 (BiRefNet architecture).
-    Weights are downloaded from HF hub on first run and cached at
-    ~/.cache/huggingface/hub — subsequent Space restarts are instant.
-    trust_remote_code=True is required by the RMBG-2.0 model card.
-    """
-    m = AutoModelForImageSegmentation.from_pretrained(
-        "briaai/RMBG-2.0",
-        trust_remote_code=True,
-    )
+def build_model_copy():
+    """Build one independent copy of the model."""
+    m = BriaRMBG()
+    m.load_state_dict(state_dict)
     m.eval()
     m.to(DEVICE)
+    # torch.compile — compiles model to native machine code
+    # First inference is slower (compile happens then)
+    # Every subsequent call is 20-40% faster
     try:
         m = torch.compile(m)
     except Exception:
-        pass  # safe to skip if g++ / inductor unavailable
+        pass  # compile not available on all platforms, safe to skip
     return m
 
 # ── Ready flag — Cause 3 fix ──────────────────────────────────────
@@ -119,29 +110,29 @@ _model_ready = False
 @app.on_event("startup")
 async def startup_event():
     global _model_ready
-    dummy = torch.zeros(1, 3, INPUT_SIZE[0], INPUT_SIZE[1]).to(DEVICE)
+    dummy = torch.zeros(1, 3, 1024, 1024).to(DEVICE)
     for i in range(POOL_SIZE):
         m = build_model_copy()
+        # Warmup in thread — event loop stays free, but pool is still empty
+        # so router cannot send requests here even if it tries
         def _run_warmup(model=m):
             with torch.inference_mode():
                 model(dummy)
         await asyncio.to_thread(_run_warmup)
-        print(f"  Copy {i + 1}/{POOL_SIZE} warmed up (RMBG-2.0 / BiRefNet)")
+        print(f"  Copy {i + 1}/{POOL_SIZE} warmed up and compiled")
         model_pool.put_nowait(m)  # slot added AFTER warmup — not before
     del dummy
-    _model_ready = True
-    print(f"Pool ready — {POOL_SIZE} BiRefNet copies loaded, accepting requests.")
+    _model_ready = True           # only NOW does root() report real pool_available
+    print(f"Pool ready — {POOL_SIZE} copies loaded, warmed up, accepting requests.")
 
 # ── Preprocessing ──────────────────────────────────────────────────
-def preprocess(image: Image.Image) -> torch.Tensor:
-    """
-    Resize to 1024×1024 and apply ImageNet normalisation.
-    BOX filter is fastest for downscaling (18ms vs 27ms for BILINEAR).
-    """
-    img    = image.convert("RGB").resize(INPUT_SIZE, Image.BOX)
+def preprocess(image: Image.Image, size=(1024, 1024)):
+    # All images resize to 1024×1024 here regardless of original size
+    # This is why the model RAM usage is always predictable
+    img    = image.convert("RGB").resize(size, Image.BOX)  # BOX=18ms vs BILINEAR=27ms for downscale
     arr    = np.array(img, dtype=np.float32) / 255.0
     tensor = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).to(DEVICE)
-    tensor = normalize(tensor, BIREFNET_MEAN, BIREFNET_STD)
+    tensor = normalize(tensor, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0])
     return tensor
 
 # ── Routes ─────────────────────────────────────────────────────────
@@ -149,11 +140,12 @@ def preprocess(image: Image.Image) -> torch.Tensor:
 @app.head("/")
 def root():
     return {
-        "status"        : "Lampdra API is running",
-        "model"         : "RMBG-2.0 (BiRefNet)",
-        "device"        : DEVICE,
-        "pool_size"     : POOL_SIZE,
-        "pool_available": model_pool.qsize() if _model_ready else 0,
+        "status"          : "Lampdra API is running",
+        "device"          : DEVICE,
+        "pool_size"       : POOL_SIZE,
+        # Returns 0 until warmup completes — prevents router sending
+        # requests to a Space that is still loading or compiling
+        "pool_available"  : model_pool.qsize() if _model_ready else 0,
     }
 
 @app.post("/remove-background")
@@ -258,8 +250,7 @@ async def remove_background(
         return Response(content=output_buffer.getvalue(), media_type="image/webp")
 
     finally:
-        # Always return the model to the pool — even on exceptions.
-        # put_nowait() is correct here: queue has no maxsize so it never
-        # raises QueueFull, and it skips coroutine overhead on every request.
-        # Consistent with startup_event which also uses put_nowait.
-        model_pool.put_nowait(model)
+        # ALWAYS return the model copy to the pool
+        # This runs even if the request crashes with any Python-level error
+        # Pool never permanently loses a slot due to soft crashes
+        await model_pool.put(model)