yamura4/RMBG-2.0-WebGPU

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 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).