RMBG-2.0 — WebGPU
A WebGPU-compatible ONNX export of briaai/RMBG-2.0 background removal model, built for ONNX Runtime 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)
<script type="module">
import { InferenceSession, Tensor } from "https://cdn.jsdelivr.net/npm/onnxruntime-web@latest/dist/ort.js";
const session = await InferenceSession.create(
"https://huggingface.co/yamura4/RMBG-2.0-WebGPU/resolve/main/onnx/model_fp16.onnx",
{ executionProviders: ["webgpu"] }
);
const input = new Tensor("float32", new Float32Array(3 * 1024 * 1024), [1, 3, 1024, 1024]);
const { alphas } = await session.run({ pixel_values: input });
// alphas.data is a Float32Array of shape (1, 1, 1024, 1024)
</script>
Usage (Node.js / Python ONNX Runtime)
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 modelconfig.json,preprocessor_config.json— model configBiRefNet_config.py— custom config class (required by HF Hub)
License
CC-BY-NC-4.0 (same as briaai/RMBG-2.0).