jessolmstead/BiRefNet_lite-ONNX-webgpu

BiRefNet lite, ONNX, rewritten for WebGPU

A modified copy of onnx-community/BiRefNet_lite-ONNX (onnx/model_fp16.onnx), the fp16 ONNX export of BiRefNet by Zheng Peng et al., rewritten so that ONNX Runtime's WebGPU execution provider runs every node on the GPU. It is what greycard, a RAW editor, downloads for its Subject and Background masks on a machine with WebGPU; the editor fetches it on first use, and nothing here needs to be fetched by hand.

The weights are untouched. Only the graph is rewritten, and on the CPU provider the rewritten graph answers exactly as the original does (max difference 0 on the matte).

What was changed, and why

The original export has three things the WebGPU provider will not run on the card, so 50 nodes failed at run time and 420 more fell back to the CPU with a memory copy each way:

1. The decoder cuts its input into a grid of patches with Split nodes of 16 and 32 outputs. A shader binds at most 16 storage buffers per stage under Dawn, and the provider's Split binds its input and every output in one shader, so such a Split fails ("Too many storage buffers in shader. Current: 17, Max is 16"). Each becomes one Slice per output: same axis, same offsets, same output names.
2. The twenty deformable convolutions sum four fp16 terms with a Sum, which the provider has no kernel for. Each becomes a chain of Adds.
3. The same convolutions floor their sampling coordinates, cast them to int64, and Slice, Add, Clip, Reshape and Concat them as int64, none of which the provider takes. The chain stays in fp16 and the cast moves to just before the GatherND that uses it. Every value in the chain is an integer bounded by the Clip, so the fp16 arithmetic is exact.

After the rewrite all 5,595 nodes are placed on WebGPU. On an RTX 5070 Ti through Dawn (Vulkan), a 1024x1024 matte takes about 0.16 s where the original took 2.9 s on a Ryzen 9 9950X3D's CPU and 1.6 s on WebGPU with the failing nodes forced to the CPU. The GPU's fp16 result differs from the CPU's along the outline by one to two pixels (mean absolute difference 4e-4); the original export shows the same on WebGPU, so it is the card's arithmetic and not the rewrite.

Reproducing the file

The rewrite is a script in the greycard repository, tools/ai/birefnet_webgpu.py, with pinned versions in tools/ai/requirements.txt:

python birefnet_webgpu.py rewrite model_fp16.onnx model_fp16_webgpu.onnx
python birefnet_webgpu.py check model_fp16_webgpu.onnx

The output is byte for byte the same on every run with those versions.

File	Bytes	SHA-256
model_fp16_webgpu.onnx	113,778,088	0a019d6ba73c9cedc9a251f8c9390b196ff6399acd281a2872692861abbd78c2

License

MIT, as BiRefNet is. The LICENSE file here is BiRefNet's, verbatim (Copyright (c) 2024 ZhengPeng). This is a modified copy of the onnx-community export; the modifications are the three graph rewrites above, made for greycard, and carry the same license.