#!/usr/bin/env python3 """One-shot: a UVR MDX-Net .onnx -> a GPU/ANE-ready fp16 CoreML .mlpackage for iOS. Independent STFT-outside-the-graph export: the STFT/iSTFT stay OUTSIDE the graph (in the app's own DSP), so the .mlpackage is the learned core only, with the ONNX's NCHW I/O kept intact — input [1, 4, dim_f, 256] complex-as-channels [L_re, L_im, R_re, R_im] -> output same shape = the vocal spectrogram. Pipeline (each step gated so a bad export fails loudly, not silently): 1. onnx2torch(onnx) -> nn.Module 2. SNR-gate the torch module vs ONNX Runtime (must be > --min-snr, default 100 dB) 3. torch.jit.trace -> coremltools.convert(mlprogram, fp16, compute_units, iOS target) -> .mlpackage 4. verify: input shape [1,4,dim_f,256] from the saved spec; on macOS also a CoreML-vs-ONNX predict SNR. Why no op-level surgery. CoreML's converter and its Metal-GPU / Apple-Neural-Engine backends lower `ConvTranspose + ReLU` natively, so targeting `mlprogram` is a single clean convert — no op-version patching, no post-conversion fixes. fp16 (`compute_precision=FLOAT16`) is the ~2x lever and is MDX-safe (peak activation << the 65504 fp16 ceiling; MDX has no whole-tensor reduction). The Apple Neural Engine is fp16-native. Compute units are a LOAD-TIME choice on-device (Swift `MLModelConfiguration.computeUnits`). Passing `--compute-units` here only sets the default the Python `predict` uses on macOS; the shipped iOS app picks the accelerator itself. Default is ALL (ANE + GPU + CPU) — see README. Env (pinned, see requirements.txt): coremltools 9.0, onnx2torch 1.5.15, torch 2.9.1+cpu, onnx, onnxruntime, numpy, Python 3.11. NOTE: converting runs anywhere; `predict()` (the numeric CoreML check in step 4) is macOS-only — on Linux the .mlpackage is produced and structurally checked, and the CoreML-vs-ONNX SNR is reported as macOS-pending. Usage: python export_mdx_coreml.py UVR_MDXNET_9482.onnx UVR_MDXNET_9482.mlpackage python export_mdx_coreml.py UVR-MDX-NET-Voc_FT.onnx UVR-MDX-NET-Voc_FT.mlpackage """ import argparse import os import platform import sys import types import warnings warnings.filterwarnings("ignore") os.environ.setdefault("GRPC_VERBOSITY", "NONE") # onnx2torch eagerly imports torchvision; MDX uses no vision ops, so stub it out when the real package # is missing or ABI-broken in the conversion env. try: import torchvision # noqa: F401 except Exception: _tv = types.ModuleType("torchvision") _tv.__version__ = "0.0.0" _ops = types.ModuleType("torchvision.ops") def _stub(*a, **k): raise NotImplementedError("torchvision op unavailable (stub)") for _n in ("nms", "batched_nms", "roi_align", "RoIAlign", "DeformConv2d", "deform_conv2d"): setattr(_ops, _n, _stub) _tv.ops = _ops sys.modules["torchvision"] = _tv sys.modules["torchvision.ops"] = _ops import numpy as np import onnx import onnx2torch import onnxruntime as ort import torch import coremltools as ct DIM_T = 256 # native trained segment (~5.92 s @ hop 1024); static in the shipped .mlpackage # --compute-units / --deployment-target string -> coremltools enum. COMPUTE_UNITS = { "all": ct.ComputeUnit.ALL, # ANE + GPU + CPU (default; best on Apple silicon) "cpuAndGPU": ct.ComputeUnit.CPU_AND_GPU, "cpuAndNE": ct.ComputeUnit.CPU_AND_NE, "cpuOnly": ct.ComputeUnit.CPU_ONLY, } DEPLOY_TARGETS = { "iOS15": ct.target.iOS15, "iOS16": ct.target.iOS16, "iOS17": ct.target.iOS17, "iOS18": ct.target.iOS18, } IN_NAME = "input" # CaC spectrogram [1, 4, dim_f, 256] (Swift MLFeatureProvider key) OUT_NAME = "output" # vocal spectrogram [1, 4, dim_f, 256] def onnx_dim_f(onnx_path): """MDX input is [batch, 4, dim_f, dim_t]; return dim_f (2048 for 9482, 3072 for Voc FT).""" dims = onnx.load(onnx_path).graph.input[0].type.tensor_type.shape.dim return int(dims[2].dim_value) def sample(dim_f): return np.random.randn(1, 4, dim_f, DIM_T).astype(np.float32) def snr(ref, got): den = np.sqrt((ref ** 2).mean()) num = np.sqrt(((ref - got) ** 2).mean()) return 20.0 * np.log10(den / (num + 1e-12)) def module_snr(mod, onnx_path, dim_f, trials=2): """Fidelity of the onnx2torch module vs ONNX Runtime — the CoreML conversion INPUT. This gate runs fully on Linux (no CoreML runtime needed).""" sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"]) iname = sess.get_inputs()[0].name worst = 1e9 for _ in range(trials): x = sample(dim_f) ref = sess.run(None, {iname: x})[0] with torch.no_grad(): got = mod(torch.from_numpy(x)).numpy() worst = min(worst, snr(ref, got)) return float(worst) def dir_size_mb(path): total = 0 for root, _, files in os.walk(path): for f in files: total += os.path.getsize(os.path.join(root, f)) return total / 1e6 def coreml_predict_snr(mlmodel, onnx_path, dim_f, trials=2): """macOS-only: run the compiled CoreML model and compare vs ONNX Runtime. Raises on non-macOS (no CoreML runtime), which the caller catches and reports as macOS-pending.""" sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"]) iname = sess.get_inputs()[0].name worst, oshape = 1e9, None for _ in range(trials): x = sample(dim_f) # fp32 reference input ref = sess.run(None, {iname: x})[0] # fp32 ONNX golden # Model I/O is fp16: feed the fp16 view so SNR reflects the true deployed path (fp16 I/O + # fp16 compute) vs the fp32 ONNX golden. got = mlmodel.predict({IN_NAME: x.astype(np.float16)})[OUT_NAME] oshape = tuple(np.asarray(got).shape) worst = min(worst, snr(ref, np.asarray(got, dtype=np.float32))) return float(worst), oshape def main(): ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) ap.add_argument("onnx", help="source UVR MDX-Net .onnx") ap.add_argument("out", help="output .mlpackage") ap.add_argument("--min-snr", type=float, default=100.0) ap.add_argument("--compute-units", choices=list(COMPUTE_UNITS), default="all", help="baked default for Python predict; the iOS app sets this at load time (default: all)") ap.add_argument("--deployment-target", choices=list(DEPLOY_TARGETS), default="iOS16") args = ap.parse_args() dim_f = onnx_dim_f(args.onnx) print(f"[1/4] onnx2torch: {args.onnx} (dim_f={dim_f}, dim_t={DIM_T})", flush=True) mod = onnx2torch.convert(args.onnx).eval() msnr = module_snr(mod, args.onnx, dim_f) print(f"[2/4] torch-module SNR vs ONNX Runtime = {msnr:.1f} dB", flush=True) if msnr < args.min_snr: raise SystemExit(f"FAIL: onnx2torch fidelity {msnr:.1f} dB < {args.min_snr} dB (do not ship)") print(f"[3/4] trace + coremltools.convert -> {args.out} " f"(mlprogram, fp16, compute_units={args.compute_units}, {args.deployment_target})", flush=True) example = torch.from_numpy(sample(dim_f)) with torch.no_grad(): traced = torch.jit.trace(mod, example) # NCHW I/O is kept from the ONNX so the app's CaC packing feeds it directly (no transpose). # I/O is real fp16 (dtype `np.float16`, matching the fp16 compute) — the fastest, most idiomatic # config for an fp16 ANE/GPU model: no fp32<->fp16 boundary cast, half the I/O bandwidth on the # ~2M-element tensors per chunk, and the ANE's native precision end to end. It is quality-safe # because the pipeline normalizes the mix, so the input spectrogram sits far below the 65504 fp16 # ceiling (peak activation ~560/1384 measured — fp16-safety end to end, extended to # the boundary). The Swift side feeds/reads an MLMultiArray(.float16); packing the fp32 host-DSP # spectrogram into fp16 (and back) is a trivial element-wise cast. See README "iOS integration". mlmodel = ct.convert( traced, inputs=[ct.TensorType(name=IN_NAME, shape=(1, 4, dim_f, DIM_T), dtype=np.float16)], outputs=[ct.TensorType(name=OUT_NAME, dtype=np.float16)], convert_to="mlprogram", compute_precision=ct.precision.FLOAT16, # fp16: the ~2x lever, MDX-safe; ANE is fp16-native compute_units=COMPUTE_UNITS[args.compute_units], minimum_deployment_target=DEPLOY_TARGETS[args.deployment_target], ) mlmodel.save(args.out) # [4/4] verify. Input shape comes straight from the saved spec (works on any OS). The numeric # CoreML-vs-ONNX SNR needs the CoreML runtime -> macOS only; on Linux report it as pending. # Read the ON-DISK artifact's spec (the source of truth for the shippable .mlpackage; the # in-memory convert result doesn't finalize the fp16 I/O boundary until serialized). spec = ct.models.MLModel(args.out, skip_model_load=True).get_spec() in_shape = list(spec.description.input[0].type.multiArrayType.shape) out_shape = list(spec.description.output[0].type.multiArrayType.shape) # CoreML ArrayFeatureType.ArrayDataType: FLOAT32=65568, FLOAT16=65552, DOUBLE=65600. io_dtype = {65568: "fp32", 65552: "fp16", 65600: "fp64"}.get( spec.description.input[0].type.multiArrayType.dataType, "?") size_mb = dir_size_mb(args.out) shape_ok = (in_shape == [1, 4, dim_f, DIM_T] and spec.WhichOneof("Type") == "mlProgram" and io_dtype == "fp16") csnr, cshape, verified = None, None, False try: csnr, cshape = coreml_predict_snr(mlmodel, args.onnx, dim_f) verified = True except Exception as e: # non-macOS (no CoreML runtime) or predict failure note = f"{type(e).__name__}: {str(e).splitlines()[0][:120]}" if str(e) else type(e).__name__ print(f"[4/4] verify: in={in_shape} out={out_shape or 'unspecified (shape-preserving U-Net)'} | " f"io={io_dtype} compute=fp16 | mlprogram | size={size_mb:.1f} MB", flush=True) if verified: ok = shape_ok and csnr >= args.min_snr and cshape == (1, 4, dim_f, DIM_T) print(f" CoreML-vs-ONNX predict SNR = {csnr:.1f} dB (macOS) | out={cshape}", flush=True) print(" " + ("PASS — GPU/ANE-ready fp16 .mlpackage" if ok else "FAIL — do not ship"), flush=True) raise SystemExit(0 if ok else 1) else: ok = shape_ok print(f" CoreML predict SNR: MACOS-PENDING (no CoreML runtime here — {note})", flush=True) print(" re-run this script on macOS to get the on-CoreML fidelity number", flush=True) print(" " + ("PASS (structural) — .mlpackage produced; numeric CoreML check macOS-pending" if ok else "FAIL — input shape mismatch"), flush=True) raise SystemExit(0 if ok else 1) if __name__ == "__main__": main()