README.md

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

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