"""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   # constant folding only: stays standard ONNX
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):   # store the float tensors `keep` picks as float16, cast back to float at load
    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)
# 1) the network's weights (conv/matmul, not the STFT kernels)
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())
# 2) the remaining big tensors: STFT kernels and folded embeddings; the iSTFT window envelope stays float32
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)