| """Build the WebGPU-capable LaMa used by the AI eraser (Photo Editor, Background Remover). |
| |
| Why this exists (measured 2026-11-11, onnxruntime-web 1.30, Chrome WebGPU): |
| LaMa's published ONNX exports run on WASM only (~30 s per 512 erase on one thread). On WebGPU every |
| run failed inside the Fourier unit ("Failed to run JSEP kernel": Add, then Sub, Slice, MatMul as each |
| was worked around). Three graph rewrites, each numerically equivalent, make it run on WebGPU at |
| ~1.7-2.2 s per erase with the same fill: |
| 1. fold at a fixed batch=1 (ORT basic optimisation = constant folding), so the DFT matrices the |
| export builds at run time from Range/Cos/Sin/int64 shape maths become constants; |
| 2. Einsum "ij,jk->ik" -> MatMul (216 nodes); |
| 3. negative-step Slice -> Gather with constant reversed indices (72 nodes). |
| The int8 OpenCV Zoo export gives WRONG output on WebGPU after the same rewrites (white fills), so the |
| source must be the fp32 export. |
| |
| Source: https://huggingface.co/Carve/LaMa-ONNX lama_fp32.onnx (Apache-2.0; sha256 1faef530...d68d6, |
| 208,044,816 bytes). LaMa itself: advimman/lama, Apache-2.0. The output is a derivative under the same |
| licence; keep this script and the licence with it wherever it is hosted. |
| |
| Usage (Python 3.11+, `pip install onnx onnxruntime numpy`): |
| python -I scripts/models/build-lama-webgpu.py lama_fp32.onnx lama_webgpu.onnx |
| It refuses to keep the result unless it matches the source within 1e-2 on the 0..255 scale. |
| """ |
|
|
| import collections |
| import sys |
|
|
| import numpy as np |
| import onnx |
| import onnxruntime as ort |
| from onnx import helper, numpy_helper, shape_inference |
|
|
|
|
| def fold(src: str, dst: str) -> None: |
| so = ort.SessionOptions() |
| so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_BASIC |
| so.add_free_dimension_override_by_name("batch", 1) |
| so.optimized_model_filepath = dst |
| ort.InferenceSession(src, so, providers=["CPUExecutionProvider"]) |
|
|
|
|
| def rewrite(path: str, dst: str) -> dict: |
| m = onnx.load(path) |
| mi = shape_inference.infer_shapes(m) |
| vi = {v.name: v for v in list(mi.graph.value_info) + list(mi.graph.input)} |
| init = {i.name: numpy_helper.to_array(i) for i in m.graph.initializer} |
| counts = collections.Counter() |
| nodes, inits = [], [] |
| for node in m.graph.node: |
| if node.op_type == "Einsum": |
| eq = [a.s.decode() for a in node.attribute if a.name == "equation"][0] |
| if eq != "ij,jk->ik": |
| raise SystemExit(f"unexpected Einsum equation {eq!r} in {node.name}") |
| node.op_type = "MatMul" |
| del node.attribute[:] |
| counts["einsum->matmul"] += 1 |
| nodes.append(node) |
| elif node.op_type == "Slice" and len(node.input) >= 5 and node.input[4] in init and (init[node.input[4]] < 0).any(): |
| starts, ends, axes, steps = (init[node.input[i]].tolist() for i in (1, 2, 3, 4)) |
| v = vi.get(node.input[0]) |
| d = [x.dim_value for x in v.type.tensor_type.shape.dim] if v else None |
| if not d or any(x <= 0 for x in d): |
| raise SystemExit(f"{node.name}: input shape not static after folding: {d}") |
| cur = node.input[0] |
| for k, (st, en, ax, sp) in enumerate(zip(starts, ends, axes, steps)): |
| idx = np.arange(d[ax])[slice(st if abs(st) < 2**62 else None, en if abs(en) < 2**62 else None, sp)] |
| name = f"{node.name}_gidx{k}" |
| inits.append(numpy_helper.from_array(idx.astype(np.int64), name)) |
| out = node.output[0] if k == len(axes) - 1 else f"{node.name}_g{k}" |
| nodes.append(helper.make_node("Gather", [cur, name], [out], axis=ax, name=f"{node.name}_gather{k}")) |
| cur = out |
| d = list(d) |
| d[ax] = len(idx) |
| counts["negslice->gather"] += 1 |
| else: |
| nodes.append(node) |
| del m.graph.node[:] |
| m.graph.node.extend(nodes) |
| m.graph.initializer.extend(inits) |
| onnx.checker.check_model(m) |
| onnx.save(m, dst) |
| return dict(counts) |
|
|
|
|
| def max_diff(a: str, b: str) -> float: |
| n = 512 |
| img = ((np.sin(np.arange(3 * n * n) * 0.37) + 1) / 2).astype(np.float32).reshape(1, 3, n, n) |
| msk = np.zeros((1, 1, n, n), np.float32) |
| msk[:, :, 200:300, 200:320] = 1 |
| feeds = {"image": img, "mask": msk} |
| oa = ort.InferenceSession(a, providers=["CPUExecutionProvider"]).run(None, feeds)[0] |
| ob = ort.InferenceSession(b, providers=["CPUExecutionProvider"]).run(None, feeds)[0] |
| return float(np.abs(oa - ob).max()) |
|
|
|
|
| def main() -> None: |
| src, dst = sys.argv[1], sys.argv[2] |
| folded = dst + ".folded.onnx" |
| fold(src, folded) |
| counts = rewrite(folded, dst) |
| diff = max_diff(src, dst) |
| print("rewrites:", counts, "max |diff| (0..255):", round(diff, 5)) |
| if diff > 1e-2: |
| raise SystemExit("the rewritten model does not match the source; refusing to keep it") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
