Initial release: fp16 CoreML .mlpackage MDX-Net stem separation models + model card + export tooling

UVR MDX-Net (9482, Voc FT, Inst HQ 3) converted to CoreML mlprogram fp16 for ANE/GPU on iOS. Includes export pipeline, LICENSE, example audio.

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+example.mp3 filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,30 @@
+MIT License
+
+Copyright (c) 2026 Gyoom
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+---
+
+The model weights are derived from the Ultimate Vocal Remover (UVR) MDX-Net
+models by Anjok07 (UVR-MDXNET 9482) and Kimberley Jensen (UVR-MDX-NET Voc FT),
+distributed under the MIT License via https://github.com/TRvlvr/model_repo and
+mirrored at https://huggingface.co/Politrees/UVR_resources. This repository
+redistributes format-converted (ONNX -> CoreML fp16) copies under the same
+MIT terms.

README.md

@@ -1,3 +1,117 @@
 ---
 license: mit
+tags:
+  - coreml
+  - apple-silicon
+  - apple-neural-engine
+  - audio-source-separation
+  - music-source-separation
+  - mdx-net
+  - uvr
+  - vocals
+  - instrumental
+  - fp16
+  - mlprogram
+  - on-device
+pipeline_tag: audio-to-audio
 ---
+
+# UVR MDX-Net CoreML — on-device music stem separation for iOS
+
+On-device **vocals / instrumental (music stem) separation** models, converted from the Ultimate Vocal
+Remover (UVR) MDX-Net checkpoints to **CoreML `.mlpackage`** files (fp16, `mlprogram`) for the Apple
+Neural Engine / GPU / CPU on iOS and macOS.
+
+Each `.mlpackage` is the **learned core only** — the STFT/iSTFT live in your app's DSP, so the whole
+signal pipeline stays under your control and runs entirely on-device. No cloud, no network.
+
+## Models
+
+| Model file | Source (UVR MDX-Net) | Stem | Size | `n_fft` | `hop` | `dim_f` |
+| --- | --- | --- | --- | --- | --- | --- |
+| `UVR_MDXNET_9482.mlpackage` | UVR-MDXNET 9482 | **vocals** | ~15 MB | 4096 | 1024 | 2048 |
+| `UVR-MDX-NET-Voc_FT.mlpackage` | UVR-MDX-NET Voc FT | **vocals** | ~32 MB | 6144 | 1024 | 3072 |
+| `UVR-MDX-NET-Inst_HQ_3.mlpackage` | UVR-MDX-NET Inst HQ 3 | **instrumental** | ~32 MB | 6144 | 1024 | 3072 |
+
+All three share the same I/O contract: **NCHW**, static `dim_t = 256`, input `[1, 4, dim_f, 256]`
+complex-as-channels `[L_re, L_im, R_re, R_im]` → output of the same shape = the predicted stem's
+spectrogram. Real **fp16** I/O with fp16 compute, `mlprogram`, `minimum_deployment_target = iOS16`.
+
+> ⚠️ **Inst HQ 3 inverts the residual.** It is architecturally identical to Voc FT but 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" will get the stems swapped.
+
+## Download
+
+```bash
+# huggingface_hub (Python)
+pip install huggingface_hub
+python - <<'PY'
+from huggingface_hub import snapshot_download
+snapshot_download("gyoom-sa/UVR-MDX-CoreML", local_dir="./UVR-MDX-CoreML")
+PY
+
+# or the CLI
+hf download gyoom-sa/UVR-MDX-CoreML --local-dir ./UVR-MDX-CoreML
+```
+
+The `.mlpackage` folders keep their on-disk structure in this repo, so `git clone` / `hf download`
+yields ready-to-compile packages. Add them to your Xcode target (or call
+`MLModel.compileModel(at:)` at runtime) and they compile to `.mlmodelc`.
+
+## Usage (iOS / Swift)
+
+The model operates on a **spectrogram**, not raw audio. Build the STFT in your app (Apple's
+Accelerate / vDSP), feed the complex planes to the model, then iSTFT the output. Per model:
+**9482** `n_fft 4096, hop 1024, dim_f 2048`; **Voc FT / Inst HQ 3** `n_fft 6144, hop 1024, dim_f 3072`.
+Use a periodic Hann window, `center=True` reflect padding, **unnormalized** STFT, drop the Nyquist bin
+(`dim_f = n_fft/2`), and pack planes in order `[L_re, L_im, R_re, R_im]` with flat row-major index
+`((plane*dim_f)+bin)*dim_t + frame`.
+
+```swift
+import CoreML
+
+// Accelerator is a LOAD-TIME choice. .all == ANE → GPU → CPU (the baked default).
+let config = MLModelConfiguration()
+config.computeUnits = .all               // .cpuAndGPU | .cpuOnly also available
+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 stem = out.featureValue(for: "output")!.multiArrayValue!     // fp16 [1,4,dim_f,256] → iSTFT
+```
+
+For the stem models, `vocals = model(mix)`; for **Inst HQ 3**, `instrumental = model(mix)` and
+`vocals = mix − instrumental`. The equal-power "shift trick" denoise (`0.5·model(x) − 0.5·model(−x)`)
+is optional.
+
+## Notes & limitations
+
+- **STFT is outside the graph.** These are the learned core only; the app must reproduce the host
+  DSP contract above (Hann window, unnormalized STFT, reflect padding, Nyquist bin dropped, plane
+  order, flat index). Getting the packing wrong silently degrades or swaps the stems.
+- **fp16 is quality-safe for MDX** — peak activations sit far below the 65504 fp16 ceiling (measured
+  ~560 for 9482 / ~1384 for Voc FT), and the pipeline normalizes the mix, so the fp16 output floor
+  (~73 dB) is well above audible noise.
+- **Residual polarity** differs per model (see the Inst HQ 3 warning).
+- These are **format-converted on-device artifacts**, not models loadable in `transformers` or via
+  HF Inference. `pipeline_tag: audio-to-audio` is set for task-based discoverability only.
+
+## Reproducibility
+
+The exact conversion pipeline (`export_mdx_coreml.py` + pinned `requirements.txt`) and the full
+technical write-up (precision, compute units, verification SNR, iOS DSP contract) live in the
+[`export/`](./export/README.md) folder of this repo. `example.mp3` is provided for a quick smoke test.
+
+## License & attribution
+
+The model weights are derived from the Ultimate Vocal Remover (UVR) MDX-Net models by **Anjok07**
+(UVR-MDXNET 9482) and **Kimberley Jensen** (UVR-MDX-NET Voc FT / Inst HQ 3), distributed under the
+MIT License via https://github.com/TRvlvr/model_repo and mirrored at
+https://huggingface.co/Politrees/UVR_resources. This repository redistributes format-converted
+(ONNX → CoreML fp16) copies under the same MIT terms. See [`LICENSE`](./LICENSE).

UVR-MDX-NET-Inst_HQ_3.mlpackage/Data/com.apple.CoreML/model.mlmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fe4055846ccddb4a04765feb20ec7fe6530624383323782a91d125f760f05ac4
+size 115669

UVR-MDX-NET-Inst_HQ_3.mlpackage/Data/com.apple.CoreML/weights/weight.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0a687c627a66e8dc3dc6947e31d35526164f234757208624506ff2aaf47bd1a5
+size 33377792

UVR-MDX-NET-Inst_HQ_3.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "95ff6d20-7277-4643-b3ca-76a3abbfb95a": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        },
+        "c2c2bdc2-b64a-4519-94aa-8b8cd087b7c5": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        }
+    },
+    "rootModelIdentifier": "95ff6d20-7277-4643-b3ca-76a3abbfb95a"
+}

UVR-MDX-NET-Voc_FT.mlpackage/Data/com.apple.CoreML/model.mlmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bdd6a5c0b2d3d9eae431f82a58e194c43f08cc3f9da664fdb4daa9d7c4b9bc43
+size 115669

UVR-MDX-NET-Voc_FT.mlpackage/Data/com.apple.CoreML/weights/weight.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ee8aeba2251618d8248bfc3bbad49bb5d3b3ee0b6caa47182e08f2e5e3131c22
+size 33377792

UVR-MDX-NET-Voc_FT.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "690b47fa-f6b7-4811-8dc1-b7c5bcc91664": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        },
+        "8d9d7018-0d95-46da-97b6-330fc9f0f65d": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        }
+    },
+    "rootModelIdentifier": "690b47fa-f6b7-4811-8dc1-b7c5bcc91664"
+}

UVR_MDXNET_9482.mlpackage/Data/com.apple.CoreML/model.mlmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0512ade072df20bd137e83051cb86dd21d32361bdd90535b94c23224ba569a83
+size 115552

UVR_MDXNET_9482.mlpackage/Data/com.apple.CoreML/weights/weight.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c58b40c6767e15c2e31be86eca74ddb72765cc169ca698559812daed0f273a1b
+size 14850048

UVR_MDXNET_9482.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "d57d5b90-d8e8-4c71-a9d2-987177ed2f60": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        },
+        "fa44167b-afd1-441a-9edb-1970eee79c0e": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        }
+    },
+    "rootModelIdentifier": "fa44167b-afd1-441a-9edb-1970eee79c0e"
+}

example.mp3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1b3783ea92eafc1741e6730bc0dc22b047c56dfcb0defc8af085ec9b5e937f56
+size 2402037

export/README.md

@@ -0,0 +1,159 @@
+# 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.

export/export_mdx_coreml.py

@@ -0,0 +1,228 @@
+#!/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()

export/requirements.txt

@@ -0,0 +1,23 @@
+# Pinned MDX-Net -> CoreML export toolchain for THIS repository (see README.md).
+# Self-contained: rebuild from this folder, no other repo paths needed.
+#
+# Rebuild (Python 3.11) from the export/ folder:
+#   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 (conv/bn/relu/convtranspose/matmul) cleanly regardless (verified,
+# torch-vs-ONNX SNR ~108-112 dB).
+#
+# Converting runs on Linux; `model.predict()` (the numeric CoreML-vs-ONNX check in step 4) needs the
+# CoreML runtime and is macOS-only.
+torch==2.9.1+cpu
+coremltools==9.0
+onnx2torch==1.5.15
+onnx
+onnxruntime
+numpy