| """Rebuilds htdemucs-wet.onnx (the Wet Studio HD stem split model). |
| |
| Input: htdemucs_fp16weights.onnx from huggingface.co/StemSplitio/htdemucs-onnx |
| (revision d54ed9eb60e258ea82131c6ee14578628816456a), an MIT export of Meta's |
| HT-Demucs v4. Why: that export doesn't build on WebGPU in onnxruntime-web |
| (a double-typed ConstantOfShape in its iSTFT has no kernel there), and its |
| 24,917 nodes take ~37 s to load. With the input shape fixed, everything that |
| only depends on shapes folds to constants: 1,427 standard-ONNX nodes, a 3 s |
| load on WebGPU. Big float tensors are then stored as float16 (+ a Cast), so |
| the file is 174 MB. Output matches the original to ~90 dB SDR. |
| |
| pip install onnx onnxruntime numpy |
| python make_model.py htdemucs_fp16weights.onnx htdemucs-wet.onnx |
| """ |
| import sys, onnx, numpy as np, onnxruntime as ort |
| from onnx import numpy_helper, helper, TensorProto |
|
|
| src, dst = sys.argv[1], sys.argv[2] |
| so = ort.SessionOptions() |
| so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_BASIC |
| so.optimized_model_filepath = dst + '.folded' |
| ort.InferenceSession(src, so, providers=['CPUExecutionProvider']) |
|
|
| m = onnx.load(dst + '.folded'); g = m.graph |
| users = {} |
| for n in g.node: |
| for i in n.input: users.setdefault(i, []).append(n.op_type) |
| def half(keep): |
| inits, casts = [], [] |
| for t in g.initializer: |
| if t.data_type == TensorProto.FLOAT and keep(t): |
| inits.append(numpy_helper.from_array(numpy_helper.to_array(t).astype(np.float16), t.name + '__fp16')) |
| casts.append(helper.make_node('Cast', [t.name + '__fp16'], [t.name], to=TensorProto.FLOAT, name=t.name + '__cast')) |
| else: |
| inits.append(t) |
| del g.initializer[:]; g.initializer.extend(inits) |
| nodes = list(g.node); del g.node[:]; g.node.extend(casts + nodes) |
| |
| half(lambda t: int(np.prod(t.dims)) >= 1024 and users.get(t.name) and all(u in ('Conv', 'ConvTranspose', 'MatMul') for u in users[t.name]) and 'stft' not in t.name.lower()) |
| |
| half(lambda t: int(np.prod(t.dims)) >= 100000 and 'Cast_output' not in t.name) |
| ops = [o for o in m.opset_import if o.domain == '']; del m.opset_import[:]; m.opset_import.extend(ops) |
| onnx.checker.check_model(m) |
| onnx.save(m, dst) |
|
|
