UVR-MDX-CoreML / export /README.md
ibrhr's picture
Initial release: fp16 CoreML .mlpackage MDX-Net stem separation models + model card + export tooling
f532e9a verified
|
Raw
History Blame Contribute Delete
9.18 kB

Export pipeline for the UVR-MDX-CoreML repository

This folder is the one-shot conversion pipeline used to produce the CoreML .mlpackage models shipped in this repository (gyoom-sa/UVR-MDX-CoreML). It is self-contained: it takes a UVR MDX-Net source .onnx and emits a GPU/ANE-ready fp16 CoreML .mlpackage for iOS — nothing here depends on anything outside this repo.

It reproduces the same STFT-outside-the-graph formula for all three models: the .mlpackage is the learned core only; the STFT/iSTFT stay in the app's own DSP.

Models

Source ONNX (from the TRvlvr/model_repo UVR releases):

Source ONNX Output .mlpackage Size dim_f Predicted stem
UVR_MDXNET_9482.onnx UVR_MDXNET_9482.mlpackage ~15 MB 2048 vocals
UVR-MDX-NET-Voc_FT.onnx UVR-MDX-NET-Voc_FT.mlpackage ~32 MB 3072 vocals
UVR-MDX-NET-Inst_HQ_3.onnx UVR-MDX-NET-Inst_HQ_3.mlpackage ~32 MB 3072 instrumental

All keep NCHW I/O with a static dim_t = 256: input [1, 4, dim_f, 256], complex-as-channels [L_re, L_im, R_re, R_im] → output with the same shape = the predicted stem's spectrogram. Real fp16 I/O, fp16 compute, mlprogram, minimum_deployment_target = iOS16 (~half the size of an fp32 .mlpackage because the weights are stored fp16).

Inst HQ 3 inverts the residual. It is architecturally identical to Voc FT (same 178 ops, same dim_f 3072 / n_fft 6144 / hop 1024, 16.68 M params, no embedded metadata) but it is an Inst model: the graph predicts the instrumental, so the free residual is vocals = mix − model(mix) — the opposite polarity from 9482/Voc FT. A caller that assumes "output = vocals" gets the stems swapped.

Pipeline (export_mdx_coreml.py)

onnx2torch(onnx) → nn.Module → [SNR-gate vs ONNX Runtime > 100 dB]
  → torch.jit.trace → coremltools.convert(mlprogram, fp16, compute_units, iOS16) → .mlpackage → verify
  1. onnx2torch(onnx)nn.Module (a plain conv U-Net: conv / bn / relu / convtranspose / matmul).
  2. SNR-gate the torch module vs ONNX Runtime (must clear --min-snr, default 100 dB).
  3. torch.jit.tracecoremltools.convert(..., convert_to="mlprogram", compute_precision=FLOAT16, compute_units=ALL, minimum_deployment_target=iOS16), keeping NCHW I/O and real fp16 I/O (dtype=np.float16).
  4. Verify: on-disk spec is an mlProgram with input/output [1,4,dim_f,256] fp16; on macOS also a CoreML-vs-ONNX predict SNR (see "Verification status").

Why onnx2torch (not a direct ONNX→CoreML path). coremltools dropped its ONNX front-end years ago; the supported route is Torch→CoreML. onnx2torch bridges our ONNX to a torch nn.Module, which traces and converts cleanly, keeping the conversion input numerically faithful to the ONNX.

Precision — fp16, and why it's safe

compute_precision=FLOAT16 runs the whole graph in fp16 (the ~2× lever, native to the ANE and GPU). MDX has no whole-tensor reduction — only per-channel BatchNorm that folds into the conv — so nothing overflows the 65504 fp16 ceiling (peak activation ~560 for 9482 / ~1384 for Voc FT, measured). I/O is real fp16 too, which is quality-safe for the same reason: the app's pipeline normalizes the mix, so the input spectrogram sits far below the ceiling. Real fp16 I/O is the fastest, most idiomatic config for an ANE/GPU model — no fp32↔fp16 boundary cast, half the I/O bandwidth on the ~2M-element tensors.

Compute units — a LOAD-TIME choice

--compute-units all bakes ComputeUnit.ALL (ANE + GPU + CPU) as the model's default, but on iOS the real selector is load-time: MLModelConfiguration.computeUnits. .all lets CoreML place ops on the Apple Neural Engine first (fastest + most power-efficient on iPhone/iPad), then GPU, then CPU. Use .cpuAndGPU to force the GPU and .cpuOnly as the CPU floor. The --compute-units flag here only affects the Python predict default; the shipping app chooses per its own config (see below).

Pinned environment

Python 3.11
coremltools    9.0
onnx2torch     1.5.15
torch          2.9.1+cpu   (via the pytorch cpu index)
onnx, onnxruntime, numpy

Rebuild from requirements.txt:

python3.11 -m venv .venv
.venv/bin/pip install --extra-index-url https://download.pytorch.org/whl/cpu -r requirements.txt

The --extra-index-url is REQUIRED: torch is pinned to the +cpu wheel and must resolve from the pytorch cpu index, otherwise pip pulls the multi-GB CUDA build (and its nvidia-* deps). Pinning torch in the same resolve also stops pip from silently upgrading it to a CUDA torch.

coremltools 9.0 prints "Torch 2.9.1 has not been tested (2.7.0 is newest tested)" — a generic warning; it converts the MDX conv U-Net cleanly regardless (verified, torch-vs-ONNX SNR ~108–112 dB).

Run

PY=.venv/bin/python
$PY export_mdx_coreml.py ../UVR_MDXNET_9482.onnx        output/UVR_MDXNET_9482.mlpackage
$PY export_mdx_coreml.py ../UVR-MDX-NET-Voc_FT.onnx     output/UVR-MDX-NET-Voc_FT.mlpackage
$PY export_mdx_coreml.py ../UVR-MDX-NET-Inst_HQ_3.onnx  output/UVR-MDX-NET-Inst_HQ_3.mlpackage --min-snr 90
# options: --compute-units {all,cpuAndGPU,cpuAndNE,cpuOnly}  --deployment-target {iOS15..iOS18}  --min-snr

The source .onnx files are the TRvlvr UVR releases (see the models table above); point the script at your download of each.

Why Inst HQ 3 needs --min-snr 90. Its onnx2torch fidelity is 95.4 dB, below the 100 dB default. That is not a broken conversion: the error is dense and unstructured with a scale-invariant absolute magnitude (~1.8e-4 max, unchanged from gaussian×1 to gaussian×0.05 inputs) — plain fp32 accumulation round-off. Inst HQ 3 simply carries ~4× larger activations than Voc FT (peak ~9.8 vs ~2.4), so the same relative round-off reads as a lower RMS SNR. A real op-level defect (the failure mode this gate exists to catch) lands at 10–30 dB, not 95. For context the shipped artifact is fp16, whose output-quantization floor is ~73 dB, so a 100 dB gate on the fp32 intermediate is stricter than the deliverable can ever be; 90 dB still sits ~17 dB above that floor.

Verification status

  • onnx2torch fidelity (any OS): torch-module vs ONNX Runtime SNR ≈ 108 dB (9482) / 112 dB (Voc FT) — the CoreML conversion input is faithful. Gated at 100 dB.
  • Structural (any OS): on-disk .mlpackage is an mlProgram, I/O [1,4,dim_f,256] fp16.
  • Numeric CoreML fidelity + on-device (macOS-pending): model.predict() needs the CoreML runtime, which is macOS-only — on Linux the script reports it as MACOS-PENDING. Re-run export_mdx_coreml.py on a Mac to get the CoreML-vs-ONNX predict SNR, and measure the ANE/GPU RTF on-device.

iOS integration (the contract to port)

The .mlpackage is only the learned core — an iOS app must reproduce the same host DSP the shipped Android app uses: periodic Hann window, center=True reflect padding, unnormalized STFT (torch.stft normalized=False), Nyquist bin dropped (dim_f = n_fft/2), plane order [L_re, L_im, R_re, R_im], flat row-major index ((plane*dim_f)+bin)*dim_t + frame. Per model: 9482 n_fft 4096, hop 1024, dim_f 2048; Voc FT / Inst HQ 3 n_fft 6144, hop 1024, dim_f 3072. The pipeline (chunk at 10% overlap → STFT → model → iSTFT → overlap-add; instrumental = mix − vocals for the stem models, the inverse for Inst HQ 3) is the standard MDX separation loop. Port the STFT/iSTFT to Accelerate/vDSP; the model call is:

import CoreML

// Load with the accelerator policy. `.all` == the baked default (ANE → GPU → CPU). Xcode compiles the
// .mlpackage into a .mlmodelc when it's added to the target (or MLModel.compileModel(at:) at runtime).
let config = MLModelConfiguration()
config.computeUnits = .all                    // GPU-only: .cpuAndGPU   ·   CPU floor: .cpuOnly
let model = try MLModel(contentsOf: compiledURL, configuration: config)

// Per chunk. dim_f = 2048 (9482) or 3072 (Voc FT / Inst HQ 3); dim_t = 256. Real fp16 I/O.
let input = try MLMultiArray(shape: [1, 4, dimF as NSNumber, dimT as NSNumber], dataType: .float16)
let p = input.dataPointer.bindMemory(to: Float16.self, capacity: input.count)
// pack the STFT: p[((plane*dimF)+bin)*dimT + frame] = value   (plane 0..3 = L_re,L_im,R_re,R_im)
let out = try model.prediction(from: MLDictionaryFeatureProvider(dictionary: ["input": input]))
let stems = out.featureValue(for: "output")!.multiArrayValue!   // fp16 [1,4,dim_f,256] → iSTFT

MLMultiArray is row-major for the given shape, so the flat index above matches the packing exactly. The UVR "shift trick" denoise (0.5·model(x) − 0.5·model(−x)) is optional. Remember the Inst HQ 3 residual inversion (vocals = mix − model(mix)) — see the models table.

The repository root README.md is the user-facing entry point; this document is the conversion tooling reference for re-exporting or auditing the shipped artifacts.