--- license: cc-by-nc-4.0 library_name: onnx pipeline_tag: image-segmentation tags: - background-removal - onnx - webgpu - remove-background --- # RMBG-2.0 — WebGPU A WebGPU-compatible ONNX export of [briaai/RMBG-2.0](https://huggingface.co/briaai/RMBG-2.0) background removal model, built for [ONNX Runtime Web](https://github.com/microsoft/onnxruntime-web) with WebGPU backend. ## The Problem The original ONNX model contains Split ops with 32 outputs (33 storage buffer bindings) and a Concat op with 1024 inputs (1025 bindings). WebGPU has a `maxStorageBuffersPerShaderStage` limit of **16**, so the model fails at runtime with: > `Uncaught (in promise) Error: MaxStorageBuffersPerShaderStage (16) exceeded. Created 33 bindings.` ## The Fix Large Split and Concat ops are cascaded into trees where every node has ≤8 outputs or inputs — max **9 bindings** per shader stage, well under the WebGPU limit. | Before | After | |---|---| | 50× Split(32/16 outputs) | Split(4) + 4×Split(8) / Split(2) + 2×Split(8) trees | | Concat(1024 inputs) | 4-level group-of-8 Concat tree | | Concat(256 inputs) | 3-level group-of-8 Concat tree | | Concat(64 inputs) | 2-level group-of-8 Concat tree | | Concat(16 inputs) | 2-level group-of-8 Concat tree | Weights are stored as fp16 (Cast to fp32 before compute). Output is numerically identical within fp16 tolerance. ## Model | File | Size | |---|---| | `onnx/model_fp16.onnx` | 490 MB | ## Usage (ONNX Runtime Web + WebGPU) ```html ``` ## Usage (Node.js / Python ONNX Runtime) ```python import onnxruntime as ort import numpy as np from PIL import Image from torchvision import transforms sess = ort.InferenceSession("onnx/model_fp16.onnx") img = Image.open("photo.jpg").convert("RGB") transform = transforms.Compose([ transforms.Resize((1024, 1024)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) input_tensor = transform(img).unsqueeze(0).numpy().astype(np.float32) alphas = sess.run(["alphas"], {"pixel_values": input_tensor})[0][0, 0] mask = (alphas * 255).astype(np.uint8) Image.fromarray(mask).save("mask.png") ``` ## Files - `onnx/model_fp16.onnx` — WebGPU-compatible model - `config.json`, `preprocessor_config.json` — model config - `BiRefNet_config.py` — custom config class (required by HF Hub) ## License CC-BY-NC-4.0 (same as [briaai/RMBG-2.0](https://huggingface.co/briaai/RMBG-2.0)).