trial(tts/moss-tts-nano): MOSS-TTS-Nano-100M + Audio-Tokenizer-Nano to CoreML

models/tts/moss-tts-nano/coreml/.python-version

@@ -0,0 +1 @@
+3.12

models/tts/moss-tts-nano/coreml/README.md

@@ -0,0 +1,96 @@
+# MOSS-TTS-Nano → CoreML
+
+Conversion of [OpenMOSS-Team/MOSS-TTS-Nano-100M](https://huggingface.co/OpenMOSS-Team/MOSS-TTS-Nano-100M)
+(0.1B multilingual streaming TTS with zero-shot voice cloning, 20 languages, native 48 kHz stereo) and
+its codec [MOSS-Audio-Tokenizer-Nano](https://huggingface.co/OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano)
+(22M, causal-transformer codec, 16×1024 RLFQ at 12.5 Hz) to CoreML. Apache-2.0 upstream.
+
+Published: [FluidInference/moss-tts-nano-coreml](https://huggingface.co/FluidInference/moss-tts-nano-coreml)
+(mlpackage + mlmodelc, `config.json` with template ids, tokenizer, preset voices). Swift backend:
+FluidAudio `Sources/FluidAudio/TTS/MossTtsNano` (`--backend moss-tts-nano`).
+
+## Pipeline
+
+```
+prompt wav ─CodecEncoder─► 16×T ref codes ─┐
+text ─SentencePiece─► ids ──────────────────┴─► rows [T,17] ─Prefill─► hidden + KV
+                                                                  │
+                     ┌────────────────────────────────────────────┘
+                     ▼   per 80 ms frame
+                 Frame (1-layer local transformer + 16 heads, sampling in-graph) ─► 16 codes ─► CodecStep ─► 2×3840 samples
+                     │                                                                ▲
+                 Step (global GPT-2, KV update) ◄──── row [assistant_slot, 16 codes] ──┘
+```
+
+Rows are `[text_token, code_0 … code_15]`; the audio pad id (1024) marks unused code slots. Prompt layout
+(voice-clone mode) is upstream's `build_inference_input_ids`: `<im_start>user … Reference: <audio_start>
+{ref rows with user_slot=8} <audio_end> … Text: {text} </user_inst><im_end> <im_start>assistant <audio_start>`.
+Generation stops when the text head picks `audio_end` (7) instead of `assistant_slot` (9).
+
+## Models (`build/`)
+
+| mlpackage | I/O | Notes |
+|---|---|---|
+| `MossNano-Prefill-T512-M1024-fp16` | rows [1,512,17] int32 + len → hidden [1,768], kv_k/kv_v [12,1,12,1024,64] | right-padded prompt, causal + key mask |
+| `MossNano-Step-M1024-fp16` | row [1,1,17] + kv in/out + cur_len → hidden | one-hot KV write at `cur_len` |
+| `MossNano-Frame-fp16` | hidden + text_u [1] + audio_u [1,16] + temps + top_p + rep_penalty + seen [1,16,1024] + greedy → should_continue, frame [1,16] | 17 sequential local passes unrolled; inverse-CDF sampling from host uniforms (top-k 50 text / 25 audio fixed) |
+| `MossNano-CodecStep-fp16` | codes [16,1,1] + frame_index + 24 KV caches → audio [1,2,3840] + caches | streaming; caches [1,4,{500,800,1200,1600},64] per layer |
+| `MossNano-CodecDecoder-fp16` | codes [16,1,T≤125] (RangeDim) → audio [1,2,T·3840] | batch decode ≤10 s (full T×T attention) |
+| `MossNano-CodecEncoder-fp32` | audio [1,2,S≤188·3840] (RangeDim, S % 3840 == 0) → codes [16,1,S/3840] | prompt encode; fp16 loses 34 % of codes |
+
+All targets macOS 14 / iOS 17 (no `StateType`). Weights: LM ~235 MB fp16 total, codec ~40 MB.
+
+## Commands
+
+`assets/en_2.wav` and `assets/zh_1.wav` are the upstream demo prompts (`MOSS-TTS-Nano/assets/audio/`,
+gitignored here as `*.wav`); copy them in before running.
+
+```bash
+uv sync
+uv run python verify_pytorch.py                       # upstream reference (sampled) → build/ref_pytorch.wav + tokens
+uv run python verify_pytorch.py --do-sample 0 --output build/ref_greedy.wav   # deterministic parity target
+uv run python convert_lm.py --fp16
+uv run python convert_codec.py --fp16 --skip-encoder && uv run python convert_codec.py --skip-decoder
+uv run python convert_codec_step.py --fp16
+uv run python infer_coreml.py --replay 375            # teacher-forced greedy replay vs PyTorch
+uv run python infer_coreml.py --codec-step build/codec/MossNano-CodecStep-fp16.mlpackage --text "…"
+uv run python benchmark.py
+```
+
+## Parity (M5 Pro, macOS 26.7, coremltools 9.0, torch 2.7.0)
+
+- Wrappers vs upstream fp32: prefill/step hidden max|Δ| 7e-6 / 5e-6; frame greedy tokens 16/16; codec
+  decoder SNR 234 dB; encoder 1584/1584 codes; streaming step decoder vs full decode SNR 78.9 dB.
+- CoreML fp16 vs wrappers: prefill hidden 4e-3, step hidden 1.1e-2; greedy replay of the 375-frame
+  reference **370/375 frames token-exact** (five near-tie flips); codec full decoder SNR 43.5 dB (T=57)
+  / 34.3 dB (T=125); streaming step SNR 56.6 dB on GPU, 40.6 dB CPU, 37.9 dB ANE; fp32 encoder exact.
+- Intelligibility (Parakeet ASR via `fluidaudiocli tts-asr-verify --score-only`, two English phrases,
+  voice-clone prompt `assets/en_2.wav`): CoreML chain (LM fp16 + streaming codec) macro WER 8.3 %
+  (only "riverbank" → "river bank"), upstream PyTorch fp32 10.1 % (same split + one dropped word).
+- Greedy decoding in upstream never emits the stop token (runs to max frames); sampling (defaults
+  text T=1.5, audio T=1.7 / top-p 0.8 / top-k 25) is the quality mode, greedy is only a parity oracle.
+
+## Latency (warm, ms per call)
+
+| model | ALL | ANE | GPU | CPU |
+|---|---|---|---|---|
+| Frame | 5.3 | 5.7 | 5.5 | 4.1 |
+| Prefill T512 | 10.9 | 117 | 11.2 | 52.9 |
+| Step M1024 | 7.6 | 15.7 | 7.8 | 58.8 |
+| CodecStep | 5.2 | 7.8 | 5.3 | 5.5 |
+| CodecDecoder (57 frames) | 6.8 | 241 | 7.6 | 89.8 |
+| CodecEncoder fp32 (7.9 s prompt) | 22.5 | 589 | 32.5 | 247 |
+
+Per 80 ms frame the LM costs ≈ 13 ms (step + frame) and streaming codec ≈ 5 ms, i.e. ≈ 4.4× real time
+with ≈ 30 ms compute to first audio after a ≈ 11 ms prefill. The Python driver measures ≈ 9 + 9 ms
+(feed construction + 38 MB KV round trip per step); a stateful (`StateType`, iOS 18) step would remove
+that copy. Prefill/Step/CodecDecoder fail ANE compilation (`ANECCompile FAILED`) and fall back — GPU is
+the intended unit for those; Frame and CodecStep run on any unit.
+
+## Known gaps / follow-ups
+
+- Text normalisation: upstream runs WeTextProcessing + a robust normaliser before tokenising; not ported
+  (FluidAudio already ships English normalisers + NeMo ITN that can cover this).
+- Prompt encode needs fp32 (residual LFQ amplifies fp16 error in deep codebooks).
+- `nq < 16` (lower-bitrate) decoding not exposed; frame graph is fixed at 16 codebooks.
+- No stateful KV variant yet; M=1024 caps prompt+generation at 1024 rows (≈ 66 s after a 200-row prompt).

models/tts/moss-tts-nano/coreml/TRIALS.md

@@ -0,0 +1,23 @@
+# Trials
+
+## T1 — LM graphs (2026-09-13)
+- Pass-through KV following the NeuTTS layout `[L,1,H,M,D]`; one-hot blend write at `cur_len`.
+- Fused frame graph unrolls 17 local-transformer passes (no local KV — prefix ≤ 17 tokens, recompute is
+  cheaper than a cache round trip) and samples in-graph. `torch.clamp(max=)` lowers to `clip` with a
+  tensor `beta` that Core ML rejects → `torch.minimum`.
+- Sampling: exact inverse-CDF in sorted space (top-k threshold → top-p keep on exclusive cumsum →
+  softmax → count(cdf < u)). Matches upstream in distribution, greedy path bit-matches (argmax).
+- fp16 greedy replay: 370/375 exact over 30 s; mismatches are near-tie argmax flips (11/16, 10/16, …).
+
+## T2 — Codec (2026-09-13)
+- HF remote code (3.3k lines) ≠ GitHub repo code: attention masks by `input_lengths`, so wrappers must
+  pass shape-derived lengths (`ones_like(x[:,0]).sum()`), not scalar 1 → first attempt gave −10 dB.
+- Upstream `apply_rope` computes `2 / D` from a traced size → CoreML `inverse` on int32. Patched the
+  module-level function with a static-D version for export.
+- Full decoder builds T×T masks per stage (last stage 32 tokens/frame): keep RangeDim ≤ 125 frames.
+- Encoder: fp16 → 1044/1584 codes exact, degrading from 0.96 (codebook 0) to 0.5 (codebook 15); fp32 exact.
+- Streaming step: shift-append caches sized to each stage's context (500/800/1200/1600), unrotated K
+  cached, RoPE applied with constant *relative* tables (q at cap−n+j, k at i) — identical dot products
+  to absolute positions, zero runtime trig, fp16-safe. Torch step vs full decode 78.9 dB; CoreML GPU 56.6 dB.
+- Compute units: `ALL` splits the 26-input step graph across units and costs 69 ms/frame in the
+  per-frame loop vs 5–8 ms on a single unit — pin CodecStep to GPU (or ANE).

models/tts/moss-tts-nano/coreml/benchmark.py

@@ -0,0 +1,89 @@
+"""Warm per-call latency of every MOSS-TTS-Nano CoreML model across compute units (M-series host)."""
+
+from __future__ import annotations
+
+import argparse
+import time
+from pathlib import Path
+
+import coremltools as ct
+import numpy as np
+
+HERE = Path(__file__).resolve().parent
+CU = {"all": ct.ComputeUnit.ALL, "ane": ct.ComputeUnit.CPU_AND_NE, "gpu": ct.ComputeUnit.CPU_AND_GPU, "cpu": ct.ComputeUnit.CPU_ONLY}
+
+
+def feed_for(ml, seq_default: int | None = None) -> dict:
+    spec = ml.get_spec()
+    feed = {}
+    for i in spec.description.input:
+        t = i.type.multiArrayType
+        shape = []
+        for k, d in enumerate(t.shape):
+            if t.shapeRange.sizeRanges and k < len(t.shapeRange.sizeRanges):
+                r = t.shapeRange.sizeRanges[k]
+                shape.append(seq_default if (seq_default and r.upperBound != r.lowerBound) else d)
+            else:
+                shape.append(d)
+        if t.dataType == t.INT32:
+            arr = np.zeros(shape, np.int32)
+            if i.name == "input_ids":
+                arr[..., 0] = 3
+                arr[..., 1:] = 1024
+            if i.name == "input_len":
+                arr[:] = 200
+            if i.name == "cur_len":
+                arr[:] = 300
+        else:
+            arr = np.zeros(shape, np.float32)
+            if i.name in ("text_temperature", "audio_temperature", "repetition_penalty"):
+                arr[:] = 1.0
+            if i.name == "audio_top_p":
+                arr[:] = 0.8
+            if i.name in ("text_u", "audio_u"):
+                arr[:] = 0.5
+            if i.name == "audio":
+                arr = np.random.default_rng(0).standard_normal(shape).astype(np.float32) * 0.1
+        feed[i.name] = arr
+    return feed
+
+
+def bench(path: Path, units: list[str], iters: int, seq_default: int | None = None) -> None:
+    row = [path.stem]
+    for cu in units:
+        try:
+            ml = ct.models.MLModel(str(path), compute_units=CU[cu])
+            feed = feed_for(ml, seq_default)
+            ml.predict(feed)
+            ml.predict(feed)
+            t0 = time.perf_counter()
+            for _ in range(iters):
+                ml.predict(feed)
+            row.append(f"{(time.perf_counter() - t0) * 1000 / iters:7.1f}")
+        except Exception as e:  # noqa: BLE001
+            row.append(f"  err:{type(e).__name__[:8]}")
+    print(" | ".join(row))
+
+
+def main() -> None:
+    p = argparse.ArgumentParser()
+    p.add_argument("--lm-dir", default=str(HERE / "build" / "lm"))
+    p.add_argument("--codec-dir", default=str(HERE / "build" / "codec"))
+    p.add_argument("--units", default="all,ane,gpu,cpu")
+    p.add_argument("--iters", type=int, default=10)
+    args = p.parse_args()
+    units = args.units.split(",")
+    print("model | " + " | ".join(f"{u:>7}" for u in units) + "   (ms per call, warm)")
+    for path in sorted(Path(args.lm_dir).glob("*.mlpackage")):
+        bench(path, units, args.iters)
+    for path in sorted(Path(args.codec_dir).glob("*.mlpackage")):
+        seq = None
+        if "CodecDecoder" in path.name:
+            seq = 57  # frames (4.6 s)
+        if "CodecEncoder" in path.name:
+            seq = 99 * 3840  # samples (7.9 s prompt)
+        bench(path, units, args.iters if seq is None else 3, seq)
+
+
+if __name__ == "__main__":
+    main()

models/tts/moss-tts-nano/coreml/convert_codec.py

@@ -0,0 +1,148 @@
+"""Convert MOSS-Audio-Tokenizer-Nano to CoreML.
+
+    MossNano-CodecDecoder-{tag}.mlpackage   codes [16,1,T≤max] → audio [1,2,T*3840]
+    MossNano-CodecEncoder-{tag}.mlpackage   audio [1,2,S≤max*3840] → codes [16,1,S/3840]
+
+Parity: decoder vs upstream decode() on real generated tokens (SNR), encoder vs upstream
+encode() on the bundled reference clip (exact code match rate).
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+import time
+from pathlib import Path
+
+import coremltools as ct
+import numpy as np
+import soundfile as sf
+import torch
+
+HERE = Path(__file__).resolve().parent
+sys.path.insert(0, str(HERE))
+
+from src.codec_coreml import MossCodecDecoder, MossCodecEncoder  # noqa: E402
+
+CODEC_REPO = "OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano"
+FRAME = 3840
+
+
+def snr_db(got: np.ndarray, want: np.ndarray) -> float:
+    noise = np.sum((got - want) ** 2)
+    return float(10 * np.log10(np.sum(want**2) / max(noise, 1e-20)))
+
+
+def precision(fp16: bool):
+    if not fp16:
+        return ct.precision.FLOAT32
+    return ct.transform.FP16ComputePrecision(op_selector=lambda op: op.op_type not in {"softmax"})
+
+
+def load_prompt(path: Path) -> torch.Tensor:
+    wav, sr = sf.read(path, dtype="float32", always_2d=True)  # [S, C]
+    wav = torch.from_numpy(wav.T)
+    if sr != 48000:
+        import torchaudio
+
+        wav = torchaudio.functional.resample(wav, sr, 48000)
+    if wav.shape[0] == 1:
+        wav = wav.repeat(2, 1)
+    S = wav.shape[1]
+    pad = (-S) % FRAME
+    return torch.nn.functional.pad(wav, (0, pad))[None]  # [1,2,S']
+
+
+def main() -> None:
+    p = argparse.ArgumentParser()
+    p.add_argument("--output-dir", default=str(HERE / "build" / "codec"))
+    p.add_argument("--fp16", action="store_true")
+    p.add_argument("--max-frames", type=int, default=125, help="decoder RangeDim upper bound in frames (10 s)")
+    p.add_argument("--max-prompt-frames", type=int, default=188, help="encoder RangeDim upper bound in frames (15 s)")
+    p.add_argument("--skip-decoder", action="store_true")
+    p.add_argument("--skip-encoder", action="store_true")
+    p.add_argument("--tokens", default=str(HERE / "build" / "ref_audio_token_ids.npy"))
+    p.add_argument("--long-tokens", default=str(HERE / "build" / "ref_greedy_audio_token_ids.npy"))
+    p.add_argument("--prompt-audio", default=str(HERE / "assets" / "en_2.wav"))
+    args = p.parse_args()
+
+    out_dir = Path(args.output_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    tag = "fp16" if args.fp16 else "fp32"
+    target = ct.target.macOS14
+
+    from transformers import AutoModel
+
+    print(f"[0] loading {CODEC_REPO}")
+    codec = AutoModel.from_pretrained(CODEC_REPO, trust_remote_code=True).eval()
+    codes = torch.from_numpy(np.load(args.tokens)).long().T[:, None, :]  # [16,1,T]
+    long_codes = torch.from_numpy(np.load(args.long_tokens)).long().T[:, None, : args.max_frames]
+    with torch.no_grad():
+        want = codec.decode(codes, return_dict=True).audio  # [1,2,S]
+        want_long = codec.decode(long_codes, return_dict=True).audio
+
+    if not args.skip_decoder:
+        print("[1] decoder")
+        dec = MossCodecDecoder(codec).eval()
+        with torch.no_grad():
+            got = dec(codes.to(torch.int32))
+        print(f"      wrapper vs upstream: shape {tuple(got.shape)} SNR={snr_db(got.numpy(), want.numpy()):.1f} dB")
+        t0 = time.perf_counter()
+        with torch.no_grad():
+            traced = torch.jit.trace(dec, (codes.to(torch.int32),), strict=False)
+        T = ct.RangeDim(lower_bound=1, upper_bound=args.max_frames, default=codes.shape[-1])
+        ml = ct.convert(
+            traced,
+            inputs=[ct.TensorType(name="codes", shape=(16, 1, T), dtype=np.int32)],
+            outputs=[ct.TensorType(name="audio", dtype=np.float32)],
+            compute_precision=precision(args.fp16),
+            minimum_deployment_target=target,
+            convert_to="mlprogram",
+        )
+        path = out_dir / f"MossNano-CodecDecoder-{tag}.mlpackage"
+        ml.save(str(path))
+        print(f"      saved {path.name} ({time.perf_counter() - t0:.0f}s)")
+        for name, c, w in (("T=57", codes, want), (f"T={long_codes.shape[-1]}", long_codes, want_long)):
+            t1 = time.perf_counter()
+            pred = ml.predict({"codes": c.numpy().astype(np.int32)})["audio"]
+            dt = time.perf_counter() - t1
+            print(f"      coreml {name}: shape {pred.shape} SNR={snr_db(pred, w.numpy()):.1f} dB  {dt*1000:.0f} ms")
+        del ml, dec
+
+    if not args.skip_encoder:
+        print("[2] encoder")
+        audio = load_prompt(Path(args.prompt_audio))  # [1,2,S]
+        with torch.no_grad():
+            ref = codec.encode(audio, return_dict=True).audio_codes  # [16,1,T]
+        enc = MossCodecEncoder(codec).eval()
+        with torch.no_grad():
+            got = enc(audio)
+        n = ref.numel()
+        print(f"      wrapper vs upstream: {tuple(got.shape)} vs {tuple(ref.shape)}  "
+              f"exact={(got.long() == ref).sum().item()}/{n}")
+        t0 = time.perf_counter()
+        with torch.no_grad():
+            traced = torch.jit.trace(enc, (audio,), strict=False)
+        S = ct.RangeDim(lower_bound=FRAME, upper_bound=args.max_prompt_frames * FRAME, default=audio.shape[-1])
+        ml = ct.convert(
+            traced,
+            inputs=[ct.TensorType(name="audio", shape=(1, 2, S), dtype=np.float32)],
+            outputs=[ct.TensorType(name="codes", dtype=np.int32)],
+            compute_precision=precision(args.fp16),
+            minimum_deployment_target=target,
+            convert_to="mlprogram",
+        )
+        path = out_dir / f"MossNano-CodecEncoder-{tag}.mlpackage"
+        ml.save(str(path))
+        print(f"      saved {path.name} ({time.perf_counter() - t0:.0f}s)")
+        t1 = time.perf_counter()
+        pred = torch.from_numpy(ml.predict({"audio": audio.numpy()})["codes"]).long()
+        dt = time.perf_counter() - t1
+        per_cb = [(pred[i] == ref[i]).float().mean().item() for i in range(16)]
+        print(f"      coreml: exact={(pred == ref).sum().item()}/{n}  {dt*1000:.0f} ms  "
+              f"per-codebook match={[round(v, 3) for v in per_cb]}")
+    print("[done]")
+
+
+if __name__ == "__main__":
+    main()