jessolmstead's picture
BiRefNet lite fp16, rewritten for WebGPU
0333e93 verified
|
Raw History Blame Contribute Delete
3.41 kB
---
license: mit
library_name: onnx
pipeline_tag: image-segmentation
base_model: onnx-community/BiRefNet_lite-ONNX
tags:
- birefnet
- salient-object-detection
- matting
- onnx
- webgpu
- onnxruntime
---
# BiRefNet lite, ONNX, rewritten for WebGPU
A modified copy of
[onnx-community/BiRefNet_lite-ONNX](https://huggingface.co/onnx-community/BiRefNet_lite-ONNX)
(`onnx/model_fp16.onnx`), the fp16 ONNX export of
[BiRefNet](https://github.com/ZhengPeng7/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](https://github.com/jessolmstead/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 `Add`s.
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.