import os import time import glob import math from dataclasses import dataclass from typing import Dict, Optional, Tuple, List import gradio as gr import spaces import numpy as np from PIL import Image import torch from torchvision import transforms from transformers import AutoModelForImageSegmentation from huggingface_hub import hf_hub_download # Slider component from gradio_imageslider import ImageSlider # InSPyReNet wrapper from transparent_background import Remover # rembg (U2Net + IS-Net via ONNX) from rembg import new_session, remove as rembg_remove # withoutBG (4-stage ONNX pipeline) from withoutbg import WithoutBG # ---------------------------- # Utilities # ---------------------------- def pil_to_rgb(pil: Image.Image) -> Image.Image: if pil.mode != "RGB": return pil.convert("RGB") return pil def ensure_rgba(pil: Image.Image) -> Image.Image: if pil.mode != "RGBA": return pil.convert("RGBA") return pil def make_checkerboard(w: int, h: int, block: int = 16) -> Image.Image: cols = int(math.ceil(w / block)) rows = int(math.ceil(h / block)) board = np.zeros((rows * block, cols * block, 3), dtype=np.uint8) c1, c2 = np.array([235, 235, 235], dtype=np.uint8), np.array([200, 200, 200], dtype=np.uint8) for r in range(rows): for c in range(cols): color = c1 if (r + c) % 2 == 0 else c2 board[r * block:(r + 1) * block, c * block:(c + 1) * block] = color return Image.fromarray(board[:h, :w, :], mode="RGB") def rgba_on_checkerboard(rgba: Image.Image) -> Image.Image: rgba = ensure_rgba(rgba) w, h = rgba.size bg = make_checkerboard(w, h) comp = Image.alpha_composite(bg.convert("RGBA"), rgba) return comp.convert("RGB") def save_temp_png(rgba: Image.Image, out_dir: str = "output_images") -> str: os.makedirs(out_dir, exist_ok=True) path = os.path.join(out_dir, "no_bg.png") ensure_rgba(rgba).save(path, format="PNG") return path def now_ms() -> float: return time.perf_counter() * 1000.0 def get_device() -> str: """Get device at runtime (important for ZeroGPU).""" return "cuda" if torch.cuda.is_available() else "cpu" @dataclass class Timing: preprocess_ms: float inference_ms: float postprocess_ms: float total_ms: float def to_text(self) -> str: return ( f"preprocess: {self.preprocess_ms:.2f} ms\n" f"inference: {self.inference_ms:.2f} ms\n" f"postprocess: {self.postprocess_ms:.2f} ms\n" f"TOTAL: {self.total_ms:.2f} ms" ) # ---------------------------- # Model Manager # ---------------------------- class ModelManager: def __init__(self): self._inspy: Optional[Remover] = None self._withoutbg: Optional[object] = None self._withoutbg_had_gpu: bool = False # Track if withoutBG was loaded with GPU self._torch_models: Dict[str, torch.nn.Module] = {} self._torch_model_on_gpu: Optional[str] = None self._rembg_sessions: Dict[str, object] = {} self._model_load_errors: Dict[str, str] = {} self._tf_1024 = transforms.Compose([ transforms.Resize((1024, 1024)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) try: torch.set_float32_matmul_precision("high") except Exception: pass def _maybe_sync(self): if get_device() == "cuda": torch.cuda.synchronize() def _load_inspy(self) -> Remover: if self._inspy is None: self._inspy = Remover(jit=False) return self._inspy def _load_withoutbg(self, force_reload: bool = False): """ Load withoutBG model. Automatically reloads if GPU became available after initial load. """ gpu_available_now = torch.cuda.is_available() # Reload if: forced, not loaded yet, or GPU is now available but wasn't before need_reload = ( force_reload or self._withoutbg is None or (gpu_available_now and not self._withoutbg_had_gpu) ) if need_reload: self._withoutbg = WithoutBG.opensource() self._withoutbg_had_gpu = gpu_available_now return self._withoutbg def _offload_torch_models_from_gpu(self, keep_name: str): if get_device() != "cuda": return if self._torch_model_on_gpu and self._torch_model_on_gpu != keep_name: prev = self._torch_models.get(self._torch_model_on_gpu) if prev is not None: prev.to("cpu") self._torch_model_on_gpu = None torch.cuda.empty_cache() def _load_torch_model(self, key: str) -> torch.nn.Module: """Load BiRefNet or BRIA RMBG 2.0 model.""" if key in self._torch_models: return self._torch_models[key] if key in self._model_load_errors: raise RuntimeError(self._model_load_errors[key]) model_configs = { "birefnet": "ZhengPeng7/BiRefNet", "bria_rmbg_2": "briaai/RMBG-2.0", } if key not in model_configs: raise ValueError(f"Unknown model key: {key}") model_id = model_configs[key] try: m = AutoModelForImageSegmentation.from_pretrained( model_id, trust_remote_code=True ) m.eval() m.to("cpu") self._torch_models[key] = m return m except OSError as e: error_msg = str(e) if "gated" in error_msg.lower() or "401" in error_msg or "access" in error_msg.lower(): self._model_load_errors[key] = ( f"Model '{model_id}' requires license acceptance.\n" f"1. Go to https://huggingface.co/{model_id}\n" f"2. Accept the license agreement\n" f"3. Add HF_TOKEN secret to your Space settings" ) else: self._model_load_errors[key] = f"Failed to load {model_id}: {error_msg}" raise RuntimeError(self._model_load_errors[key]) except ImportError as e: self._model_load_errors[key] = ( f"Import error loading {model_id}: {e}\n" f"Make sure 'timm' is in requirements.txt" ) raise RuntimeError(self._model_load_errors[key]) def _get_rembg_session(self, name: str): if name in self._rembg_sessions: return self._rembg_sessions[name] providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] try: sess = new_session(name, providers=providers) except Exception: sess = new_session(name) self._rembg_sessions[name] = sess return sess def _run_torch_alpha_model(self, model_key: str, image_rgb: Image.Image) -> Image.Image: device = get_device() m = self._load_torch_model(model_key) if device == "cuda": self._offload_torch_models_from_gpu(keep_name=model_key) if self._torch_model_on_gpu != model_key: m.to("cuda") self._torch_model_on_gpu = model_key image_rgb = pil_to_rgb(image_rgb) orig_size = image_rgb.size x = self._tf_1024(image_rgb).unsqueeze(0).to(device) with torch.inference_mode(): if device == "cuda": with torch.autocast(device_type="cuda", dtype=torch.float16): preds = m(x)[-1].sigmoid() else: preds = m(x)[-1].sigmoid() pred = preds[0].squeeze().detach().float().cpu() alpha = transforms.ToPILImage()(pred).resize(orig_size, Image.BILINEAR) out = image_rgb.convert("RGBA") out.putalpha(alpha) return out def run(self, model_name: str, input_image: Image.Image) -> Tuple[Image.Image, Timing]: if input_image is None: raise ValueError("No input image") t0 = now_ms() # Preprocess pre0 = now_ms() img_rgb = pil_to_rgb(input_image) pre1 = now_ms() # Inference inf0 = now_ms() if model_name == "InSPyReNet": remover = self._load_inspy() mask = remover.process(input_image, type="map") if isinstance(mask, Image.Image): mask = mask.convert("L") else: mask = Image.fromarray((mask * 255).astype(np.uint8), mode="L") out = img_rgb.convert("RGBA") out.putalpha(mask) elif model_name == "BiRefNet": out = self._run_torch_alpha_model("birefnet", img_rgb) elif model_name == "U2Net": sess = self._get_rembg_session("u2net") out = rembg_remove(img_rgb, session=sess) out = ensure_rgba(out) elif model_name == "BRIA RMBG 2.0": out = self._run_torch_alpha_model("bria_rmbg_2", img_rgb) elif model_name == "IS-Net": sess = self._get_rembg_session("isnet-general-use") out = rembg_remove(img_rgb, session=sess) out = ensure_rgba(out) elif model_name == "withoutBG": # Will auto-reload if GPU became available (ZeroGPU) model = self._load_withoutbg() out = model.remove_background(img_rgb) out = ensure_rgba(out) else: raise ValueError(f"Unknown model: {model_name}") self._maybe_sync() inf1 = now_ms() # Postprocess post0 = now_ms() out = ensure_rgba(out) post1 = now_ms() t1 = now_ms() timing = Timing( preprocess_ms=pre1 - pre0, inference_ms=inf1 - inf0, postprocess_ms=post1 - post0, total_ms=t1 - t0, ) return out, timing MANAGER = ModelManager() MODEL_CHOICES = [ "InSPyReNet", "BiRefNet", "U2Net", "BRIA RMBG 2.0", "IS-Net", "withoutBG", ] # ---------------------------- # Gradio handlers # ---------------------------- @spaces.GPU def run_single(model_name: str, image: Image.Image): if image is None: return None, None, "Upload an image first.", None try: out_rgba, timing = MANAGER.run(model_name, image) preview = rgba_on_checkerboard(out_rgba) out_path = save_temp_png(out_rgba) return (image, preview), out_rgba, timing.to_text(), out_path except RuntimeError as e: return None, None, f"Error: {str(e)}", None except Exception as e: return None, None, f"Unexpected error: {str(e)}", None def list_bench_images() -> List[str]: exts = ("*.jpg", "*.jpeg", "*.png", "*.webp") files = [] for e in exts: files += glob.glob(os.path.join("bench", e)) files = sorted(files) if not files: for f in ["1.jpg", "2.jpg", "3.png", "4.webp"]: if os.path.exists(f): files.append(f) return files @spaces.GPU def run_benchmark(model_name: str, repeats: int = 1): files = list_bench_images() if not files: return [], "No benchmark images found. Add 10–15 images under bench/." try: # Warmup warm_img = Image.open(files[0]).convert("RGB") for _ in range(2): _ = MANAGER.run(model_name, warm_img) rows = [] total_ms = 0.0 n_images = 0 for f in files: img = Image.open(f).convert("RGB") for r in range(repeats): out, timing = MANAGER.run(model_name, img) rows.append([ os.path.basename(f), r + 1, round(timing.total_ms, 2), round(timing.inference_ms, 2), ]) total_ms += timing.total_ms n_images += 1 avg_ms = total_ms / max(1, n_images) ips = 1000.0 / avg_ms if avg_ms > 0 else 0.0 summary = ( f"Model: {model_name}\n" f"Images: {len(files)} (repeats={repeats}) => runs={n_images}\n" f"Avg total: {avg_ms:.2f} ms\n" f"Estimated throughput: {ips:.2f} images/sec\n" f"Device: {'GPU' if torch.cuda.is_available() else 'CPU'}" ) return rows, summary except RuntimeError as e: return [], f"Error: {str(e)}" except Exception as e: return [], f"Unexpected error: {str(e)}" # ---------------------------- # UI # ---------------------------- with gr.Blocks(title="Background Removal Benchmark") as demo: gr.Markdown( """ # Background Removal Benchmark Benchmarked models: 1. **InSPyReNet** — transparent-background library 2. **BiRefNet** — ZhengPeng7/BiRefNet (requires `timm`) 3. **U2Net** — via rembg/ONNX 4. **BRIA RMBG 2.0** — briaai/RMBG-2.0 (requires license acceptance) 5. **IS-Net** — isnet-general-use via rembg 6. **withoutBG** — 4-stage ONNX pipeline (Depth → ISNet → Matting → Refiner) **Notes** - Output is true transparent PNG (RGBA) - Slider preview shows result on checkerboard - For benchmarks, add images under `bench/` folder ⚠️ **BRIA RMBG 2.0**: Requires accepting license at [huggingface.co/briaai/RMBG-2.0](https://huggingface.co/briaai/RMBG-2.0) and adding `HF_TOKEN` secret to Space settings. """ ) with gr.Tab("Try single image"): with gr.Row(): with gr.Column(scale=1): inp = gr.Image(type="pil", label="Upload image", height=420) model = gr.Dropdown(choices=MODEL_CHOICES, value="InSPyReNet", label="Model") run_btn = gr.Button("Run", variant="primary") with gr.Column(scale=2): slider = ImageSlider(label="Before / After", type="pil") out_img = gr.Image(type="pil", label="Output (RGBA)", height=420) timing_box = gr.Textbox(label="Timing / Errors", lines=5) out_file = gr.File(label="Download PNG (transparent)") run_btn.click( fn=run_single, inputs=[model, inp], outputs=[slider, out_img, timing_box, out_file] ) with gr.Tab("Benchmark (throughput estimate)"): with gr.Row(): with gr.Column(scale=1): bench_model = gr.Dropdown(choices=MODEL_CHOICES, value="InSPyReNet", label="Model") repeats = gr.Slider(1, 5, value=1, step=1, label="Repeats per image") bench_btn = gr.Button("Run benchmark", variant="primary") with gr.Column(scale=2): bench_table = gr.Dataframe( headers=["file", "repeat", "total_ms", "inference_ms"], datatype=["str", "number", "number", "number"], interactive=False ) bench_summary = gr.Textbox(label="Summary", lines=6) bench_btn.click( fn=run_benchmark, inputs=[bench_model, repeats], outputs=[bench_table, bench_summary] ) example_files = [] for f in ["1.jpg", "2.jpg", "3.png", "4.webp"]: if os.path.exists(f): example_files.append([f, "InSPyReNet"]) if example_files: gr.Examples(examples=example_files, inputs=[inp, model], label="Examples") if __name__ == "__main__": demo.launch(show_error=True)