Add Core ML export script

scripts/export_htdemucs_coreml_core.py

@@ -0,0 +1,258 @@
+#!/usr/bin/env python3
+"""Export the real-valued HTDemucs separator core to Core ML.
+
+The full HTDemucs waveform graph cannot be exported directly because it uses
+complex STFT/ISTFT tensors. This script cuts the graph at the same boundary the
+model already uses internally when `cac=True`: Swift performs STFT/ISTFT with
+Accelerate, while Core ML runs the learned real-valued separator core.
+"""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+from pathlib import Path
+from typing import Any
+
+import coremltools as ct
+import torch
+from demucs.pretrained import get_model
+
+
+MODEL_ID = "htdemucs"
+DEMUX_VERSION = "4.0.1"
+CHECKPOINT_URL = "https://dl.fbaipublicfiles.com/demucs/hybrid_transformer/955717e8-8726e21a.th"
+
+
+class HTDemucsSeparatorCore(torch.nn.Module):
+    """Real-valued HTDemucs separator core.
+
+    Inputs:
+      mix: [B, 2, segment_samples], 44.1 kHz stereo waveform.
+      spectrogram: [B, 4, 2048, segment_frames], real/imag-as-channel STFT.
+
+    Outputs:
+      spectrogram_stems: [B, 4, 4, 2048, segment_frames]
+      waveform_stems: [B, 4, 2, segment_samples]
+
+    The source order is the upstream HTDemucs order: drums, bass, other, vocals.
+    """
+
+    def __init__(self, model: torch.nn.Module) -> None:
+        super().__init__()
+        self.model = model
+
+    def forward(self, mix: torch.Tensor, spectrogram: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        model = self.model
+        x = spectrogram
+        xt = mix
+
+        batch, _, freqs, frames = x.shape
+        length = mix.shape[-1]
+
+        mean = x.mean(dim=(1, 2, 3), keepdim=True)
+        std = x.std(dim=(1, 2, 3), keepdim=True)
+        x = (x - mean) / (1e-5 + std)
+
+        mean_t = xt.mean(dim=(1, 2), keepdim=True)
+        std_t = xt.std(dim=(1, 2), keepdim=True)
+        xt = (xt - mean_t) / (1e-5 + std_t)
+
+        saved = []
+        saved_t = []
+        lengths = []
+        lengths_t = []
+
+        for idx, encode in enumerate(model.encoder):
+            lengths.append(x.shape[-1])
+            inject = None
+            if idx < len(model.tencoder):
+                lengths_t.append(xt.shape[-1])
+                tenc = model.tencoder[idx]
+                xt = tenc(xt)
+                if not tenc.empty:
+                    saved_t.append(xt)
+                else:
+                    inject = xt
+
+            x = encode(x, inject)
+            if idx == 0 and model.freq_emb is not None:
+                freq_positions = torch.arange(x.shape[-2], device=x.device)
+                emb = model.freq_emb(freq_positions).t()[None, :, :, None].expand_as(x)
+                x = x + model.freq_emb_scale * emb
+
+            saved.append(x)
+
+        if model.crosstransformer:
+            if model.bottom_channels:
+                b, c, f, t = x.shape
+                x = x.reshape(b, c, f * t)
+                x = model.channel_upsampler(x)
+                x = x.reshape(b, -1, f, t)
+                xt = model.channel_upsampler_t(xt)
+
+            x, xt = model.crosstransformer(x, xt)
+
+            if model.bottom_channels:
+                b, c, f, t = x.shape
+                x = x.reshape(b, c, f * t)
+                x = model.channel_downsampler(x)
+                x = x.reshape(b, -1, f, t)
+                xt = model.channel_downsampler_t(xt)
+
+        for idx, decode in enumerate(model.decoder):
+            skip = saved.pop(-1)
+            x, pre = decode(x, skip, lengths.pop(-1))
+
+            offset = model.depth - len(model.tdecoder)
+            if idx >= offset:
+                tdec = model.tdecoder[idx - offset]
+                length_t = lengths_t.pop(-1)
+                if tdec.empty:
+                    pre = pre[:, :, 0]
+                    xt, _ = tdec(pre, None, length_t)
+                else:
+                    skip_t = saved_t.pop(-1)
+                    xt, _ = tdec(xt, skip_t, length_t)
+
+        sources = len(model.sources)
+        x = x.view(batch, sources, -1, freqs, frames)
+        x = x * std[:, None] + mean[:, None]
+
+        xt = xt.view(batch, sources, -1, length)
+        xt = xt * std_t[:, None] + mean_t[:, None]
+        return x, xt
+
+
+def sha256_file(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def package_files(root: Path) -> list[dict[str, Any]]:
+    files: list[dict[str, Any]] = []
+    for path in sorted(root.rglob("*")):
+        if path.is_file():
+            files.append(
+                {
+                    "path": path.relative_to(root.parent).as_posix(),
+                    "bytes": path.stat().st_size,
+                    "sha256": sha256_file(path),
+                }
+            )
+    return files
+
+
+def write_manifest(path: Path, package_path: Path, segment_samples: int, segment_frames: int) -> None:
+    payload = {
+        "format_version": 1,
+        "model_id": MODEL_ID,
+        "demucs_version": DEMUX_VERSION,
+        "checkpoint_url": CHECKPOINT_URL,
+        "coreml_package": package_path.name,
+        "sample_rate": 44100,
+        "channels": 2,
+        "n_fft": 4096,
+        "hop_length": 1024,
+        "segment_samples": segment_samples,
+        "segment_seconds": segment_samples / 44100.0,
+        "segment_frames": segment_frames,
+        "frequency_bins": 2048,
+        "sources": ["drums", "bass", "other", "vocals"],
+        "vocals_source_index": 3,
+        "inputs": {
+            "mix": [1, 2, segment_samples],
+            "spectrogram": [1, 4, 2048, segment_frames],
+        },
+        "outputs": {
+            "spectrogram_stems": [1, 4, 4, 2048, segment_frames],
+            "waveform_stems": [1, 4, 2, segment_samples],
+        },
+        "preprocessing": [
+            "Decode to mono or stereo PCM.",
+            "Resample to 44100 Hz.",
+            "Duplicate mono to stereo.",
+            "Pad each chunk to segment_samples.",
+            "Reflect-pad by 1536 samples before STFT, then run normalized Hann STFT with n_fft=4096 and hop_length=1024.",
+            "Drop the final frequency bin and keep centered segment_frames frames.",
+        ],
+        "postprocessing": [
+            "Take vocals_source_index from both Core ML outputs.",
+            "Convert spectrogram real/imag channels back to complex STFT.",
+            "Append the dropped zero frequency bin, pad two frames on both sides, and run normalized Hann ISTFT.",
+            "Remove the 1536-sample STFT pad, add waveform_stems[vocals_source_index], trim original chunk length, and overlap-add chunks.",
+            "Downmix the resulting stereo vocals to mono for voice cloning.",
+        ],
+        "files": package_files(package_path),
+    }
+    path.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--out-dir", type=Path, default=Path("demucs_experiment/exported/htdemucs_coreml_core"))
+    parser.add_argument("--package-name", default="htdemucs_separator_core.mlpackage")
+    parser.add_argument("--validate", action="store_true")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    args.out_dir.mkdir(parents=True, exist_ok=True)
+
+    bag = get_model(MODEL_ID)
+    model = bag.models[0].eval()
+    segment_samples = model.valid_length(44100)
+    hop_length = model.hop_length
+    segment_frames = int((segment_samples + hop_length - 1) // hop_length)
+
+    core = HTDemucsSeparatorCore(model).eval()
+    mix = torch.zeros(1, 2, segment_samples)
+    spectrogram = torch.zeros(1, 4, 2048, segment_frames)
+
+    if args.validate:
+        with torch.no_grad():
+            probe = torch.randn_like(mix) * 0.01
+            z = model._spec(probe)
+            mag = model._magnitude(z)
+            spec_out, wave_out = core(probe, mag)
+            reconstructed = model._ispec(model._mask(z, spec_out), segment_samples) + wave_out
+            reference = model(probe)
+            max_error = (reconstructed - reference).abs().max().item()
+            print(f"wrapper_max_abs_error={max_error:.8f}")
+
+    traced = torch.jit.trace(core, (mix, spectrogram), strict=False, check_trace=False)
+    print(f"traced segment_samples={segment_samples} segment_frames={segment_frames}")
+
+    mlmodel = ct.convert(
+        traced,
+        inputs=[
+            ct.TensorType(name="mix", shape=mix.shape),
+            ct.TensorType(name="spectrogram", shape=spectrogram.shape),
+        ],
+        outputs=[
+            ct.TensorType(name="spectrogram_stems"),
+            ct.TensorType(name="waveform_stems"),
+        ],
+        convert_to="mlprogram",
+        minimum_deployment_target=ct.target.iOS26,
+        compute_precision=ct.precision.FLOAT16,
+    )
+
+    package_path = args.out_dir / args.package_name
+    if package_path.exists():
+        import shutil
+
+        shutil.rmtree(package_path)
+    mlmodel.save(str(package_path))
+    write_manifest(args.out_dir / "manifest.json", package_path, segment_samples, segment_frames)
+    print(f"saved {package_path}")
+    print(f"saved {args.out_dir / 'manifest.json'}")
+
+
+if __name__ == "__main__":
+    main()