Replace Sum with Add chains (no WebGPU Sum kernel: std::bad_alloc in wasm)

1528ba5 verified about 6 hours ago

2.42 kB

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

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