Initial release: optimized CPU runtime and converted weights of maya2-native

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/v1/backbone_wo8.data filter=lfs diff=lfs merge=lfs -text
+assets/v1/depth_frame_fp16.onnx.data filter=lfs diff=lfs merge=lfs -text
+assets/v1/tokenizer.json filter=lfs diff=lfs merge=lfs -text

NOTICE

@@ -0,0 +1,10 @@
+maya2-native-fast
+==================
+
+Derived from maya-research/maya2-native (Maya Research), Apache-2.0.
+Includes a Mimi decoder converted from kyutai/mimi (CC-BY-4.0).
+Backbone initialised from Qwen3-1.7B-Base (Qwen team), Apache-2.0.
+Speaker conditioning uses speechbrain/spkrec-ecapa-voxceleb (SpeechBrain), Apache-2.0.
+
+Apache-2.0 full text: https://www.apache.org/licenses/LICENSE-2.0
+CC-BY-4.0 full text: https://creativecommons.org/licenses/by/4.0/

README.md

@@ -0,0 +1,127 @@
+---
+license: apache-2.0
+base_model: maya-research/maya2-native
+language:
+  - mr
+  - hi
+library_name: onnx
+pipeline_tag: text-to-speech
+tags:
+  - text-to-speech
+  - onnx
+  - int8
+  - fp16
+  - marathi
+  - hindi
+  - cpu
+---
+
+# maya2-native-fast
+
+Optimized CPU inference for [maya-research/maya2-native](https://huggingface.co/maya-research/maya2-native):
+ONNX graphs, int8 backbone, fp16 depth, a torch-free numpy runtime.
+
+**This is not a fine-tune.** The weights are converted from the upstream checkpoint. No training data was used.
+
+## What is in the repo
+
+| Path | Contents |
+|---|---|
+| `assets/v1/backbone_wo8.onnx` + `.data` | 28-layer Qwen3 backbone, int8 weights (MatMulNBits, block 32), 1.6 GB |
+| `assets/v1/depth_frame_fp16.onnx` + `.data` | Depth sampler, 15 levels, fp16 storage, Gumbel top-k sampling in-graph, 873 MB |
+| `assets/v1/mimi_decoder.onnx` | Mimi decoder, fp32, dynamic frame count, 195 MB |
+| `assets/v1/embed_bf16.bin`, `audio_in_bf16.bin` | Text and audio embedding tables, bf16, memory-mapped |
+| `assets/v1/prompt.npz`, `tokenizer.json`, `voices.json`, `controls.json` | Prompt weights, extended tokenizer, voices, sentence controls |
+| `assets/v1/manifest.json` | sha256 of every file above, runtime shapes. The runtime refuses mismatches. |
+| `src/maya_fast/` | Runtime (numpy, onnxruntime, tokenizers; no torch) |
+| `tools/maya_export/` | Export and evaluation scripts, to rebuild the assets from the upstream checkpoint |
+| `tests/` | Runtime and export tests |
+
+## Measured results (Apple M-series Mac, CPU, 8 threads)
+
+| Check | Result | Target | Status |
+|---|---|---|---|
+| Depth fp16 vs fp32 reference, teacher-forced, 75 files | top-1 >= 0.9965 every level, KL <= 6e-6 | top-1 >= 0.98, KL <= 0.02 | pass |
+| Depth fp32 vs upstream PyTorch, 20 frames | rel 1.6e-6 | rel <= 1e-4 | pass |
+| Mimi decoder vs upstream | max abs 1.5e-6 | <= 1e-4 | pass |
+| Prompt port vs upstream, 500 prompts | 0 token mismatches, max embed error 7.5e-9 | <= 1e-5 | pass |
+| Backbone int8 vs fp32, teacher-forced | top-1 0.9677, KL 0.0034 | top-1 >= 0.99 | **fail (waived)** |
+| Backbone speed, 8 threads | 48 / 58 / 82 ms per token at 32 / 256 / 768 cached tokens | <= 35 ms at 256 | **fail (waived)** |
+| End-to-end, one Marathi sentence (2.96 s audio), two runs | RTF 4.67 and 4.43 | offline RTF <= 4 | **fail** |
+| Peak RSS, end-to-end | 3.9 GB | <= 3 GB target, 5 GB limit | over target |
+
+Known limits:
+- The int8 backbone misses top-1 agreement; expect small differences in generated codes.
+- End-to-end RTF on this machine is above the offline target.
+- Hindi (`--lang hi`) works mechanically; it has not been checked for quality.
+- Not yet checked by listening tests.
+- Real-time playback is not supported.
+
+## Install
+
+```bash
+python -m venv .venv
+.venv/bin/pip install -r requirements-fast.txt
+```
+
+Requires Python 3.10 to 3.13. Peak memory is about 4 GB while a sentence is synthesized.
+
+## Use: command line
+
+Run from the repo root (so `src` is importable):
+
+```bash
+PYTHONPATH=. .venv/bin/python -m src.maya_fast "नमस्कार, मी तुमची मदत करण्यासाठी इथे आहे." \
+    --assets assets/v1 -o out.wav
+```
+
+Options:
+- `--voice NAME`: named voice from `assets/v1/voices.json` (case-insensitive). Default `Charvi`.
+- `--lang mr|hi`: language tag. Default `mr`.
+- `--seed N`: random seed. Same seed and threads give the same audio.
+- `--threads N`: ORT threads. Default min(8, cores).
+
+## Use: Python
+
+```python
+from src.maya_fast.engine import MayaFast
+
+engine = MayaFast(assets_dir="assets/v1", threads=8)
+wav, metrics = engine.synthesize("नमस्कार, मी तुमची मदत करण्यासाठी इथे आहे.", voice="Charvi", seed=0, lang="mr")
+# wav: float32, 24 kHz mono
+```
+
+Write it with `soundfile` or the standard `wave` module.
+
+Assets are verified against `manifest.json` at load. A bad or partial download fails with the file name.
+
+## Rebuild the assets from upstream (optional)
+
+Needs the upstream checkpoint (`maya-research/maya2-native`, `kyutai/mimi`, `speechbrain/spkrec-ecapa-voxceleb`) and an upstream `maya3` install in a second venv (`.venv-tts`, torch 2.8, transformers 5.6.2). The export scripts write to `$MAYA_FAST_WORK/assets/v1`, default `~/.cache/maya_fast`. See `tools/maya_export/`:
+1. `export_prompt_assets.py`, `export_backbone.py --precision wo8`, `export_depth_all.py`, `export_codec.py`
+2. `finalize_assets.py` writes the runtime manifest.
+Run each step under `tools/memguard.py`, which kills the process above a memory limit.
+
+## License and attribution
+
+- These weights and code are derived from Maya2 Native (Apache-2.0) by Maya Research.
+- The Mimi decoder is derived from Kyutai Mimi (CC-BY-4.0). Keep this attribution.
+- The text backbone is derived from Qwen3-1.7B-Base (Apache-2.0).
+- Speaker conditioning uses SpeechBrain ECAPA-TDNN (Apache-2.0).
+- Full notices in `NOTICE`.
+
+The model adds no audio watermark.
+
+## Citation
+
+Please cite the upstream model:
+
+```bibtex
+@misc{maya2native2026,
+  title  = {Maya 2 Native: Streaming Multilingual Text-to-Speech},
+  author = {{Maya Research}},
+  year   = {2026},
+  url    = {https://huggingface.co/maya-research/maya2-native},
+  note   = {Release 1.0}
+}
+```

assets/v1/audio_in_bf16.bin

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

assets/v1/backbone_wo8.data

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

assets/v1/backbone_wo8.onnx

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

assets/v1/controls.json

@@ -0,0 +1,183 @@
+{
+  "text_frames_prior": {
+    "as": {
+      "a": -1.36,
+      "b": 0.7786,
+      "resid_p99": 91.6
+    },
+    "bn": {
+      "a": 3.16,
+      "b": 0.7887,
+      "resid_p99": 76.2
+    },
+    "en": {
+      "a": 0.6,
+      "b": 0.7416,
+      "resid_p99": 75.0
+    },
+    "gu": {
+      "a": 1.74,
+      "b": 0.8834,
+      "resid_p99": 83.1
+    },
+    "hi": {
+      "a": 5.65,
+      "b": 0.7644,
+      "resid_p99": 79.8
+    },
+    "kn": {
+      "a": 1.91,
+      "b": 0.7675,
+      "resid_p99": 69.5
+    },
+    "ml": {
+      "a": 1.25,
+      "b": 0.6788,
+      "resid_p99": 105.6
+    },
+    "mr": {
+      "a": 1.12,
+      "b": 0.8235,
+      "resid_p99": 60.1
+    },
+    "or": {
+      "a": 1.06,
+      "b": 0.7936,
+      "resid_p99": 68.4
+    },
+    "pa": {
+      "a": 1.86,
+      "b": 0.8079,
+      "resid_p99": 73.5
+    },
+    "ta": {
+      "a": 2.24,
+      "b": 0.6425,
+      "resid_p99": 70.8
+    },
+    "te": {
+      "a": 1.2,
+      "b": 0.7927,
+      "resid_p99": 89.6
+    }
+  },
+  "prior_default": {
+    "a": 2.0,
+    "b": 0.72,
+    "resid_p99": 90.0
+  },
+  "text": {
+    "terminals": "।.!?;।॥",
+    "fps": 12.5,
+    "frame_s": 0.08,
+    "chunk_s": 16.0
+  },
+  "guard": {
+    "enabled": true,
+    "dur_k": 2.2,
+    "dur_margin_s": 1.2,
+    "hard_cap": 1500,
+    "dur_prior_coefs": {
+      "as": [
+        0.0,
+        0.07962,
+        0.17764,
+        0.53121,
+        0.62254
+      ],
+      "bn": [
+        0.71369,
+        0.06604,
+        0.04537,
+        0.49869,
+        0.23288
+      ],
+      "en": [
+        2.87006,
+        0.01,
+        0.04601,
+        0.40336,
+        0.06541
+      ],
+      "gu": [
+        0.54378,
+        0.07044,
+        0.01,
+        0.50786,
+        0.24873
+      ],
+      "hi": [
+        2.44169,
+        0.05329,
+        0.02187,
+        0.43341,
+        0.37052
+      ],
+      "kn": [
+        1.71201,
+        0.06651,
+        0.01,
+        0.45228,
+        0.5053
+      ],
+      "ml": [
+        0.92428,
+        0.05706,
+        0.11842,
+        0.52651,
+        0.41106
+      ],
+      "mr": [
+        1.44052,
+        0.0719,
+        0.03206,
+        0.30526,
+        0.28705
+      ],
+      "or": [
+        0.0,
+        0.08134,
+        0.12578,
+        0.43209,
+        0.23588
+      ],
+      "pa": [
+        0.57673,
+        0.07431,
+        0.01,
+        0.39355,
+        0.37283
+      ],
+      "ta": [
+        2.05273,
+        0.05244,
+        0.01,
+        0.47524,
+        0.45486
+      ],
+      "te": [
+        0.45026,
+        0.07269,
+        0.0576,
+        0.46093,
+        0.39692
+      ]
+    },
+    "dur_prior_default": [
+      1.1438125,
+      0.06297000000000001,
+      0.05539583333333333,
+      0.45169916666666676,
+      0.35033166666666665
+    ]
+  },
+  "sampling": {
+    "q_serve": 16,
+    "temp": 0.7,
+    "top_k": 64,
+    "tau": 0.35,
+    "hazard": 0.6,
+    "cfg_scale": 0.0,
+    "aeos_class": 2048
+  }
+}

assets/v1/depth_frame_fp16.onnx

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

assets/v1/depth_frame_fp16.onnx.data

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

assets/v1/embed_bf16.bin

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

assets/v1/manifest.json

@@ -0,0 +1,29 @@
+{
+  "version": 1,
+  "files": {
+    "tokenizer.json": "26c51f9b772f9084ed6bdeabc44b7c15397c12966bfa3e1c09b025e82f4fdbaf",
+    "embed_bf16.bin": "7940ea16f0a59b955ac42a9f5c035930e0aded3e034ca68de18531c45a6ff861",
+    "audio_in_bf16.bin": "2670e5583c2d2e6e963a8117136a8d07e0e892d3ffda2620d3f13cdbaddc94d0",
+    "prompt.npz": "0d6392e0cba6009a913ba57605641b48b1a1fa32a0040ffd4b7b2f840be8afc5",
+    "voices.json": "dac72e40ff1c3863f91ab65d99346965df98967cc40b2a8234b757ae7892b4cf",
+    "controls.json": "527f1e3129bf6b16df9ac7e2ec7dc813da6b74a207031cc41d959b6a6337021e",
+    "backbone_wo8.onnx": "d0326c5973d9d71708d7c39c3f34851c73addc88bceb46d0ac90c37e9041955d",
+    "backbone_wo8.data": "ce9ebfb1e911bac4e3bc58b77c6657f75b79f145f89df718f23d352d824dcb01",
+    "depth_frame_fp16.onnx": "c49a4db17d5f443371c459d75c4bdbd508abc164ba30caf51f70c36032c1cea4",
+    "depth_frame_fp16.onnx.data": "494c8e441065164eb1aa26100f5c2763b447108a116914036a72e1b124463e61",
+    "mimi_decoder.onnx": "a0e6accff030071e49c1e1824b26f3a4868c39d125f11bb8079b4a4bbd9dcb58"
+  },
+  "shapes": {
+    "d_model": 2048,
+    "vocab": 194548,
+    "audio_rows": 65536,
+    "n_q": 16,
+    "n_codes": 2048,
+    "ecapa_dim": 192,
+    "rope_theta": 1000000.0,
+    "n_layers": 28,
+    "n_heads": 8,
+    "head_dim": 128,
+    "d_mem": 1024
+  }
+}

assets/v1/mimi_decoder.onnx

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

assets/v1/prompt.npz

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

assets/v1/tokenizer.json

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

assets/v1/voices.json

@@ -0,0 +1,102 @@
+{
+  "Ahana": {
+    "mode": "named",
+    "use_icl": false,
+    "use_mem": false
+  },
+  "Vaidehi": {
+    "mode": "named",
+    "use_icl": false,
+    "use_mem": false
+  },
+  "Malini": {
+    "mode": "named",
+    "use_icl": false,
+    "use_mem": false
+  },
+  "Charvi": {
+    "mode": "named",
+    "use_icl": false,
+    "use_mem": false
+  },
+  "Jahnavi": {
+    "mode": "named",
+    "use_icl": false,
+    "use_mem": false
+  },
+  "Anirban": {
+    "mode": "named",
+    "use_icl": false,
+    "use_mem": false
+  },
+  "Nakul": {
+    "mode": "named",
+    "use_icl": false,
+    "use_mem": false
+  },
+  "Vasant": {
+    "mode": "named",
+    "use_icl": false,
+    "use_mem": false
+  },
+  "Siddhant": {
+    "mode": "named",
+    "use_icl": false,
+    "use_mem": false
+  },
+  "Satvik": {
+    "mode": "named",
+    "use_icl": false,
+    "use_mem": false
+  },
+  "Abhay": {
+    "mode": "named",
+    "use_icl": false,
+    "use_mem": false
+  },
+  "Nivedita": {
+    "mode": "named",
+    "use_icl": false,
+    "use_mem": false
+  },
+  "Ishwari": {
+    "mode": "named",
+    "use_icl": false,
+    "use_mem": false
+  },
+  "Ritvik": {
+    "mode": "named",
+    "use_icl": false,
+    "use_mem": false
+  },
+  "Vrinda": {
+    "mode": "named",
+    "use_icl": false,
+    "use_mem": false
+  },
+  "Sreekar": {
+    "mode": "named",
+    "use_icl": false,
+    "use_mem": false
+  },
+  "Kayal": {
+    "mode": "named",
+    "use_icl": false,
+    "use_mem": false
+  },
+  "Shaunak": {
+    "mode": "named",
+    "use_icl": false,
+    "use_mem": false
+  },
+  "Akshat": {
+    "mode": "named",
+    "use_icl": false,
+    "use_mem": false
+  },
+  "Tanushree": {
+    "mode": "named",
+    "use_icl": false,
+    "use_mem": false
+  }
+}

requirements-fast.txt

@@ -0,0 +1,5 @@
+numpy>=2.0
+onnxruntime==1.31.*
+tokenizers>=0.22,<0.23
+psutil
+onnx

src/maya_fast/__init__.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from .engine import MayaFast
+from .assets import PromptAssets, load_manifest, AssetError
+from .backbone import Backbone
+from .depth import Depth
+from .codec import MimiDecoder
+from .prompt import AudioIn, Prompt, build_prompt
+from .control import StopPolicy, sample_cb0
+
+__all__ = [
+    "MayaFast",
+    "PromptAssets",
+    "load_manifest",
+    "AssetError",
+    "Backbone",
+    "Depth",
+    "MimiDecoder",
+    "AudioIn",
+    "Prompt",
+    "build_prompt",
+    "StopPolicy",
+    "sample_cb0",
+]

src/maya_fast/__main__.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+import argparse
+import sys
+import wave
+from pathlib import Path
+import numpy as np
+
+from .engine import MayaFast
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Maya-Fast TTS CLI")
+    parser.add_argument("text", type=str, help="Text to speak")
+    parser.add_argument("-o", "--out", type=str, default="out.wav", help="Output WAV path")
+    parser.add_argument("--voice", type=str, default="Charvi", help="Voice name (default: Charvi)")
+    parser.add_argument("--seed", type=int, default=0, help="Random seed (default: 0)")
+    parser.add_argument("--lang", type=str, default="mr", help="Language tag: mr, hi, ... (default: mr)")
+    parser.add_argument("--threads", type=int, default=None, help="ORT threads (default: min(8, cores))")
+    parser.add_argument("--assets", type=str, default=None, help="Assets directory")
+
+    args = parser.parse_args()
+
+    engine = MayaFast(assets_dir=args.assets, threads=args.threads)
+    wav, metrics = engine.synthesize(args.text, voice=args.voice, seed=args.seed, lang=args.lang)
+
+    # Convert float32 [-1.0, 1.0] to 16-bit PCM WAV using stdlib wave
+    wav_clipped = np.clip(wav, -1.0, 1.0)
+    pcm16 = (wav_clipped * 32767.0).astype(np.int16)
+
+    out_path = Path(args.out)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    with wave.open(str(out_path), "wb") as wf:
+        wf.setnchannels(1)
+        wf.setsampwidth(2)
+        wf.setframerate(24000)
+        wf.writeframes(pcm16.tobytes())
+
+    print(
+        f"audio_s={metrics['audio_s']:.2f} "
+        f"wall_s={metrics['wall_s']:.2f} "
+        f"rtf={metrics['rtf']:.2f} "
+        f"peak_rss_mb={metrics['peak_rss_mb']:.1f}"
+    )
+
+
+if __name__ == "__main__":
+    main()

src/maya_fast/assets.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+from pathlib import Path
+from typing import Any
+import numpy as np
+
+
+class AssetError(Exception):
+    pass
+
+
+def default_work_dir() -> Path:
+    env = os.environ.get("MAYA_FAST_WORK")
+    if env:
+        return Path(env)
+    return Path.home() / ".cache" / "maya_fast"
+
+
+def default_assets_dir() -> Path:
+    return default_work_dir() / "assets" / "v1"
+
+
+def compute_file_sha256(path: Path) -> str:
+    h = hashlib.sha256()
+    with open(path, "rb") as f:
+        while chunk := f.read(65536):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def load_manifest(assets_dir: str | Path, verify: bool = True) -> dict[str, Any]:
+    assets_dir = Path(assets_dir)
+    manifest_file = assets_dir / "manifest.json"
+    if not manifest_file.exists():
+        raise AssetError(f"manifest.json missing in {assets_dir}")
+
+    with open(manifest_file, "r", encoding="utf-8") as f:
+        try:
+            manifest = json.load(f)
+        except Exception as e:
+            raise AssetError(f"Corrupt manifest.json in {assets_dir}: {e}")
+
+    if verify:
+        files = manifest.get("files", {})
+        for name, expected_sha in files.items():
+            fpath = assets_dir / name
+            if not fpath.exists():
+                raise AssetError(f"Required asset file missing: {name}")
+            actual_sha = compute_file_sha256(fpath)
+            if actual_sha != expected_sha:
+                raise AssetError(f"Asset sha256 mismatch for file: {name}")
+
+    return manifest
+
+
+class PromptAssets:
+    def __init__(self, assets_dir: str | Path, verify: bool = True) -> None:
+        self.assets_dir = Path(assets_dir)
+        self.manifest = load_manifest(self.assets_dir, verify=verify)
+        self.shapes = self.manifest.get("shapes", {})
+
+        from .embed import EmbeddingTable
+        from .prompt import AudioIn
+
+        d_model = self.shapes.get("d_model", 2048)
+        vocab = self.shapes.get("vocab", 151936)
+        n_q = self.shapes.get("n_q", 16)
+        n_codes = self.shapes.get("n_codes", 2048)
+        self.embed_table = EmbeddingTable(self.assets_dir / "embed_bf16.bin", rows=vocab, dim=d_model)
+        audio_table = EmbeddingTable(self.assets_dir / "audio_in_bf16.bin", rows=n_q * n_codes, dim=d_model)
+
+        # Controls (text_frames_prior, guard, sampling, text)
+        with open(self.assets_dir / "controls.json", "r", encoding="utf-8") as f:
+            self.controls = json.load(f)
+
+        # Speaker anchor weights and audio input norm/aeos (fp32; the audio table is bf16 on disk)
+        prompt_npz = np.load(self.assets_dir / "prompt.npz")
+        self.prompt_params = {k: prompt_npz[k] for k in prompt_npz.files}
+        self.audio_in = AudioIn(audio_table, self.prompt_params["audio_in/aeos_in"],
+                                self.prompt_params["audio_in/norm"], n_q=n_q)
+
+        # Voices: name -> {mode, use_icl, use_mem}
+        with open(self.assets_dir / "voices.json", "r", encoding="utf-8") as f:
+            self.voices_info = json.load(f)
+
+        # Tokenizer (extended upstream tokenizer, fast format)
+        from tokenizers import Tokenizer
+        self.tokenizer = Tokenizer.from_file(str(self.assets_dir / "tokenizer.json"))

src/maya_fast/backbone.py

@@ -0,0 +1,157 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+import numpy as np
+import onnxruntime as ort
+
+from .ort_opts import session_options
+
+
+class Backbone:
+    def __init__(
+        self,
+        path: str | Path,
+        threads: int | None = None,
+        manifest_path: str | Path | None = None,
+        config: dict | None = None,
+    ) -> None:
+        self.path = Path(path)
+        self.opts = session_options(threads)
+        self.session = ort.InferenceSession(str(self.path), self.opts, providers=["CPUExecutionProvider"])
+
+        # Determine shapes and hyperparameters
+        # Default real Qwen3 / Maya2 backbone parameters:
+        # n_layers=28, n_heads=8 (kv heads), head_dim=128, rope_theta=10000.0, d_model=2048
+        self.n_layers = 28
+        self.n_heads = 8
+        self.head_dim = 128
+        self.rope_theta = 10000.0
+        self.d_model = 2048
+
+        # If manifest or config provided, override:
+        if config is not None:
+            self._apply_config(config)
+        elif manifest_path is not None and Path(manifest_path).exists():
+            with open(manifest_path, "r", encoding="utf-8") as f:
+                data = json.load(f)
+                shapes = data.get("shapes", {})
+                self._apply_config(shapes)
+        else:
+            # Check if manifest.json exists next to model
+            maybe_manifest = self.path.parent / "manifest.json"
+            if maybe_manifest.exists():
+                with open(maybe_manifest, "r", encoding="utf-8") as f:
+                    data = json.load(f)
+                    shapes = data.get("shapes", {})
+                    self._apply_config(shapes)
+            else:
+                # Discover n_layers from session inputs
+                past_k_inputs = [inp.name for inp in self.session.get_inputs() if inp.name.startswith("past_k_")]
+                if past_k_inputs:
+                    self.n_layers = len(past_k_inputs)
+
+        # Inspect session input shapes if available
+        for inp in self.session.get_inputs():
+            if inp.name == "past_k_0":
+                shape = inp.shape
+                # shape: [1, n_heads, P, head_dim]
+                if len(shape) == 4:
+                    if isinstance(shape[1], int):
+                        self.n_heads = shape[1]
+                    if isinstance(shape[3], int):
+                        self.head_dim = shape[3]
+            elif inp.name == "embeds":
+                if len(inp.shape) == 3 and isinstance(inp.shape[2], int):
+                    self.d_model = inp.shape[2]
+
+        self.reset()
+
+    def _apply_config(self, cfg: dict) -> None:
+        if "n_layers" in cfg:
+            self.n_layers = int(cfg["n_layers"])
+        if "n_heads" in cfg:
+            self.n_heads = int(cfg["n_heads"])
+        if "head_dim" in cfg:
+            self.head_dim = int(cfg["head_dim"])
+        if "rope_theta" in cfg:
+            self.rope_theta = float(cfg["rope_theta"])
+        if "d_model" in cfg:
+            self.d_model = int(cfg["d_model"])
+
+    def reset(self) -> None:
+        self.past_len = 0
+        self.past_k = [np.zeros((1, self.n_heads, 0, self.head_dim), dtype=np.float32) for _ in range(self.n_layers)]
+        self.past_v = [np.zeros((1, self.n_heads, 0, self.head_dim), dtype=np.float32) for _ in range(self.n_layers)]
+
+    def _compute_rope(self, T: int, past: int) -> tuple[np.ndarray, np.ndarray]:
+        # RoPE for positions [past, past + T)
+        # dim = head_dim
+        pos = np.arange(past, past + T, dtype=np.float32)[:, None]  # (T, 1)
+        dim_indices = np.arange(0, self.head_dim, 2, dtype=np.float32)  # (head_dim/2,)
+        inv_freq = 1.0 / (self.rope_theta ** (dim_indices / self.head_dim))  # (head_dim/2,)
+        sinusoid = pos * inv_freq[None, :]  # (T, head_dim/2)
+        # duplicate across second half to match typical rotary implementation: [sinusoid, sinusoid]
+        sinusoid = np.concatenate([sinusoid, sinusoid], axis=-1)  # (T, head_dim)
+        cos = np.cos(sinusoid)[None, :, :]  # (1, T, head_dim)
+        sin = np.sin(sinusoid)[None, :, :]  # (1, T, head_dim)
+        return cos.astype(np.float32), sin.astype(np.float32)
+
+    def _compute_mask(self, T: int, past: int) -> np.ndarray:
+        # Additive causal mask (1, 1, T, past+T): query i sees keys j <= past + i
+        total_len = past + T
+        i = np.arange(T)[:, None]
+        j = np.arange(total_len)[None, :]
+        mask = np.zeros((T, total_len), dtype=np.float32)
+        mask[j > (past + i)] = -np.inf
+        return mask[None, None, :, :]
+
+    def step(self, embeds: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+        """
+        embeds: (T, d_model) or (1, T, d_model)
+        Returns:
+            hidden: (d_model,) last position
+            cb0_logits: (vocab + 1,) last position
+        """
+        if embeds.ndim == 2:
+            embeds = embeds[None, :, :]
+        elif embeds.ndim != 3:
+            raise ValueError(f"Expected embeds to have 2 or 3 dims, got shape {embeds.shape}")
+
+        embeds = embeds.astype(np.float32)
+        T = embeds.shape[1]
+        P = self.past_len
+
+        cos, sin = self._compute_rope(T, P)
+        mask = self._compute_mask(T, P)
+
+        feed = {
+            "embeds": embeds,
+            "cos": cos,
+            "sin": sin,
+            "mask": mask,
+        }
+        for i in range(self.n_layers):
+            feed[f"past_k_{i}"] = self.past_k[i]
+            feed[f"past_v_{i}"] = self.past_v[i]
+
+        outputs = self.session.run(None, feed)
+
+        # Output ordering from session:
+        # output[0] = hidden (1, T, d_model)
+        # output[1] = cb0_logits (1, T, vocab + 1)
+        # output[2..] = present_k_i, present_v_i
+        hidden_out = outputs[0]
+        cb0_logits_out = outputs[1]
+
+        idx = 2
+        for i in range(self.n_layers):
+            self.past_k[i] = outputs[idx]
+            self.past_v[i] = outputs[idx + 1]
+            idx += 2
+
+        self.past_len = P + T
+
+        last_hidden = hidden_out[0, -1, :]
+        last_cb0_logits = cb0_logits_out[0, -1, :]
+        return last_hidden, last_cb0_logits

src/maya_fast/codec.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from pathlib import Path
+import numpy as np
+import onnxruntime as ort
+
+from .ort_opts import session_options
+
+
+class MimiDecoder:
+    def __init__(self, path: str | Path, threads: int | None = None) -> None:
+        self.path = Path(path)
+        self.opts = session_options(threads)
+        self.session = ort.InferenceSession(str(self.path), self.opts, providers=["CPUExecutionProvider"])
+
+    def decode(self, codes: np.ndarray) -> np.ndarray:
+        """
+        Input:
+            codes: (16, F) int64
+        Returns:
+            wav: (F * 1920,) float32
+        """
+        codes = np.asarray(codes, dtype=np.int64)
+        if codes.ndim == 2:
+            codes_input = codes[None, :, :]  # (1, 16, F)
+        elif codes.ndim == 3:
+            codes_input = codes
+        else:
+            raise ValueError(f"Expected codes to have 2 or 3 dims, got {codes.shape}")
+
+        outputs = self.session.run(None, {"codes": codes_input})
+        wav = outputs[0]  # (1, 1, samples)
+        wav = np.squeeze(wav)
+        return wav.astype(np.float32)

src/maya_fast/control.py

@@ -0,0 +1,135 @@
+from __future__ import annotations
+
+import re
+
+import numpy as np
+
+# Port of maya3.guards (DurationPrior, text_feats) and scripts/s2_serve_controls.py (expected_frames, split_text).
+# The numbers come from controls.json, written by tools/maya_export/export_prompt_assets.py from upstream.
+PAUSE = re.compile(r"[,;:.!?।॥…\-]")
+PRIOR_DEFAULT = {"a": 2.0, "b": 0.72, "resid_p99": 90.0}  # scripts/s2_serve_controls.py
+FRAME_S = 0.08  # maya3.guards.FRAME_S (one frame = 80 ms of audio)
+
+
+def text_feats(t: str) -> np.ndarray:
+    letters = sum(1 for ch in t if ch.isalpha() or (0x900 <= ord(ch) <= 0xDFF))
+    latin = sum(1 for ch in t if ch.isascii() and ch.isalpha())
+    return np.array([1.0, letters - latin, latin, len(PAUSE.findall(t)), sum(ch.isdigit() for ch in t)])
+
+
+def expected_frames(prior: dict, lang: str, nchar: int) -> tuple[float, float]:
+    p = prior.get(lang, PRIOR_DEFAULT)
+    return p["a"] + p["b"] * nchar, p["resid_p99"]
+
+
+def min_frames(ef: float) -> int:
+    return max(4, int(0.5 * ef))
+
+
+def expected_seconds(text: str, lang: str, guard: dict) -> float:
+    coefs = guard["dur_prior_coefs"]
+    coef = np.asarray(coefs.get(lang, guard["dur_prior_default"]), dtype=np.float64)
+    return float(max(0.35, text_feats(text) @ coef))
+
+
+def frame_cap(text: str, lang: str, ef: float, er: float, controls: dict) -> int:
+    """Duration cap as in tools/maya_export/golden_light.py: min(loose, max(0.5*ef + 8, duration-prior cap))."""
+    guard = controls["guard"]
+    loose = min(int(guard["hard_cap"]), int(2.0 * (ef + er) + 40))
+    if not guard["enabled"]:
+        return loose
+    dur_cap = int((float(guard["dur_k"]) * expected_seconds(text, lang, guard)
+                   + float(guard["dur_margin_s"])) / FRAME_S)
+    return int(min(loose, max(int(ef * 0.5) + 8, dur_cap)))
+
+
+def split_text(text: str, lang: str, prior: dict, chunk_s: float = 12.0, *, fps: float = 12.5,
+               terminals: str = "।.!?;।॥") -> list[str]:
+    """Split at sentence terminals into pieces of about chunk_s seconds (prior-estimated).
+    Falls back to commas, then whitespace, for run-on text."""
+    if chunk_s <= 0:
+        return [text]
+    target = chunk_s * fps
+    pieces = [p.strip() for p in re.split(rf"(?<=[{re.escape(terminals)}])\s+", text) if p.strip()]
+    out: list[str] = []
+    cur = ""
+    for p in pieces:
+        cand = (cur + " " + p).strip() if cur else p
+        if cur and expected_frames(prior, lang, len(cand))[0] > target:
+            out.append(cur)
+            cur = p
+        else:
+            cur = cand
+    if cur:
+        out.append(cur)
+
+    final: list[str] = []
+    for piece in out:
+        if expected_frames(prior, lang, len(piece))[0] <= target * 1.3:
+            final.append(piece)
+            continue
+        words = re.split(r"(?<=[,،])\s+|\s+", piece)
+        cur = ""
+        for w in words:
+            cand = (cur + " " + w).strip() if cur else w
+            if cur and expected_frames(prior, lang, len(cand))[0] > target:
+                final.append(cur)
+                cur = w
+            else:
+                cur = cand
+        if cur:
+            final.append(cur)
+    return final
+
+
+class StopPolicy:
+    """The "controlled" stop rule of ServeSession.speak: argmax, tau and hazard tests, none before min_frames."""
+
+    def __init__(self, min_frames: int, tau: float = 0.35, hazard: float = 0.6) -> None:
+        self.min_frames = int(min_frames)
+        self.tau = float(tau)
+        self.hazard = float(hazard)
+        self.H = 0.0
+
+    def check(self, t: int, p: np.ndarray, aeos_class: int = 2048) -> str | None:
+        """
+        p: (vocab + 1,) probability distribution over cb0 classes.
+        Returns: 'aeos_argmax', 'aeos_tau', 'aeos_hazard', or None
+        """
+        if t < self.min_frames:
+            return None
+
+        p_aeos = float(p[aeos_class])
+        self.H = 1.0 - (1.0 - self.H) * (1.0 - p_aeos)
+
+        if int(np.argmax(p)) == aeos_class:
+            return "aeos_argmax"
+        if p_aeos >= self.tau:
+            return "aeos_tau"
+        if self.H >= self.hazard:
+            return "aeos_hazard"
+        return None
+
+
+def sample_cb0(p: np.ndarray, noise: np.ndarray, aeos_class: int = 2048) -> int:
+    """
+    aeos zeroed and renormalised, then Gumbel-max over the rest (same distribution as upstream multinomial).
+    """
+    q = np.array(p, dtype=np.float32, copy=True)
+    q[aeos_class] = 0.0
+    s = np.sum(q)
+    if s > 0:
+        q = q / s
+    else:
+        # Uniform over classes except aeos
+        q = np.ones_like(q)
+        q[aeos_class] = 0.0
+        q = q / np.sum(q)
+
+    # Gumbel-max on probabilities q: argmax(log(q) + noise)
+    # where q == 0, log(q) = -inf
+    with np.errstate(divide="ignore"):
+        log_q = np.log(q)
+    perturbed = log_q + noise
+    perturbed[aeos_class] = -np.inf
+    return int(np.argmax(perturbed))

src/maya_fast/depth.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+from pathlib import Path
+import numpy as np
+import onnxruntime as ort
+
+from .ort_opts import session_options
+
+
+class Depth:
+    def __init__(self, path: str | Path, threads: int | None = None) -> None:
+        self.path = Path(path)
+        self.opts = session_options(threads)
+        self.session = ort.InferenceSession(str(self.path), self.opts, providers=["CPUExecutionProvider"])
+
+    def frame(
+        self,
+        h: np.ndarray,
+        cb0: int | np.integer | np.ndarray,
+        mem: np.ndarray,
+        noise: np.ndarray,
+    ) -> np.ndarray:
+        """
+        Inputs:
+            h: (d_model,) or (1, d_model) float32
+            cb0: int or (1,) int64
+            mem: (1, R, d_mem) float32
+            noise: (15, vocab) float32
+        Returns:
+            codes: (16,) int64
+        """
+        if h.ndim == 1:
+            h = h[None, :]
+        h = h.astype(np.float32)
+
+        if isinstance(cb0, (int, np.integer)):
+            cb0 = np.array([cb0], dtype=np.int64)
+        else:
+            cb0 = np.asarray(cb0, dtype=np.int64)
+            if cb0.ndim == 0:
+                cb0 = cb0.reshape(1)
+
+        mem = np.asarray(mem, dtype=np.float32)
+        if mem.ndim == 2:
+            mem = mem[None, :, :]
+
+        noise = np.asarray(noise, dtype=np.float32)
+        if noise.ndim == 2:  # (15, n_codes) -> (1, 15, n_codes)
+            noise = noise[None]
+
+        feed = {
+            "h": h,
+            "cb0": cb0,
+            "mem": mem,
+            "noise": noise,
+        }
+
+        outputs = self.session.run(None, feed)
+        codes = outputs[0]
+        if codes.ndim > 1:
+            codes = codes.squeeze()
+        return codes.astype(np.int64)

src/maya_fast/embed.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+from pathlib import Path
+import numpy as np
+
+
+class EmbeddingTable:
+    """
+    bf16 memory map embedding lookup.
+    Lookup(ids) returns float32 (len(ids), dim).
+    Exact conversion: (memmap[ids].astype(np.uint32) << 16).view(np.float32)
+    """
+    def __init__(self, path: str | Path, rows: int, dim: int) -> None:
+        self.path = Path(path)
+        self.rows = int(rows)
+        self.dim = int(dim)
+        # raw little-endian uint16
+        self.mmap = np.memmap(self.path, dtype="<u2", mode="r", shape=(self.rows, self.dim))
+
+    def lookup(self, ids: np.ndarray | list[int]) -> np.ndarray:
+        ids_arr = np.asarray(ids, dtype=np.int64)
+        bf16_rows = self.mmap[ids_arr]  # (..., dim) uint16
+        u32 = bf16_rows.astype(np.uint32) << 16
+        fp32 = u32.view(np.float32)
+        return np.ascontiguousarray(fp32)

src/maya_fast/engine.py

@@ -0,0 +1,236 @@
+from __future__ import annotations
+
+import re
+import time
+from pathlib import Path
+from typing import Any
+import numpy as np
+import psutil
+
+from .assets import PromptAssets, default_assets_dir
+from .backbone import Backbone
+from .depth import Depth
+from .codec import MimiDecoder
+from .prompt import AudioIn, build_prompt
+from .sampling import gumbel
+from .control import (
+    expected_frames,
+    frame_cap,
+    min_frames,
+    split_text,
+    StopPolicy,
+    sample_cb0,
+)
+
+# Speakable text check: must contain at least one alphanumeric character
+ALPHANUM_PATTERN = re.compile(r"[\w\d]", re.UNICODE)
+
+
+def _softmax_topk(logits: np.ndarray, temp: float = 0.7, top_k: int = 64) -> np.ndarray:
+    logits = np.asarray(logits, dtype=np.float32)
+    k = min(top_k, logits.shape[-1])
+    topk_idx = np.argpartition(logits, -k)[-k:]
+    masked = np.full_like(logits, -np.inf)
+    masked[topk_idx] = logits[topk_idx] / temp
+    max_val = np.max(masked[topk_idx])
+    exp_vals = np.exp(masked - max_val)
+    exp_vals[masked == -np.inf] = 0.0
+    sum_exp = np.sum(exp_vals)
+    if sum_exp > 0:
+        return exp_vals / sum_exp
+    probs = np.zeros_like(logits)
+    probs[topk_idx] = 1.0 / len(topk_idx)
+    return probs
+
+
+class MayaFast:
+    def __init__(
+        self,
+        assets_dir: str | Path | None = None,
+        threads: int | None = None,
+    ) -> None:
+        if assets_dir is None:
+            self.assets_dir = default_assets_dir()
+        else:
+            self.assets_dir = Path(assets_dir)
+
+        self.threads = threads
+        self.assets = PromptAssets(self.assets_dir, verify=True)
+        self.shapes = self.assets.shapes
+        self.controls = self.assets.controls
+        self.sampling_cfg = self.controls.get("sampling", {})
+
+        # Load models
+        backbone_path = self.assets_dir / "backbone_wo8.onnx"
+        self.backbone = Backbone(
+            backbone_path,
+            threads=self.threads,
+            config=self.shapes,
+        )
+
+        depth_path = self.assets_dir / "depth_frame_fp16.onnx"
+        self.depth = Depth(depth_path, threads=self.threads)
+
+        mimi_path = self.assets_dir / "mimi_decoder.onnx"
+        self.mimi = MimiDecoder(mimi_path, threads=self.threads)
+
+        self.audio_in = self.assets.audio_in
+
+    @property
+    def voices(self) -> list[str]:
+        return sorted(list(self.assets.voices_info.keys()))
+
+    def synthesize(
+        self,
+        text: str,
+        voice: str = "Charvi",
+        seed: int = 0,
+        chunk_s: float | None = None,
+        lang: str = "mr",
+    ) -> tuple[np.ndarray, dict[str, Any]]:
+        """
+        Synthesize text into 24 kHz float32 audio.
+        Returns:
+            (wav, metrics)
+        """
+        # Review focus #2: Empty or punctuation-only text: expected ValueError("no speakable text")
+        if not text or not ALPHANUM_PATTERN.search(text):
+            raise ValueError("no speakable text")
+
+        if voice not in self.assets.voices_info:
+            raise KeyError(f"Unknown voice {voice!r}. Available: {self.voices}")
+
+        t_wall_start = time.perf_counter()
+        prior = self.controls.get("text_frames_prior", {})
+        sampling_cfg = self.sampling_cfg
+        q_serve = int(sampling_cfg.get("q_serve", 16))
+        temp = float(sampling_cfg.get("temp", 0.7))
+        top_k = int(sampling_cfg.get("top_k", 64))
+        tau = float(sampling_cfg.get("tau", 0.35))
+        hazard = float(sampling_cfg.get("hazard", 0.6))
+        aeos_class = int(self.shapes.get("n_codes", 2048))
+
+        # Split text into chunks
+        c_s = 12.0 if chunk_s is None else chunk_s
+        chunks = split_text(text, lang, prior, chunk_s=c_s)
+
+        rng = np.random.default_rng(seed)
+
+        # Speaker memory for Depth: named voices use a null memory (use_mem False)
+        d_mem = int(self.shapes.get("d_mem", 1024))
+        mem = np.zeros((1, 1, d_mem), dtype=np.float32)
+
+        gen_codes: list[np.ndarray] = []
+        reasons: list[str] = []
+        caps: list[int] = []
+        max_past_len = 0
+        max_chunk_bound = 0
+
+        backbone_time = 0.0
+        depth_time = 0.0
+        mimi_time = 0.0
+
+        for chunk_text in chunks:
+            # Review focus #1: Each piece a fresh KV session
+            self.backbone.reset()
+
+            prompt = build_prompt(chunk_text, voice, self.assets, lang=lang)
+
+            ef, er = expected_frames(prior, lang, len(chunk_text))
+            cap = frame_cap(chunk_text, lang, ef, er, self.controls)
+            min_f = min_frames(ef)
+            caps.append(cap)
+            chunk_bound = prompt.n_prompt + cap
+            if chunk_bound > max_chunk_bound:
+                max_chunk_bound = chunk_bound
+
+            # Initial backbone step on prompt
+            t0 = time.perf_counter()
+            h, cb0_logits = self.backbone.step(prompt.embeds)
+            backbone_time += time.perf_counter() - t0
+
+            if self.backbone.past_len > max_past_len:
+                max_past_len = self.backbone.past_len
+
+            stop_policy = StopPolicy(min_frames=min_f, tau=tau, hazard=hazard)
+            chunk_frames: list[np.ndarray] = []
+            reason = "cap"
+
+            for t in range(cap):
+                # 1. Softmax with top_k mask over cb0_logits
+                p = _softmax_topk(cb0_logits, temp=temp, top_k=top_k)
+
+                # 2. StopPolicy check
+                stop_reason = stop_policy.check(t, p, aeos_class=aeos_class)
+                if stop_reason is not None:
+                    reason = stop_reason
+                    break
+
+                # 3. Sample cb0 (aeos zeroed, Gumbel-max over the rest)
+                cb0_noise = gumbel(rng, (cb0_logits.shape[-1],))
+                cb0 = sample_cb0(p, cb0_noise, aeos_class=aeos_class)
+
+                # 4. Depth frame
+                depth_noise = gumbel(rng, (15, int(self.shapes.get("n_codes", 2048))))
+                t_d = time.perf_counter()
+                frame_codes = self.depth.frame(h, cb0, mem, depth_noise)
+                depth_time += time.perf_counter() - t_d
+
+                chunk_frames.append(frame_codes)
+
+                # 5. AudioIn feedback
+                fb = self.audio_in.frames(frame_codes[None, :], q_use=q_serve)
+
+                # 6. Next backbone step
+                t_b = time.perf_counter()
+                h, cb0_logits = self.backbone.step(fb)
+                backbone_time += time.perf_counter() - t_b
+
+                if self.backbone.past_len > max_past_len:
+                    max_past_len = self.backbone.past_len
+
+            reasons.append(reason)
+            # Note on upstream aeos step:
+            # Upstream does one self._step(audio_in.aeos(1)) after the loop if reason.startswith("aeos").
+            # However, because each chunk resets the KV session and no further steps read from this KV session,
+            # running the final aeos step changes nothing returned and wastes compute. We skip it here.
+
+            if chunk_frames:
+                # chunk_frames: list of (16,)
+                chunk_codes_arr = np.stack(chunk_frames, axis=1)  # (16, F_chunk)
+                gen_codes.append(chunk_codes_arr)
+
+        if gen_codes:
+            all_codes = np.concatenate(gen_codes, axis=1)  # (16, Total_F)
+            total_frames = all_codes.shape[1]
+            t_m = time.perf_counter()
+            wav = self.mimi.decode(all_codes)
+            mimi_time += time.perf_counter() - t_m
+        else:
+            total_frames = 0
+            wav = np.zeros(0, dtype=np.float32)
+
+        wall_s = time.perf_counter() - t_wall_start
+        audio_s = len(wav) / 24000.0
+        rtf = wall_s / max(1e-6, audio_s)
+
+        # Peak RSS in MB
+        process = psutil.Process()
+        peak_rss_mb = process.memory_info().rss / (1024 * 1024)
+
+        metrics = {
+            "audio_s": audio_s,
+            "wall_s": wall_s,
+            "rtf": rtf,
+            "frames": total_frames,
+            "reasons": reasons,
+            "caps": caps,
+            "backbone_s": backbone_time,
+            "depth_s": depth_time,
+            "mimi_s": mimi_time,
+            "peak_rss_mb": peak_rss_mb,
+            "max_past_len": max_past_len,
+            "max_chunk_bound": max_chunk_bound,
+        }
+
+        return wav, metrics

src/maya_fast/ort_opts.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+import os
+import onnxruntime as ort
+
+
+def session_options(threads: int | None = None) -> ort.SessionOptions:
+    opts = ort.SessionOptions()
+    if threads is None:
+        cpu_count = os.cpu_count() or 1
+        threads = min(8, cpu_count)
+    opts.intra_op_num_threads = max(1, threads)
+    opts.inter_op_num_threads = 1
+    # Global Constraints:
+    # intra_op_num_threads = requested threads (default min(8, physical cores)),
+    # session.intra_op.allow_spinning=0, enable_cpu_mem_arena=False, ORT_ENABLE_ALL
+    opts.add_session_config_entry("session.intra_op.allow_spinning", "0")
+    opts.enable_cpu_mem_arena = False
+    opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+    return opts

src/maya_fast/prompt.py

@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+import dataclasses
+import math
+
+import numpy as np
+
+from .assets import PromptAssets
+from .embed import EmbeddingTable
+
+# Port of the inference path for named voices (maya3.prompt.NamedSpeakerBuilder with V5Builder.example, the
+# packer's _span and collate_v2 truncation, TTSModelV2.compose_inputs and SpeakerAnchor in eval mode).
+# Checked token for token and value for value against tools/maya_export/golden_light.py by the golden tests.
+
+RMS_EPS = float(np.finfo(np.float32).eps)  # torch.nn.RMSNorm default eps for float32
+_erf = np.frompyfunc(math.erf, 1, 1)
+
+
+def _rms_norm(x: np.ndarray, weight: np.ndarray) -> np.ndarray:
+    x = np.asarray(x, dtype=np.float32)
+    ms = np.mean(x * x, axis=-1, keepdims=True)
+    return (x * (1.0 / np.sqrt(ms + RMS_EPS)) * weight).astype(np.float32)
+
+
+def _gelu(x: np.ndarray) -> np.ndarray:
+    # exact (erf) GELU, torch's default
+    return (0.5 * x * (1.0 + _erf(x / math.sqrt(2.0)).astype(np.float32))).astype(np.float32)
+
+
+class AudioIn:
+    """Audio code frames -> backbone input (port of tts/model/audio_io.py AudioInputEmbed, eval mode)."""
+
+    def __init__(self, table: EmbeddingTable, aeos_in: np.ndarray, norm_w: np.ndarray, n_q: int = 16) -> None:
+        self.table = table
+        self.aeos_in = np.asarray(aeos_in, dtype=np.float32)
+        self.norm_w = np.asarray(norm_w, dtype=np.float32)
+        self.n_q = int(n_q)
+        if table.rows % self.n_q:
+            raise ValueError("audio table rows must be a multiple of n_q")
+        self.n_codes = table.rows // self.n_q
+
+    def frames(self, codes: np.ndarray, q_use: int = 16) -> np.ndarray:
+        """
+        codes: (N, 16) int64 in [0, n_codes)
+        Returns: (N, d_model) float32
+        """
+        codes = np.asarray(codes, dtype=np.int64)
+        if np.any((codes[:, :q_use] < 0) | (codes[:, :q_use] >= self.n_codes)):
+            raise ValueError("audio code outside [0, n_codes) reached the input table")
+        lvl = np.arange(q_use, dtype=np.int64)[None, :]
+        e = self.table.lookup(codes[:, :q_use] + lvl * self.n_codes)  # (N, q_use, d)
+        return _rms_norm(e.sum(axis=1) / math.sqrt(q_use), self.norm_w)
+
+    def aeos(self) -> np.ndarray:
+        """Returns: (d_model,) float32"""
+        return _rms_norm(self.aeos_in[None, :], self.norm_w)[0]
+
+
+def speaker_anchor_null(p: dict[str, np.ndarray]) -> np.ndarray:
+    """SpeakerAnchor for a null speaker row: fc2(gelu(fc1(null_vec))). Returns (d_model,) float32."""
+    h = _gelu(p["anchor_null"] @ p["anchor_fc1_w"].T + p["anchor_fc1_b"])
+    return (h @ p["anchor_fc2_w"].T + p["anchor_fc2_b"]).astype(np.float32)
+
+
+@dataclasses.dataclass
+class Prompt:
+    token_ids: np.ndarray  # (n_prompt,) int64, ends with the <|audio|> sentinel
+    embeds: np.ndarray     # (n_prompt, d_model) float32
+    n_prompt: int
+
+
+def build_prompt(text: str, voice: str, assets: PromptAssets, lang: str = "mr") -> Prompt:
+    """
+    Named voice prompt:  <|lang:xx|> "{Name}: " text <|audio|>
+    The speaker anchor (null for named voices) is added at the <|audio|> sentinel, the last prompt position.
+    """
+    name = _resolve_voice(voice, assets)
+    info = assets.voices_info[name]
+    if info.get("mode") != "named":
+        raise NotImplementedError(f"voice {name!r} uses mode {info.get('mode')!r}; only named voices are supported")
+
+    rendered = "{" + name + "}: " + text.strip()
+    tok = assets.tokenizer
+    lang_id = tok.token_to_id(f"<|lang:{lang}|>")
+    audio_id = tok.token_to_id("<|audio|>")
+    if lang_id is None or audio_id is None:
+        raise KeyError(f"tokenizer lacks <|lang:{lang}|> or <|audio|>")
+
+    ids = [lang_id, *tok.encode(rendered, add_special_tokens=False).ids, audio_id]
+    token_ids = np.asarray(ids, dtype=np.int64)
+    embeds = assets.embed_table.lookup(token_ids)
+    embeds[-1] += speaker_anchor_null(assets.prompt_params)
+    return Prompt(token_ids=token_ids, embeds=embeds, n_prompt=len(token_ids))
+
+
+def _resolve_voice(voice: str, assets: PromptAssets) -> str:
+    if voice in assets.voices_info:
+        return voice
+    low = str(voice).lower()  # voices are case-insensitive
+    hits = [k for k in assets.voices_info if k.lower() == low]
+    if len(hits) == 1:
+        return hits[0]
+    raise KeyError(f"Unknown voice {voice!r}. Available: {sorted(assets.voices_info)}")

src/maya_fast/sampling.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+import numpy as np
+
+
+def gumbel(rng: np.random.Generator, shape: tuple[int, ...]) -> np.ndarray:
+    """
+    Produce float32 Gumbel noise: -log(-log(U)) with U in [1e-10, 1 - 1e-7].
+    """
+    u = rng.uniform(low=1e-10, high=1.0 - 1e-7, size=shape).astype(np.float32)
+    return (-np.log(-np.log(u))).astype(np.float32)
+
+
+def topk_gumbel_argmax(
+    logits: np.ndarray,
+    temp: float = 0.7,
+    top_k: int = 64,
+    noise: np.ndarray | None = None,
+) -> int:
+    """
+    Numpy reference of the graph's rule:
+    argmax(where(rank < top_k, logits / temp + noise, -inf))
+    """
+    logits = np.asarray(logits, dtype=np.float32)
+    if noise is None:
+        noise = np.zeros_like(logits)
+    else:
+        noise = np.asarray(noise, dtype=np.float32)
+
+    k = min(top_k, logits.shape[-1])
+    topk_idx = np.argpartition(logits, -k)[-k:]
+    masked_perturbed = np.full_like(logits, -np.inf)
+    masked_perturbed[topk_idx] = logits[topk_idx] / temp + noise[topk_idx]
+    return int(np.argmax(masked_perturbed))

tests/fixtures/tiny_assets.py

@@ -0,0 +1,286 @@
+from __future__ import annotations
+
+import hashlib
+import json
+from pathlib import Path
+import numpy as np
+import onnx
+from onnx import helper, TensorProto
+
+
+def _sha256(data: bytes) -> str:
+    return hashlib.sha256(data).hexdigest()
+
+
+def _make_tiny_backbone_onnx(out_path: Path, d: int = 16, vocab: int = 64, n_layers: int = 2, n_heads: int = 2, head_dim: int = 8) -> None:
+    # Inputs:
+    # embeds (1, T, 2048) -> here (1, T, d)
+    # cos (1, T, head_dim) -> (1, T, 8)
+    # sin (1, T, head_dim) -> (1, T, 8)
+    # mask (1, 1, T, P+T)
+    # past_k_{i}, past_v_{i} ((1, 8, P, 128)) -> ((1, n_heads, P, head_dim))
+    # Outputs:
+    # hidden (1, T, 2048) -> (1, T, d)
+    # cb0_logits (1, T, 2049) -> (1, T, vocab + 1)
+    # present_k_{i}, present_v_{i}
+
+    inputs = [
+        helper.make_tensor_value_info("embeds", TensorProto.FLOAT, [1, "T", d]),
+        helper.make_tensor_value_info("cos", TensorProto.FLOAT, [1, "T", head_dim]),
+        helper.make_tensor_value_info("sin", TensorProto.FLOAT, [1, "T", head_dim]),
+        helper.make_tensor_value_info("mask", TensorProto.FLOAT, [1, 1, "T", "total_len"]),
+    ]
+    for i in range(n_layers):
+        inputs.append(helper.make_tensor_value_info(f"past_k_{i}", TensorProto.FLOAT, [1, n_heads, "P", head_dim]))
+        inputs.append(helper.make_tensor_value_info(f"past_v_{i}", TensorProto.FLOAT, [1, n_heads, "P", head_dim]))
+
+    outputs = [
+        helper.make_tensor_value_info("hidden", TensorProto.FLOAT, [1, "T", d]),
+        helper.make_tensor_value_info("cb0_logits", TensorProto.FLOAT, [1, "T", vocab + 1]),
+    ]
+    for i in range(n_layers):
+        outputs.append(helper.make_tensor_value_info(f"present_k_{i}", TensorProto.FLOAT, [1, n_heads, "total_len", head_dim]))
+        outputs.append(helper.make_tensor_value_info(f"present_v_{i}", TensorProto.FLOAT, [1, n_heads, "total_len", head_dim]))
+
+    # For tiny mock model:
+    # hidden = Identity(embeds)
+    # cb0_logits = MatMul(hidden, W_cb0) + B_cb0
+    # present_k_i = Concat(past_k_i, k_new, axis=2)
+    # We construct a simple valid ONNX graph that outputs the right shapes.
+    nodes = []
+    initializers = []
+
+    # W_cb0: (d, vocab + 1)
+    np.random.seed(42)
+    w_cb0_val = np.random.randn(d, vocab + 1).astype(np.float32) * 0.1
+    # ensure aeos (index vocab) has sensible value
+    initializers.append(helper.make_tensor("W_cb0", TensorProto.FLOAT, [d, vocab + 1], w_cb0_val.flatten()))
+
+    nodes.append(helper.make_node("Identity", ["embeds"], ["hidden"]))
+    nodes.append(helper.make_node("MatMul", ["hidden", "W_cb0"], ["cb0_logits"]))
+
+    # For past_k, past_v: generate dummy k_new from embeds:
+    # embeds (1, T, d) -> slice or reshape to (1, n_heads, T, head_dim)
+    # In our case d = 16 = 2 * 8 = n_heads * head_dim!
+    # So we can reshape embeds -> (1, T, n_heads, head_dim) then transpose to (1, n_heads, T, head_dim)
+    initializers.append(helper.make_tensor("shape_k", TensorProto.INT64, [4], [1, -1, n_heads, head_dim]))
+    nodes.append(helper.make_node("Reshape", ["embeds", "shape_k"], ["reshaped_kv"]))
+    nodes.append(helper.make_node("Transpose", ["reshaped_kv"], ["k_new"], perm=[0, 2, 1, 3]))
+
+    for i in range(n_layers):
+        nodes.append(helper.make_node("Concat", [f"past_k_{i}", "k_new"], [f"present_k_{i}"], axis=2))
+        nodes.append(helper.make_node("Concat", [f"past_v_{i}", "k_new"], [f"present_v_{i}"], axis=2))
+
+    graph = helper.make_graph(nodes, "tiny_backbone", inputs, outputs, initializers)
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)], ir_version=10)
+    onnx.save(model, str(out_path))
+
+
+def _make_tiny_depth_onnx(out_path: Path, d: int = 16, vocab: int = 64, d_mem: int = 16) -> None:
+    # DepthFrame (batched, as the real sampled graph):
+    # Inputs:
+    # h (B, 2048) -> (B, d)
+    # cb0 (B,) int64
+    # mem (B, R, 1024) -> (B, R, d_mem)
+    # noise (B, 15, 2048) float32 -> (B, 15, vocab)
+    # Output:
+    # codes (B, 16) int64
+
+    inputs = [
+        helper.make_tensor_value_info("h", TensorProto.FLOAT, ["B", d]),
+        helper.make_tensor_value_info("cb0", TensorProto.INT64, ["B"]),
+        helper.make_tensor_value_info("mem", TensorProto.FLOAT, ["B", "R", d_mem]),
+        helper.make_tensor_value_info("noise", TensorProto.FLOAT, ["B", 15, vocab]),
+    ]
+    outputs = [
+        helper.make_tensor_value_info("codes", TensorProto.INT64, ["B", 16]),
+    ]
+
+    # For tiny test: codes[:, 0] = cb0, codes[:, 1..15] = argmax(noise[:, i]) over the vocab.
+    initializers = [helper.make_tensor("unsq_axes", TensorProto.INT64, [1], [1])]
+    nodes = [
+        helper.make_node("ArgMax", ["noise"], ["deeper_codes"], axis=-1, keepdims=0),
+        helper.make_node("Unsqueeze", ["cb0", "unsq_axes"], ["cb0_col"]),
+        helper.make_node("Concat", ["cb0_col", "deeper_codes"], ["codes"], axis=1),
+    ]
+
+    graph = helper.make_graph(nodes, "tiny_depth", inputs, outputs, initializers)
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)], ir_version=10)
+    onnx.save(model, str(out_path))
+
+
+def _make_tiny_mimi_onnx(out_path: Path) -> None:
+    # MimiDecoder:
+    # Input: codes (1, 16, F) int64
+    # Output: wav (1, 1, F * 1920) float32
+    inputs = [
+        helper.make_tensor_value_info("codes", TensorProto.INT64, [1, 16, "F"]),
+    ]
+    outputs = [
+        helper.make_tensor_value_info("wav", TensorProto.FLOAT, [1, 1, "samples"]),
+    ]
+
+    # In ONNX helper:
+    # Shape of codes -> shape (3,)
+    # Gather index 2 -> F
+    # Expand 1920 samples per frame.
+    # We can do: Cast(codes, FLOAT), ReduceMean along axis 1 -> (1, 1, F)
+    # Then Tile by [1, 1, 1920] -> (1, 1, F * 1920)
+    initializers = [
+        helper.make_tensor("tile_repeats", TensorProto.INT64, [3], [1, 1, 1920]),
+    ]
+    nodes = [
+        helper.make_node("Cast", ["codes"], ["codes_float"], to=TensorProto.FLOAT),
+        helper.make_node("ReduceMean", ["codes_float"], ["mean_frames"], axes=[1], keepdims=1),
+        helper.make_node("Tile", ["mean_frames", "tile_repeats"], ["wav"]),
+    ]
+
+    graph = helper.make_graph(nodes, "tiny_mimi", inputs, outputs, initializers)
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)], ir_version=10)
+    onnx.save(model, str(out_path))
+
+
+def make_tiny_assets(assets_dir: Path, aeos_bias: float = 0.0) -> Path:
+    assets_dir = Path(assets_dir)
+    assets_dir.mkdir(parents=True, exist_ok=True)
+
+    d = 16
+    vocab = 64
+    n_layers = 2
+    n_heads = 2
+    head_dim = 8
+    d_mem = 16
+    n_q = 16
+    n_codes = vocab
+    ecapa_dim = 16
+
+    files_sha = {}
+
+    # 1. embed_bf16.bin
+    np.random.seed(123)
+    embeds_fp32 = np.random.randn(vocab, d).astype(np.float32)
+    # Convert fp32 to bf16 uint16
+    embeds_bf16 = (embeds_fp32.view(np.uint32) >> 16).astype(np.uint16)
+    embed_bytes = embeds_bf16.tobytes()
+    (assets_dir / "embed_bf16.bin").write_bytes(embed_bytes)
+    files_sha["embed_bf16.bin"] = _sha256(embed_bytes)
+
+    # 2. prompt.npz
+    # anchor_fc1_w, anchor_fc1_b, anchor_fc2_w, anchor_fc2_b, anchor_null, aeos, plus audio_in/<name>
+    # Note: audio_in parameters: table (n_q * n_codes, d), aeos_in (d,), norm.weight (d,)
+    prompt_dict = {
+        "anchor_fc1_w": np.random.randn(32, ecapa_dim).astype(np.float32),
+        "anchor_fc1_b": np.zeros(32, dtype=np.float32),
+        "anchor_fc2_w": np.zeros((d, 32), dtype=np.float32),  # zero init as upstream
+        "anchor_fc2_b": np.zeros(d, dtype=np.float32),
+        "anchor_null": np.random.randn(ecapa_dim).astype(np.float32),
+        "aeos": np.random.randn(d).astype(np.float32),
+        "audio_in/aeos_in": np.random.randn(d).astype(np.float32),
+        "audio_in/norm": np.ones(d, dtype=np.float32),
+    }
+    np.savez(assets_dir / "prompt.npz", **prompt_dict)
+    files_sha["prompt.npz"] = _sha256((assets_dir / "prompt.npz").read_bytes())
+
+    # audio input table, bf16 raw uint16 row-major (same layout as the real export)
+    audio_f32 = np.random.randn(n_q * n_codes, d).astype(np.float32)
+    audio_bf16 = (audio_f32.view(np.uint32) >> 16).astype("<u2")
+    (assets_dir / "audio_in_bf16.bin").write_bytes(audio_bf16.tobytes())
+    files_sha["audio_in_bf16.bin"] = _sha256((assets_dir / "audio_in_bf16.bin").read_bytes())
+
+    # 3. voices.json + voices.npz
+    voices_json = {
+        "Charvi": {"mode": "named", "use_icl": False, "use_mem": False}
+    }
+    (assets_dir / "voices.json").write_text(json.dumps(voices_json, indent=2))
+    files_sha["voices.json"] = _sha256((assets_dir / "voices.json").read_bytes())
+
+
+    # 4. controls.json
+    controls_json = {
+        "text_frames_prior": {
+            "mr": {"a": 4.0, "b": 0.5, "resid_p99": 10.0}
+        },
+        "text": {"terminals": "।.!?;।॥", "fps": 12.5, "frame_s": 0.08, "chunk_s": 16.0},
+        "guard": {
+            "enabled": True,
+            "dur_k": 2.2,
+            "dur_margin_s": 1.2,
+            "hard_cap": 1500,
+            "dur_prior_coefs": {"mr": [1.0, 0.05, 0.03, 0.3, 0.3]},
+            "dur_prior_default": [1.0, 0.05, 0.03, 0.3, 0.3]
+        },
+        "sampling": {
+            "q_serve": 16,
+            "temp": 0.7,
+            "top_k": 64,
+            "tau": 0.35,
+            "hazard": 0.6
+        }
+    }
+    (assets_dir / "controls.json").write_text(json.dumps(controls_json, indent=2))
+    files_sha["controls.json"] = _sha256((assets_dir / "controls.json").read_bytes())
+
+    # 5. tokenizer.json (minimal mock tokenizer using tokenizers library)
+    from tokenizers import Tokenizer, models, pre_tokenizers
+    vocab_dict = {f"t_{i}": i for i in range(vocab - 5)}
+    vocab_dict["<|lang:mr|>"] = vocab - 5
+    vocab_dict["<|audio|>"] = vocab - 4
+    vocab_dict["<|aeos|>"] = vocab - 3
+    vocab_dict["[UNK]"] = vocab - 2
+    vocab_dict[" "] = vocab - 1
+    tok = Tokenizer(models.WordLevel(vocab_dict, unk_token="[UNK]"))
+    tok.pre_tokenizer = pre_tokenizers.WhitespaceSplit()
+    tok.save(str(assets_dir / "tokenizer.json"))
+    files_sha["tokenizer.json"] = _sha256((assets_dir / "tokenizer.json").read_bytes())
+
+    # 6. ONNX models
+    _make_tiny_backbone_onnx(assets_dir / "backbone_wo8.onnx", d=d, vocab=vocab, n_layers=n_layers, n_heads=n_heads, head_dim=head_dim)
+    # If aeos_bias is requested, we can adjust W_cb0 in backbone or create an alternate backbone
+    if aeos_bias != 0.0:
+        # Load and modify W_cb0 last column
+        m = onnx.load(str(assets_dir / "backbone_wo8.onnx"))
+        for init in m.graph.initializer:
+            if init.name == "W_cb0":
+                arr = onnx.numpy_helper.to_array(init)
+                # bias is applied to the output logits, or we add bias node.
+                # In tiny_backbone, cb0_logits = hidden @ W_cb0. If we add a large negative bias to index vocab (aeos):
+                # Let's add an explicit bias tensor:
+                bias_arr = np.zeros((1, 1, vocab + 1), dtype=np.float32)
+                bias_arr[0, 0, vocab] = aeos_bias
+                bias_init = helper.make_tensor("bias_cb0", TensorProto.FLOAT, [1, 1, vocab + 1], bias_arr.flatten())
+                m.graph.initializer.append(bias_init)
+                # change cb0_logits node
+                for node in m.graph.node:
+                    if node.name == "" and node.op_type == "MatMul":
+                        node.output[0] = "cb0_logits_pre"
+                m.graph.node.append(helper.make_node("Add", ["cb0_logits_pre", "bias_cb0"], ["cb0_logits"]))
+                break
+        onnx.save(m, str(assets_dir / "backbone_wo8.onnx"))
+
+    files_sha["backbone_wo8.onnx"] = _sha256((assets_dir / "backbone_wo8.onnx").read_bytes())
+
+    _make_tiny_depth_onnx(assets_dir / "depth_frame_fp16.onnx", d=d, vocab=vocab, d_mem=d_mem)
+    files_sha["depth_frame_fp16.onnx"] = _sha256((assets_dir / "depth_frame_fp16.onnx").read_bytes())
+
+    _make_tiny_mimi_onnx(assets_dir / "mimi_decoder.onnx")
+    files_sha["mimi_decoder.onnx"] = _sha256((assets_dir / "mimi_decoder.onnx").read_bytes())
+
+    # 7. manifest.json
+    manifest = {
+        "version": 1,
+        "files": files_sha,
+        "shapes": {
+            "d_model": d,
+            "vocab": vocab,
+            "n_q": n_q,
+            "n_codes": n_codes,
+            "n_layers": n_layers,
+            "n_heads": n_heads,
+            "head_dim": head_dim,
+            "d_mem": d_mem,
+            "rope_theta": 10000.0,
+        }
+    }
+    (assets_dir / "manifest.json").write_text(json.dumps(manifest, indent=2))
+
+    return assets_dir

tests/test_audio_utils.py

@@ -0,0 +1,76 @@
+import unittest
+import numpy as np
+import os
+import sys
+import tempfile
+import soundfile as sf
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from src.audio_utils import resample, load_and_resample, condition_audio
+
+
+def tone(freq, sr, seconds=1.0, amp=0.5):
+    t = np.arange(int(sr * seconds)) / sr
+    return (amp * np.sin(2 * np.pi * freq * t)).astype(np.float32)
+
+
+class TestResample(unittest.TestCase):
+    def test_removes_content_above_nyquist(self):
+        # 10 kHz is above the 8 kHz Nyquist of 16 kHz output; linear interp aliases it to 6 kHz
+        out = resample(tone(10000, 48000), 48000, 16000)
+        self.assertLess(np.sqrt(np.mean(out[1000:-1000] ** 2)), 0.01)
+
+    def test_preserves_speech_band_tone(self):
+        out = resample(tone(1000, 44100), 44100, 16000)
+        self.assertEqual(len(out), 16000)
+        self.assertAlmostEqual(np.sqrt(np.mean(out[1000:-1000] ** 2)), 0.5 / np.sqrt(2), delta=0.02)
+
+    def test_upsample_8k(self):
+        out = resample(tone(1000, 8000), 8000, 16000)
+        self.assertEqual(len(out), 16000)
+        self.assertEqual(out.dtype, np.float32)
+
+    def test_16k_passthrough(self):
+        x = tone(1000, 16000)
+        self.assertTrue(np.array_equal(resample(x, 16000, 16000), x))
+
+    def test_load_and_resample_uses_soxr(self):
+        stereo = np.stack([tone(10000, 48000), tone(10000, 48000)], axis=1)
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
+            path = f.name
+        try:
+            sf.write(path, stereo, 48000)
+            data, duration = load_and_resample(path)
+        finally:
+            os.remove(path)
+        self.assertEqual(data.ndim, 1)
+        self.assertEqual(data.dtype, np.float32)
+        self.assertAlmostEqual(duration, 1.0, places=3)
+        self.assertLess(np.sqrt(np.mean(data[1000:-1000] ** 2)), 0.01)
+
+
+class TestConditionAudio(unittest.TestCase):
+    def test_removes_dc_normalizes_and_pads(self):
+        x = (0.1 * np.sin(2 * np.pi * 440 * np.arange(16000) / 16000) + 0.05).astype(np.float32)
+        out = condition_audio(x)
+        self.assertEqual(len(out), 16000 + 2 * 3200)
+        self.assertTrue(np.all(out[:3200] == 0) and np.all(out[-3200:] == 0))
+        self.assertLess(abs(out[3200:-3200].mean()), 1e-3)
+        self.assertAlmostEqual(np.abs(out).max(), 0.9, delta=1e-3)
+
+    def test_all_zero_stays_zero(self):
+        out = condition_audio(np.zeros(1600, dtype=np.float32))
+        self.assertFalse(np.isnan(out).any())
+        self.assertEqual(np.abs(out).max(), 0.0)
+
+    def test_near_silent_gain_is_capped(self):
+        x = (0.001 * np.sin(2 * np.pi * 440 * np.arange(16000) / 16000)).astype(np.float32)
+        self.assertAlmostEqual(np.abs(condition_audio(x)).max(), 0.01, delta=1e-4)
+
+    def test_empty_input_returns_padding_only(self):
+        self.assertEqual(len(condition_audio(np.zeros(0, dtype=np.float32))), 6400)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_maya_fast_assets.py

@@ -0,0 +1,51 @@
+import unittest
+from pathlib import Path
+import tempfile
+import numpy as np
+
+from tests.fixtures.tiny_assets import make_tiny_assets
+from src.maya_fast.assets import load_manifest, AssetError, PromptAssets
+from src.maya_fast.embed import EmbeddingTable
+
+
+class TestMayaFastAssets(unittest.TestCase):
+    def setUp(self):
+        self.temp_dir = tempfile.TemporaryDirectory()
+        self.assets_dir = make_tiny_assets(Path(self.temp_dir.name))
+
+    def tearDown(self):
+        self.temp_dir.cleanup()
+
+    def test_manifest_sha_mismatch_raises(self):
+        # corrupt a file
+        corrupt_target = self.assets_dir / "embed_bf16.bin"
+        with open(corrupt_target, "ab") as f:
+            f.write(b"\x00\x00\x00\x00")
+
+        with self.assertRaises(AssetError) as ctx:
+            load_manifest(self.assets_dir, verify=True)
+        self.assertIn("embed_bf16.bin", str(ctx.exception))
+
+    def test_manifest_missing_file_raises(self):
+        (self.assets_dir / "prompt.npz").unlink()
+        with self.assertRaises(AssetError) as ctx:
+            load_manifest(self.assets_dir, verify=True)
+        self.assertIn("prompt.npz", str(ctx.exception))
+
+    def test_bf16_lookup_exact(self):
+        # write 3 known float32 values that are exact in bf16, look up -> equal
+        with tempfile.TemporaryDirectory() as td:
+            p = Path(td) / "test_embed.bin"
+            # 1.0, -2.5, 0.125 are exactly representable in bf16
+            vals = np.array([[1.0, -2.5, 0.125], [0.0, 4.0, -0.5]], dtype=np.float32)
+            u32 = vals.view(np.uint32)
+            bf16 = (u32 >> 16).astype(np.uint16)
+            p.write_bytes(bf16.tobytes())
+
+            table = EmbeddingTable(p, rows=2, dim=3)
+            res = table.lookup(np.array([0, 1]))
+            np.testing.assert_array_equal(res, vals)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_maya_fast_backbone.py

@@ -0,0 +1,50 @@
+import unittest
+from pathlib import Path
+import tempfile
+import numpy as np
+
+from tests.fixtures.tiny_assets import make_tiny_assets
+from src.maya_fast.backbone import Backbone
+
+
+class TestMayaFastBackbone(unittest.TestCase):
+    def setUp(self):
+        self.temp_dir = tempfile.TemporaryDirectory()
+        self.assets_dir = make_tiny_assets(Path(self.temp_dir.name))
+
+    def tearDown(self):
+        self.temp_dir.cleanup()
+
+    def test_backbone_step_kv_grows(self):
+        # runtime: tiny synthetic 2-layer graph built with onnx.helper in tests/fixtures/tiny_assets.py;
+        # past_len goes 5 -> 6 -> 7 and reset() returns it to 0
+        backbone_path = self.assets_dir / "backbone_wo8.onnx"
+        bb = Backbone(backbone_path, threads=1, manifest_path=self.assets_dir / "manifest.json")
+        self.assertEqual(bb.past_len, 0)
+
+        # Prefill 5 tokens
+        # d_model = 16 for tiny
+        d = 16
+        prompt_embeds = np.random.randn(5, d).astype(np.float32)
+        h, cb0 = bb.step(prompt_embeds)
+        self.assertEqual(bb.past_len, 5)
+        self.assertEqual(h.shape, (d,))
+        self.assertEqual(cb0.shape, (65,))  # vocab + 1
+
+        # Step 6
+        tok1 = np.random.randn(1, d).astype(np.float32)
+        h, cb0 = bb.step(tok1)
+        self.assertEqual(bb.past_len, 6)
+
+        # Step 7
+        tok2 = np.random.randn(1, d).astype(np.float32)
+        h, cb0 = bb.step(tok2)
+        self.assertEqual(bb.past_len, 7)
+
+        # Reset
+        bb.reset()
+        self.assertEqual(bb.past_len, 0)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_maya_fast_codec.py

@@ -0,0 +1,45 @@
+import unittest
+from pathlib import Path
+import tempfile
+import numpy as np
+
+from tests.fixtures.tiny_assets import make_tiny_assets
+from src.maya_fast.assets import default_assets_dir
+from src.maya_fast.codec import MimiDecoder
+
+
+class TestMayaFastCodec(unittest.TestCase):
+    def setUp(self):
+        self.temp_dir = tempfile.TemporaryDirectory()
+        self.assets_dir = make_tiny_assets(Path(self.temp_dir.name))
+
+    def tearDown(self):
+        self.temp_dir.cleanup()
+
+    def test_codec_length(self):
+        # tiny assets: F = 5 -> 9600 samples
+        codec_path = self.assets_dir / "mimi_decoder.onnx"
+        mimi = MimiDecoder(codec_path, threads=1)
+
+        F = 5
+        codes = np.random.randint(0, 64, size=(16, F), dtype=np.int64)
+        wav = mimi.decode(codes)
+
+        self.assertEqual(wav.shape, (F * 1920,))
+        self.assertEqual(wav.dtype, np.float32)
+
+
+class TestMayaFastCodecReal(unittest.TestCase):
+    def test_real_decoder_length(self):
+        path = default_assets_dir() / "mimi_decoder.onnx"
+        if not path.exists():
+            self.skipTest(f"real Mimi decoder not found at {path}")
+        mimi = MimiDecoder(path, threads=1)
+        codes = np.random.default_rng(0).integers(0, 2048, size=(16, 5), dtype=np.int64)
+        wav = mimi.decode(codes)
+        self.assertEqual(wav.shape, (5 * 1920,))
+        self.assertTrue(np.all(np.isfinite(wav)))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_maya_fast_control.py

@@ -0,0 +1,123 @@
+import json
+import unittest
+from pathlib import Path
+import numpy as np
+
+from src.maya_fast.assets import PromptAssets, default_assets_dir
+from src.maya_fast.control import (
+    expected_frames,
+    frame_cap,
+    min_frames,
+    split_text,
+    StopPolicy,
+    sample_cb0,
+)
+
+
+class TestMayaFastControl(unittest.TestCase):
+    def test_controls_basic(self):
+        prior = {"mr": {"a": 4.0, "b": 0.5, "resid_p99": 10.0}}
+        ef, er = expected_frames(prior, "mr", 20)
+        self.assertEqual(ef, 14.0)
+        self.assertEqual(er, 10.0)
+
+        controls = {
+            "guard": {
+                "enabled": True,
+                "dur_k": 2.2,
+                "dur_margin_s": 1.2,
+                "hard_cap": 1500,
+                "dur_prior_coefs": {"mr": [1.0, 0.05, 0.03, 0.3, 0.3]},
+                "dur_prior_default": [1.0, 0.05, 0.03, 0.3, 0.3],
+            }
+        }
+        cap = frame_cap("नमस्कार", "mr", ef, er, controls)
+        self.assertGreater(cap, 0)
+        self.assertLessEqual(cap, 1500)
+
+        mf = min_frames(ef)
+        self.assertGreaterEqual(mf, 4)
+
+        # test split_text
+        text = "हा पहिला भाग आहे. हा दुसरा भाग आहे."
+        chunks = split_text(text, "mr", prior, chunk_s=1.0)
+        self.assertGreater(len(chunks), 1)
+
+    def test_stop_policy_reasons(self):
+        # test_stop_policy_reasons: hand-built p vectors trigger each reason at the right t; nothing fires before min_frames
+        policy = StopPolicy(min_frames=5, tau=0.35, hazard=0.6)
+        aeos_class = 2048
+        p = np.zeros(2049, dtype=np.float32)
+
+        # Before min_frames, none should fire even if aeos is 1.0
+        p[aeos_class] = 1.0
+        for t in range(5):
+            self.assertIsNone(policy.check(t, p))
+
+        # At t = 5, argmax aeos triggers aeos_argmax
+        self.assertEqual(policy.check(5, p), "aeos_argmax")
+
+        # Test tau trigger when argmax is another token
+        policy = StopPolicy(min_frames=5, tau=0.35, hazard=0.6)
+        p = np.zeros(2049, dtype=np.float32)
+        p[0] = 0.5
+        p[aeos_class] = 0.4  # >= 0.35 but argmax is 0
+        self.assertEqual(policy.check(5, p), "aeos_tau")
+
+        # Test hazard trigger: accumulated hazard 1 - prod(1 - p_aeos) >= 0.6
+        # If p_aeos = 0.2 at each step after min_frames:
+        # t=5: 1 - 0.8 = 0.2
+        # t=6: 1 - 0.8^2 = 0.36
+        # t=7: 1 - 0.8^3 = 0.488
+        # t=8: 1 - 0.8^4 = 0.5904
+        # t=9: 1 - 0.8^5 = 0.67232 >= 0.6 -> aeos_hazard!
+        policy = StopPolicy(min_frames=5, tau=0.35, hazard=0.6)
+        p = np.zeros(2049, dtype=np.float32)
+        p[0] = 0.7
+        p[aeos_class] = 0.2
+        for t in range(5, 9):
+            self.assertIsNone(policy.check(t, p))
+        self.assertEqual(policy.check(9, p), "aeos_hazard")
+
+    def test_sample_cb0(self):
+        # aeos zeroed, Gumbel-max over the rest
+        aeos_class = 64
+        p = np.zeros(65, dtype=np.float32)
+        p[aeos_class] = 0.99
+        p[10] = 0.01
+
+        noise = np.zeros(65, dtype=np.float32)
+        # Even though aeos had 0.99, sample_cb0 masks aeos out, so token 10 must be chosen
+        chosen = sample_cb0(p, noise, aeos_class=aeos_class)
+        self.assertEqual(chosen, 10)
+
+
+GOLDEN = Path.home() / ".cache" / "maya_fast" / "golden"
+
+
+class TestMayaFastControlGolden(unittest.TestCase):
+    """Sentence controls (ef, er, cap, minimum, split_text) must equal the upstream values in light.json."""
+
+    @classmethod
+    def setUpClass(cls):
+        assets_dir = default_assets_dir()
+        if not (GOLDEN / "light.json").exists() or not (assets_dir / "manifest.json").exists():
+            raise unittest.SkipTest(f"goldens at {GOLDEN} or assets at {assets_dir} not found")
+        with open(GOLDEN / "light.json", "r", encoding="utf-8") as f:
+            cls.gj = json.load(f)
+        cls.controls = PromptAssets(assets_dir).controls
+
+    def test_controls_match_golden(self):
+        prior = self.controls["text_frames_prior"]
+        for s_id, item in self.gj["sentences"].items():
+            text = item["text"]
+            ef, er = expected_frames(prior, "mr", len(text))
+            self.assertAlmostEqual(ef, item["ef"], places=9, msg=s_id)
+            self.assertAlmostEqual(er, item["er"], places=9, msg=s_id)
+            self.assertEqual(frame_cap(text, "mr", ef, er, self.controls), item["cap"], s_id)
+            self.assertEqual(min_frames(ef), item["minimum"], s_id)
+            self.assertEqual(split_text(text, "mr", prior, chunk_s=16.0), item["split_text"], s_id)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_maya_fast_depth.py

@@ -0,0 +1,41 @@
+import unittest
+from pathlib import Path
+import tempfile
+import numpy as np
+
+from tests.fixtures.tiny_assets import make_tiny_assets
+from src.maya_fast.depth import Depth
+
+
+class TestMayaFastDepth(unittest.TestCase):
+    def setUp(self):
+        self.temp_dir = tempfile.TemporaryDirectory()
+        self.assets_dir = make_tiny_assets(Path(self.temp_dir.name))
+
+    def tearDown(self):
+        self.temp_dir.cleanup()
+
+    def test_depth_runtime_shapes(self):
+        # tiny assets: frame() returns 16 codes in [0, 2048), codes[0] == cb0
+        depth_path = self.assets_dir / "depth_frame_fp16.onnx"
+        depth = Depth(depth_path, threads=1)
+
+        d = 16
+        vocab = 64
+        d_mem = 16
+
+        h = np.random.randn(d).astype(np.float32)
+        cb0 = 17
+        mem = np.random.randn(1, 1, d_mem).astype(np.float32)
+        noise = np.random.randn(15, vocab).astype(np.float32)
+
+        codes = depth.frame(h, cb0, mem, noise)
+
+        self.assertEqual(codes.shape, (16,))
+        self.assertEqual(codes.dtype, np.int64)
+        self.assertEqual(codes[0], cb0)
+        self.assertTrue(np.all((codes >= 0) & (codes < 2048)))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_maya_fast_engine.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+import ast
+from pathlib import Path
+import tempfile
+import unittest
+import numpy as np
+
+from tests.fixtures.tiny_assets import make_tiny_assets
+from src.maya_fast.engine import MayaFast
+from src.maya_fast.control import split_text
+
+
+class TestMayaFastEngine(unittest.TestCase):
+    def setUp(self):
+        self.temp_dir = tempfile.TemporaryDirectory()
+        self.assets_dir = make_tiny_assets(Path(self.temp_dir.name))
+
+    def tearDown(self):
+        self.temp_dir.cleanup()
+
+    def test_synthesize_returns_audio(self):
+        # test_synthesize_returns_audio (length = frames * 1920)
+        engine = MayaFast(assets_dir=self.assets_dir, threads=1)
+        wav, metrics = engine.synthesize("t_1 t_2", voice="Charvi", seed=42)
+        self.assertEqual(len(wav), metrics["frames"] * 1920)
+        self.assertEqual(wav.dtype, np.float32)
+        self.assertIn("audio_s", metrics)
+        self.assertIn("wall_s", metrics)
+        self.assertIn("rtf", metrics)
+        self.assertIn("reasons", metrics)
+
+    def test_same_seed_same_audio(self):
+        engine = MayaFast(assets_dir=self.assets_dir, threads=1)
+        wav1, _ = engine.synthesize("t_1 t_2", voice="Charvi", seed=123)
+        wav2, _ = engine.synthesize("t_1 t_2", voice="Charvi", seed=123)
+        np.testing.assert_array_equal(wav1, wav2)
+
+    def test_empty_text_raises(self):
+        engine = MayaFast(assets_dir=self.assets_dir, threads=1)
+        with self.assertRaises(ValueError):
+            engine.synthesize("", voice="Charvi")
+        with self.assertRaises(ValueError):
+            engine.synthesize("।,.", voice="Charvi")
+
+    def test_cap_stops_generation(self):
+        # tiny assets built with an aeos logit bias of -100 -> reasons == ["cap"], frames == cap
+        with tempfile.TemporaryDirectory() as td:
+            biased_assets_dir = make_tiny_assets(Path(td), aeos_bias=-100.0)
+            engine = MayaFast(assets_dir=biased_assets_dir, threads=1)
+            wav, metrics = engine.synthesize("t_1", voice="Charvi", seed=0)
+            self.assertEqual(metrics["reasons"], ["cap"])
+
+    def test_long_text_is_split_and_kv_bounded(self):
+        # 3-sentence text -> len(reasons) == len(split_text(...)), and Backbone.past_len never exceeds the largest piece's prompt + cap
+        text = "t_1 t_2. t_1 t_3. t_2 t_3."
+        engine = MayaFast(assets_dir=self.assets_dir, threads=1)
+        prior = engine.controls.get("text_frames_prior", {})
+        expected_chunks = split_text(text, "mr", prior, chunk_s=0.5)
+
+        wav, metrics = engine.synthesize(text, voice="Charvi", seed=0, chunk_s=0.5)
+        self.assertEqual(len(metrics["reasons"]), len(expected_chunks))
+        self.assertLessEqual(metrics["max_past_len"], metrics["max_chunk_bound"])
+
+    def test_thread_count_does_not_change_output(self):
+        # threads 1 and 16 -> identical wav
+        engine1 = MayaFast(assets_dir=self.assets_dir, threads=1)
+        engine2 = MayaFast(assets_dir=self.assets_dir, threads=16)
+        wav1, _ = engine1.synthesize("t_1 t_2", voice="Charvi", seed=7)
+        wav2, _ = engine2.synthesize("t_1 t_2", voice="Charvi", seed=7)
+        np.testing.assert_array_equal(wav1, wav2)
+
+    def test_unknown_voice_lists_voices(self):
+        engine = MayaFast(assets_dir=self.assets_dir, threads=1)
+        with self.assertRaises(KeyError) as ctx:
+            engine.synthesize("t_1", voice="UnknownVoice")
+        self.assertIn("Charvi", str(ctx.exception))
+
+    def test_runtime_has_no_torch(self):
+        # ast scan of every file under src/maya_fast/: no torch import
+        pkg_dir = Path(__file__).resolve().parent.parent / "src" / "maya_fast"
+        for py_file in pkg_dir.glob("**/*.py"):
+            with open(py_file, "r", encoding="utf-8") as f:
+                tree = ast.parse(f.read(), filename=str(py_file))
+            for node in ast.walk(tree):
+                if isinstance(node, ast.Import):
+                    for alias in node.names:
+                        self.assertNotEqual(alias.name.split(".")[0], "torch", f"torch imported in {py_file}")
+                elif isinstance(node, ast.ImportFrom):
+                    if node.module:
+                        self.assertNotEqual(node.module.split(".")[0], "torch", f"torch imported in {py_file}")
+
+
+if __name__ == "__main__":
+    unittest.main()