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 (MIT) for Transformers.js, rebuilt so it runs on any WebGPU device.

The ONNX export in 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)

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.