BiRefNet rebuilt for WebGPU, fp32 and fp16

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 ZhengPeng
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,94 @@
+---
+license: mit
+library_name: transformers.js
+base_model: ZhengPeng7/BiRefNet
+pipeline_tag: image-segmentation
+tags:
+  - background-removal
+  - onnx
+  - webgpu
+---
+
+# BiRefNet, ONNX for WebGPU
+
+[BiRefNet](https://huggingface.co/ZhengPeng7/BiRefNet) for background removal, rebuilt so it runs on WebGPU in the browser. This is the full model with the Swin-L backbone. For the smaller one see [wonfy/BiRefNet_lite-ONNX-webgpu](https://huggingface.co/wonfy/BiRefNet_lite-ONNX-webgpu). The weights are unchanged. Only the graph was rewritten, starting from [onnx-community/BiRefNet-ONNX](https://huggingface.co/onnx-community/BiRefNet-ONNX).
+
+The stock export doesn't run on WebGPU. Some of its shaders need more storage buffers than a GPU allows (WebGPU only promises 8 per shader), and its int64 index maths and `Sum` nodes have no WebGPU kernels, so they fall back to the CPU with very large tensors.
+
+| File | Precision | Size |
+|---|---|---|
+| `onnx/model.onnx` | fp32 | 1038 MB |
+| `onnx/model_fp16.onnx` | fp16 weights, float32 input and output | 520 MB |
+
+Input `input_image`: float32 `[1, 3, 1024, 1024]`, RGB resized to 1024 x 1024, scaled to 0..1, then ImageNet mean `[0.485, 0.456, 0.406]` and std `[0.229, 0.224, 0.225]`.
+Output `output_image`: float32 `[1, 1, 1024, 1024]` logits. Apply a sigmoid to get the alpha mask.
+
+## What was changed
+
+1. **Constant folding** with ONNX Runtime's basic optimizer, which only produces standard ONNX ops. It takes the graph from 21,968 to 5,271 nodes and removes every shape op, `ScatterND`, `Where` and `ConstantOfShape`.
+2. **4 `Concat` nodes** with 16, 64, 256 and 1024 inputs became trees with at most 6 inputs each.
+3. **59 `Split` nodes** with up to 32 outputs became trees with at most 6 outputs each.
+4. **20 `Sum` nodes** became chains of `Add` in the same order, so the result stays bit-identical.
+5. **80 `GatherND` nodes** (the sampling in the deformable convolutions) became a plain `Gather` on the data reshaped to `[H*W, C]`. Their index maths, which was int64, now runs in float (the values are whole numbers), then `y * W + x` is taken in int32. No int64 tensor is left in the graph. Doing the last step in int32 also keeps the index exact in the fp16 model.
+6. **fp16:** converted with `onnxconverter-common` and `keep_io_types`, so input and output stay float32.
+
+Every node now uses at most 7 storage buffers, under the WebGPU minimum of 8.
+
+## Results
+
+On the CPU (ONNX Runtime 1.31, three test images: hair, a car and two cats), compared with the original onnx-community fp32 model:
+
+| File | Logits | Alpha mask |
+|---|---|---|
+| `model.onnx` | bit-identical, max difference 0.0 | identical |
+| `model_fp16.onnx` | max difference 0.06 to 1.25 | mean difference under 1e-5, at most 0.02% of pixels move by more than 1/255 |
+
+On WebGPU (Chrome 155, AMD Radeon RX 9070 XT, Transformers.js 4.3.1 with ONNX Runtime Web 1.31.0-dev.20260914), `model_fp16.onnx` on the same three inputs:
+
+- About 800 ms per image once warm, about 4.3 s to load.
+- 0.5% to 1.3% of pixels move by more than 1/255 from the fp32 CPU result, and 0.01% to 0.02% of pixels change between inside and outside, all of them on edges.
+
+The fp32 file is about 1 GB. ONNX Runtime Web reads the whole model into WebAssembly memory before it moves it to the GPU, so the fp16 file is the safer choice in a browser.
+
+## Use it
+
+With Transformers.js:
+
+```js
+import { AutoModel, AutoProcessor, RawImage } from "@huggingface/transformers";
+
+const model = await AutoModel.from_pretrained("wonfy/BiRefNet-ONNX-webgpu", {
+  device: "webgpu",
+  dtype: "fp32",
+  model_file_name: "model_fp16", // leave this out for the fp32 file
+});
+const processor = await AutoProcessor.from_pretrained("wonfy/BiRefNet-ONNX-webgpu");
+
+const image = await RawImage.fromURL("photo.jpg");
+const { pixel_values } = await processor(image);
+const { output_image } = await model({ input_image: pixel_values });
+// alpha = sigmoid(output_image.data[i]), at 1024 x 1024
+```
+
+`dtype: "fp32"` with `model_file_name: "model_fp16"` is the combination that was tested. It loads the fp16 file as it is.
+
+With ONNX Runtime Web directly, create an `InferenceSession` from either file with `executionProviders: ["webgpu"]`.
+
+## Rebuild it
+
+The scripts are in `scripts/`. They were run with Python 3.13, onnx 1.23.2, onnxruntime 1.31.0, onnxconverter-common 1.16.0, numpy and Pillow.
+
+```
+python scripts/fold.py model.onnx folded.onnx
+python scripts/patch.py folded.onnx onnx/model.onnx
+python scripts/to_fp16.py onnx/model.onnx onnx/model_fp16.onnx
+python scripts/verify.py samples/ results/ model.onnx onnx/model.onnx onnx/model_fp16.onnx
+```
+
+`model.onnx` in the first line is the fp32 file from onnx-community/BiRefNet-ONNX. `patch.py` prints an audit at the end: the most storage buffers any node uses, and any op still producing int64.
+
+## Licence
+
+MIT, the same as BiRefNet. See `LICENSE`. BiRefNet is by Peng Zheng et al., [paper](https://arxiv.org/abs/2401.03407) and [code](https://github.com/ZhengPeng7/BiRefNet).
+
+BiRefNet's training data includes DIS5K, and DIS5K's own terms restrict commercial use of the dataset. The MIT licence covers the model and code, not the training data. Check what that means for your use.

config.json

@@ -0,0 +1,6 @@
+{
+  "model_type": "swin",
+  "transformers.js_config": {
+    "dtype": "fp32"
+  }
+}

onnx/model.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:90a8e190dcc73f0f553f83df14aa0bf7ad888c3cfe8f2180d9ab474ca86462d4
+size 1038412111

onnx/model_fp16.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f5a69905dd90e6592a8030d3f7d1701d3c8c6bb840515691b2f7aab2e3d4bfd6
+size 520175713

preprocessor_config.json

@@ -0,0 +1,12 @@
+{
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "feature_extractor_type": "ViTFeatureExtractor",
+  "image_mean": [0.485, 0.456, 0.406],
+  "image_processor_type": "ViTFeatureExtractor",
+  "image_std": [0.229, 0.224, 0.225],
+  "resample": 2,
+  "rescale_factor": 0.00392156862745098,
+  "size": { "height": 1024, "width": 1024 }
+}

scripts/fold.py

@@ -0,0 +1,12 @@
+import sys
+import time
+
+import onnxruntime as ort
+
+source, target = sys.argv[1], sys.argv[2]
+options = ort.SessionOptions()
+options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_BASIC
+options.optimized_model_filepath = target
+started = time.time()
+ort.InferenceSession(source, options, providers=["CPUExecutionProvider"])
+print(f"folded in {time.time() - started:.1f} s -> {target}")

scripts/patch.py

@@ -0,0 +1,197 @@
+"""Rewrites a constant-folded BiRefNet graph so every WebGPU shader fits in 8 storage buffers and nothing falls back to the CPU."""
+import math
+import sys
+from collections import Counter, defaultdict
+
+import numpy as np
+import onnx
+from onnx import TensorProto, helper, numpy_helper, shape_inference
+
+FAN = 6
+
+source, target = sys.argv[1], sys.argv[2]
+model = shape_inference.infer_shapes(onnx.load(source))
+graph = model.graph
+shapes = {}
+for value in list(graph.value_info) + list(graph.input) + list(graph.output):
+    t = value.type.tensor_type
+    shapes[value.name] = ([d.dim_value for d in t.shape.dim], t.elem_type)
+inits = {i.name: i for i in graph.initializer}
+for name, init in inits.items():
+    shapes[name] = (list(init.dims), init.data_type)
+
+counter = Counter()
+
+
+def fresh(base):
+    counter[base] += 1
+    return f"wonfy/{base}/{counter[base]}"
+
+
+new_inits = []
+
+
+def const(base, array):
+    name = fresh(base)
+    new_inits.append(numpy_helper.from_array(np.asarray(array), name))
+    return name
+
+
+def attr(node, name, default=None):
+    for a in node.attribute:
+        if a.name == name:
+            return helper.get_attribute_value(a)
+    return default
+
+
+consumers = defaultdict(list)
+for n in graph.node:
+    for x in n.input:
+        consumers[x].append(n)
+producers = {out: n for n in graph.node for out in n.output}
+
+# The int64 index chain of each deformable-conv sample, done in float instead (the values are whole numbers).
+cast_nodes = [n for n in graph.node if n.op_type == "Cast" and attr(n, "to") == TensorProto.INT64]
+region = set()
+frontier = [n.output[0] for n in cast_nodes]
+while frontier:
+    name = frontier.pop()
+    if name in region:
+        continue
+    region.add(name)
+    for c in consumers[name]:
+        if c.op_type in ("Slice", "Add", "Clip", "Reshape"):
+            frontier.append(c.output[0])
+        elif c.op_type == "Concat":
+            region.add(c.output[0])
+        elif c.op_type != "GatherND":
+            raise SystemExit(f"unexpected consumer {c.op_type} of int64 tensor {name}")
+cast_source = {n.output[0]: n.input[0] for n in cast_nodes}
+
+drop = set(id(n) for n in cast_nodes)
+renames = {}
+for n in graph.node:
+    if n.op_type in ("Add", "Clip") and n.output[0] in region:
+        for i, x in enumerate(n.input):
+            if x in inits and inits[x].data_type == TensorProto.INT64:
+                n.input[i] = const("float_const", numpy_helper.to_array(inits[x]).astype(np.float32))
+for n in graph.node:
+    for i, x in enumerate(n.input):
+        if x in cast_source:
+            n.input[i] = cast_source[x]
+
+out_nodes = []
+stats = Counter()
+
+
+def emit(op, inputs, outputs, **attrs):
+    out_nodes.append(helper.make_node(op, inputs, outputs, name=fresh(op.lower()), **attrs))
+
+
+def concat_tree(inputs, output, axis):
+    while len(inputs) > FAN:
+        grouped = []
+        for start in range(0, len(inputs), FAN):
+            group = inputs[start:start + FAN]
+            if len(group) == 1:
+                grouped.append(group[0])
+                continue
+            name = fresh("concat")
+            emit("Concat", group, [name], axis=axis)
+            grouped.append(name)
+        inputs = grouped
+    emit("Concat", inputs, [output], axis=axis)
+
+
+def split_tree(source_name, sizes, outputs, axis):
+    if len(outputs) <= FAN:
+        emit("Split", [source_name, const("split_sizes", np.array(sizes, dtype=np.int64))], outputs, axis=axis)
+        return
+    size = math.ceil(len(outputs) / FAN)
+    parts, part_sizes = [], []
+    for start in range(0, len(outputs), size):
+        group_outputs, group_sizes = outputs[start:start + size], sizes[start:start + size]
+        name = group_outputs[0] if len(group_outputs) == 1 else fresh("split")
+        parts.append((name, group_sizes, group_outputs))
+        part_sizes.append(sum(group_sizes))
+    split_tree(source_name, part_sizes, [p[0] for p in parts], axis)
+    for name, group_sizes, group_outputs in parts:
+        if len(group_outputs) > 1:
+            split_tree(name, group_sizes, group_outputs, axis)
+
+
+concat_inputs = {n.output[0]: list(n.input) for n in graph.node if n.op_type == "Concat" and n.output[0] in region}
+
+for n in graph.node:
+    if id(n) in drop:
+        continue
+    if n.op_type == "Concat" and n.output[0] in region:
+        stats["index concat removed"] += 1
+        continue
+    if n.op_type == "Concat" and len(n.input) > FAN:
+        concat_tree(list(n.input), n.output[0], attr(n, "axis"))
+        stats["concat"] += 1
+        continue
+    if n.op_type == "Split" and len(n.output) > FAN:
+        sizes = numpy_helper.to_array(inits[n.input[1]]).tolist()
+        axis = attr(n, "axis", 0)
+        split_tree(n.input[0], sizes, list(n.output), axis)
+        stats["split"] += 1
+        continue
+    if n.op_type == "Sum":
+        total = n.input[0]
+        for i, x in enumerate(n.input[1:], start=1):
+            name = n.output[0] if i == len(n.input) - 1 else fresh("sum")
+            emit("Add", [total, x], [name])
+            total = name
+        stats["sum"] += 1
+        continue
+    if n.op_type == "GatherND":
+        assert attr(n, "batch_dims") == 2, "only batch_dims 2 is handled"
+        data, index = n.input
+        data_shape = shapes[data][0]
+        assert data_shape[:2] == [1, 1], data_shape
+        _, _, height, width, channels = data_shape
+        y, x = concat_inputs[index]
+        count = shapes[n.output[0]][0][2]
+        y32, x32, row, flat, flat1, table, picked = (fresh(k) for k in ("y32", "x32", "row", "flat", "flat1", "table", "picked"))
+        emit("Cast", [y], [y32], to=TensorProto.INT32)
+        emit("Cast", [x], [x32], to=TensorProto.INT32)
+        emit("Mul", [y32, const("width", np.array(width, dtype=np.int32))], [row])
+        emit("Add", [row, x32], [flat])
+        emit("Reshape", [flat, const("shape", np.array([count], dtype=np.int64))], [flat1])
+        emit("Reshape", [data, const("shape", np.array([height * width, channels], dtype=np.int64))], [table])
+        emit("Gather", [table, flat1], [picked], axis=0)
+        emit("Reshape", [picked, const("shape", np.array([1, 1, count, channels], dtype=np.int64))], [n.output[0]])
+        stats["gathernd"] += 1
+        continue
+    out_nodes.append(n)
+
+del graph.node[:]
+graph.node.extend(out_nodes)
+graph.initializer.extend(new_inits)
+del graph.value_info[:]
+used = {x for n in graph.node for x in n.input}
+kept = [i for i in graph.initializer if i.name in used]
+del graph.initializer[:]
+graph.initializer.extend(kept)
+del model.opset_import[:]
+model.opset_import.extend([helper.make_opsetid("", 17)])
+
+model = shape_inference.infer_shapes(model, strict_mode=True)
+onnx.checker.check_model(model)
+onnx.save(model, target)
+print("rewrote", dict(stats))
+
+# Audit: buffers per node, and anything WebGPU would run on the CPU.
+graph = model.graph
+types = {v.name: v.type.tensor_type.elem_type for v in list(graph.value_info) + list(graph.output) + list(graph.input)}
+cpu_inputs = {("Reshape", 1), ("Split", 1), ("Slice", 1), ("Slice", 2), ("Slice", 3), ("Slice", 4), ("Clip", 1), ("Clip", 2), ("Pad", 1), ("Pad", 2), ("Resize", 1), ("Resize", 2), ("Resize", 3)}
+worst = Counter()
+for n in graph.node:
+    buffers = sum(1 for i, x in enumerate(n.input) if x and (n.op_type, i) not in cpu_inputs) + len(n.output)
+    worst[n.op_type] = max(worst[n.op_type], buffers)
+print("most buffers per op:", dict(sorted(worst.items(), key=lambda kv: -kv[1])))
+int64 = Counter(n.op_type for n in graph.node if any(types.get(o) == TensorProto.INT64 for o in n.output))
+print("ops still making int64:", dict(int64) or "none")
+print("ops:", dict(Counter(n.op_type for n in graph.node).most_common()))

scripts/to_fp16.py

@@ -0,0 +1,14 @@
+import sys
+from collections import Counter
+
+import onnx
+from onnxconverter_common import float16
+
+source, target = sys.argv[1], sys.argv[2]
+model = float16.convert_float_to_float16(onnx.load(source), keep_io_types=True, disable_shape_infer=False)
+onnx.checker.check_model(model)
+onnx.save(model, target)
+graph = model.graph
+float_inits = Counter(onnx.TensorProto.DataType.Name(i.data_type) for i in graph.initializer)
+print("initializer types", dict(float_inits))
+print("casts", sum(1 for n in graph.node if n.op_type == "Cast"), "of", len(graph.node), "nodes")

scripts/verify.py

@@ -0,0 +1,52 @@
+"""Runs each model on the sample images on the CPU and compares its output with the first model's."""
+import os
+import sys
+import time
+
+import numpy as np
+import onnxruntime as ort
+from PIL import Image
+
+samples_dir, out_dir, *models = sys.argv[1:]
+os.makedirs(out_dir, exist_ok=True)
+MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
+STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
+
+
+def pixels(path):
+    image = Image.open(path).convert("RGB")
+    resized = np.asarray(image.resize((1024, 1024), Image.BILINEAR), dtype=np.float32) / 255
+    return image, ((resized - MEAN) / STD).transpose(2, 0, 1)[None]
+
+
+def sigmoid(x):
+    return 1 / (1 + np.exp(-x))
+
+
+samples = [(name, *pixels(os.path.join(samples_dir, name))) for name in sorted(os.listdir(samples_dir))]
+reference = {}
+for index, path in enumerate(models):
+    label = os.path.splitext(os.path.basename(path))[0]
+    options = ort.SessionOptions()
+    options.log_severity_level = 3
+    started = time.time()
+    session = ort.InferenceSession(path, options, providers=["CPUExecutionProvider"])
+    print(f"{label}: loaded in {time.time() - started:.1f} s, {os.path.getsize(path) / 1e6:.0f} MB")
+    input_name = session.get_inputs()[0].name
+    for name, image, x in samples:
+        started = time.time()
+        logits = session.run(None, {input_name: x})[0].astype(np.float32)[0, 0]
+        took = time.time() - started
+        alpha = sigmoid(logits)
+        Image.fromarray((alpha * 255).round().astype(np.uint8)).resize(image.size, Image.BILINEAR).save(os.path.join(out_dir, f"{os.path.splitext(name)[0]}-{label}.png"))
+        if index == 0:
+            reference[name] = (logits, alpha)
+            print(f"  {name}: {took:.1f} s, logits {logits.min():.2f}..{logits.max():.2f}, {100 * (alpha > 0.5).mean():.1f}% foreground")
+            continue
+        ref_logits, ref_alpha = reference[name]
+        diff = np.abs(alpha - ref_alpha)
+        print(
+            f"  {name}: {took:.1f} s, max logit diff {np.abs(logits - ref_logits).max():.2e}, "
+            f"alpha diff max {diff.max():.2e} mean {diff.mean():.2e}, "
+            f"pixels off by more than 1/255: {100 * (diff > 1 / 255).mean():.3f}%"
+        )