Initial release: HTDemucs to Core ML conversion (FP32 + FP16)

ATTRIBUTION.md

@@ -0,0 +1,61 @@
+# Attribution
+
+The Core ML model produced by this repository is a converted form of
+**Hybrid Transformer Demucs (HTDemucs)** by Meta AI / Meta Platforms, Inc.
+
+When you ship the resulting `.mlpackage` in an application, you must keep the
+following attribution available to users (e.g., in an "About" / "Legal"
+screen, in your privacy policy, or in a `THIRD_PARTY_LICENSES` file):
+
+---
+
+> This product uses Hybrid Transformer Demucs by Meta Platforms, Inc.,
+> licensed under the MIT License.
+>
+> Source: https://github.com/facebookresearch/demucs
+>
+> Reference papers:
+>   - Défossez, A. (2021). *Hybrid Spectrogram and Waveform Source Separation.*
+>     Proceedings of the ISMIR 2021 Workshop on Music Source Separation.
+>   - Rouard, S., Massa, F., & Défossez, A. (2023).
+>     *Hybrid Transformers for Music Source Separation.* ICASSP 2023.
+
+---
+
+## Demucs upstream license
+
+```
+MIT License
+
+Copyright (c) Meta Platforms, Inc. and affiliates.
+
+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 pre-trained HTDemucs weights are released under the same MIT license by
+Meta and are downloaded automatically by the `demucs` Python package the
+first time `convert.py` runs.
+
+## Not affiliated
+
+This repository is an independent open-source project. It is **not
+affiliated with, endorsed by, or sponsored by Apple, Meta, or the Demucs
+authors**. Apple, Core ML, the Neural Engine, and any other Apple
+trademarks are property of Apple Inc. and used here for descriptive
+purposes only.

CONVERSION_NOTES.md

@@ -0,0 +1,190 @@
+# Conversion notes — the three workarounds
+
+This document explains the three non-trivial pieces in `convert.py`. Each one
+is a Core ML / `coremltools` quirk that won't show up until you're already
+several layers deep into a conversion attempt.
+
+The same patterns apply to **any audio model** that uses STFT + Transformer
+(MDX-Net, Spleeter, OpenUnmix, BS-RoFormer, …), so they're reusable beyond
+HTDemucs.
+
+---
+
+## 1. `complex64` is not supported in Core ML
+
+### Problem
+
+`torch.stft(..., return_complex=True)` produces `complex64`. `coremltools`
+cannot represent complex tensors, so any graph that carries them through
+will fail to convert (or, worse, "convert" but produce nonsense).
+
+The naive workaround — call `torch.view_as_real` once at the boundary —
+breaks again the moment you need `torch.istft`, because `istft` requires a
+complex input.
+
+### Solution
+
+Replace the entire STFT/ISTFT pair with **purely real-valued operations**
+(`RealSTFT` and `RealISTFT` in `convert.py`).
+
+**`RealSTFT`** is straightforward — keep `torch.stft` (still emits complex
+internally) but immediately call `view_as_real` so the *traced graph* never
+carries a complex tensor:
+
+```python
+z = torch.stft(x_flat, ..., return_complex=True)
+z_ri = torch.view_as_real(z)   # (..., freqs, frames, 2)
+```
+
+**`RealISTFT`** is the harder half. We rebuild ISTFT from scratch using a
+matrix-form IDFT plus an explicit overlap-add:
+
+1. Pre-compute `cos`/`sin` basis matrices for a one-sided spectrum, with
+   correct DC/Nyquist scaling and the `normalized=True` factor folded in.
+2. Per-frame IDFT becomes a single `matmul`:
+   `signal = real @ cos_basis - imag @ sin_basis`.
+3. Apply the synthesis (Hann) window.
+4. Overlap-add into an output buffer.
+5. Divide by the pre-computed sum-of-squared-windows for correct
+   reconstruction normalization.
+6. Strip the `n_fft // 2` center padding.
+
+This is more code than `torch.istft`, but it converts cleanly to Core ML
+and has no `complex` operations anywhere.
+
+### Why not use `coremltools`'s built-in STFT op?
+
+`coremltools` has gained more audio ops over time, but as of writing:
+- The op coverage for STFT/ISTFT around HTDemucs's specific window size,
+  hop length, `normalized=True`, and `center=True` combination is brittle.
+- HTDemucs's own STFT pre/post-processing (the `_pad1d` reflect padding,
+  the `_spec` trim of the last freq bin, the `+2 / -2` frame trim) needs
+  to be reproduced bit-for-bit. It's easier to keep the whole STFT
+  pipeline as plain PyTorch tensor ops.
+
+---
+
+## 2. `nn.MultiheadAttention` cannot be traced
+
+### Problem
+
+`coremltools` can't convert `_native_multi_head_attention`, the fused C++
+op that PyTorch dispatches to inside `nn.MultiheadAttention`. You'll get
+something like:
+
+```
+PyTorch convert function for op '_native_multi_head_attention' not implemented.
+```
+
+### Solution
+
+Replace every `nn.MultiheadAttention` instance with a hand-written
+`ManualMHA` module that decomposes attention into the primitive ops
+`coremltools` *does* support: `linear`, `matmul`, `softmax`.
+
+The substitution is in-place via `_replace_mha_recursive`, which walks the
+HTDemucs `crosstransformer` and swaps modules. This preserves all
+pre-trained weights — `ManualMHA` uses the same `in_proj_weight`,
+`in_proj_bias`, and `out_proj` tensors as the original module.
+
+`ManualMHA` handles two cases:
+- **Self-attention** (`query == key == value` by pointer equality):
+  single `in_proj` then `chunk(3)`.
+- **Cross-attention**: split the `in_proj` weight into Q/K/V slices and
+  apply each linear separately.
+
+Trade-off: a few percent slower than the fused op on CPU/GPU. Negligible
+for the once-per-10-seconds inference cadence in stem separation.
+
+---
+
+## 3. Core ML's 1D `scatter_add` is fragile
+
+### Problem
+
+A natural way to write overlap-add inside an ISTFT is something like:
+
+```python
+output = torch.zeros(batch, out_length)
+output.scatter_add_(1, ola_indices, frames_signal_flat)
+```
+
+`coremltools` *will* convert this, but for some shape/index combinations
+the resulting Core ML graph mis-compiles silently — outputs come out
+slightly wrong or full-on garbage on the GPU backend.
+
+### Solution
+
+Pre-compute the OLA index tensor **once at module init time** and store it
+as a registered buffer:
+
+```python
+frame_offsets = torch.arange(num_frames) * hop_length
+local_offsets = torch.arange(n_fft)
+ola_indices = (frame_offsets.unsqueeze(1) + local_offsets.unsqueeze(0)).reshape(-1)
+self.register_buffer("ola_indices", ola_indices.long())
+```
+
+At forward-time, `expand` it to the batch dimension and call the canonical
+`scatter_add_` with constant indices. This sidesteps the buggy code path
+because the indices are now part of the constant graph rather than a
+runtime computation.
+
+A pre-computed `win_sum` buffer (sum of squared windows over all frames)
+takes care of normalization the same way — done once, not recomputed
+inside the model.
+
+This is why the converter's `RealISTFT.__init__` takes a fixed
+`num_frames` parameter: the converter bakes a specific segment length
+into the model. If you want a different segment length, re-run
+`convert.py --segment N`.
+
+---
+
+## Compute units
+
+`compute_units=ct.ComputeUnit.CPU_AND_GPU` is the default in `convert.py`
+and the right choice for HTDemucs.
+
+`ct.ComputeUnit.ALL` (which lets the runtime route ops to the **Apple
+Neural Engine**) produces incorrect output on some HTDemucs shapes —
+likely a mismatch between ANE's quantization assumptions and the
+pre-/post-processing math around the network. The validation step at the
+end of `convert.py` will report a large max-diff if this happens. Do not
+ignore it.
+
+If you really need ANE inference, you'd have to: split the model so that
+only the encoder/transformer/decoder runs on ANE, and keep the STFT/ISTFT
+on CPU/GPU. That's a larger restructure — out of scope for this repo.
+
+---
+
+## Validation step
+
+After conversion, `convert.py` runs the same dummy input through both the
+PyTorch wrapper and the saved Core ML model and reports max/mean
+differences:
+
+| Precision | Expected max diff | Threshold |
+|---|---|---|
+| FP32 | ~1e-4 to 1e-2 | < 0.1 |
+| FP16 | ~1e-2 to 1e-1 | < 0.2 |
+
+If you see drift larger than the threshold, something has gone wrong —
+typically a `coremltools` version mismatch or an ANE routing bug.
+
+---
+
+## Reproducing on a fresh machine
+
+```bash
+python3 -m venv venv && source venv/bin/activate
+pip install -r requirements.txt
+python convert.py --fp16    # or default FP32
+```
+
+First run downloads ~300 MB of HTDemucs weights to `~/.cache/torch/hub/`.
+Conversion takes 2–5 minutes on an Apple Silicon Mac with 16 GB RAM.
+
+If `coremltools` complains about Python or `torch` versions, see the
+pinned bounds in `requirements.txt`.

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

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

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

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

HTDemucs_CoreML.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "70E43B23-A4BB-4678-9897-9F7F13EFF6F4": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        },
+        "FDA2CD24-2727-44BF-B7A5-1AEBBA959083": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        }
+    },
+    "rootModelIdentifier": "70E43B23-A4BB-4678-9897-9F7F13EFF6F4"
+}

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

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

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

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

HTDemucs_CoreML_FP16.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "02CB79C0-A986-4E3F-ABF8-796A39543460": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        },
+        "5472320A-6037-4CD3-AA89-B348A89A8D53": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        }
+    },
+    "rootModelIdentifier": "5472320A-6037-4CD3-AA89-B348A89A8D53"
+}

LICENSE

@@ -0,0 +1,27 @@
+MIT License
+
+Copyright (c) 2026 Dejan Nikolic
+
+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.
+
+---
+
+This project converts the Demucs (Hybrid Transformer Demucs) model from
+PyTorch to Apple Core ML format. The original Demucs project is licensed
+under MIT as well; see ATTRIBUTION.md for the full upstream notice.

README.md

@@ -0,0 +1,126 @@
+---
+license: mit
+library_name: coreml
+tags:
+  - audio
+  - music-source-separation
+  - stem-separation
+  - demucs
+  - htdemucs
+  - coreml
+  - apple-silicon
+  - swift
+  - macos
+  - ios
+pipeline_tag: audio-to-audio
+---
+
+# HTDemucs → Core ML
+
+Convert Meta's [Hybrid Transformer Demucs](https://github.com/facebookresearch/demucs) (HTDemucs) into a Core ML `.mlpackage` you can drop into a macOS or iOS app and run with `MLModel`.
+
+The hard part of converting HTDemucs to Core ML is not the network itself — it is the STFT/ISTFT and the multi-head attention around it. This repo contains a **single-file converter** (`convert.py`, ~600 LoC) that solves the three blockers you hit otherwise:
+
+1. Core ML doesn't support `complex64` → real-valued STFT/ISTFT.
+2. `coremltools` can't trace `nn.MultiheadAttention` → manual decomposition.
+3. Core ML's 1D `scatter_add` is fragile → pre-computed OLA index buffer.
+
+The result is a **stand-alone `.mlpackage`** that takes raw stereo audio and outputs four stems (vocals, drums, bass, other) at 44.1 kHz.
+
+## Why another conversion?
+
+There is one prior public Core ML conversion of HTDemucs by [john-rocky/CoreML-Models](https://github.com/john-rocky/CoreML-Models) at 7.8 s segments / 80 MB. This repo offers:
+
+- **Longer segments (10 s by default)** → fewer overlap-add boundaries on long files.
+- **CLI flags** for segment length, FP16 quantization, compute-unit selection.
+- **Source order** reordered to `[vocals, drums, bass, other]` (DJ/UI convention).
+- **Documented workarounds** so you can reproduce or adapt the pipeline for other audio models (Spleeter, OpenUnmix, MDX-Net).
+
+## Quick start
+
+```bash
+git clone https://github.com/dexxdean/htdemucs-coreml
+cd htdemucs-coreml
+python3 -m venv venv && source venv/bin/activate
+pip install -r requirements.txt
+
+# default: 10 s segments, FP32, ~400 MB
+python convert.py
+
+# half size, ~200 MB, slight numerical drift but inaudible in practice
+python convert.py --fp16
+
+# shorter segments if you want lower latency / smaller buffers
+python convert.py --segment 7
+```
+
+The output is `HTDemucs_CoreML.mlpackage` (or `HTDemucs_CoreML_FP16.mlpackage`).
+
+## Usage in Swift
+
+```swift
+import CoreML
+import AVFoundation
+
+// 1. Load the model.
+let url = Bundle.main.url(forResource: "HTDemucs_CoreML", withExtension: "mlpackage")!
+let config = MLModelConfiguration()
+config.computeUnits = .cpuAndGPU   // see "Compute units" below
+let model = try MLModel(contentsOf: url, configuration: config)
+
+// 2. Feed a (1, 2, 441000) Float32 MLMultiArray named "audio".
+//    Output is a (1, 4, 2, 441000) Float32 array named "sources",
+//    in the order [vocals, drums, bass, other].
+```
+
+A more complete example with chunking, overlap-add, and `AVAudioEngine` playback is in [`examples/swift/StemSeparator.swift`](examples/swift/StemSeparator.swift).
+
+## Model I/O
+
+| | |
+|---|---|
+| Input name | `audio` |
+| Input shape | `(1, 2, segment_samples)` Float32 |
+| Output name | `sources` |
+| Output shape | `(1, 4, 2, segment_samples)` Float32 |
+| Output order | `vocals, drums, bass, other` |
+| Sample rate | 44 100 Hz, stereo |
+| Default segment | 441 000 samples (10 s) |
+| Min. deployment | macOS 14 / iOS 17 |
+
+## Compute units
+
+HTDemucs is **not stable on the Apple Neural Engine**. Use `.cpuAndGPU` (the default baked into the model). Forcing `.all` or `.cpuAndNeuralEngine` may produce silent garbage on some shapes — the validation step in `convert.py` will warn if numerical drift is large.
+
+## File sizes
+
+| Variant | Size | Notes |
+|---|---|---|
+| FP32, 10 s | ~400 MB | Default, full reference quality. |
+| FP16, 10 s | ~200 MB | Inaudible quality difference for music separation. |
+| FP32, 7.8 s | ~310 MB | Closer to john-rocky's segment length. |
+
+## How it works
+
+See [`CONVERSION_NOTES.md`](CONVERSION_NOTES.md) for the technical deep-dive on the three workarounds (real STFT, manual MHA, OLA scatter).
+
+## License & attribution
+
+This repo is **MIT-licensed** — see [`LICENSE`](LICENSE).
+
+The converted model derives from [facebookresearch/demucs](https://github.com/facebookresearch/demucs) © Meta Platforms, Inc., MIT-licensed. The pre-trained HTDemucs weights are downloaded by the `demucs` Python package at conversion time from Meta's official release. **You must comply with Demucs' MIT license when redistributing the resulting `.mlpackage`** — keep the attribution in `ATTRIBUTION.md` alongside the model and in your app's about/legal screen.
+
+This project is **not affiliated with Apple, Meta, or Demucs**. The package name `HTDemucs_CoreML.mlpackage` was chosen to avoid any confusion with Apple-internal model names (e.g., `MusicSourceSeparation`).
+
+## Citation
+
+If you use this in academic work, please cite the original Demucs papers:
+
+```bibtex
+@inproceedings{rouard2023hybrid,
+  title={Hybrid Transformers for Music Source Separation},
+  author={Rouard, Simon and Massa, Francisco and D{\'e}fossez, Alexandre},
+  booktitle={ICASSP 2023},
+  year={2023}
+}
+```

convert.py

@@ -0,0 +1,658 @@
+#!/usr/bin/env python3
+"""
+convert.py — Convert Demucs (Hybrid Transformer) to Core ML.
+
+Core ML does not support complex64 tensors. This script wraps HTDemucs with
+a real-valued STFT/ISTFT implementation (rfft -> view_as_real for STFT,
+matrix IDFT + overlap-add for ISTFT) while keeping the neural network
+(encoder/transformer/decoder) unchanged.
+
+Default output: HTDemucs_CoreML.mlpackage
+
+Prerequisites:
+    python3 -m venv venv && source venv/bin/activate
+    pip install -r requirements.txt
+
+Usage:
+    python convert.py                    # FP32, ~400 MB
+    python convert.py --fp16             # FP16, ~200 MB
+    python convert.py --segment 7        # 7-second segments instead of 10
+    python convert.py --output Foo.mlpackage
+"""
+
+import argparse
+import math
+import warnings
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+# ---------------------------------------------------------------------------
+# Defaults (override via CLI args)
+# ---------------------------------------------------------------------------
+MODEL_NAME = "htdemucs"
+SAMPLE_RATE = 44100
+SEGMENT_SAMPLES = 441000               # 10s @ 44.1 kHz
+NUM_CHANNELS = 2
+NUM_SOURCES = 4
+DEFAULT_OUTPUT = "HTDemucs_CoreML.mlpackage"
+
+# Demucs internal source order: drums(0), bass(1), other(2), vocals(3)
+# We reorder to vocals, drums, bass, other (typical UI / DJ convention).
+SOURCE_REORDER = [3, 0, 1, 2]
+SOURCE_NAMES = ["vocals", "drums", "bass", "other"]
+
+
+# ---------------------------------------------------------------------------
+# ManualMHA: replaces nn.MultiheadAttention.
+# coremltools cannot convert the fused _native_multi_head_attention op,
+# so we decompose attention into matmul + softmax explicitly.
+# ---------------------------------------------------------------------------
+class ManualMHA(nn.Module):
+    """Drop-in für nn.MultiheadAttention, dekomponiert in matmul+softmax."""
+
+    def __init__(self, mha: nn.MultiheadAttention):
+        super().__init__()
+        self.embed_dim = mha.embed_dim
+        self.num_heads = mha.num_heads
+        self.head_dim = mha.embed_dim // mha.num_heads
+        self.in_proj_weight = mha.in_proj_weight
+        self.in_proj_bias = mha.in_proj_bias
+        self.out_proj = mha.out_proj
+        # Cross-attention: separate k/v projections
+        self.kdim = mha.kdim
+        self.vdim = mha.vdim
+        self._qkv_same_embed_dim = mha._qkv_same_embed_dim
+
+    def forward(self, query, key, value, need_weights=False, **kwargs):
+        B, T, E = query.shape
+        S = key.shape[1]
+
+        if self._qkv_same_embed_dim and query.data_ptr() == key.data_ptr():
+            # Self-attention: single in_proj for Q, K, V.
+            qkv = F.linear(query, self.in_proj_weight, self.in_proj_bias)
+            q, k, v = qkv.chunk(3, dim=-1)
+        else:
+            # Cross-attention or different inputs.
+            w_q, w_k, w_v = self.in_proj_weight.chunk(3, dim=0)
+            b_q, b_k, b_v = (self.in_proj_bias.chunk(3, dim=0)
+                              if self.in_proj_bias is not None
+                              else (None, None, None))
+            q = F.linear(query, w_q, b_q)
+            k = F.linear(key, w_k, b_k)
+            v = F.linear(value, w_v, b_v)
+
+        q = q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
+        v = v.view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
+
+        scale = self.head_dim ** -0.5
+        attn = torch.matmul(q, k.transpose(-2, -1)) * scale
+        attn = F.softmax(attn, dim=-1)
+        out = torch.matmul(attn, v)
+
+        out = out.transpose(1, 2).contiguous().view(B, T, E)
+        out = self.out_proj(out)
+        return out, None
+
+
+def _replace_mha_recursive(module: nn.Module) -> None:
+    """Replace all nn.MultiheadAttention submodules with ManualMHA, in place."""
+    for name, child in module.named_children():
+        if isinstance(child, nn.MultiheadAttention):
+            setattr(module, name, ManualMHA(child))
+        else:
+            _replace_mha_recursive(child)
+
+
+# ---------------------------------------------------------------------------
+# 1D reflect-pad helper (mirrors demucs.hdemucs.pad1d).
+# ---------------------------------------------------------------------------
+def _pad1d(x: torch.Tensor, paddings: tuple, mode: str = "reflect"):
+    """Reflect-pad along the last dim, with a fallback for very short signals."""
+    pl, pr = paddings
+    length = x.shape[-1]
+    max_pad = max(pl, pr)
+    if length <= max_pad:
+        extra_pad = max_pad - length + 1
+        x = F.pad(x, (0, extra_pad))
+        padded = F.pad(x, (pl, pr), mode=mode)
+        end = padded.shape[-1] - extra_pad
+        return padded[..., :end]
+    return F.pad(x, (pl, pr), mode=mode)
+
+
+# ---------------------------------------------------------------------------
+# RealSTFT: real-valued STFT via rfft -> view_as_real.
+# Produces (..., freqs, frames, 2) so no complex64 leaks into the traced graph.
+# ---------------------------------------------------------------------------
+class RealSTFT(nn.Module):
+    """STFT that returns only real tensors."""
+
+    def __init__(self, n_fft: int, hop_length: int):
+        super().__init__()
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        self.register_buffer("window", torch.hann_window(n_fft))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Input:  (B, C, T)
+        Output: (B, C, freqs, frames, 2)   -- [real, imag]
+        """
+        B, C, T = x.shape
+        x_flat = x.reshape(B * C, T)
+
+        # torch.stft -> complex -> immediately view_as_real.
+        z = torch.stft(
+            x_flat, self.n_fft, self.hop_length,
+            window=self.window, win_length=self.n_fft,
+            normalized=True, center=True, return_complex=True,
+        )
+        # z: (B*C, freqs, frames) complex64.
+        z_ri = torch.view_as_real(z)  # (B*C, freqs, frames, 2) float32.
+        _, Fr, Fm, _ = z_ri.shape
+        return z_ri.view(B, C, Fr, Fm, 2)
+
+
+# ---------------------------------------------------------------------------
+# RealISTFT: real-valued ISTFT via matrix IDFT + overlap-add.
+# Avoids view_as_complex (not supported by coremltools).
+# ---------------------------------------------------------------------------
+class RealISTFT(nn.Module):
+    """Pure real-valued ISTFT (matrix IDFT + OLA)."""
+
+    def __init__(self, n_fft: int, hop_length: int, num_frames: int):
+        super().__init__()
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        freqs = n_fft // 2 + 1
+
+        # Synthesis window.
+        window = torch.hann_window(n_fft)
+        self.register_buffer("window", window)
+
+        # IDFT basis matrices: cos / sin for a one-sided spectrum.
+        n = torch.arange(n_fft, dtype=torch.float32).unsqueeze(0)      # (1, N)
+        k = torch.arange(freqs, dtype=torch.float32).unsqueeze(1)      # (freqs, 1)
+        angles = 2.0 * math.pi * k * n / n_fft                         # (freqs, N)
+
+        cos_basis = torch.cos(angles)
+        sin_basis = torch.sin(angles)
+
+        # Scaling: DC and Nyquist single, rest double (one-sided spectrum).
+        # Normalization: /N * sqrt(N) because the forward STFT used normalized=True.
+        norm = math.sqrt(n_fft)
+        scale = torch.ones(freqs, 1) * (2.0 / n_fft * norm)
+        scale[0] = 1.0 / n_fft * norm
+        scale[-1] = 1.0 / n_fft * norm
+
+        self.register_buffer("cos_basis", cos_basis * scale)   # (freqs, N)
+        self.register_buffer("sin_basis", sin_basis * scale)   # (freqs, N)
+
+        # Pre-compute OLA indices and window-sum buffer.
+        # Core ML's 1D scatter_add can mis-compile for some shapes; using a
+        # pre-built index tensor + the canonical scatter_add_ call sidesteps it.
+        out_length = (num_frames - 1) * hop_length + n_fft
+        frame_offsets = torch.arange(num_frames) * hop_length
+        local_offsets = torch.arange(n_fft)
+        ola_indices = (frame_offsets.unsqueeze(1) + local_offsets.unsqueeze(0)).reshape(-1)
+        self.register_buffer("ola_indices", ola_indices.long())
+
+        window_sq = window * window
+        win_sum = torch.zeros(out_length)
+        for i in range(num_frames):
+            start = i * hop_length
+            win_sum[start:start + n_fft] += window_sq
+        win_sum = win_sum.clamp(min=1e-8)
+        self.register_buffer("win_sum", win_sum)
+        self.out_length = out_length
+
+    def forward(self, z_ri: torch.Tensor, length: int) -> torch.Tensor:
+        """
+        Input:  z_ri (batch, freqs, frames, 2)
+        Output: (batch, length)
+        """
+        real = z_ri[..., 0]  # (batch, freqs, frames)
+        imag = z_ri[..., 1]
+
+        # Per-frame IDFT: (batch, frames, freqs) @ (freqs, N) -> (batch, frames, N)
+        real_t = real.transpose(-2, -1)
+        imag_t = imag.transpose(-2, -1)
+
+        frames_signal = (
+            torch.matmul(real_t, self.cos_basis)
+            - torch.matmul(imag_t, self.sin_basis)
+        )
+
+        # Apply synthesis window.
+        frames_signal = frames_signal * self.window.unsqueeze(0).unsqueeze(0)
+
+        # --- Overlap-add via scatter_add ---
+        batch = frames_signal.shape[0]
+        idx = self.ola_indices.unsqueeze(0).expand(batch, -1)
+
+        flat = frames_signal.reshape(batch, -1)
+        output = torch.zeros(batch, self.out_length, device=z_ri.device)
+        output.scatter_add_(1, idx, flat)
+
+        # Window normalization (pre-computed buffer).
+        output = output / self.win_sum.unsqueeze(0)
+
+        # Strip center padding.
+        pad = self.n_fft // 2
+        output = output[:, pad:pad + length]
+
+        return output
+
+
+# ---------------------------------------------------------------------------
+# RealValuedHTDemucs: wrapper that swaps STFT/ISTFT for real-valued versions
+# while keeping the actual network (encoder / transformer / decoder) intact.
+# ---------------------------------------------------------------------------
+class RealValuedHTDemucs(nn.Module):
+    """
+    Wraps HTDemucs with real-valued STFT/ISTFT.
+
+    Data flow:
+    1. RealSTFT -> (B, C, Fr, T, 2)             [real]
+    2. Spec trimming (real instead of complex)
+    3. _magnitude (cac=True): permute+reshape -> (B, C*2, Fr, T)   [real]
+    4. Encoder / CrossTransformer / Decoder                        [all real]
+    5. _mask (cac=True): reshape+permute -> (B, S, C, Fr, T, 2)    [real]
+    6. RealISTFT -> waveform                                       [real]
+    7. + time branch (denormalized)                                [real]
+    """
+
+    def __init__(self, model: nn.Module, segment_samples: int):
+        super().__init__()
+        self.segment_samples = segment_samples
+
+        # Adopt network submodules from the loaded HTDemucs.
+        self.encoder = model.encoder
+        self.tencoder = model.tencoder
+        self.decoder = model.decoder
+        self.tdecoder = model.tdecoder
+        self.crosstransformer = model.crosstransformer
+        self.freq_emb = model.freq_emb
+        self.freq_emb_scale = model.freq_emb_scale
+        self.sources = model.sources
+        self.depth = model.depth
+
+        # Bottom-channel projection (present in some HTDemucs variants).
+        self.bottom_channels = model.bottom_channels
+        if self.bottom_channels:
+            self.channel_upsampler = model.channel_upsampler
+            self.channel_downsampler = model.channel_downsampler
+            self.channel_upsampler_t = model.channel_upsampler_t
+            self.channel_downsampler_t = model.channel_downsampler_t
+
+        # STFT / ISTFT parameters.
+        self.nfft = model.nfft
+        self.hop_length = model.hop_length
+
+        # Real-valued STFT / ISTFT modules.
+        self.real_stft = RealSTFT(model.nfft, model.hop_length)
+
+        # Frame count for ISTFT (fixed for the chosen segment size).
+        le = int(math.ceil(segment_samples / model.hop_length))
+        num_frames_istft = le + 4  # after padding inside _real_ispec
+        self.real_istft = RealISTFT(model.nfft, model.hop_length, num_frames_istft)
+
+        # nn.MultiheadAttention -> ManualMHA (fused op not supported by coremltools).
+        _replace_mha_recursive(self)
+
+    def _real_spec(self, mix: torch.Tensor) -> torch.Tensor:
+        """
+        Real-valued STFT + trim.
+        Input:  (B, C, T)
+        Output: (B, C, Fr, le, 2) -- trimmed, real
+        """
+        hl = self.hop_length
+        length = mix.shape[-1]
+
+        le = int(math.ceil(length / hl))
+        pad = hl // 2 * 3
+        x = _pad1d(mix, (pad, pad + le * hl - length), mode="reflect")
+
+        z_ri = self.real_stft(x)  # (B, C, Fr, frames, 2)
+
+        # Trim: drop the last freq bin, keep frames [2 : 2+le].
+        z_ri = z_ri[:, :, :-1, :, :]
+        z_ri = z_ri[:, :, :, 2:2 + le, :]
+        return z_ri
+
+    def _real_magnitude(self, z_ri: torch.Tensor) -> torch.Tensor:
+        """
+        cac=True: real/imag channels.
+        Input:  (B, C, Fr, T, 2)
+        Output: (B, C*2, Fr, T)
+        """
+        # Move the (..., 2) dim into the channel axis:
+        # (B, C, Fr, T, 2) -> (B, C, 2, Fr, T) -> (B, C*2, Fr, T).
+        B, C, Fr, T, _ = z_ri.shape
+        m = z_ri.permute(0, 1, 4, 2, 3)
+        m = m.reshape(B, C * 2, Fr, T)
+        return m
+
+    def _real_mask(self, m: torch.Tensor) -> torch.Tensor:
+        """
+        cac=True: network output -> real/imag tensor.
+        Input:  (B, S, C*2, Fr, T) -- denormalized network output
+        Output: (B*S*C, Fr, T, 2)  -- ready for RealISTFT
+        """
+        B, S, _, Fr, T = m.shape
+        out = m.view(B, S, -1, 2, Fr, T).permute(0, 1, 2, 4, 5, 3)
+        out = out.reshape(B * S * (out.shape[2]), Fr, T, 2)
+        return out
+
+    def _real_ispec(self, z_ri: torch.Tensor, length: int) -> torch.Tensor:
+        """
+        Real-valued ISTFT.
+        Input:  (batch, Fr, T, 2)
+        Output: (batch, length)
+        """
+        hl = self.hop_length
+        # Pad freq: add 1 bin at the end.
+        z_ri = F.pad(z_ri, (0, 0, 0, 0, 0, 1))
+        # Pad frames: add 2 on each side.
+        z_ri = F.pad(z_ri, (0, 0, 2, 2))
+
+        pad = hl // 2 * 3
+        le = hl * int(math.ceil(length / hl)) + 2 * pad
+        x = self.real_istft(z_ri, le)
+        x = x[:, pad:pad + length]
+        return x
+
+    def forward(self, mix: torch.Tensor) -> torch.Tensor:
+        """
+        Input:  (1, 2, segment_samples)
+        Output: (1, 4, 2, segment_samples)  -- [vocals, drums, bass, other]
+        """
+        length = mix.shape[-1]
+
+        # --- Frequency branch: real-valued STFT ---
+        z_ri = self._real_spec(mix)         # (B, C, Fr, T, 2)
+        mag = self._real_magnitude(z_ri)    # (B, C*2, Fr, T) float
+        x = mag
+
+        B, C_mag, Fq, T = x.shape
+
+        # Normalize.
+        mean = x.mean(dim=(1, 2, 3), keepdim=True)
+        std = x.std(dim=(1, 2, 3), keepdim=True)
+        x = (x - mean) / (1e-5 + std)
+
+        # --- Time branch ---
+        xt = mix
+        meant = xt.mean(dim=(1, 2), keepdim=True)
+        stdt = xt.std(dim=(1, 2), keepdim=True)
+        xt = (xt - meant) / (1e-5 + stdt)
+
+        # --- Encoder ---
+        saved = []
+        saved_t = []
+        lengths = []
+        lengths_t = []
+
+        for idx, encode in enumerate(self.encoder):
+            lengths.append(x.shape[-1])
+            inject = None
+            if idx < len(self.tencoder):
+                lengths_t.append(xt.shape[-1])
+                tenc = self.tencoder[idx]
+                xt = tenc(xt)
+                if not tenc.empty:
+                    saved_t.append(xt)
+                else:
+                    inject = xt
+            x = encode(x, inject)
+            if idx == 0 and self.freq_emb is not None:
+                frs = torch.arange(x.shape[-2], device=x.device)
+                emb = self.freq_emb(frs).t()[None, :, :, None].expand_as(x)
+                x = x + self.freq_emb_scale * emb
+            saved.append(x)
+
+        # --- Cross-Transformer ---
+        if self.crosstransformer:
+            if self.bottom_channels:
+                b, c, f, t = x.shape
+                from einops import rearrange
+                x = rearrange(x, "b c f t-> b c (f t)")
+                x = self.channel_upsampler(x)
+                x = rearrange(x, "b c (f t)-> b c f t", f=f)
+                xt = self.channel_upsampler_t(xt)
+
+            x, xt = self.crosstransformer(x, xt)
+
+            if self.bottom_channels:
+                x = rearrange(x, "b c f t-> b c (f t)")
+                x = self.channel_downsampler(x)
+                x = rearrange(x, "b c (f t)-> b c f t", f=f)
+                xt = self.channel_downsampler_t(xt)
+
+        # --- Decoder ---
+        for idx, decode in enumerate(self.decoder):
+            skip = saved.pop(-1)
+            x, pre = decode(x, skip, lengths.pop(-1))
+
+            offset = self.depth - len(self.tdecoder)
+            if idx >= offset:
+                tdec = self.tdecoder[idx - offset]
+                length_t = lengths_t.pop(-1)
+                if tdec.empty:
+                    pre = pre[:, :, 0]
+                    xt, _ = tdec(pre, None, length_t)
+                else:
+                    skip = saved_t.pop(-1)
+                    xt, _ = tdec(xt, skip, length_t)
+
+        # --- Frequency branch: denormalize + mask ---
+        S = len(self.sources)
+        x = x.view(B, S, -1, Fq, T)
+        x = x * std[:, None] + mean[:, None]
+
+        # _real_mask -> (B*S*C, Fr, T, 2)
+        zout_ri = self._real_mask(x)
+
+        # Real-valued ISTFT.
+        x_freq = self._real_ispec(zout_ri, length)
+        # x_freq: (B*S*C, length) -> (B, S, C, length)
+        C_orig = NUM_CHANNELS
+        x_freq = x_freq.view(B, S, C_orig, length)
+
+        # --- Time branch: denormalize ---
+        xt = xt.view(B, S, -1, length)
+        xt = xt * stdt[:, None] + meant[:, None]
+
+        # --- Combine ---
+        x_out = x_freq + xt
+
+        # Reorder sources: drums,bass,other,vocals -> vocals,drums,bass,other.
+        x_out = x_out[:, SOURCE_REORDER, :, :]
+
+        return x_out
+
+
+# ---------------------------------------------------------------------------
+# Metadata
+# ---------------------------------------------------------------------------
+def _add_metadata(mlmodel, segment_samples: int) -> None:
+    mlmodel.author = "HTDemucs CoreML conversion"
+    mlmodel.license = (
+        "MIT. Original Demucs: Copyright (c) Meta Platforms, Inc. and "
+        "affiliates, MIT License. See LICENSE and ATTRIBUTION."
+    )
+    mlmodel.short_description = (
+        f"Hybrid Transformer Demucs (HTDemucs) -- music source separation "
+        f"into {', '.join(SOURCE_NAMES)} at {SAMPLE_RATE} Hz."
+    )
+    mlmodel.input_description["audio"] = (
+        f"Stereo audio. Shape (1, 2, {segment_samples}), Float32, {SAMPLE_RATE} Hz."
+    )
+    mlmodel.output_description["sources"] = (
+        f"Separated stems. Shape (1, 4, 2, {segment_samples}). "
+        f"Order: [{', '.join(SOURCE_NAMES)}]."
+    )
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(
+        description="Convert Demucs (HTDemucs) to Core ML mlpackage."
+    )
+    p.add_argument(
+        "--segment", type=float, default=10.0,
+        help="Segment length in seconds (default: 10.0).",
+    )
+    p.add_argument(
+        "--fp16", action="store_true",
+        help="Quantize to FP16 (~half the file size, minor accuracy loss).",
+    )
+    p.add_argument(
+        "--output", type=str, default=None,
+        help="Output mlpackage path (default: HTDemucs_CoreML[_FP16].mlpackage).",
+    )
+    p.add_argument(
+        "--compute-units", choices=["cpu_and_gpu", "all", "cpu_only"],
+        default="cpu_and_gpu",
+        help="Default ComputeUnit baked into the model (default: cpu_and_gpu). "
+             "HTDemucs is unstable on the Neural Engine -- keep 'cpu_and_gpu' "
+             "unless you have specifically validated 'all'.",
+    )
+    return p.parse_args()
+
+
+def main() -> None:
+    import coremltools as ct
+
+    warnings.filterwarnings("ignore", category=UserWarning)
+    warnings.filterwarnings("ignore", category=FutureWarning)
+
+    args = parse_args()
+
+    segment_samples = int(round(args.segment * SAMPLE_RATE))
+    output_path = args.output or (
+        "HTDemucs_CoreML_FP16.mlpackage" if args.fp16 else DEFAULT_OUTPUT
+    )
+    precision = ct.precision.FLOAT16 if args.fp16 else ct.precision.FLOAT32
+    compute_units = {
+        "cpu_and_gpu": ct.ComputeUnit.CPU_AND_GPU,
+        "all": ct.ComputeUnit.ALL,
+        "cpu_only": ct.ComputeUnit.CPU_ONLY,
+    }[args.compute_units]
+
+    print("=" * 60)
+    print("  HTDemucs -> Core ML Converter")
+    print("  (real-valued STFT / ISTFT wrapper)")
+    print("=" * 60)
+    print(f"  Model:        {MODEL_NAME}")
+    print(f"  Sample rate:  {SAMPLE_RATE} Hz")
+    print(f"  Segment:      {segment_samples} samples ({args.segment:.1f}s)")
+    print(f"  Stems:        {', '.join(SOURCE_NAMES)}")
+    print(f"  Precision:    {'FP16' if args.fp16 else 'FP32'}")
+    print(f"  Compute:      {args.compute_units}")
+    print(f"  Output:       {output_path}")
+    print("=" * 60)
+
+    # --- Load model ---
+    print(f"\n[1/5] Loading Demucs '{MODEL_NAME}' ...")
+    from demucs.pretrained import get_model
+    bag = get_model(MODEL_NAME)
+    model = bag.models[0]
+    model.eval()
+    model.use_train_segment = False
+    num_params = sum(p.numel() for p in model.parameters()) / 1e6
+    print(f"       {num_params:.1f}M parameters loaded.")
+
+    # --- Build wrapper ---
+    print("\n[2/5] Building real-valued wrapper ...")
+    wrapper = RealValuedHTDemucs(model, segment_samples=segment_samples)
+    wrapper.eval()
+
+    dummy = torch.randn(1, NUM_CHANNELS, segment_samples)
+
+    # --- PyTorch sanity check ---
+    print("\n[3/5] PyTorch forward pass ...")
+    with torch.no_grad():
+        out_wrapper = wrapper(dummy)
+
+    print(f"       Output shape: {out_wrapper.shape}")
+    expected = (1, NUM_SOURCES, NUM_CHANNELS, segment_samples)
+    assert out_wrapper.shape == expected, f"Shape {out_wrapper.shape} != {expected}"
+    print("       OK.")
+
+    # --- Trace ---
+    print("\n[4/5] torch.jit.trace ...")
+    with torch.no_grad():
+        traced = torch.jit.trace(wrapper, dummy, strict=False)
+    print("       Trace OK.")
+
+    # --- Core ML conversion ---
+    print("\n[5/5] Core ML conversion ...")
+    mlmodel = ct.convert(
+        traced,
+        inputs=[
+            ct.TensorType(
+                name="audio",
+                shape=(1, NUM_CHANNELS, segment_samples),
+                dtype=np.float32,
+            )
+        ],
+        outputs=[ct.TensorType(name="sources")],
+        convert_to="mlprogram",
+        compute_units=compute_units,
+        compute_precision=precision,
+        minimum_deployment_target=ct.target.macOS14,
+    )
+
+    _add_metadata(mlmodel, segment_samples)
+    mlmodel.save(output_path)
+
+    # --- Validation ---
+    # Important: reload with the SAME compute_units we converted for.
+    # MLModel(path) without a config defaults to ComputeUnit.ALL, which on
+    # HTDemucs may dispatch to ANE and crash with E5RT errors -- exactly
+    # the bug we baked the CPU_AND_GPU default into the model to avoid.
+    print("\n[Val] Validating Core ML vs. PyTorch reference ...")
+    try:
+        val_config = ct.ComputeUnit.CPU_AND_GPU
+        mlmodel_loaded = ct.models.MLModel(output_path, compute_units=val_config)
+        with torch.no_grad():
+            ref = wrapper(dummy).numpy()
+        pred = mlmodel_loaded.predict({"audio": dummy.numpy()})
+        cml_out = pred["sources"]
+
+        assert ref.shape == cml_out.shape, f"Shape mismatch: {ref.shape} vs {cml_out.shape}"
+        max_diff = float(np.max(np.abs(ref - cml_out)))
+        mean_diff = float(np.mean(np.abs(ref - cml_out)))
+        print(f"      Max diff:  {max_diff:.6f}")
+        print(f"      Mean diff: {mean_diff:.6f}")
+        threshold = 0.2 if args.fp16 else 0.1
+        if max_diff < threshold:
+            print("      Validation OK.")
+        else:
+            print("      Large numerical drift (expected for FP16 on ANE).")
+    except Exception as e:
+        print(f"      Validation skipped: {e}")
+
+    # --- Summary ---
+    size_mb = sum(
+        f.stat().st_size for f in Path(output_path).rglob("*") if f.is_file()
+    ) / (1024 * 1024)
+
+    print("\n" + "=" * 60)
+    print(f"  Done: {output_path} ({size_mb:.0f} MB)")
+    print()
+    print("  Next step: drag the .mlpackage into your Xcode project")
+    print("  and load it via MLModel(contentsOf: ...). See examples/swift/.")
+    print("=" * 60)
+
+
+if __name__ == "__main__":
+    main()

examples/swift/StemSeparator.swift

@@ -0,0 +1,258 @@
+// StemSeparator.swift
+//
+// Minimal example showing how to load HTDemucs_CoreML.mlpackage in Swift,
+// chunk a stereo audio file, run inference per chunk, and reassemble the
+// four stems via overlap-add.
+//
+// Drop the .mlpackage into your Xcode target's Resources, then call:
+//
+//     let separator = try StemSeparator()
+//     let stems = try await separator.separate(fileURL: someURL) { progress in
+//         print("\(Int(progress * 100))%")
+//     }
+//     // stems[.vocals], stems[.drums], stems[.bass], stems[.other]
+//
+// This is a pared-down reference. Real apps usually want:
+//   - Resampling to 44.1 kHz before chunking (AVAudioConverter).
+//   - Triangular overlap-add windowing (this file uses linear taper for brevity).
+//   - Cancellation / progress on a Task.
+//
+// License: MIT (same as the rest of this repo).
+
+import AVFoundation
+import CoreML
+
+public enum StemKind: Int, CaseIterable {
+    case vocals = 0
+    case drums  = 1
+    case bass   = 2
+    case other  = 3
+}
+
+public enum StemSeparatorError: LocalizedError {
+    case modelNotFound
+    case unsupportedFormat
+    case inferenceFailed(String)
+
+    public var errorDescription: String? {
+        switch self {
+        case .modelNotFound:        return "HTDemucs_CoreML.mlpackage not found in bundle."
+        case .unsupportedFormat:    return "Audio must be 44.1 kHz stereo Float32."
+        case .inferenceFailed(let m): return "Inference failed: \(m)"
+        }
+    }
+}
+
+public final class StemSeparator {
+
+    // The converter bakes a fixed segment length into the model. Defaults
+    // match `python convert.py` (10 s @ 44.1 kHz). If you converted with
+    // --segment 7, change segmentSamples to 308700, etc.
+    private let segmentSamples = 441_000
+    private let overlapSamples = 44_100   // 1 s overlap-add
+    private let sampleRate     = 44_100.0
+
+    private let model: MLModel
+    private let inferenceLock = NSLock()  // MLModel.prediction is not thread-safe.
+
+    public init() throws {
+        guard let url = Bundle.main.url(
+            forResource: "HTDemucs_CoreML", withExtension: "mlpackage"
+        ) ?? Bundle.main.url(
+            forResource: "HTDemucs_CoreML", withExtension: "mlmodelc"
+        ) else {
+            throw StemSeparatorError.modelNotFound
+        }
+
+        let config = MLModelConfiguration()
+        config.computeUnits = .cpuAndGPU   // do NOT use .all -- ANE is unstable here
+        self.model = try MLModel(contentsOf: url, configuration: config)
+    }
+
+    /// Separate a file into four stems, returning each as an
+    /// AVAudioPCMBuffer at 44.1 kHz / stereo / Float32.
+    public func separate(
+        fileURL: URL,
+        progress: @Sendable @escaping (Double) -> Void = { _ in }
+    ) async throws -> [StemKind: AVAudioPCMBuffer] {
+        let mix = try loadAndResample(fileURL: fileURL)
+        return try await Task.detached(priority: .userInitiated) {
+            try self.runChunked(mix: mix, progress: progress)
+        }.value
+    }
+
+    // MARK: - Loading & resampling
+
+    private func loadAndResample(fileURL: URL) throws -> AVAudioPCMBuffer {
+        let file = try AVAudioFile(forReading: fileURL)
+        let target = AVAudioFormat(
+            commonFormat: .pcmFormatFloat32,
+            sampleRate: sampleRate,
+            channels: 2,
+            interleaved: false
+        )!
+        guard let converter = AVAudioConverter(from: file.processingFormat, to: target) else {
+            throw StemSeparatorError.unsupportedFormat
+        }
+
+        let outFrames = AVAudioFrameCount(
+            Double(file.length) * sampleRate / file.processingFormat.sampleRate
+        )
+        guard let out = AVAudioPCMBuffer(pcmFormat: target, frameCapacity: outFrames) else {
+            throw StemSeparatorError.unsupportedFormat
+        }
+
+        let input = AVAudioPCMBuffer(pcmFormat: file.processingFormat,
+                                     frameCapacity: AVAudioFrameCount(file.length))!
+        try file.read(into: input)
+
+        var consumed = false
+        var error: NSError?
+        converter.convert(to: out, error: &error) { _, status in
+            if consumed { status.pointee = .endOfStream; return nil }
+            consumed = true
+            status.pointee = .haveData
+            return input
+        }
+        if let error { throw StemSeparatorError.inferenceFailed(error.localizedDescription) }
+        return out
+    }
+
+    // MARK: - Chunked inference + overlap-add
+
+    private func runChunked(
+        mix: AVAudioPCMBuffer,
+        progress: @Sendable @escaping (Double) -> Void
+    ) throws -> [StemKind: AVAudioPCMBuffer] {
+        let total = Int(mix.frameLength)
+        let stride = segmentSamples - overlapSamples
+        let chunks = max(1, Int(ceil(Double(total) / Double(stride))))
+
+        // Output buffers (one per stem).
+        var outputs: [StemKind: UnsafeMutablePointer<Float>] = [:]
+        var weights = [Float](repeating: 0, count: total)
+        defer { outputs.values.forEach { $0.deallocate() } }
+
+        for kind in StemKind.allCases {
+            outputs[kind] = UnsafeMutablePointer<Float>.allocate(capacity: total * 2)
+            outputs[kind]!.initialize(repeating: 0, count: total * 2)
+        }
+
+        // Triangular fade window for overlap-add.
+        var window = [Float](repeating: 1, count: segmentSamples)
+        for i in 0..<overlapSamples {
+            let w = Float(i) / Float(overlapSamples)
+            window[i] = w
+            window[segmentSamples - 1 - i] = w
+        }
+
+        for c in 0..<chunks {
+            let start = c * stride
+            let chunk = sliceChunk(mix: mix, start: start)
+            let stemChunks = try predict(chunk: chunk)
+
+            for kind in StemKind.allCases {
+                guard let dst = outputs[kind] else { continue }
+                let src = stemChunks[kind]!
+                for f in 0..<segmentSamples {
+                    let global = start + f
+                    if global >= total { break }
+                    let w = window[f]
+                    dst[global * 2 + 0] += src[f * 2 + 0] * w
+                    dst[global * 2 + 1] += src[f * 2 + 1] * w
+                    if kind == .vocals {  // accumulate window only once
+                        weights[global] += w
+                    }
+                }
+            }
+            progress(Double(c + 1) / Double(chunks))
+        }
+
+        // Normalize by accumulated window weights.
+        for kind in StemKind.allCases {
+            guard let dst = outputs[kind] else { continue }
+            for f in 0..<total {
+                let w = max(weights[f], 1e-6)
+                dst[f * 2 + 0] /= w
+                dst[f * 2 + 1] /= w
+            }
+        }
+
+        // Wrap into AVAudioPCMBuffers.
+        let outFormat = AVAudioFormat(
+            commonFormat: .pcmFormatFloat32,
+            sampleRate: sampleRate,
+            channels: 2,
+            interleaved: false
+        )!
+        var result: [StemKind: AVAudioPCMBuffer] = [:]
+        for kind in StemKind.allCases {
+            let buf = AVAudioPCMBuffer(pcmFormat: outFormat,
+                                       frameCapacity: AVAudioFrameCount(total))!
+            buf.frameLength = AVAudioFrameCount(total)
+            let src = outputs[kind]!
+            for f in 0..<total {
+                buf.floatChannelData![0][f] = src[f * 2 + 0]
+                buf.floatChannelData![1][f] = src[f * 2 + 1]
+            }
+            result[kind] = buf
+        }
+        return result
+    }
+
+    // MARK: - Single-chunk prediction
+
+    private func predict(chunk: [Float]) throws -> [StemKind: [Float]] {
+        precondition(chunk.count == segmentSamples * 2,
+                     "chunk must be \(segmentSamples * 2) interleaved floats")
+
+        let array = try MLMultiArray(
+            shape: [1, 2, NSNumber(value: segmentSamples)],
+            dataType: .float32
+        )
+        // Deinterleave into (1, 2, N).
+        let ptr = array.dataPointer.bindMemory(to: Float.self, capacity: array.count)
+        for f in 0..<segmentSamples {
+            ptr[0 * segmentSamples + f] = chunk[f * 2 + 0]
+            ptr[1 * segmentSamples + f] = chunk[f * 2 + 1]
+        }
+
+        inferenceLock.lock()
+        defer { inferenceLock.unlock() }
+
+        let provider = try MLDictionaryFeatureProvider(dictionary: ["audio": array])
+        let result = try model.prediction(from: provider)
+        guard let out = result.featureValue(for: "sources")?.multiArrayValue else {
+            throw StemSeparatorError.inferenceFailed("missing 'sources' output")
+        }
+
+        // Output shape: (1, 4, 2, segmentSamples). Order: [vocals, drums, bass, other].
+        var stems: [StemKind: [Float]] = [:]
+        let outPtr = out.dataPointer.bindMemory(to: Float.self, capacity: out.count)
+        for kind in StemKind.allCases {
+            var samples = [Float](repeating: 0, count: segmentSamples * 2)
+            let stemBase = kind.rawValue * 2 * segmentSamples
+            for f in 0..<segmentSamples {
+                samples[f * 2 + 0] = outPtr[stemBase + 0 * segmentSamples + f]
+                samples[f * 2 + 1] = outPtr[stemBase + 1 * segmentSamples + f]
+            }
+            stems[kind] = samples
+        }
+        return stems
+    }
+
+    private func sliceChunk(mix: AVAudioPCMBuffer, start: Int) -> [Float] {
+        let total = Int(mix.frameLength)
+        var out = [Float](repeating: 0, count: segmentSamples * 2)
+        let l = mix.floatChannelData![0]
+        let r = mix.floatChannelData![1]
+        for f in 0..<segmentSamples {
+            let g = start + f
+            if g < total {
+                out[f * 2 + 0] = l[g]
+                out[f * 2 + 1] = r[g]
+            }
+        }
+        return out
+    }
+}

requirements.txt

@@ -0,0 +1,16 @@
+# Tested combinations (any of these should work):
+#
+#   Verified set A (used to produce the released mlpackage):
+#     torch 2.8.0, torchaudio 2.8.0, coremltools 9.0,
+#     demucs 4.0.1, numpy 2.0.2, einops 0.8.2, Python 3.9
+#
+# Lower bounds reflect the API surface the converter relies on; upper
+# bounds are kept open since coremltools 9.x and torch 2.8.x are known
+# to convert this model cleanly.
+
+torch>=2.1
+torchaudio>=2.1
+demucs>=4.0,<5.0
+coremltools>=7.2
+numpy>=1.24
+einops>=0.7