feat(tts/zipvoice): LuxTTS/ZipVoice CoreML conversion + ANE-canonical decoder rewrite

models/tts/zipvoice/.gitignore

@@ -0,0 +1,5 @@
+.venv/
+build/
+upstream/
+uv.lock
+__pycache__/

models/tts/zipvoice/README.md

@@ -0,0 +1,55 @@
+# ZipVoice / LuxTTS → CoreML
+
+CoreML conversion for [LuxTTS](https://huggingface.co/YatharthS/LuxTTS) (48 kHz
+zero-shot voice cloning), a distilled [ZipVoice](https://github.com/k2-fsa/ZipVoice)
+(k2-fsa) with a 4-step flow-matching solver and a dual-head 48 kHz Vocos vocoder.
+The same pipeline applies to upstream ZipVoice checkpoints (en+zh, dialog) —
+identical graphs, different weights/vocoder config.
+
+Requested in FluidAudio issue #49 ([comment](https://github.com/FluidInference/FluidAudio/issues/49#issuecomment-4762131530)).
+
+## Status
+
+| Component | CoreML | Parity (vs torch) | Notes |
+|---|---|---|---|
+| TextEncoder (Zipformer, 4M) | ✅ fp16, 8.3 MB | cos 0.999997 | fixed 256-token bucket + mask |
+| FmDecoder (Zipformer, 119M) | ✅ fp16, 238 MB | cos ≥0.9987/step | fixed 1024-frame bucket + mask |
+| Duration expansion + solver | host (Python/Swift) | exact | anchor-Euler, 4 steps, t_shift 0.5 |
+| Vocoder (dual-head Vocos 48k) | ⬜ torch for now | — | ISTFT×2 + Linkwitz-Riley crossover; needs DFT-matmul ISTFT |
+| Prompt transcription | out of scope | — | upstream uses Whisper; FluidAudio would use Parakeet |
+| End-to-end audio | — | log-mel cos 0.99925, RMS match | phase-only waveform divergence |
+
+## Performance (M5 Pro, macOS 26.5)
+
+| Model | CPU | GPU | ANE (98–99% resident) |
+|---|---|---|---|
+| FmDecoder (per step) | 155 ms | **14.1 ms** | 286 ms |
+| TextEncoder | 5.5 ms | 1.4 ms | 4.8 ms |
+
+`all` compute units picks GPU. ~8 s utterance ≈ 5 + 4×14 ≈ **~60 ms** for the
+core (≈130× realtime) before vocoder. ANE residency is high but latency is
+poor (seq-first layouts / rel-pos attention transposes) — see trials.md.
+
+## Layout
+
+- `coreml/convert_coreml.py` — export TextEncoder + FmDecoder (component split
+  mirrors upstream `zipvoice/bin/onnx_export.py`; duration expansion stays host-side)
+- `coreml/parity.py` — per-component + end-to-end parity vs torch oracle
+- `scripts/reference_infer.py` — pure-torch reference (parity oracle)
+- `upstream/` — clones of ysharma3501/LuxTTS and k2-fsa/ZipVoice (gitignored)
+
+## Usage
+
+```bash
+uv venv --python 3.11 .venv
+# install deps (see pyproject; universal uv sync fights the cu126 pins upstream)
+.venv/bin/python scripts/reference_infer.py --prompt-audio <ref.wav>
+.venv/bin/python coreml/convert_coreml.py
+.venv/bin/python -m coreml.parity
+```
+
+## Licenses
+
+LuxTTS & ZipVoice: Apache-2.0. Frontend uses piper-phonemize (espeak-ng
+data — GPL-3.0); a FluidAudio integration should reuse the in-house G2P
+frontend instead.

models/tts/zipvoice/coreml/ane/__init__.py

@@ -0,0 +1,5 @@
+"""ANE-canonical (B, C, 1, S) rewrite of Zipformer2EncoderLayer.
+
+See layer.py for the module, parity.py for fp32 gates, bench.py for CoreML
+conversion + ANE/GPU benchmarks.
+"""

models/tts/zipvoice/coreml/ane/bench.py

@@ -0,0 +1,181 @@
+"""CoreML benchmark: original Zipformer2EncoderLayer vs AneZipformerLayer.
+
+Converts both to CoreML (iOS17, fp16, mlprogram, fixed S=1024), times each
+under CPU_AND_NE and CPU_AND_GPU (3 warmup + 10 timed), checks the ANE
+layer's fp16 output against the torch fp32 reference, and runs
+coreml-cli --fallback on the compiled ANE package.
+
+Run: .venv/bin/python -m coreml.ane.bench
+"""
+
+import shutil
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+sys.path.insert(0, ".")
+
+import coremltools as ct
+import numpy as np
+import torch
+from torch import Tensor, nn
+
+from coreml.ane.layer import AneZipformerLayer, ane_to_tbc, tbc_to_ane
+from coreml.convert_coreml import (
+    load_model,
+    patch_coremltools_int,
+    patch_simple_downsample,
+)
+
+SEQ_LEN = 1024
+EMBED_DIM = 512
+OUT_DIR = Path("build/coreml/ane_trial")
+COREML_CLI = Path("/Users/hanweng/Documents/mobius-zipvoice/tools/coreml-cli")
+
+
+class OrigLayerWrapper(nn.Module):
+    """(S, 1, C) seq-first original layer with frozen pos_emb constant."""
+
+    def __init__(self, layer, pos_emb: Tensor):
+        super().__init__()
+        self.layer = layer
+        self.register_buffer("pos_emb", pos_emb)
+
+    def forward(self, src: Tensor, time_emb: Tensor) -> Tensor:
+        return self.layer(src, self.pos_emb, time_emb=time_emb)
+
+
+def convert(traced, inputs, name):
+    mlmodel = ct.convert(
+        traced,
+        inputs=inputs,
+        outputs=[ct.TensorType(name="out", dtype=np.float32)],
+        minimum_deployment_target=ct.target.iOS17,
+        compute_precision=ct.precision.FLOAT16,
+        convert_to="mlprogram",
+    )
+    path = OUT_DIR / f"{name}.mlpackage"
+    if path.exists():
+        shutil.rmtree(path)
+    mlmodel.save(str(path))
+    print(f"saved {path}")
+    return path
+
+
+def bench(path, feeds, units, warmup=3, runs=10):
+    model = ct.models.MLModel(str(path), compute_units=units)
+    for _ in range(warmup):
+        out = model.predict(feeds)
+    times = []
+    for _ in range(runs):
+        t0 = time.perf_counter()
+        out = model.predict(feeds)
+        times.append((time.perf_counter() - t0) * 1e3)
+    return float(np.mean(times)), out["out"]
+
+
+def cosine(a, b):
+    a = np.asarray(a, dtype=np.float64).ravel()
+    b = np.asarray(b, dtype=np.float64).ravel()
+    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))
+
+
+def fallback_check(mlpackage: Path):
+    compiled = ct.models.utils.compile_model(str(mlpackage))
+    dest = mlpackage.with_suffix(".mlmodelc")
+    if dest.exists():
+        shutil.rmtree(dest)
+    shutil.copytree(compiled, dest)
+    print(f"\n--- coreml-cli --fallback: {dest.name} ---")
+    try:
+        res = subprocess.run(
+            ["uv", "run", "coreml-cli", str(dest.resolve()), "--fallback"],
+            cwd=COREML_CLI,
+            capture_output=True,
+            text=True,
+            timeout=600,
+        )
+        print(res.stdout)
+        if res.returncode != 0:
+            print(res.stderr[-2000:])
+    except Exception as e:  # noqa: BLE001 — report and continue
+        print(f"coreml-cli failed: {e}")
+
+
+def main():
+    torch.manual_seed(0)
+    OUT_DIR.mkdir(parents=True, exist_ok=True)
+    patch_coremltools_int()
+    patch_simple_downsample()
+    model, _ = load_model()
+    model.eval()
+
+    enc0 = model.fm_decoder.encoders[0]
+    layer = enc0.layers[0] if hasattr(enc0, "layers") else enc0.encoder.layers[0]
+    with torch.no_grad():
+        pos_emb = enc0.encoder_pos(torch.randn(SEQ_LEN, 1, EMBED_DIM)).detach()
+
+    src = torch.randn(SEQ_LEN, 1, EMBED_DIM)
+    time_emb = torch.randn(1, EMBED_DIM)
+
+    # torch fp32 reference
+    with torch.no_grad():
+        ref = layer(src, pos_emb, time_emb=time_emb).numpy()
+
+    # --- original layer, seq-first ---
+    orig = OrigLayerWrapper(layer, pos_emb).eval()
+    with torch.no_grad():
+        traced_orig = torch.jit.trace(orig, (src, time_emb))
+    orig_path = convert(
+        traced_orig,
+        [
+            ct.TensorType(name="src", shape=(SEQ_LEN, 1, EMBED_DIM), dtype=np.float32),
+            ct.TensorType(name="time_emb", shape=(1, EMBED_DIM), dtype=np.float32),
+        ],
+        "OrigLayer",
+    )
+
+    # --- ANE-canonical layer ---
+    ane = AneZipformerLayer(layer, pos_emb, SEQ_LEN).eval()
+    x_ane = tbc_to_ane(src)
+    t_ane = time_emb.reshape(1, EMBED_DIM, 1, 1)
+    with torch.no_grad():
+        traced_ane = torch.jit.trace(ane, (x_ane, t_ane))
+    ane_path = convert(
+        traced_ane,
+        [
+            ct.TensorType(name="x", shape=(1, EMBED_DIM, 1, SEQ_LEN), dtype=np.float32),
+            ct.TensorType(name="time_emb", shape=(1, EMBED_DIM, 1, 1), dtype=np.float32),
+        ],
+        "AneLayer",
+    )
+
+    orig_feeds = {"src": src.numpy(), "time_emb": time_emb.numpy()}
+    ane_feeds = {"x": x_ane.numpy(), "time_emb": t_ane.numpy()}
+
+    print(f"\n{'model':<12} {'units':<12} {'mean ms':>10}")
+    results = {}
+    for name, path, feeds in (
+        ("orig", orig_path, orig_feeds),
+        ("ane", ane_path, ane_feeds),
+    ):
+        for units in (ct.ComputeUnit.CPU_AND_NE, ct.ComputeUnit.CPU_AND_GPU):
+            ms, out = bench(path, feeds, units)
+            results[(name, units.name)] = (ms, out)
+            print(f"{name:<12} {units.name:<12} {ms:>10.2f}")
+
+    # fp16 CoreML (ANE) vs torch fp32 accuracy
+    ane_ne_out = results[("ane", "CPU_AND_NE")][1]  # (1, C, 1, S)
+    ane_out_tbc = ane_to_tbc(torch.from_numpy(ane_ne_out)).numpy()
+    orig_ne_out = results[("orig", "CPU_AND_NE")][1]
+    print(f"\nane CoreML fp16 (NE) vs torch fp32: cos={cosine(ane_out_tbc, ref):.6f}, "
+          f"max_abs={np.abs(ane_out_tbc - ref).max():.4f}")
+    print(f"orig CoreML fp16 (NE) vs torch fp32: cos={cosine(orig_ne_out, ref):.6f}")
+
+    fallback_check(ane_path)
+    fallback_check(orig_path)
+
+
+if __name__ == "__main__":
+    main()

models/tts/zipvoice/coreml/ane/convert_decoder.py

@@ -0,0 +1,121 @@
+"""Convert AneFmDecoder to CoreML (iOS17 fp16, fixed S=1024).
+
+Same I/O contract as the original FmDecoder (t, x, text_condition,
+speech_condition, guidance_scale, padding_mask -> v), so coreml/parity.py
+feeds and swift/RssBench.swift work unchanged. Saves
+build/coreml-ane/AneFmDecoder.mlpackage, copies TextEncoder.mlpackage from
+build/coreml, and compiles both to .mlmodelc (the decoder as
+FmDecoder.mlmodelc so rss_bench finds it).
+
+Run: .venv/bin/python -m coreml.ane.convert_decoder
+"""
+
+import argparse
+import shutil
+import sys
+from pathlib import Path
+
+sys.path.insert(0, ".")
+
+import coremltools as ct
+import numpy as np
+import torch
+
+from coreml.ane.decoder import AneFmDecoder, AneFmDecoderIO
+from coreml.convert_coreml import FEAT_DIM, load_model, patch_coremltools_int
+
+SEQ_LEN = 1024
+SRC_DIR = Path("build/coreml")
+
+
+def compile_to(mlpackage: Path, dest: Path):
+    compiled = ct.models.utils.compile_model(str(mlpackage))
+    if dest.exists():
+        shutil.rmtree(dest)
+    shutil.copytree(compiled, dest)
+    print(f"compiled {dest}")
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--out-dir", default="build/coreml-ane")
+    parser.add_argument(
+        "--deployment-target", default="iOS17", choices=["iOS17", "iOS18"],
+        help="iOS18 enables grouped-channel palettization downstream",
+    )
+    args = parser.parse_args()
+    out_dir = Path(args.out_dir)
+
+    out_dir.mkdir(parents=True, exist_ok=True)
+    patch_coremltools_int()
+    model, _ = load_model()
+
+    ane = AneFmDecoderIO(AneFmDecoder(model.fm_decoder, SEQ_LEN)).eval()
+    for sl, err in ane.core.basis_errs.items():
+        print(f"pos basis S={sl}: rel_max_err={err:.2e}")
+
+    t = torch.tensor([0.5])
+    x = torch.randn(1, SEQ_LEN, FEAT_DIM)
+    text = torch.randn(1, SEQ_LEN, FEAT_DIM)
+    speech = torch.randn(1, SEQ_LEN, FEAT_DIM)
+    g = torch.tensor([3.0])
+    mask = torch.zeros(1, SEQ_LEN)
+    mask[0, SEQ_LEN - 124 :] = 1.0
+
+    with torch.no_grad():
+        ref = ane(t, x, text, speech, g, mask)
+        traced = torch.jit.trace(ane, (t, x, text, speech, g, mask))
+
+    mlmodel = ct.convert(
+        traced,
+        inputs=[
+            ct.TensorType(name="t", shape=(1,), dtype=np.float32),
+            ct.TensorType(name="x", shape=(1, SEQ_LEN, FEAT_DIM), dtype=np.float32),
+            ct.TensorType(
+                name="text_condition", shape=(1, SEQ_LEN, FEAT_DIM), dtype=np.float32
+            ),
+            ct.TensorType(
+                name="speech_condition", shape=(1, SEQ_LEN, FEAT_DIM), dtype=np.float32
+            ),
+            ct.TensorType(name="guidance_scale", shape=(1,), dtype=np.float32),
+            ct.TensorType(name="padding_mask", shape=(1, SEQ_LEN), dtype=np.float32),
+        ],
+        outputs=[ct.TensorType(name="v", dtype=np.float32)],
+        minimum_deployment_target=getattr(ct.target, args.deployment_target),
+        compute_precision=ct.precision.FLOAT16,
+        convert_to="mlprogram",
+    )
+    pkg = out_dir / "AneFmDecoder.mlpackage"
+    if pkg.exists():
+        shutil.rmtree(pkg)
+    mlmodel.save(str(pkg))
+    print(f"saved {pkg}")
+
+    # Sanity: fp16 CoreML (CPU) vs torch fp32 on the trace inputs.
+    cm = ct.models.MLModel(str(pkg), compute_units=ct.ComputeUnit.CPU_ONLY)
+    out = cm.predict(
+        {
+            "t": t.numpy(),
+            "x": x.numpy(),
+            "text_condition": text.numpy(),
+            "speech_condition": speech.numpy(),
+            "guidance_scale": g.numpy(),
+            "padding_mask": mask.numpy(),
+        }
+    )["v"]
+    n = SEQ_LEN - 124
+    a, b = ref.numpy()[:, :n].ravel(), out[:, :n].ravel()
+    c = float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))
+    print(f"CoreML fp16 (CPU) vs torch fp32: cos={c:.6f} max_abs={np.abs(a - b).max():.4f}")
+
+    # TextEncoder + compiled models.
+    te_dst = out_dir / "TextEncoder.mlpackage"
+    if not te_dst.exists():
+        shutil.copytree(SRC_DIR / "TextEncoder.mlpackage", te_dst)
+        print(f"copied {te_dst}")
+    compile_to(pkg, out_dir / "FmDecoder.mlmodelc")
+    compile_to(te_dst, out_dir / "TextEncoder.mlmodelc")
+
+
+if __name__ == "__main__":
+    main()