diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..aba9b11b8a2ca632a2b54e83bb440b048ff787e1 100644 --- a/.gitattributes +++ b/.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 diff --git a/NOTICE b/NOTICE new file mode 100644 index 0000000000000000000000000000000000000000..718c30264b27ec9dffcd1b9085e3e2819b869daa --- /dev/null +++ b/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/ diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..cbaf0e18389d09b239859e50d06bad8c1ae6e88f --- /dev/null +++ b/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} +} +``` diff --git a/assets/v1/audio_in_bf16.bin b/assets/v1/audio_in_bf16.bin new file mode 100644 index 0000000000000000000000000000000000000000..19ee58a290fcbc073e96429aba7824b704d3037d --- /dev/null +++ b/assets/v1/audio_in_bf16.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2670e5583c2d2e6e963a8117136a8d07e0e892d3ffda2620d3f13cdbaddc94d0 +size 268435456 diff --git a/assets/v1/backbone_wo8.data b/assets/v1/backbone_wo8.data new file mode 100644 index 0000000000000000000000000000000000000000..5539441c2e11b87eacea6cdafee174bb7e86f948 --- /dev/null +++ b/assets/v1/backbone_wo8.data @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ce9ebfb1e911bac4e3bc58b77c6657f75b79f145f89df718f23d352d824dcb01 +size 1602727936 diff --git a/assets/v1/backbone_wo8.onnx b/assets/v1/backbone_wo8.onnx new file mode 100644 index 0000000000000000000000000000000000000000..215b6d8daa76fdc5d999a8a324a1b30453b326ca --- /dev/null +++ b/assets/v1/backbone_wo8.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d0326c5973d9d71708d7c39c3f34851c73addc88bceb46d0ac90c37e9041955d +size 1121189 diff --git a/assets/v1/controls.json b/assets/v1/controls.json new file mode 100644 index 0000000000000000000000000000000000000000..c4003982e90902e6a9577f5673de915d564aacb6 --- /dev/null +++ b/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 + } +} \ No newline at end of file diff --git a/assets/v1/depth_frame_fp16.onnx b/assets/v1/depth_frame_fp16.onnx new file mode 100644 index 0000000000000000000000000000000000000000..99b36b6a93465797b4d15bce6954ea2348579673 --- /dev/null +++ b/assets/v1/depth_frame_fp16.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c49a4db17d5f443371c459d75c4bdbd508abc164ba30caf51f70c36032c1cea4 +size 4507646 diff --git a/assets/v1/depth_frame_fp16.onnx.data b/assets/v1/depth_frame_fp16.onnx.data new file mode 100644 index 0000000000000000000000000000000000000000..c9bc40966be009b38d2f7cd197d8d7a21d8a99e0 --- /dev/null +++ b/assets/v1/depth_frame_fp16.onnx.data @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:494c8e441065164eb1aa26100f5c2763b447108a116914036a72e1b124463e61 +size 872850432 diff --git a/assets/v1/embed_bf16.bin b/assets/v1/embed_bf16.bin new file mode 100644 index 0000000000000000000000000000000000000000..96feb2df4e82c973c376b3f74cec2b3e22bc2883 --- /dev/null +++ b/assets/v1/embed_bf16.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7940ea16f0a59b955ac42a9f5c035930e0aded3e034ca68de18531c45a6ff861 +size 796868608 diff --git a/assets/v1/manifest.json b/assets/v1/manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..0336e1c019372a9879f6d9be08ea2c8792d460b7 --- /dev/null +++ b/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 + } +} \ No newline at end of file diff --git a/assets/v1/mimi_decoder.onnx b/assets/v1/mimi_decoder.onnx new file mode 100644 index 0000000000000000000000000000000000000000..252d9745c41063d4c4efd95babfa04fdc99c2b65 --- /dev/null +++ b/assets/v1/mimi_decoder.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a0e6accff030071e49c1e1824b26f3a4868c39d125f11bb8079b4a4bbd9dcb58 +size 194722647 diff --git a/assets/v1/prompt.npz b/assets/v1/prompt.npz new file mode 100644 index 0000000000000000000000000000000000000000..20cb1e278ee5ad9475d5a558232d6d94dc4f76e9 --- /dev/null +++ b/assets/v1/prompt.npz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0d6392e0cba6009a913ba57605641b48b1a1fa32a0040ffd4b7b2f840be8afc5 +size 4616734 diff --git a/assets/v1/tokenizer.json b/assets/v1/tokenizer.json new file mode 100644 index 0000000000000000000000000000000000000000..664e46d978d0001f75a8dfb0228a7916c6fd398f --- /dev/null +++ b/assets/v1/tokenizer.json @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:26c51f9b772f9084ed6bdeabc44b7c15397c12966bfa3e1c09b025e82f4fdbaf +size 16103051 diff --git a/assets/v1/voices.json b/assets/v1/voices.json new file mode 100644 index 0000000000000000000000000000000000000000..ef4fcfdf131e1766ed4dd8c4b4702fb6419c1658 --- /dev/null +++ b/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 + } +} \ No newline at end of file diff --git a/requirements-fast.txt b/requirements-fast.txt new file mode 100644 index 0000000000000000000000000000000000000000..4bfbfbe816975ba0d104d3eb0c3ee763a8a3ad6f --- /dev/null +++ b/requirements-fast.txt @@ -0,0 +1,5 @@ +numpy>=2.0 +onnxruntime==1.31.* +tokenizers>=0.22,<0.23 +psutil +onnx diff --git a/src/maya_fast/__init__.py b/src/maya_fast/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6a7457e4001c2243c22afa13ce51345f3d1ad1be --- /dev/null +++ b/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", +] diff --git a/src/maya_fast/__main__.py b/src/maya_fast/__main__.py new file mode 100644 index 0000000000000000000000000000000000000000..da771d018c5fd719d03debdd83f713324e3bd26c --- /dev/null +++ b/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() diff --git a/src/maya_fast/__pycache__/__init__.cpython-313.pyc b/src/maya_fast/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..62d5d5ee234cfb45ca06bd4e240d4feff5487fd0 Binary files /dev/null and b/src/maya_fast/__pycache__/__init__.cpython-313.pyc differ diff --git a/src/maya_fast/__pycache__/__main__.cpython-313.pyc b/src/maya_fast/__pycache__/__main__.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..63d75bc8a77cf52113cae9cee3c03f950659c298 Binary files /dev/null and b/src/maya_fast/__pycache__/__main__.cpython-313.pyc differ diff --git a/src/maya_fast/__pycache__/assets.cpython-313.pyc b/src/maya_fast/__pycache__/assets.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ac738adfe4e2b4457575e40acec615c9e5ae4627 Binary files /dev/null and b/src/maya_fast/__pycache__/assets.cpython-313.pyc differ diff --git a/src/maya_fast/__pycache__/backbone.cpython-313.pyc b/src/maya_fast/__pycache__/backbone.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f2f900e44187e6864141f63f540f891f6d105da9 Binary files /dev/null and b/src/maya_fast/__pycache__/backbone.cpython-313.pyc differ diff --git a/src/maya_fast/__pycache__/codec.cpython-313.pyc b/src/maya_fast/__pycache__/codec.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ce612a2160cbb174f0d45104ccb64ff06b277654 Binary files /dev/null and b/src/maya_fast/__pycache__/codec.cpython-313.pyc differ diff --git a/src/maya_fast/__pycache__/control.cpython-313.pyc b/src/maya_fast/__pycache__/control.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cd3c64e92e12264adea15075450d6f2c1734a3ea Binary files /dev/null and b/src/maya_fast/__pycache__/control.cpython-313.pyc differ diff --git a/src/maya_fast/__pycache__/depth.cpython-313.pyc b/src/maya_fast/__pycache__/depth.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5c8f89974ca8d585d370a5f99ecae26f108db5e5 Binary files /dev/null and b/src/maya_fast/__pycache__/depth.cpython-313.pyc differ diff --git a/src/maya_fast/__pycache__/embed.cpython-313.pyc b/src/maya_fast/__pycache__/embed.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d946d0e0a6edab39c5dd5c699f07babd151eb907 Binary files /dev/null and b/src/maya_fast/__pycache__/embed.cpython-313.pyc differ diff --git a/src/maya_fast/__pycache__/engine.cpython-313.pyc b/src/maya_fast/__pycache__/engine.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6a3799c4b27b5235cb7e88701afaefdd31f4b073 Binary files /dev/null and b/src/maya_fast/__pycache__/engine.cpython-313.pyc differ diff --git a/src/maya_fast/__pycache__/ort_opts.cpython-313.pyc b/src/maya_fast/__pycache__/ort_opts.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..20cdb93508b43ef8ec76e83af2a9fdfe7c72b04b Binary files /dev/null and b/src/maya_fast/__pycache__/ort_opts.cpython-313.pyc differ diff --git a/src/maya_fast/__pycache__/prompt.cpython-313.pyc b/src/maya_fast/__pycache__/prompt.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0eaee448a3cec81cb07bcd8e71cf541c2ee81317 Binary files /dev/null and b/src/maya_fast/__pycache__/prompt.cpython-313.pyc differ diff --git a/src/maya_fast/__pycache__/sampling.cpython-313.pyc b/src/maya_fast/__pycache__/sampling.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f73ca153ae2f35de5c4fa72326c27fa7c12096ba Binary files /dev/null and b/src/maya_fast/__pycache__/sampling.cpython-313.pyc differ diff --git a/src/maya_fast/assets.py b/src/maya_fast/assets.py new file mode 100644 index 0000000000000000000000000000000000000000..7c072e355947c387fe09978d7838d75168e14f2a --- /dev/null +++ b/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")) diff --git a/src/maya_fast/backbone.py b/src/maya_fast/backbone.py new file mode 100644 index 0000000000000000000000000000000000000000..b55eb1b2861e6a25a4af57a131ff1eb6bba238d6 --- /dev/null +++ b/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 diff --git a/src/maya_fast/codec.py b/src/maya_fast/codec.py new file mode 100644 index 0000000000000000000000000000000000000000..e336ecbc6a241e9dba4318e5fa46d68b0014eaaf --- /dev/null +++ b/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) diff --git a/src/maya_fast/control.py b/src/maya_fast/control.py new file mode 100644 index 0000000000000000000000000000000000000000..bae2118815047c37d54540358d0520a38348872f --- /dev/null +++ b/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)) diff --git a/src/maya_fast/depth.py b/src/maya_fast/depth.py new file mode 100644 index 0000000000000000000000000000000000000000..6a447adc4783a30f54b7f275b2d55ccb210ad89c --- /dev/null +++ b/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) diff --git a/src/maya_fast/embed.py b/src/maya_fast/embed.py new file mode 100644 index 0000000000000000000000000000000000000000..2daa3478ec222c9a2f5702d51248380b0c3cc70b --- /dev/null +++ b/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=" 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) diff --git a/src/maya_fast/engine.py b/src/maya_fast/engine.py new file mode 100644 index 0000000000000000000000000000000000000000..0948eb3b0ed56b18ce2b356cb449e396d4a2a87a --- /dev/null +++ b/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 diff --git a/src/maya_fast/ort_opts.py b/src/maya_fast/ort_opts.py new file mode 100644 index 0000000000000000000000000000000000000000..8267b5bb33f2d47d53d39782137a026ead5866f5 --- /dev/null +++ b/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 diff --git a/src/maya_fast/prompt.py b/src/maya_fast/prompt.py new file mode 100644 index 0000000000000000000000000000000000000000..d17bc100181d110f999f389b90ded3b479e7ec56 --- /dev/null +++ b/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)}") diff --git a/src/maya_fast/sampling.py b/src/maya_fast/sampling.py new file mode 100644 index 0000000000000000000000000000000000000000..b224bac9dedf36b861e3e0608f0dd78c2f007f54 --- /dev/null +++ b/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)) diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/tests/fixtures/__init__.py b/tests/fixtures/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/tests/fixtures/tiny_assets.py b/tests/fixtures/tiny_assets.py new file mode 100644 index 0000000000000000000000000000000000000000..29bfe5bb854504b089d21aadcaa777183407ce5f --- /dev/null +++ b/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/ + # 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(""] = 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 diff --git a/tests/test_audio_utils.py b/tests/test_audio_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..4ae708af5ee87526d22eae1c7448d77f8da28396 --- /dev/null +++ b/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() diff --git a/tests/test_maya_fast_assets.py b/tests/test_maya_fast_assets.py new file mode 100644 index 0000000000000000000000000000000000000000..fdbe243c55670b7f6ed53597832d5ceb8a1271f2 --- /dev/null +++ b/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() diff --git a/tests/test_maya_fast_backbone.py b/tests/test_maya_fast_backbone.py new file mode 100644 index 0000000000000000000000000000000000000000..e3b168ae5118339347dc5865756994cfc1422792 --- /dev/null +++ b/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() diff --git a/tests/test_maya_fast_codec.py b/tests/test_maya_fast_codec.py new file mode 100644 index 0000000000000000000000000000000000000000..0bd7598151c3383cbb909a89376e6f33bfdbf98d --- /dev/null +++ b/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() diff --git a/tests/test_maya_fast_control.py b/tests/test_maya_fast_control.py new file mode 100644 index 0000000000000000000000000000000000000000..ef193e3a81476e7c485df88e7b7bfa29fd12bcc1 --- /dev/null +++ b/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() diff --git a/tests/test_maya_fast_depth.py b/tests/test_maya_fast_depth.py new file mode 100644 index 0000000000000000000000000000000000000000..cea7e82595659d5ad99091d4050c341dcffaed1d --- /dev/null +++ b/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() diff --git a/tests/test_maya_fast_engine.py b/tests/test_maya_fast_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..936b2b28146e23a606d4bfceef8d8709fcf80a9e --- /dev/null +++ b/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() diff --git a/tests/test_maya_fast_prompt.py b/tests/test_maya_fast_prompt.py new file mode 100644 index 0000000000000000000000000000000000000000..07ffddd02facdf30f557507dd8b16f12fc195c92 --- /dev/null +++ b/tests/test_maya_fast_prompt.py @@ -0,0 +1,75 @@ +import json +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 PromptAssets, default_assets_dir +from src.maya_fast.prompt import AudioIn, build_prompt + +GOLDEN = Path.home() / ".cache" / "maya_fast" / "golden" +TOL = 1e-5 + + +class TestMayaFastPromptTiny(unittest.TestCase): + def setUp(self): + self.temp_dir = tempfile.TemporaryDirectory() + self.assets_dir = make_tiny_assets(Path(self.temp_dir.name)) + self.assets = PromptAssets(self.assets_dir) + + def tearDown(self): + self.temp_dir.cleanup() + + def test_audio_in_shapes(self): + audio_in = self.assets.audio_in + codes = np.zeros((4, 16), dtype=np.int64) + frames = audio_in.frames(codes) + self.assertEqual(frames.shape, (4, self.assets.shapes["d_model"])) + aeos = audio_in.aeos() + self.assertEqual(aeos.shape, (self.assets.shapes["d_model"],)) + + def test_build_prompt(self): + p = build_prompt("t_1 t_2", "Charvi", self.assets) + self.assertGreater(p.n_prompt, 0) + self.assertEqual(len(p.token_ids), p.n_prompt) + self.assertEqual(p.embeds.shape, (p.n_prompt, self.assets.shapes["d_model"])) + + def test_unknown_voice_lists_voices(self): + with self.assertRaisesRegex(KeyError, "Charvi"): + build_prompt("t_1", "nobody", self.assets) + + +class TestMayaFastPromptGolden(unittest.TestCase): + """Compare the numpy port against the upstream PyTorch goldens (tools/maya_export/golden_light.py).""" + + @classmethod + def setUpClass(cls): + assets_dir = default_assets_dir() + if not (GOLDEN / "light.npz").exists() or not (assets_dir / "manifest.json").exists(): + raise unittest.SkipTest(f"goldens at {GOLDEN} or assets at {assets_dir} not found") + cls.g = np.load(GOLDEN / "light.npz") + with open(GOLDEN / "light.json", "r", encoding="utf-8") as f: + cls.gj = json.load(f) + cls.assets = PromptAssets(assets_dir) + + def test_prompt_matches_golden(self): + for s_id, item in self.gj["sentences"].items(): + for voice in self.gj["voices"]: + key = f"prompt/{s_id}/{voice}" + n = int(self.g[f"{key}/n_prompt"]) + p = build_prompt(item["text"], voice, self.assets, lang="mr") + self.assertEqual(p.n_prompt, n, key) + np.testing.assert_array_equal(p.token_ids, self.g[f"{key}/token_ids"][:n], err_msg=key) + err = float(np.max(np.abs(p.embeds - self.g[f"{key}/prompt_emb"]))) + self.assertLessEqual(err, TOL, key) + + def test_audio_in_matches_golden(self): + audio_in = self.assets.audio_in + frames = audio_in.frames(self.g["audio_in_test_codes"]) + self.assertLessEqual(float(np.max(np.abs(frames - self.g["audio_in_frames"]))), TOL) + self.assertLessEqual(float(np.max(np.abs(audio_in.aeos() - self.g["aeos"]))), TOL) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_maya_fast_sampling.py b/tests/test_maya_fast_sampling.py new file mode 100644 index 0000000000000000000000000000000000000000..d945ef22d75edcdb092b6f5b523b9f33a88b9b9c --- /dev/null +++ b/tests/test_maya_fast_sampling.py @@ -0,0 +1,81 @@ +import math +import unittest +import numpy as np + +from src.maya_fast.sampling import gumbel, topk_gumbel_argmax + + +def _gammp(a: float, x: float) -> float: + if x <= 0: + return 0.0 + if x < a + 1.0: + ap = a + sum_val = 1.0 / a + del_val = sum_val + for _ in range(1, 300): + ap += 1.0 + del_val *= x / ap + sum_val += del_val + if abs(del_val) < abs(sum_val) * 1e-12: + break + return sum_val * math.exp(-x + a * math.log(x) - math.lgamma(a)) + else: + b = x + 1.0 - a + c = 1.0 / 1e-30 + d = 1.0 / b + h = d + for i in range(1, 300): + an = -i * (i - a) + b += 2.0 + d = an * d + b + if abs(d) < 1e-30: + d = 1e-30 + c = b + an / c + if abs(c) < 1e-30: + c = 1e-30 + d = 1.0 / d + del_h = d * c + h *= del_h + if abs(del_h - 1.0) < 1e-12: + break + q = math.exp(-x + a * math.log(x) - math.lgamma(a)) * h + return 1.0 - q + + +def _chi2_pvalue(stat: float, df: int) -> float: + return max(0.0, min(1.0, 1.0 - _gammp(df / 2.0, stat / 2.0))) + + +class TestMayaFastSampling(unittest.TestCase): + def test_topk_gumbel_matches_multinomial(self): + # numpy: fixed 2048-logit vector, 40 000 draws, chi-square p > 0.01 against softmax(topk_masked(logits) / 0.7) + np.random.seed(42) + vocab_size = 2048 + logits = np.random.randn(vocab_size).astype(np.float32) + temp = 0.7 + top_k = 64 + + topk_indices = np.argpartition(logits, -top_k)[-top_k:] + scaled = logits[topk_indices] / temp + exp_scaled = np.exp(scaled - np.max(scaled)) + probs = exp_scaled / np.sum(exp_scaled) + + n_draws = 40000 + rng = np.random.default_rng(12345) + counts = np.zeros(vocab_size, dtype=np.int64) + + for _ in range(n_draws): + noise = gumbel(rng, (vocab_size,)) + sample = topk_gumbel_argmax(logits, temp=temp, top_k=top_k, noise=noise) + counts[sample] += 1 + + obs = counts[topk_indices].astype(np.float64) + exp = (probs * n_draws).astype(np.float64) + + chi2_stat = float(np.sum((obs - exp) ** 2 / exp)) + p_val = _chi2_pvalue(chi2_stat, df=top_k - 1) + self.assertGreater(p_val, 0.01) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_memguard.py b/tests/test_memguard.py new file mode 100644 index 0000000000000000000000000000000000000000..17b716e76e6501eb68b2ff2f8f133d273157d68a --- /dev/null +++ b/tests/test_memguard.py @@ -0,0 +1,128 @@ +import os +import sys +import tempfile +import threading +import unittest +from pathlib import Path + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import psutil + +from tools.memguard import GuardResult, run_guarded + + +class TestMemguard(unittest.TestCase): + def setUp(self): + self._tmp = tempfile.mkdtemp(prefix="memguard_test_") + self._saved_work = os.environ.get("MAYA_FAST_WORK") + os.environ["MAYA_FAST_WORK"] = self._tmp + + def tearDown(self): + if self._saved_work is None: + os.environ.pop("MAYA_FAST_WORK", None) + else: + os.environ["MAYA_FAST_WORK"] = self._saved_work + + def test_kills_on_rss_limit(self): + child = ( + "import time\n" + "b = bytearray(300 * 1024 * 1024)\n" + "for i in range(0, len(b), 4096):\n" + " b[i] = 1\n" + "time.sleep(10)\n" + ) + r = run_guarded( + [sys.executable, "-c", child], + rss_limit_mb=100, + min_free_mb=0, + poll_s=0.1, + ) + self.assertTrue(r.killed, f"expected kill, got {r}") + self.assertEqual(r.reason, "rss limit") + self.assertIsInstance(r, GuardResult) + + def test_refuses_when_not_enough_free(self): + r = run_guarded( + [sys.executable, "-c", "pass"], + rss_limit_mb=500, + min_free_mb=int(psutil.virtual_memory().total / (1024 * 1024)) + 1, + ) + self.assertEqual(r.returncode, -1) + self.assertEqual(r.reason, "insufficient free memory") + self.assertFalse(r.killed) + + def test_clean_run_reports_peak(self): + child = ( + "import time\n" + "b = bytearray(50 * 1024 * 1024)\n" + "for i in range(0, len(b), 4096):\n" + " b[i] = 1\n" + "time.sleep(1.5)\n" + ) + r = run_guarded( + [sys.executable, "-c", child], + rss_limit_mb=500, + min_free_mb=0, + poll_s=0.1, + ) + self.assertEqual(r.returncode, 0) + self.assertIsNone(r.reason) + self.assertFalse(r.killed) + self.assertGreaterEqual(r.peak_rss_mb, 40) + + def test_lock_serializes_runs(self): + marker = Path(self._tmp) / "markers.txt" + child = ( + "import time\n" + f"with open({str(marker)!r}, 'a') as f:\n" + " f.write('S\\n')\n" + " f.flush()\n" + "time.sleep(0.7)\n" + f"with open({str(marker)!r}, 'a') as f:\n" + " f.write('E\\n')\n" + " f.flush()\n" + ) + barrier = threading.Barrier(2) + results = [] + errors = [] + + def run(): + try: + barrier.wait(timeout=10) + results.append( + run_guarded( + [sys.executable, "-c", child], + rss_limit_mb=500, + min_free_mb=0, + poll_s=0.1, + ) + ) + except Exception as exc: # pragma: no cover - surfaced below + errors.append(exc) + + threads = [threading.Thread(target=run) for _ in range(2)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=60) + + self.assertEqual(errors, []) + self.assertEqual(len(results), 2, "both guarded runs must finish") + for r in results: + self.assertEqual(r.returncode, 0, f"run failed: {r}") + self.assertIsNone(r.reason) + + events = marker.read_text().split() + self.assertEqual(events.count("S"), 2) + self.assertEqual(events.count("E"), 2) + depth = 0 + for ev in events: + depth += 1 if ev == "S" else -1 + self.assertGreaterEqual(depth, 0, f"overlapping runs: {events}") + self.assertLessEqual(depth, 1, f"overlapping runs: {events}") + self.assertEqual(depth, 0, f"unbalanced run markers: {events}") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_metrics.py b/tests/test_metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..3762676558cca8f0bcf725142fe86d778978a13b --- /dev/null +++ b/tests/test_metrics.py @@ -0,0 +1,28 @@ +import unittest +import os +import sys + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from src.metrics import score_transcript + + +class TestScoreTranscript(unittest.TestCase): + def test_identical_is_zero(self): + self.assertEqual(score_transcript("जे कोणी उंच भागात", "जे कोणी उंच भागात")["wer"], 0.0) + + def test_lenient_forgives_vowel_length(self): + s = score_transcript("पोलिस अधीक्षक", "पोलीस अधीक्षक") + self.assertEqual(s["wer"], 0.5) + self.assertEqual(s["wer_lenient"], 0.0) + + def test_empty_prediction_scores_one(self): + s = score_transcript("जे कोणी उंच भागात", "") + self.assertEqual((s["wer"], s["cer"], s["wer_lenient"], s["cer_lenient"]), (1.0, 1.0, 1.0, 1.0)) + + def test_empty_reference_scores_zero(self): + self.assertEqual(score_transcript("।", "काहीतरी")["wer"], 0.0) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_models.py b/tests/test_models.py new file mode 100644 index 0000000000000000000000000000000000000000..efae7d57f4f567f972aab15febcba42815f0e356 --- /dev/null +++ b/tests/test_models.py @@ -0,0 +1,44 @@ +import unittest +import numpy as np +import os +import sys + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from src.models import IndicConformerONNXASR + + +class FakeStream: + def __init__(self, recorder): + self._recorder = recorder + self.result = type("Result", (), {"text": " नमस्कार "})() + + def accept_waveform(self, sr, samples): + self._recorder["sr"] = sr + self._recorder["samples"] = samples + + +class FakeRecognizer: + def __init__(self): + self.recorded = {} + + def create_stream(self): + return FakeStream(self.recorded) + + def decode_stream(self, stream): + pass + + +class TestIndicConformerTranscribe(unittest.TestCase): + def test_indicconformer_transcribe_conditions_audio(self): + engine = IndicConformerONNXASR() + fake = FakeRecognizer() + engine._recognizer = fake # skip load(): no model download + text = engine.transcribe(np.full(16000, 0.05, dtype=np.float32)) + self.assertEqual(len(fake.recorded["samples"]), 22400) + self.assertEqual(fake.recorded["sr"], 16000) + self.assertEqual(text, "नमस्कार") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_normalizer.py b/tests/test_normalizer.py new file mode 100644 index 0000000000000000000000000000000000000000..bd4185df864e16d191866468b828de0cedb48e6d --- /dev/null +++ b/tests/test_normalizer.py @@ -0,0 +1,26 @@ +import unittest +import os +import sys + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from src.normalizer import normalize_marathi_text + +class TestNormalizer(unittest.TestCase): + def test_strips_zero_width_joiners(self): + self.assertEqual(normalize_marathi_text("क्‍ष ‌आणि"), "क्ष आणि") + + def test_folds_nukta(self): + # U+095B (ज़) and decomposed ज + U+093C both become ज + self.assertEqual(normalize_marathi_text("ज़ ज़"), "ज ज") + + def test_strict_keeps_vowel_length(self): + self.assertNotEqual(normalize_marathi_text("पोलिस"), normalize_marathi_text("पोलीस")) + + def test_lenient_folds_vowel_length_and_chandrabindu(self): + self.assertEqual(normalize_marathi_text("पोलिस", lenient=True), normalize_marathi_text("पोलीस", lenient=True)) + self.assertEqual(normalize_marathi_text("सुर", lenient=True), normalize_marathi_text("सूर", lenient=True)) + self.assertEqual(normalize_marathi_text("हँ", lenient=True), normalize_marathi_text("हं", lenient=True)) + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..85ad624049fb4fbe7eb5630722c90523506e313d --- /dev/null +++ b/tests/test_pipeline.py @@ -0,0 +1,33 @@ +import unittest +import numpy as np +import os +import sys + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from src.normalizer import normalize_marathi_text +from src.audio_utils import load_and_resample +from src.models import get_model + +class TestMarathiASR(unittest.TestCase): + def test_devanagari_normalization(self): + sample = "पोलीस अधीक्षक, 'चंद्रशेखर'! आणि सोळंकी।॥" + normalized = normalize_marathi_text(sample) + self.assertEqual(normalized, "पोलीस अधीक्षक चंद्रशेखर आणि सोळंकी") + + def test_audio_loader_and_resampler(self): + # Test on existing benchmark sample + audio_path = "benchmarks/audio/clean_sample_0.wav" + if os.path.exists(audio_path): + data, duration = load_and_resample(audio_path, target_sr=16000) + self.assertEqual(data.dtype, np.float32) + self.assertTrue(duration > 0) + self.assertEqual(data.ndim, 1) + + def test_indicconformer_instantiation(self): + engine = get_model("indicconformer-mr-int8") + self.assertEqual(engine.model_id, "indicconformer-mr-int8") + self.assertTrue(engine.disk_size_mb > 0) + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_ptt_buffer.py b/tests/test_ptt_buffer.py new file mode 100644 index 0000000000000000000000000000000000000000..5f302e16ced525d5fb3947a36b09af0b2b5165dc --- /dev/null +++ b/tests/test_ptt_buffer.py @@ -0,0 +1,88 @@ +import unittest +import numpy as np +import os +import sys + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from src.ptt_buffer import ClipAssembler + + +def chunk(value, n=160): + return np.full(n, value, dtype=np.float32) + + +class TestClipAssembler(unittest.TestCase): + def test_clip_includes_pre_and_post_roll(self): + a = ClipAssembler() + for _ in range(100): + self.assertIsNone(a.feed(chunk(1))) + self.assertIsNone(a.start()) + for _ in range(50): + self.assertIsNone(a.feed(chunk(2))) + a.stop() + outs = [a.feed(chunk(3)) for _ in range(25)] + self.assertTrue(all(o is None for o in outs[:-1])) + clip = outs[-1] + self.assertEqual(len(clip), 4800 + 8000 + 4000) + self.assertTrue(np.all(clip[:4800] == 1) and np.all(clip[4800:12800] == 2) and np.all(clip[12800:] == 3)) + + def test_pre_roll_limited_to_available_history(self): + a = ClipAssembler() + for _ in range(10): + a.feed(chunk(1)) + self.assertIsNone(a.start()) + a.feed(chunk(2)) + a.stop() + clip = [a.feed(chunk(3)) for _ in range(25)][-1] + self.assertEqual(len(clip), 1600 + 160 + 4000) + self.assertTrue(np.all(clip[:1600] == 1)) + self.assertTrue(np.all(clip[1600:1760] == 2)) + + def test_post_roll_trimmed_exactly_with_odd_chunk_size(self): + a = ClipAssembler() + for _ in range(20): + a.feed(chunk(1, 512)) + a.start() + for _ in range(10): + a.feed(chunk(2, 512)) + a.stop() + outs = [a.feed(chunk(3, 512)) for _ in range(8)] + self.assertTrue(all(o is None for o in outs[:-1])) + clip = outs[-1] + held = 10 * 512 + # pre-roll is capped at 4800 samples even though 10240 were available + self.assertEqual(len(clip), 4800 + held + 4000) + self.assertTrue(np.all(clip[-4000:] == 3)) + + def test_restart_during_post_roll_flushes_without_duplication(self): + a = ClipAssembler() + a.start() + for _ in range(10): + a.feed(chunk(2)) + a.stop() + for _ in range(5): + a.feed(chunk(3)) + first = a.start() + self.assertEqual(len(first), 1600 + 800) + for _ in range(10): + a.feed(chunk(4)) + a.stop() + second = [a.feed(chunk(5)) for _ in range(25)][-1] + self.assertTrue(np.all(second[:1600] == 4)) + self.assertEqual(len(second), 1600 + 4000) + + def test_is_recording_tracks_hold_only(self): + a = ClipAssembler() + self.assertFalse(a.is_recording) + a.start() + self.assertTrue(a.is_recording) + a.stop() + self.assertFalse(a.is_recording) + for _ in range(5): + a.feed(chunk(3)) + self.assertFalse(a.is_recording) + + +if __name__ == "__main__": + unittest.main() diff --git a/tools/maya_export/__init__.py b/tools/maya_export/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/tools/maya_export/common.py b/tools/maya_export/common.py new file mode 100644 index 0000000000000000000000000000000000000000..79353014c2fe285a6f8b8d06fbf72ff2ffb4b086 --- /dev/null +++ b/tools/maya_export/common.py @@ -0,0 +1,152 @@ +from __future__ import annotations + +import gc +import os +from pathlib import Path +from typing import Iterator +import torch +from safetensors import safe_open + + +def work_dir() -> Path: + env = os.environ.get("MAYA_FAST_WORK") + if env: + return Path(env) + return Path.home() / ".cache" / "maya_fast" + + +def ckpt_paths() -> dict: + from maya3 import hub + return hub.resolve("maya-research/maya2-native") + + +def iter_ckpt_tensors( + skip_prefixes: tuple[str, ...] = (), + files: list[str] | None = None, +) -> Iterator[tuple[str, torch.Tensor]]: + """Yield one tensor at a time from checkpoint safetensors, stripping model. prefix. + If files is None, uses ckpt_paths()["ckpt"]. + """ + from tts.model import load_s3 as L + + if files is None: + paths = ckpt_paths() + src = os.path.abspath(str(paths["ckpt"])) + _, files, _ = L._safetensors_files(src) + + seen = set() + for f in files: + with safe_open(f, "pt") as fh: + for k in fh.keys(): + name = k[len(L.MODEL_PREFIX):] if k.startswith(L.MODEL_PREFIX) else k + if any(name.startswith(p) for p in skip_prefixes): + continue + if name in seen: + continue + seen.add(name) + t = fh.get_tensor(k) + yield name, t + del t + + +def _stub_lm(path: str, dtype: torch.dtype, attn: str): + from transformers import AutoConfig, AutoModelForCausalLM + cfg = AutoConfig.from_pretrained(path) + cfg.num_hidden_layers = 0 + cfg.layer_types = [] + with torch.device("meta"): + lm = AutoModelForCausalLM.from_config(cfg, dtype=torch.float32, attn_implementation="sdpa") + for m in lm.modules(): + if hasattr(m, "inv_freq") and hasattr(m, "rope_type"): + from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS + fn = m.rope_init_fn if getattr(m, "rope_init_fn", None) is not None else ( + m.compute_default_rope_parameters if m.rope_type == "default" else ROPE_INIT_FUNCTIONS[m.rope_type]) + with torch.device("cpu"): + inv, scale = fn(m.config, torch.device("cpu")) + m.inv_freq = inv + if hasattr(m, "original_inv_freq"): + m.original_inv_freq = inv + if hasattr(m, "attention_scaling"): + m.attention_scaling = scale + lm.tie_weights() + return lm + + +class _UpcastFp32(torch.nn.Module): + """Parametrization: compute on an fp32 copy of a bf16 weight, made per call.""" + + def forward(self, w: torch.Tensor) -> torch.Tensor: + return w.float() + + +class _Bf16Embedding(torch.nn.Embedding): + def forward(self, idx: torch.Tensor) -> torch.Tensor: + # Row lookup on the bf16 table, then upcast only the rows read. bf16 to fp32 is exact. + return torch.nn.functional.embedding(idx, self.weight).float() + + +def _keep_bf16_params_fp32_compute(model: torch.nn.Module) -> None: + """Parameters stored as bf16 are upcast to fp32 at use, so weights stay half-size in RAM. + Embedding tables are read by row. Other modules get a parametrization that upcasts per call. + Upcasting bf16 is exact, so results match an fp32 copy of the same weights.""" + from torch.nn.utils import parametrize + + embed_weights = set() + for mod in model.modules(): + if isinstance(mod, torch.nn.Embedding) and mod.weight.dtype == torch.bfloat16: + embed_weights.add(id(mod.weight)) + mod.__class__ = _Bf16Embedding + for mod in list(model.modules()): # snapshot: registering adds child modules + for pname, p in list(mod._parameters.items()): + if p is None or p.dtype != torch.bfloat16 or id(p) in embed_weights: + continue + if isinstance(mod, torch.nn.Embedding): + continue + parametrize.register_parametrization(mod, pname, _UpcastFp32(), unsafe=True) + + +def build_light_model(): + """Build light TTSModelV2 on meta device with 0 backbone decoder layers, + and populate non-backbone layer weights via iter_ckpt_tensors.""" + import tts.model.backbone as BB + import tts.model.load_s3 as L + + orig_stub = BB._lm_from_config + try: + BB._lm_from_config = _stub_lm + paths = ckpt_paths() + src = str(paths["ckpt"]) + model_cfg, cfg_path = L.resolve_model_cfg(src) + model_cfg["attn"] = "sdpa" + + anchor_cfg = model_cfg.get("speaker_anchor") + depth_grouped = bool(model_cfg.get("depth_grouped", True)) + + with torch.device("meta"): + model = L.build_eval_model( + model_cfg, + anchor_cfg=anchor_cfg, + amortize_frac=1.0, + depth_grouped=depth_grouped, + ) + + mods = dict(model.named_modules()) + for name, t in iter_ckpt_tensors(skip_prefixes=("backbone.lm.model.layers.",)): + parent_name, _, leaf = name.rpartition(".") + parent = mods.get(parent_name) + if parent is None or (leaf not in parent._parameters and leaf not in parent._buffers): + continue + is_param = leaf in parent._parameters + if is_param: + # Keep the checkpoint dtype (bf16) to halve resident memory. + parent._parameters[leaf] = torch.nn.Parameter(t, requires_grad=False) + else: + parent._buffers[leaf] = t.float() + del t + + model.backbone.lm.tie_weights() + _keep_bf16_params_fp32_compute(model) + gc.collect() + return model.eval() + finally: + BB._lm_from_config = orig_stub diff --git a/tools/maya_export/depth_all.py b/tools/maya_export/depth_all.py new file mode 100644 index 0000000000000000000000000000000000000000..8f41bdda3f37e7b4b1a3c2b907d6903a88cdac96 --- /dev/null +++ b/tools/maya_export/depth_all.py @@ -0,0 +1,113 @@ +"""Single depth graph: levels 1..15, teacher-forced, weights stored in fp16 and computed in fp32. + +Same computation as tts/model/depth_v2.py for levels 1..15: each level's prefix (_prefix), the level's weight group +(GROUP_OF_LEVEL: groups 0..7 private, levels 9..15 share group 8), the 4 shared cross-attention blocks, LayerNorm and +the level head. Group 8 is one set of weights used by seven levels, so it is stored once in the graph. + +Storage: every weight is an fp16 initializer; each use is a Cast to fp32 (parametrization in torch, Cast node in ONNX). +Built one tensor at a time from the checkpoint, so the fp32 copy is never resident. +""" +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from tts.model.depth_v2 import D, GROUP_OF_LEVEL, LAYERS, RANK, _CrossBlock, _SelfBlock +from tts.tokens import N_CODES, N_Q + +LEVELS = 15 # serving q_serve = 16 -> levels 1..15 +TEMP = 0.7 +TOP_K = 64 +N_PRIVATE_GROUPS = 9 # groups 0..8 are used by levels 1..15 + + +class _UpcastFp32(nn.Module): + def forward(self, w: torch.Tensor) -> torch.Tensor: + return w.float() + + +class _Fp16Embedding(nn.Embedding): + def forward(self, idx: torch.Tensor) -> torch.Tensor: + return F.embedding(idx, self.weight).float() + + +class DepthAll(nn.Module): + def __init__(self, n_codes: int = N_CODES, d_backbone: int = 2048) -> None: + super().__init__() + self.n_codes = n_codes + self.in_proj = nn.Linear(d_backbone, D, bias=False) + self.code_lo = nn.Embedding((N_Q - 1) * n_codes, RANK) + self.code_up = nn.Linear(RANK, D, bias=False) + self.level_emb = nn.Embedding(N_Q, D) + self.groups = nn.ModuleList( + nn.ModuleList(_SelfBlock() for _ in range(LAYERS)) for _ in range(N_PRIVATE_GROUPS)) + self.cross = nn.ModuleList(_CrossBlock() for _ in range(LAYERS)) + self.norm = nn.LayerNorm(D) + self.heads = nn.ModuleList(nn.Linear(D, n_codes, bias=False) for _ in range(LEVELS)) + + def _prefix(self, h: torch.Tensor, codes: torch.Tensor, level: int) -> torch.Tensor: + ks = torch.arange(level, device=h.device) + flat = codes[:, :level] + ks * self.n_codes + e = self.code_up(self.code_lo(flat)) + self.level_emb(ks + 1)[None] + return torch.cat([self.in_proj(h)[:, None] + self.level_emb.weight[:1], e], 1) + + def _level_logits(self, h: torch.Tensor, codes_prev: torch.Tensor, level: int, mem: torch.Tensor) -> torch.Tensor: + """Logits of one level given the codes of the levels before it (codes_prev has `level` columns).""" + x = self._prefix(h, codes_prev, level) + T = x.shape[1] + mask = torch.ones(T, T, dtype=torch.bool, device=x.device).tril()[None, None] + group = self.groups[GROUP_OF_LEVEL[level]] + for li in range(LAYERS): + x = group[li](x, mask) + x = self.cross[li](x, mem, None) + return self.heads[level - 1](self.norm(x[:, -1])) + + def forward(self, h: torch.Tensor, codes: torch.Tensor, mem: torch.Tensor) -> torch.Tensor: + """Teacher-forced: h (B, 2048); codes (B, Q_SERVE=16) int64; mem (B, 1, 1024). Returns logits (B, 15, 2048).""" + return torch.stack([self._level_logits(h, codes[:, :level], level, mem) for level in range(1, LEVELS + 1)], 1) + + def sample(self, h: torch.Tensor, cb0: torch.Tensor, mem: torch.Tensor, noise: torch.Tensor) -> torch.Tensor: + """Sampled frame. h (B, 2048); cb0 (B,) int64; mem (B, 1, 1024); noise (B, 15, 2048) Gumbel. + Level L: argmax(where(rank < top_k, logits / TEMP + noise, -inf)) over logits of the previous levels' codes. + Returns codes (B, 16) int64, cb0 first. Gumbel-max over the top-k softmax, same distribution as upstream + depth_generate (torch.multinomial over the top-k softmax of logits / TEMP).""" + codes = [cb0] + for level in range(1, LEVELS + 1): + prev = torch.stack(codes, 1) # (B, level) + scaled = self._level_logits(h, prev, level, mem) / TEMP + kth = scaled.topk(TOP_K, -1).values[:, -1:] + masked = scaled.masked_fill(scaled < kth, float("-inf")) + codes.append((masked + noise[:, level - 1]).argmax(-1)) + return torch.stack(codes, 1) + + +class DepthSample(nn.Module): + """Export wrapper: the sampled-frame graph (inputs h, cb0, mem, noise; output codes).""" + + def __init__(self, depth: DepthAll) -> None: + super().__init__() + self.depth = depth + + def forward(self, h, cb0, mem, noise): + return self.depth.sample(h, cb0, mem, noise) + + +def load_fp16(model: DepthAll, get) -> None: + """Load every parameter from `get(ckpt_name) -> tensor` as fp16, and mark it for an fp32 upcast on use. + `get` returns the bf16/fp32 checkpoint tensor; mapping from module names to checkpoint names is in export.""" + names = {} + for name, p in model.named_parameters(): + names[name] = get(name) + for name, t in names.items(): + mod_name, _, leaf = name.rpartition(".") + mod = model.get_submodule(mod_name) + mod._parameters[leaf] = nn.Parameter(t.to(torch.float16), requires_grad=False) + for mod_name, mod in list(model.named_modules()): + for leaf, p in list(mod._parameters.items()): + if p is None or p.dtype != torch.float16: + continue + if isinstance(mod, nn.Embedding): + mod.__class__ = _Fp16Embedding + continue + torch.nn.utils.parametrize.register_parametrization(mod, leaf, _UpcastFp32(), unsafe=True) diff --git a/tools/maya_export/depth_check.py b/tools/maya_export/depth_check.py new file mode 100644 index 0000000000000000000000000000000000000000..6f385da166698240e0dd4de1e5bfba0c785a07c8 --- /dev/null +++ b/tools/maya_export/depth_check.py @@ -0,0 +1,20 @@ +"""fp32 depth graphs vs upstream PyTorch depth (light model, exact fp32 weights) on 20 golden frames.""" +import numpy as np, onnxruntime as ort, torch +from tools.maya_export.common import build_light_model, work_dir +from tts.model.depth_v2 import GROUP_OF_LEVEL +B = work_dir() / 'depth_build' / 'fp32' +z = np.load(work_dir() / 'golden' / 'ref' / 'conv_01_s0.npz') +codes = z['codes'][:, :20].T.astype(np.int64); h = z['h_frames'][:20].astype(np.float32) +model = build_light_model(); d = model.depth.eval() +mem_t = d.null_mem.detach().float().expand(20, -1, -1).contiguous() +worst_rel = 0.0; worst_abs = 0.0 +with torch.inference_mode(): + for L in range(1, 16): + x = d._prefix(torch.from_numpy(h), torch.from_numpy(codes), L) + ref = d.heads[L - 1](d._run(x, GROUP_OF_LEVEL[L], mem_t, None)).float().numpy() + s = ort.InferenceSession(str(B / f'level_{L:02d}.onnx'), providers=['CPUExecutionProvider']) + got = s.run(None, {'h': h, 'codes': codes, 'mem': mem_t.numpy()})[0] + del s + err = np.abs(got - ref) + worst_abs = max(worst_abs, float(err.max())); worst_rel = max(worst_rel, float(err.max() / np.abs(ref).max())) +print(f'fp32 vs upstream PyTorch depth: WORST abs {worst_abs:.3e} rel {worst_rel:.3e}') diff --git a/tools/maya_export/depth_frame.py b/tools/maya_export/depth_frame.py new file mode 100644 index 0000000000000000000000000000000000000000..e0570a53b8620e3ec6533e63a072b22c3c4e9d5f --- /dev/null +++ b/tools/maya_export/depth_frame.py @@ -0,0 +1,55 @@ +"""Per-level teacher-forced depth graph (one ONNX graph per level 1..15). + +Each graph is the upstream DepthV2 computation for one level, with the same classes (_SelfBlock, _CrossBlock) and the +same prefix (_prefix) and last-position read-out (_run) as tts/model/depth_v2.py, mask None (the upstream kmask is +all True for named voices, so padding is never masked). It holds only the weights that level needs: its weight group +(4 self-attention blocks), the 4 shared cross-attention blocks, the code embeddings, the input projection and the +level's head. The 15 graphs are chained by the runtime; the group's weights are not shared between levels 9..15 here +(each level graph carries group 8), which costs bandwidth but keeps every part small. +""" +from __future__ import annotations + +import torch +import torch.nn as nn + +from tts.model.depth_v2 import D, GROUP_OF_LEVEL, LAYERS, RANK, _CrossBlock, _SelfBlock +from tts.tokens import N_CODES, N_Q + +Q_SERVE = 16 # serving codebooks per frame (cb0..cb15); levels 1..15 read columns 0..L-1 of these + + +class DepthLevel(nn.Module): + def __init__(self, level: int, n_codes: int = N_CODES, d_backbone: int = 2048) -> None: + super().__init__() + if not 1 <= level <= N_Q - 1: + raise ValueError(f"level {level} outside 1..{N_Q - 1}") + self.level = level + self.group = GROUP_OF_LEVEL[level] + self.n_codes = n_codes + self.in_proj = nn.Linear(d_backbone, D, bias=False) + self.code_lo = nn.Embedding((N_Q - 1) * n_codes, RANK) + self.code_up = nn.Linear(RANK, D, bias=False) + self.level_emb = nn.Embedding(N_Q, D) + self.trunk = nn.ModuleList(_SelfBlock() for _ in range(LAYERS)) + self.cross = nn.ModuleList(_CrossBlock() for _ in range(LAYERS)) + self.norm = nn.LayerNorm(D) + self.head = nn.Linear(D, n_codes, bias=False) + + def forward(self, h: torch.Tensor, codes: torch.Tensor, mem: torch.Tensor) -> torch.Tensor: + """h (B, 2048) fp32 backbone hidden; codes (B, Q_SERVE) int64 with columns cb0..cb_{L-1} used; mem (B, 1, 1024). + Returns the level's logits (B, 2048).""" + L = self.level + ks = torch.arange(L, device=h.device) + flat = codes[:, :L] + ks * self.n_codes + e = self.code_up(self.code_lo(flat)) + self.level_emb(ks + 1)[None] + x = torch.cat([self.in_proj(h)[:, None] + self.level_emb.weight[:1], e], 1) # (B, L+1, D) + T = x.shape[1] + mask = torch.ones(T, T, dtype=torch.bool, device=x.device).tril()[None, None] + for li in range(LAYERS): + x = self.trunk[li](x, mask) + x = self.cross[li](x, mem, None) + return self.head(self.norm(x[:, -1])) + + +def level_group(level: int) -> int: + return GROUP_OF_LEVEL[level] diff --git a/tools/maya_export/eval_teacher_forced.py b/tools/maya_export/eval_teacher_forced.py new file mode 100644 index 0000000000000000000000000000000000000000..bb4daf412a0b6264c76f613ec8a683b896f71064 --- /dev/null +++ b/tools/maya_export/eval_teacher_forced.py @@ -0,0 +1,207 @@ +"""Teacher-forced quality gate for the fused backbone. + +Reference: fp32 per-layer graphs, one layer resident at a time, run over each reference `seq_emb` +(work_dir()/golden/ref/*.npz, from golden_reference.py), then the fp32 head. +Candidate: the fused graph (backbone_{candidate}.onnx), run over the same sequences. +Compared at the generating positions only: the frame-i prediction is made at position n_p - 1 + i. + +Metrics: top-1 agreement of cb0 argmax, and mean KL(ref || cand) over the distribution softmax(logits / 0.7) +(the sampling temperature), in nats. + + .venv/bin/python tools/memguard.py --rss-limit-mb 4000 -- \ + .venv-tts/bin/python -m tools.maya_export.eval_teacher_forced backbone --candidate wo8 +""" +from __future__ import annotations + +import argparse +import glob +import json +from pathlib import Path + +import numpy as np +import onnxruntime as ort +import torch +from transformers import AutoConfig +from transformers.models.qwen3.modeling_qwen3 import Qwen3RotaryEmbedding + +from tools.maya_export import export_backbone as EB +from tools.maya_export.common import ckpt_paths, work_dir +from tools.maya_export.export_depth import N_LEVELS as DEPTH_LEVELS +from tools.maya_export.streamed_backbone import _Weights, shard_files + +N_LAYERS = EB.N_LAYERS +TEMP = 0.7 +TOP1_GATE = 0.99 +KL_GATE = 0.01 +DEPTH_TOP1_GATE = 0.98 +DEPTH_KL_GATE = 0.02 + + +def _sess(path: Path, threads: int = 8) -> ort.InferenceSession: + so = ort.SessionOptions() + so.intra_op_num_threads = threads + so.add_session_config_entry("session.intra_op.allow_spinning", "0") + so.enable_cpu_mem_arena = False + return ort.InferenceSession(str(path), so, providers=["CPUExecutionProvider"]) + + +def _rope_table(max_len: int) -> tuple[np.ndarray, np.ndarray]: + cfg = AutoConfig.from_pretrained(ckpt_paths()["backbone_dir"]) + cfg._attn_implementation = "sdpa" + rot = Qwen3RotaryEmbedding(config=cfg).eval() + cos, sin = EB.rope(rot, max_len, 0) + return cos.numpy(), sin.numpy() + + +def _load_refs(ref_dir: Path, n_prompt: dict[str, int]): + items = [] + for path in sorted(glob.glob(str(ref_dir / "*_s*.npz"))): + name = Path(path).stem + with np.load(path) as z: + seq = z["seq_emb"].astype(np.float32) + F = int(z["codes"].shape[1]) + s_id = name.rsplit("_s", 1)[0] + n_p = n_prompt[s_id] + assert seq.shape[0] == n_p + F, (name, seq.shape, n_p, F) + items.append({"name": name, "seq": seq, "n_p": n_p, "F": F}) + return items + + +def _gen_positions(n_p: int, F: int) -> np.ndarray: + return n_p - 1 + np.arange(F) + + +def _run_candidate(backbone_path: Path, items, cos_tab, sin_tab): + sess = _sess(backbone_path) + out = [] + for it in items: + T = it["seq"].shape[0] + zeros = np.zeros((1, EB.N_KV_HEADS, 0, EB.HEAD_DIM), np.float32) + feed = {"embeds": it["seq"][None], "cos": cos_tab[:, :T], "sin": sin_tab[:, :T], + "mask": EB.masks(T, 0).numpy()} + for i in range(N_LAYERS): + feed[f"past_k_{i}"] = zeros + feed[f"past_v_{i}"] = zeros + res = sess.run(None, feed) + cb0 = res[1][0] # (T, 2049) + out.append(cb0[_gen_positions(it["n_p"], it["F"])]) + del sess + return out + + +def _run_reference(build_dir: Path, items, cos_tab, sin_tab): + """One fp32 layer at a time over every sequence; the hidden state of every item is kept between layers.""" + hidden = [it["seq"][None].copy() for it in items] + for i in range(N_LAYERS): + sess = _sess(build_dir / f"layer_{i:02d}.onnx", threads=4) + for k, it in enumerate(items): + T = it["seq"].shape[0] + zeros = np.zeros((1, EB.N_KV_HEADS, 0, EB.HEAD_DIM), np.float32) + out, _, _ = sess.run(None, {"h": hidden[k], "cos": cos_tab[:, :T], "sin": sin_tab[:, :T], + "mask": EB.masks(T, 0).numpy(), "past_k": zeros, "past_v": zeros}) + hidden[k] = out + del sess + head = _sess(build_dir / "head.onnx", threads=4) + out = [] + for k, it in enumerate(items): + _, cb0 = head.run(None, {"hid": hidden[k]}) + out.append(cb0[0][_gen_positions(it["n_p"], it["F"])]) + return out + + +def _softmax(x: np.ndarray, t: float) -> np.ndarray: + z = x.astype(np.float64) / t + z = z - z.max(-1, keepdims=True) + e = np.exp(z) + return e / e.sum(-1, keepdims=True) + + +def backbone(candidate: str) -> dict: + ref_dir = work_dir() / "golden" / "ref" + light = json.loads((work_dir() / "golden" / "light.json").read_text(encoding="utf-8")) + voice = "Charvi" + golden = np.load(work_dir() / "golden" / "light.npz") + n_prompt = {s_id: int(golden[f"prompt/{s_id}/{voice}/n_prompt"]) for s_id in light["sentences"]} + + items = _load_refs(ref_dir, n_prompt) + max_len = max(it["seq"].shape[0] for it in items) + cos_tab, sin_tab = _rope_table(max_len) + + cand = _run_candidate(work_dir() / "assets" / "v1" / f"backbone_{candidate}.onnx", items, cos_tab, sin_tab) + ref = _run_reference(work_dir() / "backbone_build" / "fp32", items, cos_tab, sin_tab) + + top1_hits, total, kls = 0, 0, [] + for c, r in zip(cand, ref): + top1_hits += int((c.argmax(-1) == r.argmax(-1)).sum()) + total += r.shape[0] + p = _softmax(r, TEMP) + q = _softmax(c, TEMP) + kls.append(float(np.sum(p * (np.log(p + 1e-30) - np.log(q + 1e-30)), axis=-1).sum())) + result = {"top1": top1_hits / total, "kl_mean": sum(kls) / total, "frames": total, + "files": len(items), "gate": {"top1>=": TOP1_GATE, "kl_mean<=": KL_GATE}} + result["pass"] = bool(result["top1"] >= TOP1_GATE and result["kl_mean"] <= KL_GATE) + return result + + +def depth(candidate: str) -> dict: + """Per-level depth gate. Candidate: per-level graphs (depth_build/{candidate}); one level resident at a time, + every file run through it. Reference: the stored depth_logits (fp16) of the corpus.""" + ref_dir = work_dir() / "golden" / "ref" + items = [] + for path in sorted(glob.glob(str(ref_dir / "*_s*.npz"))): + with np.load(path) as z: + items.append({"codes": z["codes"].T.astype(np.int64), "h": z["h_frames"].astype(np.float32), + "ref": z["depth_logits"].astype(np.float32)}) + mem_w = _Weights(shard_files(ckpt_paths()["ckpt"])) + null_mem = mem_w.get("depth.null_mem").float().numpy().reshape(1, 1, 1024) + build = work_dir() / "depth_build" / candidate + single = build / "depth_all.onnx" + per_top1, per_kl = [], [] + if single.exists(): + # one graph, all levels: run each file once and read the 15 level outputs + sess = _sess(single, threads=8) + outs = [] + for it in items: + F = it["h"].shape[0] + outs.append(sess.run(None, {"h": it["h"], "codes": it["codes"], + "mem": np.repeat(null_mem, F, axis=0)})[0]) + del sess + levels = [(level, [o[:, level - 1] for o in outs]) for level in range(1, DEPTH_LEVELS + 1)] + else: + levels = None + for level in range(1, DEPTH_LEVELS + 1): + if levels is None: + sess = _sess(build / f"level_{level:02d}.onnx", threads=8) + got_per_item = [] + for it in items: + F = it["h"].shape[0] + mem = np.repeat(null_mem, F, axis=0) + got_per_item.append(sess.run(None, {"h": it["h"], "codes": it["codes"], "mem": mem})[0]) + del sess + else: + got_per_item = levels[level - 1][1] + hits, total, kl_sum = 0, 0, 0.0 + for it, got in zip(items, got_per_item): + F = it["h"].shape[0] + want = it["ref"][:, level - 1] + hits += int((got.argmax(-1) == want.argmax(-1)).sum()) + total += F + p = _softmax(want, TEMP) + q = _softmax(got, TEMP) + kl_sum += float(np.sum(p * (np.log(p + 1e-30) - np.log(q + 1e-30)))) + per_top1.append(hits / total) + per_kl.append(kl_sum / total) + passes = [t >= DEPTH_TOP1_GATE and k <= DEPTH_KL_GATE for t, k in zip(per_top1, per_kl)] + return {"candidate": candidate, "frames": total, "files": len(items), + "top1_per_level": per_top1, "kl_per_level": per_kl, + "gate": {"top1>=": DEPTH_TOP1_GATE, "kl<=": DEPTH_KL_GATE}, + "pass_per_level": passes, "pass": all(passes)} + + +if __name__ == "__main__": + ap = argparse.ArgumentParser() + ap.add_argument("target", choices=["backbone", "depth"]) + ap.add_argument("--candidate", choices=["wo8", "wo8_fph", "fp32", "fp16"], required=True) + a = ap.parse_args() + res = backbone(a.candidate) if a.target == "backbone" else depth(a.candidate) + print(json.dumps(res, indent=2)) diff --git a/tools/maya_export/export_backbone.py b/tools/maya_export/export_backbone.py new file mode 100644 index 0000000000000000000000000000000000000000..6851e39de6b7a8636b91454fedbd93156c2e6a7b --- /dev/null +++ b/tools/maya_export/export_backbone.py @@ -0,0 +1,279 @@ +"""Fused backbone ONNX: 28 Qwen3 decoder layers + final RMSNorm + cb0 head, with KV cache in/out. + +Build (one layer resident at a time, so it fits the 4 GB limit): + 1. each layer: transformers Qwen3DecoderLayer loaded from the checkpoint (tools/maya_export/streamed_backbone.py), + exported to ONNX (opset 17, dynamo=False), saved with external data; + for wo8 quantized with MatMulNBitsQuantizer(bits=8, block_size=32, is_symmetric=True, accuracy_level=4). + 2. head: final RMSNorm + cb0_head (fp32). + 3. merge: graph structure only (weights stay in the per-part data files), tensors renamed, then the weight bytes + are copied into one data file next to the merged model. + +Output: work_dir()/assets/v1/backbone_{precision}.onnx + backbone_{precision}.data +Per-part files: work_dir()/backbone_build/{precision}/ (kept for the layer-by-layer reference in the tests). + + .venv/bin/python tools/memguard.py --rss-limit-mb 4000 --min-free-mb 5000 -- \ + .venv-tts/bin/python -m tools.maya_export.export_backbone --precision wo8 +""" +from __future__ import annotations + +import argparse +import shutil +from pathlib import Path + +import numpy as np +import onnx +import torch +import torch.nn as nn +from onnx import TensorProto, helper + +from tools.maya_export.common import ckpt_paths, work_dir +from tools.maya_export.streamed_backbone import StreamedBackbone, shard_files + +N_LAYERS = 28 +N_KV_HEADS = 8 +HEAD_DIM = 128 +D_MODEL = 2048 +N_CB0 = 2049 # 2048 codes + +OPSET = 17 +PREFILL = 8 # trace length for the export; dynamic axes cover any T +COPY_CHUNK = 1 << 24 + + +class LayerKV(nn.Module): + """One decoder layer, KV cache as flat tensors, additive attention mask passed in.""" + + def __init__(self, layer): + super().__init__() + self.layer = layer + + def forward(self, h, cos, sin, mask, past_k, past_v): + from transformers import DynamicCache + + cache = DynamicCache() + cache.update(past_k, past_v, 0) + out = self.layer(h, attention_mask=mask, position_ids=None, past_key_values=cache, + use_cache=True, cache_position=None, position_embeddings=(cos, sin)) + if isinstance(out, tuple): + out = out[0] + return out, cache.layers[0].keys, cache.layers[0].values + + +class Head(nn.Module): + def __init__(self, norm: nn.Module, cb0_head: nn.Linear) -> None: + super().__init__() + self.norm = norm + self.cb0_head = cb0_head + + def forward(self, hid): + x = self.norm(hid) + return x, self.cb0_head(x) + + +def masks(T: int, past: int) -> torch.Tensor: + """Additive causal mask (1, 1, T, past+T): query i sees keys j <= past + i.""" + i = torch.arange(T)[:, None] + j = torch.arange(past + T)[None, :] + return torch.zeros(T, past + T).masked_fill(j > past + i, float("-inf"))[None, None] + + +def rope(rot, T: int, past: int) -> tuple[torch.Tensor, torch.Tensor]: + pos = torch.arange(past, past + T)[None] + return rot(torch.zeros(1, T, 1), pos) + + +def _export(module: nn.Module, args: tuple, path: Path, **kw) -> None: + with torch.inference_mode(): + torch.onnx.export(module, args, str(path), opset_version=OPSET, dynamo=False, + do_constant_folding=True, **kw) + + +def _save_external(model: onnx.ModelProto, path: Path) -> None: + # onnx appends to an existing data file; remove it so a rebuild does not leave stale bytes behind + path.with_name(path.name + ".data").unlink(missing_ok=True) + onnx.save_model(model, str(path), save_as_external_data=True, all_tensors_to_one_file=True, + location=path.name + ".data", size_threshold=1024) # small tensors inline: ORT shape inference needs them + + +def _quantize_wo8(model: onnx.ModelProto, exclude: list[str] | None = None) -> onnx.ModelProto: + from onnxruntime.quantization.matmul_nbits_quantizer import MatMulNBitsQuantizer + + q = MatMulNBitsQuantizer(model, bits=8, block_size=32, is_symmetric=True, accuracy_level=4, + nodes_to_exclude=exclude or []) + q.process() + return q.model.model + + +def export_layer(sb: StreamedBackbone, i: int, precision: str, build: Path) -> Path: + sb._load_layer(i) + sb.layer.self_attn.layer_idx = 0 # LayerKV owns a one-slot cache + wrap = LayerKV(sb.layer).eval() + h0 = torch.randn(1, PREFILL, D_MODEL) * 0.02 + cos, sin = rope(sb.rotary, PREFILL, 0) + pk = torch.zeros(1, N_KV_HEADS, 0, HEAD_DIM) + pv = torch.zeros(1, N_KV_HEADS, 0, HEAD_DIM) + tmp = build / f"layer_{i:02d}_raw.onnx" + _export(wrap, (h0, cos, sin, masks(PREFILL, 0), pk, pv), tmp, + input_names=["h", "cos", "sin", "mask", "past_k", "past_v"], + output_names=["out", "new_k", "new_v"], + dynamic_axes={"h": {1: "t"}, "cos": {1: "t"}, "sin": {1: "t"}, "mask": {2: "t", 3: "total"}, + "past_k": {2: "past"}, "past_v": {2: "past"}, + "out": {1: "t"}, "new_k": {2: "total"}, "new_v": {2: "total"}}) + model = onnx.load(str(tmp)) + if precision == "wo8": + model = _quantize_wo8(model) + out = build / f"layer_{i:02d}.onnx" + _save_external(model, out) + tmp.unlink() + tmp.with_suffix(".onnx.data").unlink(missing_ok=True) + return out + + +def export_head(sb: StreamedBackbone, build: Path) -> Path: + cb0 = nn.Linear(D_MODEL, N_CB0, bias=False) + cb0.weight.data.copy_(sb.weights.get("backbone.cb0_head.weight").float()) + head = Head(sb.norm, cb0).eval() + tmp = build / "head_raw.onnx" + _export(head, (torch.randn(1, PREFILL, D_MODEL),), tmp, input_names=["hid"], + output_names=["hidden", "cb0_logits"], + dynamic_axes={"hid": {1: "t"}, "hidden": {1: "t"}, "cb0_logits": {1: "t"}}) + out = build / "head.onnx" + _save_external(onnx.load(str(tmp)), out) + tmp.unlink() + tmp.with_suffix(".onnx.data").unlink(missing_ok=True) + return out + + +def _rename(model: onnx.ModelProto, prefix: str, mapping: dict[str, str], keep: set[str]): + """Rename every tensor and node of a part. Names in `mapping` are replaced; names in `keep` are shared + across parts (cos/sin/mask); everything else gets `prefix`.""" + def new(name: str) -> str: + if name == "": + return name + if name in mapping: + return mapping[name] + if name in keep: + return name + return prefix + name + + g = model.graph + nodes = [] + for node in g.node: + n = onnx.NodeProto() + n.CopyFrom(node) + del n.input[:] + n.input.extend(new(x) for x in node.input) + del n.output[:] + n.output.extend(new(x) for x in node.output) + n.name = prefix + node.name if node.name else "" + nodes.append(n) + inits = [] + for init in g.initializer: + t = TensorProto() + t.CopyFrom(init) + t.name = new(init.name) + inits.append(t) + return nodes, inits + + +def merge(layer_paths: list[Path], head_path: Path, out_path: Path, precision: str) -> None: + all_nodes, all_inits, opsets = [], [], {} + for i, path in enumerate(layer_paths): + m = onnx.load(str(path), load_external_data=False) + src_in = "embeds" if i == 0 else f"hid_{i - 1}" + mapping = {"h": src_in, "past_k": f"past_k_{i}", "past_v": f"past_v_{i}", + "out": f"hid_{i}", "new_k": f"present_k_{i}", "new_v": f"present_v_{i}"} + nodes, inits = _rename(m, f"L{i}_", mapping, keep={"cos", "sin", "mask"}) + all_nodes += nodes + all_inits += inits + for op in m.opset_import: + opsets[op.domain] = max(opsets.get(op.domain, 0), op.version) + hm = onnx.load(str(head_path), load_external_data=False) + nodes, inits = _rename(hm, "H_", {"hid": f"hid_{len(layer_paths) - 1}", "hidden": "hidden", + "cb0_logits": "cb0_logits"}, keep=set()) + all_nodes += nodes + all_inits += inits + for op in hm.opset_import: + opsets[op.domain] = max(opsets.get(op.domain, 0), op.version) + + data_name = out_path.with_suffix(".data").name + _consolidate(all_inits, layer_paths + [head_path], out_path.parent, data_name) + + inputs = [helper.make_tensor_value_info("embeds", TensorProto.FLOAT, [1, "T", D_MODEL]), + 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"])] + outputs = [helper.make_tensor_value_info("hidden", TensorProto.FLOAT, [1, "T", D_MODEL]), + helper.make_tensor_value_info("cb0_logits", TensorProto.FLOAT, [1, "T", N_CB0])] + n = len(layer_paths) + for i in range(n): + inputs.append(helper.make_tensor_value_info(f"past_k_{i}", TensorProto.FLOAT, [1, N_KV_HEADS, "P", HEAD_DIM])) + inputs.append(helper.make_tensor_value_info(f"past_v_{i}", TensorProto.FLOAT, [1, N_KV_HEADS, "P", HEAD_DIM])) + outputs.append(helper.make_tensor_value_info(f"present_k_{i}", TensorProto.FLOAT, [1, N_KV_HEADS, "total", HEAD_DIM])) + outputs.append(helper.make_tensor_value_info(f"present_v_{i}", TensorProto.FLOAT, [1, N_KV_HEADS, "total", HEAD_DIM])) + + graph = helper.make_graph(all_nodes, f"maya2_backbone_{precision}", inputs, outputs, all_inits) + opset_imports = [helper.make_opsetid(d, v) for d, v in sorted(opsets.items())] + ir_version = onnx.load(str(layer_paths[0]), load_external_data=False).ir_version + model = helper.make_model(graph, opset_imports=opset_imports, ir_version=ir_version) + onnx.save_model(model, str(out_path)) + + +def _consolidate(inits: list[TensorProto], parts: list[Path], out_dir: Path, data_name: str) -> None: + """Copy every external tensor's bytes into one data file, updating offset and location in place.""" + offset = 0 + with open(out_dir / data_name, "wb") as dst: + for t in inits: + if t.data_location != TensorProto.EXTERNAL: + continue + entries = {e.key: e.value for e in t.external_data} + loc = entries["location"] + src = _source_file(loc, parts) + length = int(entries["length"]) + start = int(entries.get("offset", 0)) + with open(src, "rb") as f: + f.seek(start) + remaining = length + while remaining: + chunk = f.read(min(COPY_CHUNK, remaining)) + if not chunk: + raise IOError(f"short read for {t.name} in {src}") + dst.write(chunk) + remaining -= len(chunk) + del t.external_data[:] + for k, v in (("location", data_name), ("offset", str(offset)), ("length", str(length))): + e = t.external_data.add() + e.key, e.value = k, v + offset += length + + +def _source_file(location: str, parts: list[Path]) -> Path: + for p in parts: + if (p.parent / location).exists(): + return p.parent / location + raise FileNotFoundError(location) + + +def build(precision: str) -> Path: + if precision not in ("fp32", "wo8"): + raise ValueError(precision) + build_dir = work_dir() / "backbone_build" / precision + build_dir.mkdir(parents=True, exist_ok=True) + paths = ckpt_paths() + sb = StreamedBackbone(paths["backbone_dir"], shard_files(paths["ckpt"]), n_layers=N_LAYERS) + layer_paths = [] + for i in range(N_LAYERS): + layer_paths.append(export_layer(sb, i, precision, build_dir)) + print(f"layer {i:02d} {precision} ok", flush=True) + head = export_head(sb, build_dir) + out_dir = work_dir() / "assets" / "v1" + out_dir.mkdir(parents=True, exist_ok=True) + out = out_dir / f"backbone_{precision}.onnx" + merge(layer_paths, head, out, precision) + print(f"wrote {out} and {out.with_suffix('.data').name}", flush=True) + return out + + +if __name__ == "__main__": + ap = argparse.ArgumentParser() + ap.add_argument("--precision", choices=("fp32", "wo8"), required=True) + build(ap.parse_args().precision) diff --git a/tools/maya_export/export_codec.py b/tools/maya_export/export_codec.py new file mode 100644 index 0000000000000000000000000000000000000000..4f94b90fabaac6344ba5eea65b33dc91be2effb1 --- /dev/null +++ b/tools/maya_export/export_codec.py @@ -0,0 +1,137 @@ +"""Export the Mimi decoder (codes -> 24 kHz wav) to ONNX, fp32, dynamic frame count, and check parity against +transformers' MimiModel.decode on random codes. Run guarded: + + .venv/bin/python tools/memguard.py --rss-limit-mb 2000 -- \ + .venv-tts/bin/python -m tools.maya_export.export_codec +""" +from __future__ import annotations + +import time +from pathlib import Path + +import numpy as np +import onnxruntime as ort +import torch +from transformers import MimiModel + +from tools.maya_export.common import ckpt_paths, work_dir + +N_Q = 16 +N_CODES = 2048 +SAMPLES_PER_FRAME = 1920 # 24 kHz / 12.5 Hz +PARITY_TOL = 1e-4 + + +class _Decode(torch.nn.Module): + def __init__(self, model: MimiModel) -> None: + super().__init__() + self.model = model + + def forward(self, codes: torch.Tensor) -> torch.Tensor: + return self.model.decode(codes)[0] # (1, 1, F * 1920) + + +def _mimi_dir() -> Path: + from maya3 import hub + return Path(hub.resolve("maya-research/maya2-native")["mimi_dir"]) + + +def _traceable_sliding_mask(config, inputs_embeds, **_): + """Explicit causal sliding-window additive mask. transformers' mask helper fails under torch.onnx tracing; + check_parity compares the exported graph with the unpatched model, so a wrong mask shows up there.""" + t = inputs_embeds.shape[1] + i = torch.arange(t)[:, None] + j = torch.arange(t)[None, :] + allowed = (j <= i) & (j > i - int(config.sliding_window)) + return torch.where(allowed, 0.0, float("-inf"))[None, None].to(inputs_embeds.dtype) + + +def check_mask(model: MimiModel, lengths=(5, 300)) -> None: + """The traced graph uses _traceable_sliding_mask. Compare it with transformers' own mask on lengths beyond + the sliding window (250 frames), where the window actually cuts attention.""" + import transformers.models.mimi.modeling_mimi as M + + for t in lengths: + emb = torch.zeros(1, t, model.config.hidden_size) + pos = torch.arange(t)[None] + ref = M.create_sliding_window_causal_mask(config=model.config, inputs_embeds=emb, attention_mask=None, + past_key_values=None, position_ids=pos) + ref = ref.reshape(t, t) + ref_allowed = ref if ref.dtype == torch.bool else ref == 0 + mine = _traceable_sliding_mask(model.config, emb).reshape(t, t) == 0 + if not torch.equal(ref_allowed, mine): + raise SystemExit(f"sliding-window mask differs from transformers at T={t}") + print(f"sliding-window mask matches transformers at T={list(lengths)}") + + +def export(out_path: Path) -> MimiModel: + import transformers.models.mimi.modeling_mimi as M + + model = MimiModel.from_pretrained(str(_mimi_dir()), dtype=torch.float32, attn_implementation="eager").eval() + wrapper = _Decode(model).eval() + dummy = torch.randint(0, N_CODES, (1, N_Q, 5), dtype=torch.long) + saved = M.create_sliding_window_causal_mask + M.create_sliding_window_causal_mask = _traceable_sliding_mask + try: + _trace(wrapper, dummy, out_path) + finally: + M.create_sliding_window_causal_mask = saved + return model + + +def _trace(wrapper: _Decode, dummy: torch.Tensor, out_path: Path) -> None: + with torch.inference_mode(): + torch.onnx.export( + wrapper, + (dummy,), + str(out_path), + input_names=["codes"], + output_names=["wav"], + dynamic_axes={"codes": {2: "F"}, "wav": {2: "samples"}}, + opset_version=17, + dynamo=False, + ) + + +def check_parity(out_path: Path, model: MimiModel, frames=(5, 17, 50)) -> float: + sess = ort.InferenceSession(str(out_path), providers=["CPUExecutionProvider"]) + worst = 0.0 + rng = np.random.default_rng(0) + for F in frames: + codes = rng.integers(0, N_CODES, size=(1, N_Q, F), dtype=np.int64) + with torch.inference_mode(): + ref = model.decode(torch.from_numpy(codes))[0].float().numpy() + got = sess.run(None, {"codes": codes})[0] + assert got.shape == ref.shape == (1, 1, F * SAMPLES_PER_FRAME), (got.shape, ref.shape) + worst = max(worst, float(np.max(np.abs(got - ref)))) + return worst + + +def speed(out_path: Path, F: int = 50, reps: int = 20) -> float: + sess = ort.InferenceSession(str(out_path), providers=["CPUExecutionProvider"]) + codes = np.random.default_rng(1).integers(0, N_CODES, size=(1, N_Q, F), dtype=np.int64) + sess.run(None, {"codes": codes}) + times = [] + for _ in range(reps): + t0 = time.perf_counter() + sess.run(None, {"codes": codes}) + times.append(time.perf_counter() - t0) + return float(np.median(times)) / F * 1000.0 # ms per frame + + +def main() -> None: + out = work_dir() / "assets" / "v1" / "mimi_decoder.onnx" + out.parent.mkdir(parents=True, exist_ok=True) + model = export(out) + print(f"exported {out} ({out.stat().st_size / 2**20:.1f} MiB)") + check_mask(model) + err = check_parity(out, model) + print(f"parity max abs error vs MimiModel.decode: {err:.3e} (tol {PARITY_TOL:.0e})") + ms = speed(out) + print(f"decode median: {ms:.2f} ms per frame at F=50 (gate <= 8 ms)") + if err > PARITY_TOL: + raise SystemExit("parity gate FAILED") + + +if __name__ == "__main__": + main() diff --git a/tools/maya_export/export_depth.py b/tools/maya_export/export_depth.py new file mode 100644 index 0000000000000000000000000000000000000000..c9c2a6f6f8edd377d4a449846a5d00297a3232e2 --- /dev/null +++ b/tools/maya_export/export_depth.py @@ -0,0 +1,83 @@ +"""Export the 15 per-level depth graphs (fp32) to work_dir()/depth_build/fp32/level_XX.onnx. + +Each level's weights are read from the checkpoint and loaded into a fresh DepthLevel; one level at a time. + + .venv/bin/python tools/memguard.py --rss-limit-mb 4000 --min-free-mb 4000 -- \ + .venv-tts/bin/python -m tools.maya_export.export_depth +""" +from __future__ import annotations + +import torch + +from tts.model.depth_v2 import GROUP_OF_LEVEL +from tts.tokens import N_Q + +from tools.maya_export.common import ckpt_paths, work_dir +from tools.maya_export.depth_frame import DepthLevel, Q_SERVE +from tools.maya_export.export_backbone import D_MODEL, _export, _quantize_wo8, _save_external +from tools.maya_export.streamed_backbone import _Weights, shard_files + +D_DEPTH = 1024 +N_LEVELS = 15 # serving q_serve = 16: levels 1..15 (the checkpoint has up to 31) + + +def ckpt_name(name: str, level: int) -> str: + """DepthLevel parameter name -> checkpoint key (without the model. prefix).""" + if name.startswith("trunk."): + _, li, rest = name.split(".", 2) + return f"depth.groups.{GROUP_OF_LEVEL[level]}.{li}.{rest}" + if name.startswith("head."): + return f"depth.heads.{level - 1}.weight" + return f"depth.{name}" + + +def _head_nodes(model) -> list[str]: + """Names of the node(s) producing the `logits` output (the level head's MatMul): kept in fp32 for wo8_fph.""" + names = [] + for node in model.graph.node: + if "logits" in node.output: + names.append(node.name) + return names + + +def export_level(level: int, weights: _Weights, build, precision: str) -> None: + mod = DepthLevel(level).eval() + state = {name: weights.get(ckpt_name(name, level)).float() for name, _ in mod.named_parameters()} + mod.load_state_dict(state, strict=True) + del state + h = torch.randn(1, D_MODEL) * 0.02 + codes = torch.zeros(1, Q_SERVE, dtype=torch.long) + mem = torch.zeros(1, 1, D_DEPTH) + tmp = build / f"level_{level:02d}_raw.onnx" + _export(mod, (h, codes, mem), tmp, input_names=["h", "codes", "mem"], output_names=["logits"], + dynamic_axes={"h": {0: "B"}, "codes": {0: "B"}, "mem": {0: "B"}, "logits": {0: "B"}}) + import onnx + + model = onnx.load(str(tmp)) + if precision == "wo8": + model = _quantize_wo8(model) + elif precision == "wo8_fph": # int8 trunks, fp32 head + model = _quantize_wo8(model, exclude=_head_nodes(model)) + _save_external(model, build / f"level_{level:02d}.onnx") + tmp.unlink() + tmp.with_suffix(".onnx.data").unlink(missing_ok=True) + del mod + + +def main() -> None: + import argparse + + ap = argparse.ArgumentParser() + ap.add_argument("--precision", choices=("fp32", "wo8", "wo8_fph"), default="fp32") + precision = ap.parse_args().precision + build = work_dir() / "depth_build" / precision + build.mkdir(parents=True, exist_ok=True) + paths = ckpt_paths() + weights = _Weights(shard_files(paths["ckpt"])) + for level in range(1, N_LEVELS + 1): + export_level(level, weights, build, precision) + print(f"level {level:02d} fp32 ok", flush=True) + + +if __name__ == "__main__": + main() diff --git a/tools/maya_export/export_depth_all.py b/tools/maya_export/export_depth_all.py new file mode 100644 index 0000000000000000000000000000000000000000..8c0dcbcb6bae32c95315ebdc75b90a5e68a5e6ab --- /dev/null +++ b/tools/maya_export/export_depth_all.py @@ -0,0 +1,72 @@ +"""Export the single fp16-storage depth graph (levels 1..15) to work_dir()/depth_build/fp16/depth_all.onnx. + + .venv/bin/python tools/memguard.py --rss-limit-mb 4000 --min-free-mb 4000 -- \ + .venv-tts/bin/python -m tools.maya_export.export_depth_all +""" +from __future__ import annotations + +import onnx +import torch + +from tts.tokens import N_CODES + +from tools.maya_export.common import ckpt_paths, work_dir +from tools.maya_export.depth_all import LEVELS, DepthAll, DepthSample, load_fp16 +from tools.maya_export.export_backbone import D_MODEL, _save_external +from tools.maya_export.streamed_backbone import _Weights, shard_files + +Q_SERVE = 16 + + +def ckpt_name(name: str) -> str: + """DepthAll parameter name -> checkpoint key (without the model. prefix).""" + if name.startswith("groups."): + _, g, li, rest = name.split(".", 3) + return f"depth.groups.{g}.{li}.{rest}" + return f"depth.{name}" + + +def main() -> None: + out_dir = work_dir() / "depth_build" / "fp16" + out_dir.mkdir(parents=True, exist_ok=True) + weights = _Weights(shard_files(ckpt_paths()["ckpt"])) + with torch.device("meta"): + model = DepthAll() + load_fp16(model, lambda n: weights.get(ckpt_name(n))) + model.eval() + h = torch.randn(1, D_MODEL) * 0.02 + codes = torch.zeros(1, Q_SERVE, dtype=torch.long) + mem = torch.zeros(1, 1, 1024) + tmp = out_dir / "depth_all_raw.onnx" + # no constant folding: it would fold each Cast(fp16 weight) into an fp32 constant and double the stored weights + with torch.inference_mode(): + torch.onnx.export(model, (h, codes, mem), str(tmp), input_names=["h", "codes", "mem"], + output_names=["logits"], opset_version=17, dynamo=False, do_constant_folding=False, + dynamic_axes={"h": {0: "B"}, "codes": {0: "B"}, "mem": {0: "B"}, "logits": {0: "B"}}) + _save_external(onnx.load(str(tmp)), out_dir / "depth_all.onnx") + tmp.unlink() + tmp.with_suffix(".onnx.data").unlink(missing_ok=True) + print(f"wrote {out_dir / 'depth_all.onnx'} ({LEVELS} levels)", flush=True) + + # sampled frame graph (the runtime's Depth.frame): inputs h, cb0, mem, noise; output codes (B, 16) + sample = DepthSample(model).eval() + # example batch 2: a batch of 1 would specialise the dynamic batch dimension + h = torch.randn(2, D_MODEL) * 0.02 + cb0 = torch.zeros(2, dtype=torch.long) + mem = torch.zeros(2, 1, 1024) + noise = torch.zeros(2, LEVELS, N_CODES) + tmp = out_dir / "depth_frame_raw.onnx" + # dynamo exporter: traces symbolically, so the per-use fp32 weight copies are not held in memory + B = torch.export.Dim("B") + with torch.inference_mode(): + prog = torch.export.export(sample, (h, cb0, mem, noise), + dynamic_shapes={"h": {0: B}, "cb0": {0: B}, "mem": {0: B}, "noise": {0: B}}) + torch.onnx.export(prog, (), str(tmp), input_names=["h", "cb0", "mem", "noise"], output_names=["codes"]) + _save_external(onnx.load(str(tmp)), out_dir / "depth_frame_fp16.onnx") + tmp.unlink() + tmp.with_suffix(".onnx.data").unlink(missing_ok=True) + print(f"wrote {out_dir / 'depth_frame_fp16.onnx'}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/tools/maya_export/export_prompt_assets.py b/tools/maya_export/export_prompt_assets.py new file mode 100644 index 0000000000000000000000000000000000000000..6c1ccf3d4ef2e8c82ed17dbc8eecee598fc09ccc --- /dev/null +++ b/tools/maya_export/export_prompt_assets.py @@ -0,0 +1,115 @@ +"""Export the prompt-side assets (tokenizer, bf16 embedding tables, speaker anchor, audio input, voices, +sentence controls) to MAYA_FAST_WORK/assets/v1 for the numpy runtime. Run guarded: + + .venv/bin/python tools/memguard.py --rss-limit-mb 4000 --min-free-mb 4000 -- \ + .venv-tts/bin/python -m tools.maya_export.export_prompt_assets +""" +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import torch + +from maya3 import guards as G +from maya3.engine import ROOT +from maya3.prompt import MODES +from scripts import s2_serve_controls as S + +from tools.maya_export.common import build_light_model, ckpt_paths, work_dir +from tools.maya_export.manifest import write_manifest + +N_Q = 16 +N_CODES = 2048 +ECAPA_DIM = 192 +HARD_CAP = 1500 # frame cap ceiling used by the goldens (golden_light.py: min(1500, ...)) +CHUNK_S = 16.0 # upstream default in s2_serve_controls.py (golden_light.py: chunk_s) + + +def _bf16_to_file(t: torch.Tensor, path: Path) -> int: + """Write a bf16 tensor as raw little-endian uint16 (row-major). Returns the row count.""" + raw = t.detach().contiguous().view(torch.int16).numpy().astype(" dict: + coef_model = G.DurationPrior() + g = G.GuardConfig() + prior = S.load_prior(prior_path) + return { + "text_frames_prior": prior, + "prior_default": {"a": S.PRIOR_DEFAULT["a"], "b": S.PRIOR_DEFAULT["b"], + "resid_p99": S.PRIOR_DEFAULT["resid_p99"]}, + "text": {"terminals": S.TERMINALS, "fps": S.FPS, "frame_s": G.FRAME_S, "chunk_s": CHUNK_S}, + "guard": { + "enabled": bool(g.enabled), + "dur_k": float(g.dur_k), + "dur_margin_s": float(g.dur_margin_s), + "hard_cap": HARD_CAP, + "dur_prior_coefs": {lang: [float(x) for x in c] for lang, c in coef_model.coef.items()}, + "dur_prior_default": [float(x) for x in coef_model.default], + }, + "sampling": {"q_serve": 16, "temp": 0.7, "top_k": 64, "tau": 0.35, "hazard": 0.6, + "cfg_scale": 0.0, "aeos_class": N_CODES}, + } + + +def main() -> None: + out = work_dir() / "assets" / "v1" + out.mkdir(parents=True, exist_ok=True) + + model = build_light_model() + with torch.inference_mode(): + # tokenizer: the extended tokenizer the model was trained with (backend, fast format) + model.backbone.tokenizer.backend_tokenizer.save(str(out / "tokenizer.json")) + + embed_rows = _bf16_to_file(model.backbone.lm.model.embed_tokens.weight, out / "embed_bf16.bin") + audio_rows = _bf16_to_file(model.audio_in.table.weight, out / "audio_in_bf16.bin") + + anchor = model.speaker_anchor + prompt = { + "anchor_fc1_w": anchor.fc1.weight.detach().float().numpy(), + "anchor_fc1_b": anchor.fc1.bias.detach().float().numpy(), + "anchor_fc2_w": anchor.fc2.weight.detach().float().numpy(), + "anchor_fc2_b": anchor.fc2.bias.detach().float().numpy(), + "anchor_null": anchor.null_vec.detach().float().numpy(), + "audio_in/aeos_in": model.audio_in.aeos_in.detach().float().numpy(), + "audio_in/norm": model.audio_in.norm.weight.detach().float().numpy(), + } + np.savez(out / "prompt.npz", **prompt) + + # voices: every trained named voice, with the conditioning flags the upstream builder uses + paths = ckpt_paths() + with open(paths["voices_json"], "r", encoding="utf-8") as f: + raw = json.load(f) + items = raw["voices"] if isinstance(raw, dict) and "voices" in raw else raw + voices = {} + for it in items: + name = it.get("name") + if not name: + continue + mode = it.get("mode", "named") + use_icl, use_mem = MODES[mode] + voices[name] = {"mode": mode, "use_icl": bool(use_icl), "use_mem": bool(use_mem)} + (out / "voices.json").write_text(json.dumps(voices, indent=2, ensure_ascii=False), encoding="utf-8") + + controls = _controls(str(Path(ROOT) / "configs" / "text_frames_prior.json")) + (out / "controls.json").write_text(json.dumps(controls, indent=2, ensure_ascii=False), encoding="utf-8") + + names = ["tokenizer.json", "embed_bf16.bin", "audio_in_bf16.bin", "prompt.npz", + "voices.json", "controls.json"] + write_manifest(out, names, shapes={ + "d_model": int(model.audio_in.aeos_in.shape[0]), + "vocab": embed_rows, + "audio_rows": audio_rows, + "n_q": N_Q, + "n_codes": N_CODES, + "ecapa_dim": ECAPA_DIM, + }) + print(f"wrote prompt assets to {out}") + + +if __name__ == "__main__": + main() diff --git a/tools/maya_export/finalize_assets.py b/tools/maya_export/finalize_assets.py new file mode 100644 index 0000000000000000000000000000000000000000..d122bfec1be47b7e3f6fa3d4a744f18aba131be7 --- /dev/null +++ b/tools/maya_export/finalize_assets.py @@ -0,0 +1,31 @@ +"""Rewrite work_dir()/assets/v1/manifest.json with every runtime file and the runtime shapes. + +Shapes used by the runtime: rope_theta (Qwen3 backbone), n_layers / n_heads (KV heads) / head_dim, d_mem (depth +speaker memory width). Run after the backbone, depth and Mimi exports. +""" +from __future__ import annotations + +import json + +from tools.maya_export.common import work_dir +from tools.maya_export.manifest import write_manifest + +RUNTIME_FILES = [ + "tokenizer.json", "embed_bf16.bin", "audio_in_bf16.bin", "prompt.npz", "voices.json", "controls.json", + "backbone_wo8.onnx", "backbone_wo8.data", + "depth_frame_fp16.onnx", "depth_frame_fp16.onnx.data", + "mimi_decoder.onnx", +] + + +def main() -> None: + out = work_dir() / "assets" / "v1" + prev = json.loads((out / "manifest.json").read_text(encoding="utf-8")) + shapes = dict(prev.get("shapes", {})) + shapes.update({"rope_theta": 1000000.0, "n_layers": 28, "n_heads": 8, "head_dim": 128, "d_mem": 1024}) + write_manifest(out, RUNTIME_FILES, shapes=shapes) + print(f"wrote {out / 'manifest.json'} ({len(RUNTIME_FILES)} files)", flush=True) + + +if __name__ == "__main__": + main() diff --git a/tools/maya_export/golden_light.py b/tools/maya_export/golden_light.py new file mode 100644 index 0000000000000000000000000000000000000000..555db37c3fd52600fa4fe3c165e4a61ff465cbe0 --- /dev/null +++ b/tools/maya_export/golden_light.py @@ -0,0 +1,138 @@ +from __future__ import annotations + +import json +from pathlib import Path +import numpy as np +import torch + +from maya3 import guards as G +from maya3.engine import ROOT +from maya3.prompt import DEFAULT_BUILDER, MODES, VoiceRegistry, load_named_voices +from scripts import s2_serve_controls as S +from tts.data.packing_v2 import PackerV2 + +from tools.maya_export.common import build_light_model, ckpt_paths, work_dir + + +def main() -> None: + out_dir = work_dir() / "golden" + out_dir.mkdir(parents=True, exist_ok=True) + + testset_path = Path("benchmarks/tts_testset_mr.json") + if not testset_path.exists(): + testset_path = Path("/Users/anujbarve/Documents/marathi-helper/benchmarks/tts_testset_mr.json") + with open(testset_path, "r", encoding="utf-8") as f: + test_sentences = json.load(f) + + paths = ckpt_paths() + prior_path = Path(ROOT) / "configs" / "text_frames_prior.json" + prior = S.load_prior(str(prior_path)) + + registry = VoiceRegistry() + load_named_voices(registry, paths["voices_json"]) + # Also load the voices.json file to get the list of voice names + with open(paths["voices_json"], "r", encoding="utf-8") as f: + raw_v = json.load(f) + v_items = raw_v["voices"] if isinstance(raw_v, dict) and "voices" in raw_v else raw_v + voice_names = [v["name"] for v in v_items if isinstance(v, dict) and "name" in v] + + model = build_light_model() + packer = PackerV2(model.backbone.tokenizer, model.backbone.token_ids, max_tokens=8192) + + g = G.GuardConfig() + dur = G.DurationPrior() + chunk_s = 16.0 # upstream default in s2_serve_controls.py + + npz_data = {} + json_data = { + "voices": {}, + "sentences": {}, + } + + # 1. Per voice: mem, ecapa + print(f"Generating voice goldens for {len(voice_names)} voices...") + for vname in voice_names: + voice = registry.get(vname) + use_icl, use_mem = MODES[voice.mode] + builder = voice.builder or DEFAULT_BUILDER + voice.prefix = builder.prefix(packer, voice) + voice.spk_refs = ( + np.asarray(voice.ecapa, np.float32).reshape(1, 192) + if voice.ecapa is not None + else np.zeros((1, 192), np.float32) + ) + refs_t = torch.as_tensor(voice.spk_refs, device="cpu")[None] + with torch.inference_mode(): + mem = model.depth.spk_memory(refs_t, "ecapa", null_mask=torch.tensor([not use_mem], device="cpu")) + + npz_data[f"voice/{vname}/mem"] = mem.cpu().float().numpy() + npz_data[f"voice/{vname}/ecapa"] = voice.spk_refs + + json_data["voices"][vname] = { + "mode": voice.mode, + "use_mem": bool(use_mem), + "use_icl": bool(use_icl), + } + + # 2. Per sentence × every voice + print(f"Generating prompt goldens for {len(test_sentences)} sentences × {len(voice_names)} voices...") + lang = "mr" + for s_idx, item in enumerate(test_sentences): + s_id = item["id"] + text = item["text"] + + # Sentence-level controls + ef, er = S.expected_frames(prior, lang, len(text)) + loose = min(1500, int(2.0 * (ef + er) + 40)) + cap = loose + if g.enabled: + cap = min(cap, max(int(ef * 0.5) + 8, dur.cap_frames(text, lang, g.dur_k, g.dur_margin_s))) + minimum = max(4, int(0.5 * ef)) + split_chunks = S.split_text(text, lang, prior, chunk_s) + + json_data["sentences"][s_id] = { + "text": text, + "ef": float(ef), + "er": float(er), + "cap": int(cap), + "minimum": int(minimum), + "split_text": split_chunks, + } + + for vname in voice_names: + voice = registry.get(vname) + builder = voice.builder or DEFAULT_BUILDER + with torch.inference_mode(): + ex = builder.example(packer, voice, text, lang) + prompt_emb, n_p = S.prompt_from_example(ex, model, "cpu", torch.float32, 16) + + prefix_key = f"prompt/{s_id}/{vname}" + npz_data[f"{prefix_key}/token_ids"] = np.asarray(ex["tokens"], dtype=np.int64) + npz_data[f"{prefix_key}/kinds"] = np.asarray(ex["kind"], dtype=np.int64) + npz_data[f"{prefix_key}/prompt_emb"] = prompt_emb[0].cpu().float().numpy() + npz_data[f"{prefix_key}/n_prompt"] = np.int64(n_p) + + # 3. Fixed random code rows (seed 0) audio_in frames and aeos + print("Generating audio_in and aeos goldens...") + torch.manual_seed(0) + test_codes = torch.randint(0, 2048, (8, 16), dtype=torch.int64) + with torch.inference_mode(): + audio_in_frames = model.audio_in.frames(test_codes, 16) + aeos = model.audio_in.aeos(1)[0] + + npz_data["audio_in_test_codes"] = test_codes.numpy() + npz_data["audio_in_frames"] = audio_in_frames.cpu().float().numpy() + npz_data["aeos"] = aeos.cpu().float().numpy() + + npz_path = out_dir / "light.npz" + json_path = out_dir / "light.json" + + np.savez(npz_path, **npz_data) + with open(json_path, "w", encoding="utf-8") as f: + json.dump(json_data, f, indent=2, ensure_ascii=False) + + print(f"Successfully wrote {npz_path} and {json_path}") + + +if __name__ == "__main__": + main() diff --git a/tools/maya_export/golden_reference.py b/tools/maya_export/golden_reference.py new file mode 100644 index 0000000000000000000000000000000000000000..f51cbb645f1008b958a6e50eadb48401c5b00e7b --- /dev/null +++ b/tools/maya_export/golden_reference.py @@ -0,0 +1,208 @@ +"""Heavy reference corpus: upstream serving path (controlled stop, temp 0.7, top-k 64, depth sampler) with the +backbone streamed one layer at a time (tools/maya_export/streamed_backbone.py), so it fits 4 GB. + +Per sentence the three seeds run as one batch (same prompt, no padding). Each row has its own torch.Generator for +the cb0 draw; the depth draws use the global RNG, seeded per step from the sentence index, so a rerun reproduces. + +Per (sentence, seed) writes work_dir()/golden/ref/{id}_s{seed}.npz with: + codes (16, F) int64 generated codes (cb0 first) + wav (F*1920,) fp32 Mimi decode of codes, 24 kHz + reason str stop reason ("cap", "aeos_argmax", "aeos_tau", "aeos_hazard") + h_frames (F, 2048) fp32 backbone hidden fed to depth for each frame + cb0_logits (F, 2049) fp16 raw cb0 head logits per frame + depth_logits (F, 15, 2048) fp16 teacher-forced depth head logits for levels 1..15, from h_frames and codes + seq_emb (n_p + F, 2048) fp32 backbone input sequence: prompt embeddings, then the fed-back frame embeddings + +Run guarded (one model process at a time): + .venv/bin/python tools/memguard.py --rss-limit-mb 4000 --min-free-mb 4000 -- \ + .venv-tts/bin/python -m tools.maya_export.golden_reference +""" +from __future__ import annotations + +import json +import os +from pathlib import Path + +import numpy as np +import torch + +from maya3.engine import ROOT +from maya3.prompt import DEFAULT_BUILDER, MODES, VoiceRegistry, load_named_voices +from scripts import s2_serve_controls as S +from tts.data.packing_v2 import PackerV2 +from tts.model.depth_v2 import GROUP_OF_LEVEL +from tts.tokens import AEOS_CLASS + +from tools.maya_export.common import build_light_model, ckpt_paths, work_dir +from tools.maya_export.streamed_backbone import StreamedBackbone, shard_files + +VOICE = "Charvi" +SEEDS = (0, 1, 2) +TEMP = 0.7 +TOP_K = 64 +TAU = 0.35 +HAZARD = 0.6 +Q_SERVE = 16 +N_LEVELS = Q_SERVE - 1 + + +def _cb0_probs(logits: torch.Tensor) -> torch.Tensor: + """Upstream ServeSession._cb0_probs for a batch: temperature, top-k mask, softmax.""" + lg = logits.float() / TEMP + kth = lg.topk(TOP_K, -1).values[:, -1:] + lg = lg.masked_fill(lg < kth, -float("inf")) + return lg.softmax(-1) + + +def _voice_mem(model, voice) -> torch.Tensor: + """Depth speaker memory for the voice (named voices: null memory, as the golden prompts).""" + _, use_mem = MODES[voice.mode] + refs = np.asarray(voice.ecapa, np.float32).reshape(1, 192) if voice.ecapa is not None else np.zeros((1, 192), np.float32) + refs_t = torch.as_tensor(refs, device="cpu")[None] + return model.depth.spk_memory(refs_t, "ecapa", null_mask=torch.tensor([not use_mem], device="cpu")) + + +def _teacher_forced_depth(model, h_frames: torch.Tensor, codes: torch.Tensor, mem: torch.Tensor) -> torch.Tensor: + """Logits of every depth head for levels 1..15 on the generated codes. h_frames (F, D), codes (16, F).""" + depth = model.depth + F = h_frames.shape[0] + mem_f = mem.expand(F, -1, -1) + ref_mask = torch.ones(F, 1, dtype=torch.bool) + out = [] + codes_t = codes.T.contiguous() # (F, 16) + for level in range(1, Q_SERVE): + x = depth._prefix(h_frames, codes_t[:, :level], level) + h = depth._run(x, GROUP_OF_LEVEL[level], mem_f, ref_mask) + out.append(depth.heads[level - 1](h).float()) + return torch.stack(out, 1) # (F, 15, 2048) + + +def main() -> None: + out_dir = work_dir() / "golden" / "ref" + out_dir.mkdir(parents=True, exist_ok=True) + golden = json.loads((work_dir() / "golden" / "light.json").read_text(encoding="utf-8")) + sentences = golden["sentences"] + limit = int(os.environ.get("MAYA_REF_LIMIT", "0")) # smoke runs only; 0 = whole corpus + if limit: + sentences = dict(list(sentences.items())[:limit]) + + model = build_light_model() + paths = ckpt_paths() + registry = VoiceRegistry() + load_named_voices(registry, paths["voices_json"]) + voice = registry.get(VOICE) + builder = voice.builder or DEFAULT_BUILDER + packer = PackerV2(model.backbone.tokenizer, model.backbone.token_ids, max_tokens=8192) + voice.prefix = builder.prefix(packer, voice) + voice.spk_refs = np.zeros((1, 192), np.float32) + mem = _voice_mem(model, voice) + audio_in = model.audio_in + cb0_head = model.backbone.cb0_logits + + stream = StreamedBackbone(paths["backbone_dir"], shard_files(paths["ckpt"])) + B = len(SEEDS) + + for s_idx, (s_id, item) in enumerate(sentences.items()): + text, cap, min_f = item["text"], int(item["cap"]), int(item["minimum"]) + ex = builder.example(packer, voice, text, "mr") + with torch.inference_mode(): + prompt_emb, n_p = S.prompt_from_example(ex, model, "cpu", torch.float32, Q_SERVE) + prompt_emb = prompt_emb.float() # (1, n_p, D) + + gens = [torch.Generator().manual_seed(seed) for seed in SEEDS] + stream.reset() + with torch.inference_mode(): + h_all = stream.forward(prompt_emb.expand(B, -1, -1)) + h = h_all[:, -1] + H = np.zeros(B) + reason = [None] * B + frames = [[] for _ in range(B)] + h_rec = [[] for _ in range(B)] + logit_rec = [[] for _ in range(B)] + fed = [[] for _ in range(B)] + + for t in range(cap): + logits = cb0_head(h).float() + p = _cb0_probs(logits) + p_aeos = p[:, AEOS_CLASS].numpy() + active = [i for i in range(B) if reason[i] is None] + if not active: + break + for i in active: + if t >= min_f: + H[i] = 1.0 - (1.0 - H[i]) * (1.0 - p_aeos[i]) + if int(p[i].argmax()) == AEOS_CLASS: + reason[i] = "aeos_argmax" + elif p_aeos[i] >= TAU: + reason[i] = "aeos_tau" + elif H[i] >= HAZARD: + reason[i] = "aeos_hazard" + cb0 = torch.zeros(B, dtype=torch.long) + for i in range(B): + if reason[i] is not None: + continue + q = p[i].clone() + q[AEOS_CLASS] = 0.0 + q = q / q.sum() + cb0[i] = int(torch.multinomial(q, 1, generator=gens[i])) + for i in active: + if reason[i] is None: + logit_rec[i].append(logits[i].half()) + h_rec[i].append(h[i].float()) + + torch.manual_seed(1000 * s_idx + t) + mem_b = mem.expand(B, -1, -1) + frame = S.depth_generate(model.depth, h.float(), cb0, mem_b, Q_SERVE, TEMP, TOP_K, 0.0, + torch.ones(B, 1, dtype=torch.bool)) + fb = audio_in.frames(frame, Q_SERVE) # (B, D) + for i in active: + if reason[i] is None: + frames[i].append(frame[i]) + fed[i].append(fb[i].clone()) + h = stream.forward(fb[:, None, :])[:, -1] + + for i in range(B): + if reason[i] is None: + reason[i] = "cap" + + for i, seed in enumerate(SEEDS): + F = len(frames[i]) + if F == 0: + raise RuntimeError(f"{s_id} seed {seed}: no frames generated (reason {reason[i]})") + codes = torch.stack(frames[i], -1) # (16, F) + h_frames = torch.stack(h_rec[i], 0) # (F, D) + depth_logits = _teacher_forced_depth(model, h_frames, codes, mem).half() + seq = torch.cat([prompt_emb[0], torch.stack(fed[i], 0)], 0) + np.savez( + out_dir / f"{s_id}_s{seed}.npz", + codes=codes.numpy().astype(np.int64), + reason=np.array(reason[i]), + h_frames=h_frames.numpy().astype(np.float32), + cb0_logits=torch.stack(logit_rec[i], 0).numpy().astype(np.float16), + depth_logits=depth_logits.numpy(), + seq_emb=seq.numpy().astype(np.float32), + ) + print(f"{s_id}: frames {[len(f) for f in frames]} reasons {reason} cap {cap} min {min_f}", flush=True) + + _decode_wav(out_dir, sentences.keys()) + + +def _decode_wav(out_dir: Path, s_ids) -> None: + """Mimi decode of every saved codes array, written into the same npz as `wav`.""" + from transformers import MimiModel + + mimi = MimiModel.from_pretrained(ckpt_paths()["mimi_dir"], dtype=torch.float32).eval() + for s_id in s_ids: + for seed in SEEDS: + path = out_dir / f"{s_id}_s{seed}.npz" + with np.load(path) as z: + arrays = {k: z[k] for k in z.files} + with torch.inference_mode(): + wav = mimi.decode(torch.from_numpy(arrays["codes"])[None])[0][0, 0].float().numpy() + arrays["wav"] = wav + np.savez(path, **arrays) + print(f"wrote wav for {len(s_ids) * len(SEEDS)} files", flush=True) + + +if __name__ == "__main__": + main() diff --git a/tools/maya_export/manifest.py b/tools/maya_export/manifest.py new file mode 100644 index 0000000000000000000000000000000000000000..3630aeb00c0462955e74f7bfe45c829724f94dec --- /dev/null +++ b/tools/maya_export/manifest.py @@ -0,0 +1,26 @@ +from __future__ import annotations + +import hashlib +import json +from pathlib import Path + + +def sha256_file(path: Path) -> str: + h = hashlib.sha256() + with open(path, "rb") as f: + while chunk := f.read(1 << 20): + h.update(chunk) + return h.hexdigest() + + +def write_manifest(assets_dir: Path, names: list[str], shapes: dict, version: int = 1) -> dict: + """Hash every named file in assets_dir and write manifest.json next to them. + The runtime refuses any asset whose sha256 differs (src/maya_fast/assets.py).""" + manifest = { + "version": version, + "files": {name: sha256_file(assets_dir / name) for name in names}, + "shapes": shapes, + } + with open(assets_dir / "manifest.json", "w", encoding="utf-8") as f: + json.dump(manifest, f, indent=2, ensure_ascii=False) + return manifest diff --git a/tools/maya_export/streamed_backbone.py b/tools/maya_export/streamed_backbone.py new file mode 100644 index 0000000000000000000000000000000000000000..68ecede344532e93b77d21792c53c90fec6c6b5b --- /dev/null +++ b/tools/maya_export/streamed_backbone.py @@ -0,0 +1,91 @@ +"""Reference backbone that streams one decoder layer at a time from the checkpoint. + +The upstream Qwen3 backbone is 4.7 GiB in bf16. Holding it resident breaks the 4 GB RAM limit, so this runs the +same transformers Qwen3DecoderLayer / Qwen3RMSNorm / Qwen3RotaryEmbedding code, one layer at a time: read the layer's +weights from safetensors, upcast to fp32, run it, move on. The KV cache is a transformers DynamicCache, so the +prefill / one-token-per-step path is the upstream one. Costs: one read of ~2.6 GiB of layer weights per call of +`forward` (per generated token), so it is slow; it is used only for the reference corpus. +""" +from __future__ import annotations + +from pathlib import Path + +import torch +from safetensors import safe_open +from transformers import AutoConfig, DynamicCache +from transformers.models.qwen3.modeling_qwen3 import Qwen3DecoderLayer, Qwen3RMSNorm, Qwen3RotaryEmbedding + +LAYER_PREFIX = "backbone.lm.model.layers.{i}." +NORM_KEY = "backbone.lm.model.norm.weight" + + +class _Weights: + """Name -> tensor lookup over the checkpoint shards, reading one tensor at a time.""" + + def __init__(self, shard_files: list[str], model_prefix: str = "model.") -> None: + self.index: dict[str, str] = {} + for f in shard_files: + with safe_open(f, "pt") as fh: + for k in fh.keys(): + self.index[k] = f + self.model_prefix = model_prefix + + def get(self, name: str) -> torch.Tensor: + """name: checkpoint key without the leading `model.` (e.g. backbone.lm.model.norm.weight).""" + key = name if name in self.index else self.model_prefix + name + with safe_open(self.index[key], "pt") as fh: + return fh.get_tensor(key) + + +class StreamedBackbone: + def __init__(self, backbone_cfg_path: str, shard_files: list[str], n_layers: int | None = None) -> None: + self.cfg = AutoConfig.from_pretrained(backbone_cfg_path) + self.cfg._attn_implementation = "sdpa" # sdpa applies the causal mask when mask is None and T > 1 + self.n_layers = int(n_layers if n_layers is not None else self.cfg.num_hidden_layers) + self.weights = _Weights(shard_files) + self.rotary = Qwen3RotaryEmbedding(self.cfg) + self.norm = Qwen3RMSNorm(self.cfg.hidden_size, eps=self.cfg.rms_norm_eps) + self.norm.weight.data.copy_(self.weights.get(NORM_KEY).float()) + self.layer = Qwen3DecoderLayer(self.cfg, layer_idx=0).eval() + self.cache = DynamicCache(config=self.cfg) + self.past_len = 0 + + def reset(self) -> None: + self.cache = DynamicCache(config=self.cfg) + self.past_len = 0 + + def _load_layer(self, i: int) -> None: + prefix = LAYER_PREFIX.format(i=i) + state = {} + for name, _ in self.layer.named_parameters(): + state[name] = self.weights.get(prefix + name).float() + with torch.no_grad(): + for name, param in self.layer.named_parameters(): + param.data.copy_(state[name]) + del state + self.layer.self_attn.layer_idx = i + + @torch.inference_mode() + def forward(self, embeds: torch.Tensor) -> torch.Tensor: + """embeds: (1, T, hidden) fp32, the next T positions after the cached ones. Returns (1, T, hidden) after the + final norm (the upstream `last_hidden_state`).""" + t = embeds.shape[1] + position_ids = torch.arange(self.past_len, self.past_len + t)[None] + position_embeddings = self.rotary(embeds, position_ids) + h = embeds + for i in range(self.n_layers): + self._load_layer(i) + h = self.layer( + h, + attention_mask=None, # sdpa: causal when T > 1 and no cache, none needed for one-token steps + position_ids=position_ids, + position_embeddings=position_embeddings, + past_key_values=self.cache, + use_cache=True, + ) + self.past_len += t + return self.norm(h) + + +def shard_files(ckpt_dir: str | Path) -> list[str]: + return sorted(str(p) for p in Path(ckpt_dir).glob("*.safetensors")) diff --git a/tools/memguard.py b/tools/memguard.py new file mode 100644 index 0000000000000000000000000000000000000000..74c842cebb66ef07aaaff389573a15c3a05839ac --- /dev/null +++ b/tools/memguard.py @@ -0,0 +1,191 @@ +"""Memory guard: run one command at a time under RAM and swap limits. + +Every command that loads model weights runs under this guard, per the plan's +Global Constraints: + +* an exclusive lock on ``$MAYA_FAST_WORK/model.lock`` (default + ``~/.cache/maya_fast/model.lock``) is held for the whole run, so only one + model process exists at a time across agents; +* the run is refused when free memory is below ``min_free_mb``; +* the child tree is killed when its RSS exceeds ``rss_limit_mb`` or when swap + grew by more than ``max_swap_growth_mb`` during the run. + +CLI: + python tools/memguard.py --rss-limit-mb N --min-free-mb M -- +prints one line ``MEMGUARD peak_rss_mb=... swap_delta_mb=... reason=...`` +and exits non-zero on any violation. + +MB means MiB (1024**2) throughout. +""" + +from __future__ import annotations + +import argparse +import fcntl +import os +import subprocess +import sys +import time +from dataclasses import dataclass +from pathlib import Path + +import psutil + +MB = 1024 * 1024 + + +@dataclass +class GuardResult: + returncode: int + peak_rss_mb: float + swap_before_mb: float + swap_after_mb: float + killed: bool + reason: str | None + + +def lock_path() -> Path: + """Shared model lock: one guarded model run at a time, across processes.""" + work = os.environ.get("MAYA_FAST_WORK") or str(Path.home() / ".cache" / "maya_fast") + return Path(work) / "model.lock" + + +def _acquire_lock(path: Path): + path.parent.mkdir(parents=True, exist_ok=True) + fh = open(path, "a") + try: + fcntl.flock(fh.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + except (BlockingIOError, OSError): + print(f"MEMGUARD waiting for model lock {path}", file=sys.stderr, flush=True) + fcntl.flock(fh.fileno(), fcntl.LOCK_EX) + return fh + + +def _tree_rss(proc: subprocess.Popen) -> int: + try: + ps_proc = psutil.Process(proc.pid) + processes = [ps_proc] + ps_proc.children(recursive=True) + except psutil.NoSuchProcess: + return 0 + rss = 0 + for p in processes: + try: + rss += p.memory_info().rss + except psutil.NoSuchProcess: + continue + return rss + + +def _kill_tree(proc: subprocess.Popen) -> None: + try: + ps_proc = psutil.Process(proc.pid) + for p in ps_proc.children(recursive=True) + [ps_proc]: + try: + p.kill() + except psutil.NoSuchProcess: + pass + except psutil.NoSuchProcess: + pass + try: + proc.wait(timeout=30) + except subprocess.TimeoutExpired: # pragma: no cover - children ignored SIGKILL + pass + + +def run_guarded( + cmd: list[str], + *, + rss_limit_mb: int, + min_free_mb: int, + max_swap_growth_mb: int = 64, + poll_s: float = 0.2, +) -> GuardResult: + """Run ``cmd`` while holding the shared model lock and enforcing limits. + + The lock is taken *before* the free-memory check and held until the child + exits, so a refused run still serializes against other model runs and a + finished run frees its memory before the next one is checked. + """ + lock = _acquire_lock(lock_path()) + try: + swap_before = psutil.swap_memory().used + if psutil.virtual_memory().available < min_free_mb * MB: + return GuardResult( + returncode=-1, + peak_rss_mb=0.0, + swap_before_mb=swap_before / MB, + swap_after_mb=swap_before / MB, + killed=False, + reason="insufficient free memory", + ) + + proc = subprocess.Popen(cmd) + peak_rss = 0 + killed = False + reason: str | None = None + while proc.poll() is None: + rss = _tree_rss(proc) + peak_rss = max(peak_rss, rss) + if rss > rss_limit_mb * MB: + reason, killed = "rss limit", True + _kill_tree(proc) + break + if psutil.swap_memory().used - swap_before > max_swap_growth_mb * MB: + reason, killed = "swap growth", True + _kill_tree(proc) + break + time.sleep(poll_s) + + swap_after = psutil.swap_memory().used + returncode = proc.returncode if proc.returncode is not None else -1 + if reason is not None and returncode == 0: + returncode = -1 + return GuardResult( + returncode=returncode, + peak_rss_mb=peak_rss / MB, + swap_before_mb=swap_before / MB, + swap_after_mb=swap_after / MB, + killed=killed, + reason=reason, + ) + finally: + fcntl.flock(lock.fileno(), fcntl.LOCK_UN) + lock.close() + + +def main(argv: list[str] | None = None) -> int: + argv = list(sys.argv[1:] if argv is None else argv) + if "--" in argv: + split = argv.index("--") + flag_args, cmd = argv[:split], argv[split + 1 :] + else: + flag_args, cmd = argv, [] + parser = argparse.ArgumentParser(prog="memguard", add_help=True) + parser.add_argument("--rss-limit-mb", type=int, required=True) + parser.add_argument("--min-free-mb", type=int, required=True) + parser.add_argument("--max-swap-growth-mb", type=int, default=64) + parser.add_argument("--poll-s", type=float, default=0.2) + args = parser.parse_args(flag_args) + if not cmd: + parser.error("no command given; use: -- ") + + result = run_guarded( + cmd, + rss_limit_mb=args.rss_limit_mb, + min_free_mb=args.min_free_mb, + max_swap_growth_mb=args.max_swap_growth_mb, + poll_s=args.poll_s, + ) + print( + f"MEMGUARD peak_rss_mb={result.peak_rss_mb:.1f} " + f"swap_delta_mb={result.swap_after_mb - result.swap_before_mb:.1f} " + f"reason={result.reason}", + flush=True, + ) + if result.reason is not None: + return 1 + return result.returncode if result.returncode >= 0 else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) \ No newline at end of file