BiRefNet-lite ONNX rebuilt for WebGPU (≤7 storage buffers per node)

README.md

@@ -0,0 +1,45 @@
+---
+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). 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.

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:a9e7aebec91bad0cf3949015e8ddc379602db6a2b7a06918e81340f4e0ce3716
+size 224182130

onnx/model_fp16.onnx

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

preprocessor_config.json

@@ -0,0 +1,23 @@
+{
+  "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
+  }
+}

rebuild.py

@@ -0,0 +1,162 @@
+"""
+Rebuild BiRefNet-lite ONNX so it runs on any WebGPU device.
+
+The official export has Concat nodes with up to 1024 inputs and Split nodes
+with up to 32 outputs. ONNX Runtime's WebGPU backend binds every input and
+output of a node as a storage buffer, and GPUs cap that per shader
+(maxStorageBuffersPerShaderStage: 8 minimum, 10 or 16 on many devices), so
+those nodes fail with "Too many storage buffers in shader".
+
+This replaces them with trees of small nodes that compute exactly the same
+values (concatenation and splitting are associative along one axis):
+  - Concat: at most CONCAT_FAN_IN inputs per node (+1 output)
+  - Split:  at most SPLIT_FAN_OUT outputs per node (+ data + sizes inputs)
+Every node then uses at most MAX_BUFFERS storage buffers.
+
+Usage: python rebuild.py <in.onnx> <out.onnx>
+"""
+
+import sys
+
+import numpy as np
+import onnx
+from onnx import helper, numpy_helper
+
+CONCAT_FAN_IN = 6
+SPLIT_FAN_OUT = 5
+MAX_BUFFERS = 7
+
+
+def buffers(node: onnx.NodeProto) -> int:
+    return len([i for i in node.input if i]) + len([o for o in node.output if o])
+
+
+class Rewriter:
+    def __init__(self, model: onnx.ModelProto):
+        self.model = model
+        self.graph = model.graph
+        self.counter = 0
+        self.initializers = {t.name: t for t in self.graph.initializer}
+        self.constants = {
+            n.output[0]: n for n in self.graph.node if n.op_type == "Constant" and len(n.output) == 1
+        }
+        self.names = {o for n in self.graph.node for o in n.output} | set(self.initializers)
+        self.names |= {i.name for i in self.graph.input}
+
+    def fresh(self, base: str) -> str:
+        while True:
+            self.counter += 1
+            name = f"{base}/wgpu_{self.counter}"
+            if name not in self.names:
+                self.names.add(name)
+                return name
+
+    def const_value(self, name: str) -> np.ndarray:
+        if name in self.initializers:
+            return numpy_helper.to_array(self.initializers[name])
+        if name in self.constants:
+            attr = next(a for a in self.constants[name].attribute if a.name == "value")
+            return numpy_helper.to_array(attr.t)
+        raise ValueError(f"split sizes '{name}' are not constant")
+
+    def int64_initializer(self, base: str, values: list[int]) -> str:
+        name = self.fresh(base + "/sizes")
+        tensor = numpy_helper.from_array(np.array(values, dtype=np.int64), name)
+        self.graph.initializer.append(tensor)
+        self.initializers[name] = tensor
+        return name
+
+    def concat_tree(self, node: onnx.NodeProto) -> list[onnx.NodeProto]:
+        axis = next(a.i for a in node.attribute if a.name == "axis")
+        inputs = list(node.input)
+        made: list[onnx.NodeProto] = []
+        while len(inputs) > CONCAT_FAN_IN:
+            grouped = []
+            for start in range(0, len(inputs), CONCAT_FAN_IN):
+                chunk = inputs[start : start + CONCAT_FAN_IN]
+                if len(chunk) == 1:
+                    grouped.append(chunk[0])
+                    continue
+                out = self.fresh(node.name)
+                made.append(helper.make_node("Concat", chunk, [out], name=self.fresh(node.name), axis=axis))
+                grouped.append(out)
+            inputs = grouped
+        made.append(helper.make_node("Concat", inputs, list(node.output), name=node.name, axis=axis))
+        return made
+
+    def split_tree(self, node: onnx.NodeProto) -> list[onnx.NodeProto]:
+        axis = next((a.i for a in node.attribute if a.name == "axis"), 0)
+        if len(node.input) < 2 or not node.input[1]:
+            raise ValueError(f"{node.name}: Split without explicit sizes isn't handled")
+        sizes = [int(v) for v in self.const_value(node.input[1]).tolist()]
+        outputs = list(node.output)
+        assert len(sizes) == len(outputs)
+        made: list[onnx.NodeProto] = []
+
+        def split(data: str, parts: list[tuple[int, str]], base: str) -> None:
+            # parts: (size, final output name)
+            if len(parts) <= SPLIT_FAN_OUT:
+                made.append(
+                    helper.make_node(
+                        "Split",
+                        [data, self.int64_initializer(base, [s for s, _ in parts])],
+                        [o for _, o in parts],
+                        name=self.fresh(base),
+                        axis=axis,
+                    )
+                )
+                return
+            # Groups of roughly equal count, each split again below.
+            per_group = -(-len(parts) // SPLIT_FAN_OUT)
+            groups = [parts[i : i + per_group] for i in range(0, len(parts), per_group)]
+            group_outputs = []
+            for group in groups:
+                group_outputs.append(group[0][1] if len(group) == 1 else self.fresh(base))
+            made.append(
+                helper.make_node(
+                    "Split",
+                    [data, self.int64_initializer(base, [sum(s for s, _ in g) for g in groups])],
+                    group_outputs,
+                    name=self.fresh(base),
+                    axis=axis,
+                )
+            )
+            for group, out in zip(groups, group_outputs):
+                if len(group) > 1:
+                    split(out, group, base)
+
+        split(node.input[0], list(zip(sizes, outputs)), node.name)
+        return made
+
+    def run(self) -> dict[str, int]:
+        stats = {"Concat": 0, "Split": 0}
+        new_nodes: list[onnx.NodeProto] = []
+        for node in self.graph.node:
+            if buffers(node) > MAX_BUFFERS and node.op_type == "Concat":
+                new_nodes.extend(self.concat_tree(node))
+                stats["Concat"] += 1
+            elif buffers(node) > MAX_BUFFERS and node.op_type == "Split":
+                new_nodes.extend(self.split_tree(node))
+                stats["Split"] += 1
+            else:
+                new_nodes.append(node)
+        del self.graph.node[:]
+        self.graph.node.extend(new_nodes)
+        return stats
+
+
+def main() -> None:
+    source, target = sys.argv[1], sys.argv[2]
+    model = onnx.load(source)
+    stats = Rewriter(model).run()
+    too_wide = [(n.op_type, n.name, buffers(n)) for n in model.graph.node if buffers(n) > MAX_BUFFERS]
+    if too_wide:
+        raise SystemExit(f"still too wide: {too_wide[:10]}")
+    onnx.checker.check_model(model)
+    onnx.save(model, target)
+    widest = max(buffers(n) for n in model.graph.node)
+    print(f"{source} -> {target}: rewrote {stats}, nodes {len(model.graph.node)}, widest node {widest} buffers")
+
+
+if __name__ == "__main__":
+    main()