# nemo128.onnx with its STFT in float32 instead of float64: onnxruntime-web's WebGPU build has no float64 Cast
import sys, glob, onnx, numpy as np, onnxruntime as ort, soundfile as sf
from onnx import numpy_helper
src, out = sys.argv[1], sys.argv[2]
m = onnx.load(src)
g = m.graph
cast_in = next(n for n in g.node if n.name == 'node_Cast_6')
cast_out = next(n for n in g.node if n.name == 'node_Cast_7')
stft = next(n for n in g.node if n.op_type == 'STFT')
# STFT reads the padded float32 signal directly; its window becomes float32; the result needs no cast back
stft.input[0] = cast_in.input[0]
for init in g.initializer:
    if init.name in stft.input and init.data_type == onnx.TensorProto.DOUBLE:
        init.CopyFrom(numpy_helper.from_array(numpy_helper.to_array(init).astype(np.float32), init.name))
for n in g.node:
    n.input[:] = [cast_out.input[0] if i == cast_out.output[0] else i for i in n.input]
g.node.remove(cast_in); g.node.remove(cast_out)
# drop stale value_info for the removed or retyped tensors
keep = [v for v in g.value_info if v.name not in {'tmp_16', 'image', 'spectrogram', 'tmp_17'}]
del g.value_info[:]; g.value_info.extend(keep)
onnx.checker.check_model(m)
onnx.save(m, out)
worst = 0
for f in sys.argv[3:]:
    a, _ = sf.read(f, dtype='float32')
    feeds = {'waveforms': a[None], 'waveforms_lens': np.array([len(a)], dtype=np.int64)}
    r64 = ort.InferenceSession(src).run(None, feeds)[0]
    r32 = ort.InferenceSession(out).run(None, feeds)[0]
    worst = max(worst, float(np.abs(r64 - r32).max()))
print('max feature difference fp64 vs fp32 STFT:', worst)