File size: 4,959 Bytes
2c110d2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | """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()
|