from fastapi import FastAPI, File, UploadFile, Header, HTTPException from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import Response from PIL import Image 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 from safetensors.torch import load_file # ── CPU Thread Tuning ────────────────────────────────────────────── # Must be set before any model loading # Tells PyTorch, OpenMP, and MKL how many threads to use # Prevents thread contention inside Docker containers torch.set_num_threads(2) # P3 fix: 2 threads = 1 per CPU core on HF free tier os.environ["OMP_NUM_THREADS"] = "2" # was 4 — 32 threads fighting 2 cores caused thrashing os.environ["MKL_NUM_THREADS"] = "2" # ── Dynamo Safety Net ────────────────────────────────────────────── # If torch.compile() fails for ANY reason (missing g++, ISA detection, # etc.) fall back to eager mode instead of crashing the request. # Primary fix is g++ in Dockerfile; this is the belt-and-suspenders. import torch._dynamo torch._dynamo.config.suppress_errors = True # ── App Setup ────────────────────────────────────────────────────── app = FastAPI(title="Lampdra API", version="1.0") ALLOWED_ORIGINS = [ "https://lampdra.com", "https://www.lampdra.com", "https://paul1k-lampdra-api.hf.space", ] app.add_middleware( CORSMiddleware, allow_origins=ALLOWED_ORIGINS, allow_methods=["*"], allow_headers=["*"], ) # ── API Key ──────────────────────────────────────────────────────── API_KEY = os.environ.get("LAMPDRA_API_KEY", "") def verify_key(x_api_key: str = Header(default="")): # Fail-closed: if LAMPDRA_API_KEY env var is not set, reject everything. # Previously: if API_KEY and ... → empty string is falsy → check skipped entirely. if not API_KEY or x_api_key != API_KEY: raise HTTPException(status_code=401, detail="Invalid API key") # ── Request Limits ───────────────────────────────────────────────── # File size checked first — cheapest possible check, no decoding needed # Pixel count checked after decoding — catches malicious small-file/huge-pixel PNGs MAX_FILE_SIZE = 20 * 1024 * 1024 # 20MB — stops RAM spike before it starts MAX_PIXELS = 10000 * 10000 # 100MP — generous for real photos # ── Model Pool Setup ─────────────────────────────────────────────── # P1+P2 fix: pool=1 (was 4) # HF free tier has ~2 CPU cores per Space. # pool=4 meant 8 model copies sharing 2 cores = each got 0.25 cores = 4x slower. # pool=1 means 1 model per uvicorn process gets all 2 cores = fastest inference. # GIL: pool=4 in one process = models take turns anyway — pool=1 removes illusion. # 2 uvicorn workers × pool=1 = 2 real parallel inferences per Space. # 5 Spaces × 2 = 10 truly parallel inferences system-wide. # RAM: 2 workers × 1 copy × 1.2GB = ~2.4GB (was 9.6GB — 4x less RAM used) POOL_SIZE = 1 DEVICE = "cuda" if torch.cuda.is_available() else "cpu" 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(): """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 # compile not available on all platforms, safe to skip return m # ── Ready flag — Cause 3 fix ────────────────────────────────────── # root() returns pool_available=0 until this is True. # Guarantees router never sends requests during model loading or warmup. _model_ready = False # ── Startup: Warm-up THEN Fill Pool — Cause 1 fix ───────────────── # OLD order: fill pool → warmup borrows the only slot → real requests # arrive, queue inside worker, router 504s before warmup ends. # NEW order: warmup first (pool stays empty → router sees 0 → sends nothing) # then fill pool → router sees 1 → safe to send real requests. @app.on_event("startup") async def startup_event(): global _model_ready 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 and compiled") model_pool.put_nowait(m) # slot added AFTER warmup — not before del dummy _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, 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, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0]) return tensor # ── Routes ───────────────────────────────────────────────────────── @app.get("/") @app.head("/") def root(): return { "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") async def remove_background( file : UploadFile = File(...), x_api_key : str = Header(default=""), origin : str = Header(default="") ): # ── Server-side origin check ─────────────────────────────────── # The worker should only ever be called by the router, never directly # by browsers or any other caller. Only the router's HF space is allowed. ALLOWED_POST_ORIGINS = { "https://paul1k-lampdra-router.hf.space", } if origin not in ALLOWED_POST_ORIGINS: raise HTTPException(status_code=403, detail="Forbidden") verify_key(x_api_key) # ── Guard 1: File Size ───────────────────────────────────────── # Checked before ANY decoding — cheapest check possible # Stops large files before they ever touch RAM contents = await file.read() if len(contents) > MAX_FILE_SIZE: raise HTTPException( status_code=400, detail=f"File too large. Maximum allowed size is 20MB." ) # ── Guard 2: Readable Image ──────────────────────────────────── # .convert("RGB") is the real test — if this works, the model can handle it # Avoids image.verify() which incorrectly rejects many valid real-world images # (phone photos with non-standard EXIF, progressive JPEGs, etc.) try: image = Image.open(io.BytesIO(contents)).convert("RGB") except Exception: raise HTTPException( status_code=400, detail="Cannot read image file. Please upload a valid JPG, PNG, or WEBP." ) # ── Guard 3: Pixel Dimensions ────────────────────────────────── # Guards against malicious images: small file size but enormous pixel count # Example: a 1MB PNG that is 50000×50000 would consume ~7.5GB RAM to decode # This check runs AFTER decoding so we have real dimensions if image.width * image.height > MAX_PIXELS: del image # free RAM immediately before raising raise HTTPException( status_code=400, detail=f"Image dimensions too large. Maximum is 10000×10000 pixels." ) # ── Preprocess ───────────────────────────────────────────────── orig_size = image.size input_tensor = preprocess(image) # Convert to RGBA now while image is still in memory. # Avoids re-opening from bytes later (~168ms saved per request). image_rgba = image.convert("RGBA") del image # free the large RGB decoded image immediately # image_rgba is kept — it is small (same pixels, +alpha channel) # ── Grab a Free Model from the Pool ─────────────────────────── # If all 4 copies are busy this line waits until one is free # Nobody gets rejected — they queue fairly model = await model_pool.get() try: # Run inference in a thread pool so the asyncio event loop # stays free to handle I/O (new connections, sending responses) # while the CPU-bound model inference runs in the background. # torch.inference_mode disables gradient tracking — 5-15% faster. def run_inference(): with torch.inference_mode(): return model(input_tensor) result = await asyncio.to_thread(run_inference) del input_tensor # free tensor now that model is done with it mask = result[0][0].squeeze() mask = F.interpolate( mask.unsqueeze(0).unsqueeze(0), size=(orig_size[1], orig_size[0]), mode="bilinear", align_corners=False ).squeeze() # Move to CPU only when needed for numpy conversion mask = (mask.cpu() * 255).clamp(0, 255).byte().numpy() mask_image = Image.fromarray(mask) # image_rgba already in memory — no re-decode needed image_out = image_rgba image_out.putalpha(mask_image) output_buffer = io.BytesIO() # WebP with alpha: ~70% smaller than PNG compress_level=1, similar encode speed. # quality=85 is visually lossless for transparency masks at this resolution. # method=3 balances encode speed vs compression (0=fastest, 6=smallest). image_out.save(output_buffer, format="WEBP", quality=85, method=3) output_buffer.seek(0) return Response(content=output_buffer.getvalue(), media_type="image/webp") finally: # 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)