export/README.md

# 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](https://github.com/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.trace` → `coremltools.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`](requirements.txt):

```bash
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

```bash
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:

```swift
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.