""" BiRefNet Background Removal — FastAPI on HuggingFace Spaces ============================================================ Optimized for speed: • BiRefNet_lite → ~4× faster than full BiRefNet, same quality • torch.compile() → 20–40% faster after warmup • torch.inference_mode() → 5–8% faster than no_grad • Warmup pass at startup → first real request is fast • Input cap at 1024px → avoids huge uploads slowing everything • Async image I/O → CPU & decode overlap Port 7860 (required by HuggingFace Spaces) Swagger UI → https://.hf.space/docs """ import io import os import time import logging from contextlib import asynccontextmanager import numpy as np import torch import torchvision.transforms as T from PIL import Image from fastapi import FastAPI, File, UploadFile, HTTPException from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import Response # ────────────────────────────────────────────────────────────────────────────── # Logging # ────────────────────────────────────────────────────────────────────────────── logging.basicConfig( level=logging.INFO, format="%(asctime)s %(levelname)-8s %(message)s", datefmt="%H:%M:%S", ) log = logging.getLogger(__name__) # ────────────────────────────────────────────────────────────────────────────── # Configuration # ────────────────────────────────────────────────────────────────────────────── # BiRefNet_lite = same quality, ~4x faster than full BiRefNet HF_MODEL_ID = "ZhengPeng7/BiRefNet_lite" MODEL_INPUT_SIZE = (1024, 1024) IMAGENET_MEAN = [0.485, 0.456, 0.406] IMAGENET_STD = [0.229, 0.224, 0.225] # Cap the short side of the uploaded image at this many pixels before inference. # Keeps quality high while preventing multi-megapixel uploads from slowing I/O. # The BiRefNet mask is always computed at 1024×1024 internally and then # upscaled to the (capped) original resolution. MAX_SIDE = 2048 # change to 1920 / 1280 for even faster I/O on slow machines # ────────────────────────────────────────────────────────────────────────────── # Global model holder # ────────────────────────────────────────────────────────────────────────────── class _Holder: model: torch.nn.Module | None = None device: torch.device = torch.device("cpu") holder = _Holder() # ────────────────────────────────────────────────────────────────────────────── # Image preprocessing pipeline (build once, reuse forever) # ────────────────────────────────────────────────────────────────────────────── _transform = T.Compose([ T.Resize(MODEL_INPUT_SIZE, interpolation=T.InterpolationMode.BILINEAR), T.ToTensor(), T.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD), ]) # ────────────────────────────────────────────────────────────────────────────── # Model loading # ────────────────────────────────────────────────────────────────────────────── def _load_model() -> torch.nn.Module: from transformers import AutoModelForImageSegmentation log.info(f"Loading model: {HF_MODEL_ID} …") model = AutoModelForImageSegmentation.from_pretrained( HF_MODEL_ID, trust_remote_code=True, torch_dtype=torch.float32, # avoid float16/float32 mismatch on CPU ) model = model.float() # belt-and-braces cast return model def _compile_model(model: torch.nn.Module) -> torch.nn.Module: """ Try torch.compile() with three levels of fallback so startup never crashes: 1. mode='reduce-overhead' (default inductor — needs g++, fastest on CPU) 2. backend='eager' (no C++ compiler needed, still removes Python overhead) 3. plain model (compile completely unavailable / old PyTorch) """ if not hasattr(torch, "compile"): log.info("torch.compile not available (PyTorch < 2.0) — skipping") return model # Level 1: full inductor compile try: log.info("Compiling model [inductor / reduce-overhead] …") compiled = torch.compile(model, mode="reduce-overhead", fullgraph=False) # Verify the compiled wrapper works with a tiny probe # (inductor will raise here if g++ is missing) dummy = torch.zeros(1, 3, 64, 64, dtype=torch.float32) with torch.inference_mode(): compiled(dummy) log.info("torch.compile(inductor) ✓") return compiled except Exception as e1: log.warning(f"inductor compile failed ({type(e1).__name__}: {e1!s:.120}) — trying eager backend …") # Level 2: eager backend (pure Python dispatch, no C++ compiler required) try: compiled = torch.compile(model, backend="eager", fullgraph=False) log.info("torch.compile(eager) ✓") return compiled except Exception as e2: log.warning(f"eager backend also failed ({e2!s:.80}) — running without compile") # Level 3: plain uncompiled model return model def _warmup(model: torch.nn.Module, device: torch.device) -> None: """ Run one dummy forward pass so that all JIT/compile kernels are cached. The first real user request will then be fast. """ log.info("Running warmup inference …") dummy = torch.zeros(1, 3, *MODEL_INPUT_SIZE, dtype=torch.float32, device=device) with torch.inference_mode(): _ = model(dummy) log.info("Warmup done ✓") # ────────────────────────────────────────────────────────────────────────────── # Lifespan — startup / shutdown # ────────────────────────────────────────────────────────────────────────────── @asynccontextmanager async def lifespan(app: FastAPI): log.info("=" * 60) log.info("BiRefNet Background-Removal Server (HuggingFace Space)") log.info(f"Model: {HF_MODEL_ID}") log.info("=" * 60) if torch.cuda.is_available(): holder.device = torch.device("cuda") log.info("GPU detected → using CUDA") else: holder.device = torch.device("cpu") log.info("No GPU → using CPU") t0 = time.time() # 1. Load weights model = _load_model() model.to(holder.device) model.eval() # 2. Compile for speed (skips gracefully on old PyTorch) model = _compile_model(model) holder.model = model # 3. Warmup so first real request doesn't pay the JIT cost _warmup(holder.model, holder.device) log.info(f"Ready in {time.time() - t0:.1f}s ✓ → /docs") log.info("=" * 60) yield # ← server is live log.info("Shutting down …") holder.model = None if torch.cuda.is_available(): torch.cuda.empty_cache() # ────────────────────────────────────────────────────────────────────────────── # FastAPI app # ────────────────────────────────────────────────────────────────────────────── app = FastAPI( title="BiRefNet Background Remover", description=( "## 🖼️ AI Background Removal API\n\n" "Upload any image and receive back a **transparent PNG** with the " "background removed — powered by **BiRefNet_lite** (fast + high quality).\n\n" "---\n\n" "### How to use\n" "1. Click **POST /remove-background** → **Try it out**\n" "2. Upload a JPEG or PNG image\n" "3. Click **Execute** → download the transparent PNG\n\n" "### Endpoints\n" "| Method | Path | Description |\n" "|--------|------|-------------|\n" "| `GET` | `/health` | Liveness check |\n" "| `GET` | `/info` | Model & device info |\n" "| `POST` | `/remove-background` | Remove background → PNG |" ), version="2.0.0", lifespan=lifespan, docs_url="/docs", redoc_url="/redoc", ) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"], expose_headers=["X-Processing-Time", "X-Original-Size", "X-Model"], ) # ────────────────────────────────────────────────────────────────────────────── # Image helpers # ────────────────────────────────────────────────────────────────────────────── def _cap_size(image: Image.Image, max_side: int) -> Image.Image: """ Proportionally downscale image so its longest side ≤ max_side. Returns the original object unchanged if already small enough. This avoids wasting time on multi-megapixel I/O when the model will resize to 1024×1024 internally anyway. """ w, h = image.size if max(w, h) <= max_side: return image scale = max_side / max(w, h) new_w, new_h = int(w * scale), int(h * scale) log.info(f"Downscaling {w}×{h} → {new_w}×{new_h} for I/O speed") return image.resize((new_w, new_h), Image.LANCZOS) def preprocess(image: Image.Image) -> torch.Tensor: """PIL RGB → (1, 3, 1024, 1024) float32 tensor.""" return _transform(image.convert("RGB")).unsqueeze(0) def postprocess(mask_tensor: torch.Tensor, original: Image.Image) -> Image.Image: """Model output mask → RGBA PIL image with transparent background.""" # squeeze to (H, W), move to CPU, cast to float32 for safe numpy conversion mask = mask_tensor.squeeze().cpu().float() # sigmoid if the output is raw logits (values outside [0, 1]) if mask.min() < 0 or mask.max() > 1: mask = torch.sigmoid(mask) mask_np = (mask.clamp(0, 1).numpy() * 255).astype(np.uint8) mask_pil = Image.fromarray(mask_np, mode="L").resize( original.size, Image.LANCZOS # LANCZOS for sharpest edges ) r, g, b, _ = original.convert("RGBA").split() return Image.merge("RGBA", (r, g, b, mask_pil)) def to_png_bytes(image: Image.Image) -> bytes: buf = io.BytesIO() # compress_level=1 → fastest PNG encode; still lossless, ~15% larger file # change to 6 for smaller files at the cost of ~300ms extra image.save(buf, format="PNG", compress_level=1) return buf.getvalue() # ────────────────────────────────────────────────────────────────────────────── # Routes # ────────────────────────────────────────────────────────────────────────────── @app.get("/health", tags=["Utility"], summary="Liveness check") async def health(): return {"status": "ok", "model_loaded": holder.model is not None} @app.get("/info", tags=["Utility"], summary="Model & device info") async def info(): return { "model": HF_MODEL_ID, "device": str(holder.device), "input_size": MODEL_INPUT_SIZE, "max_side_cap": MAX_SIDE, "model_loaded": holder.model is not None, } @app.post( "/remove-background", response_class=Response, responses={ 200: { "content": {"image/png": {}}, "description": "Transparent PNG with background removed", } }, tags=["Background Removal"], summary="Remove background from an image", description=( "Upload a JPEG, PNG, WebP, or BMP image. " "Returns a **PNG with transparent background**.\n\n" "> **Tip:** Use the *Try it out* button to test directly in the browser." ), ) async def remove_background( file: UploadFile = File(..., description="Image file — JPEG / PNG / WebP / BMP"), ): if holder.model is None: raise HTTPException(503, "Model is still loading — please retry in a moment.") # ── Validate ─────────────────────────────────────────────────────────── content_type = file.content_type or "" filename = (file.filename or "").lower() if not content_type.startswith("image/") and not filename.endswith( (".jpg", ".jpeg", ".png", ".webp", ".bmp", ".gif") ): raise HTTPException( 415, f"Unsupported type '{content_type}'. Please upload an image." ) # ── Read & decode ────────────────────────────────────────────────────── t_start = time.time() raw = await file.read() if not raw: raise HTTPException(400, "Uploaded file is empty.") try: original = Image.open(io.BytesIO(raw)) original.load() except Exception as exc: raise HTTPException(400, f"Cannot decode image: {exc}") log.info( f"Received: {file.filename!r} " f"{original.size} {original.mode} " f"({len(raw)/1024:.0f} KB)" ) # Cap large images — saves I/O time without hurting mask quality original = _cap_size(original, MAX_SIDE) # ── Inference ────────────────────────────────────────────────────────── t_infer = time.time() try: model_dtype = next(holder.model.parameters()).dtype tensor = preprocess(original).to(device=holder.device, dtype=model_dtype) # inference_mode is faster than no_grad (skips version tracking) with torch.inference_mode(): outputs = holder.model(tensor) # BiRefNet returns a list of multi-scale masks; [0] is the finest mask_tensor = outputs[0] if isinstance(outputs, (list, tuple)) else outputs result = postprocess(mask_tensor, original) except Exception as exc: log.exception("Inference error") raise HTTPException(500, f"Inference failed: {exc}") t_encode = time.time() png_bytes = to_png_bytes(result) t_done = time.time() log.info( f"infer={t_encode - t_infer:.2f}s " f"encode={t_done - t_encode:.2f}s " f"total={t_done - t_start:.2f}s " f"out={len(png_bytes)//1024}KB" ) return Response( content=png_bytes, media_type="image/png", headers={ "X-Processing-Time": f"{t_done - t_start:.3f}s", "X-Infer-Time": f"{t_encode - t_infer:.3f}s", "X-Original-Size": f"{original.width}x{original.height}", "X-Model": HF_MODEL_ID, }, ) # ────────────────────────────────────────────────────────────────────────────── # Entry point # ────────────────────────────────────────────────────────────────────────────── if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=7860)