File size: 2,416 Bytes
20f6e53 1528ba5 20f6e53 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 | ---
library_name: transformers.js
tags:
- background-removal
- image-segmentation
- webgpu
- onnx
pipeline_tag: image-segmentation
license: mit
base_model:
- ZhengPeng7/BiRefNet_lite
- onnx-community/BiRefNet_lite-ONNX
---
# BiRefNet-lite ONNX, rebuilt for WebGPU
[BiRefNet-lite](https://github.com/ZhengPeng7/BiRefNet) (MIT) for Transformers.js, rebuilt so it runs on **any WebGPU device**.
The ONNX export in [onnx-community/BiRefNet_lite-ONNX](https://huggingface.co/onnx-community/BiRefNet_lite-ONNX) has Concat nodes with up to 1024 inputs and Split nodes with up to 32 outputs. ONNX Runtime's WebGPU backend binds each input and output as a storage buffer, and GPUs limit that per shader (`maxStorageBuffersPerShaderStage`: 8 minimum, often 10 or 16), so those nodes fail with:
```
Too many storage buffers in shader. Current: 17, Max is 16
```
This repo replaces those 4 Concat and 59 Split nodes with trees of smaller ones (at most 6 inputs per Concat, 5 outputs per Split, so every node uses ≤ 7 storage buffers).
It also replaces the 20 `Sum` nodes with chains of `Add` (same order). ONNX Runtime's WebGPU backend has no Sum kernel, so Sum ran on the CPU, inside WebAssembly's 4 GB heap; in the deformable convolutions it adds four 392 MiB tensors (784 MiB each after the casts to fp32 the CPU needs), which fails with `std::bad_alloc`.
Nothing else changes: the outputs are **bit-for-bit identical** to the original on ONNX Runtime CPU, for both files.
| File | Precision | Size | Needs |
|---|---|---|---|
| `onnx/model_fp16.onnx` | fp16 | 115 MB | WebGPU with `shader-f16` |
| `onnx/model.onnx` | fp32 | 224 MB | WebGPU |
Reproduce with `rebuild.py` (included): `python rebuild.py model.onnx out.onnx` on the files from `onnx-community/BiRefNet_lite-ONNX` (source SHA-256: fp32 `5600024376f572a557870a5eb0afb1e5961636bef4e1e22132025467d0f03333`, fp16 `d39b897ceb16ae654c1731f3dba0cf9b368d9cae74b5a57459b455cc8bfec402`).
## Usage (Transformers.js)
```js
import { pipeline } from "@huggingface/transformers";
const segmenter = await pipeline("image-segmentation", "edelfrade/BiRefNet_lite-webgpu", { device: "webgpu", dtype: "fp16" });
const [{ mask }] = await segmenter(image);
```
## License
MIT, same as the original model. Credit to Peng Zheng et al., *Bilateral Reference for High-Resolution Dichotomous Image Segmentation* (CAAI AIR 2024), and to onnx-community for the ONNX export.
|