Add files using upload-large-folder tool

.gitattributes

@@ -33,3 +33,4 @@ 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
+cuda_fp32/talker_cache.onnx.data filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,13 @@
+# Downloaded checkpoint + intermediate/build artifacts
+model/
+cache/
+*.onnx
+*.onnx.data
+*.safetensors
+_genai_*/
+__pycache__/
+*.pyc
+
+codes/**/__pycache__/
+_load_test.py
+build_tok.log

README.md

@@ -0,0 +1,103 @@
+---
+license: apache-2.0
+pipeline_tag: text-to-speech
+language:
+- zh
+- en
+- ja
+- ko
+- de
+- fr
+- ru
+- pt
+- es
+- it
+tags:
+- tts
+- qwen
+- audio
+- onnx
+- onnxruntime-genai
+- multilingual
+base_model: 
+- Qwen/Qwen3-TTS-12Hz-0.6B-VoiceDesign
+arxiv: 2601.15621
+---
+
+# Qwen3-TTS-12Hz-0.6B-CustomVoice
+
+[Qwen3-TTS](https://github.com/QwenLM/Qwen3-TTS) is a series of advanced multilingual, controllable, robust, and streaming text-to-speech models developed by the Qwen team. 
+
+This specific checkpoint is the **0.6B CustomVoice** variant, based on the **12Hz** tokenizer. It supports 9 premium timbres and allows for fine-grained style control over target voices via natural language instructions across 10 major languages.
+
+- **Paper:** [Qwen3-TTS Technical Report](https://huggingface.co/papers/2601.15621)
+- **GitHub:** [QwenLM/Qwen3-TTS](https://github.com/QwenLM/Qwen3-TTS)
+- **Demo:** [Hugging Face Spaces](https://huggingface.co/spaces/Qwen/Qwen3-TTS)
+
+## Key Features
+* **Multilingual Synthesis**: Supports Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, and Italian.
+* **Intelligent Control**: Adapts tone, rhythm, and emotional expression based on natural language instructions (e.g., "Speak in a very happy tone").
+* **Low Latency**: Optimized for streaming generation with the Qwen3-TTS-Tokenizer-12Hz, achieving end-to-end synthesis latency as low as 97ms.
+
+## Quickstart
+
+To use Qwen3-TTS, you can install the `qwen-tts` package:
+
+```bash
+pip install -U qwen-tts
+```
+
+### Sample Usage
+
+```python
+import torch
+import soundfile as sf
+from qwen_tts import Qwen3TTSModel
+
+# Load the model
+model = Qwen3TTSModel.from_pretrained(
+    "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice",
+    device_map="cuda:0",
+    dtype=torch.bfloat16,
+    attn_implementation="flash_attention_2",
+)
+
+# Generate speech with specific instructions
+wavs, sr = model.generate_custom_voice(
+    text="其实我真的有发现，我是一个特别善于观察别人情绪的人。",
+    language="Chinese", 
+    speaker="Vivian",
+    instruct="用特别愤怒的语气说", 
+)
+
+# Save the generated audio
+sf.write("output_custom_voice.wav", wavs[0], sr)
+```
+
+## Supported Speakers
+
+For `Qwen3-TTS-12Hz-0.6B-CustomVoice`, the following speakers are supported. We recommend using each speaker’s native language for the best results:
+
+| Speaker | Voice Description | Native Language |
+| --- | --- | --- |
+| Vivian | Bright young female voice. | Chinese |
+| Serena | Warm, gentle young female voice. | Chinese |
+| Uncle_Fu | Seasoned male voice, mellow timbre. | Chinese |
+| Dylan | Youthful Beijing male voice. | Chinese (Beijing) |
+| Eric | Lively Chengdu male voice. | Chinese (Sichuan) |
+| Ryan | Dynamic male voice with rhythm. | English |
+| Aiden | Sunny American male voice. | English |
+| Ono_Anna | Playful Japanese female voice. | Japanese |
+| Sohee | Warm Korean female voice. | Korean |
+
+## Citation
+If you find Qwen3-TTS useful for your research, please consider citing:
+
+```bibtex
+@article{Qwen3-TTS,
+  title={Qwen3-TTS Technical Report},
+  author={Hangrui Hu and Xinfa Zhu and Ting He and Dake Guo and Bin Zhang and Xiong Wang and Zhifang Guo and Ziyue Jiang and Hongkun Hao and Zishan Guo and Xinyu Zhang and Pei Zhang and Baosong Yang and Jin Xu and Jingren Zhou and Junyang Lin},
+  journal={arXiv preprint arXiv:2601.15621},
+  year={2026}
+}
+```

STATUS.md

@@ -0,0 +1,258 @@
+# Qwen3-TTS — Current Status
+
+_Updated: 2026-06-19_
+
+## ✅ De-duplicated talker — ship talker_cache only (~870 MB/dir saved)
+`talker.onnx` (no-cache) and `talker_cache.onnx` (KV-cache) hold the SAME transformer weights
+→ shipping both duplicated ~870 MB per dir. Fix: default export now builds **only
+`talker_cache`** (it does prefill+decode, faster O(n), and is what inference auto-uses); plain
+`talker` is still buildable via explicit `--components talker`. Dropped the redundant
+`talker.onnx` (+ manifest entry) from all dirs that had both → **freed ~22.5 GB** across the
+8 built variants. inference.py guards for "neither talker present". Verified generation still
+works on talker_cache-only dirs (KV-cache path, audio produced).
+
+## ✅ Config-driven dims — 0.6B now exports; 1.7B undisturbed
+`user_script.py` had 1.7B dims hardcoded (hidden 2048, 28 layers, etc.) in the talker /
+talker_cache / code_predictor / residual_embed io+dummy funcs → 0.6B (hidden **1024**) would
+export wrong. Fixed with a module `_DIMS` populated by `_load_tts()` from the loaded model's
+`talker_config` (Olive calls model_loader before io/dummy). **Defaults equal the 1.7B values**,
+so 1.7B exports are byte-identical (verified: voicedesign/customvoice/base17 dirs untouched).
+**Validated**: full `base/0.6B` int4 export — all 9 components incl. talker/talker_cache/
+speaker_encoder — `onnx/base06/cpu_int4/`, talker `inputs_embeds [batch,seq,1024]` ✓.
+(0.6B differs from 1.7B only in hidden_size; layers/kv/head_dim/groups are identical.)
+
+## ✅ EXPORT OOM RESOLVED — `optimize.py` isolates each component in a subprocess
+The Windows pagefile/OOM kill on full-run exports is fixed: `main()` now re-invokes itself once
+per component in a fresh subprocess (`--_child`), so memory is fully reclaimed between builds and
+the heavy talker/talker_cache run alone. Default behavior; `--no-isolate` forces single-process.
+`speaker_encoder` auto-dropped from defaults for non-base models. **Validated**: a full
+**cpu_fp16** voicedesign export (the heaviest case — 5.6 GB fp16 talker + talker_cache) completed
+in ONE command, all 8 components, no OOM → `onnx/voicedesign/cpu_fp16/`.
+
+## ✅ All 3 model types exported + validated with their feature
+| dir | model_type | feature | validated |
+|---|---|---|---|
+| `onnx/voicedesign/cpu_int4/` | voice_design | `--instruct` | ✅ generates |
+| `onnx/customvoice/cpu_int4/`  | custom_voice | `--speaker ryan` | ✅ generates (27 frames) |
+| `onnx/base17_cpu_int4/` | base | clone `--ref-audio/--ref-text` | ✅ runs e2e; speaker_encoder parity 1.0 |
+Model-type gating enforces the right feature per model (wrong flag → clear error).
+Pending exports: 0.6B Base, other precisions/devices (fp16/fp32, cuda) — mechanical repeats
+(component-by-component due to the pagefile cap). KV-cache talker still blocked.
+
+## TL;DR (current session)
+- ✅ Full text→speech `generate()` implemented + **100% greedy parity vs PyTorch** (all-16
+  codebooks 432/432). Needed talker→(logits,hidden) + `residual_embed.onnx` + the
+  `dynamic_shapes` seq fix — all done & validated.
+- ✅ **Model-type feature gating in inference.py** — `model_type` read from config;
+  `_check_features` enforces: voice_design→`--instruct`, custom_voice→`--speaker`,
+  base→clone (`--ref-audio/--ref-text`). VALIDATED both ways: voice_design+instruct generates;
+  voice_design+speaker raises a clear ValueError. `--text` prints `model_type=…`.
+- ✅ **voicedesign `onnx/voicedesign/cpu_int4/` restored** (all 7 components) and generation
+  re-validated (model_type=voice_design, 24-frame wav). Re-export had to be done
+  **component-by-component** (talker solo) — a single full-run OOM/pagefile-killed the talker
+  (`OSError 1455: paging file too small`). The lighter parts export together.
+- ✅ **Env fix:** `optimize.py` + `inference.py` PEP-723 now pin `numba>=0.60 / llvmlite>=0.43`
+  — librosa otherwise pulls numba 0.53.1 → llvmlite 0.36 which won't build on Python 3.12.
+- ✅ Base checkpoints downloaded (`base/1.7B`, `base/0.6B`); **Base-1.7B int4 exported** →
+  `onnx/base17_cpu_int4/` (7 components incl. `residual_embed` + 2-output dynamic talker).
+- ✅ **Base voice cloning — IMPLEMENTED & runs end-to-end.** `generate(ref_audio, ref_text)`
+  → `_generate_clone`. Pieces:
+  • `speaker_encoder.onnx` exported (ECAPA + **inline mel/STFT** front-end, dynamic audio len,
+    fp32, dynamo; `audio[B,T]→x-vector[B,2048]`). `SpeakerEncoderWrapper` reimplements
+    mel_spectrogram WITHOUT its `if torch.min(y)<-1` debug branch (broke torch.export).
+  • ICL prefill mirrors `generate_icl_prompt`: ref_text+text+eos / codec_bos + **per-frame
+    ref-code sum (= `step_embed`/residual_embed!)**, with the x-vector injected in the codec
+    prefix. Ref audio→codes via `tok_encoder` in 1 s windows (`encode_chunked`).
+  • AR loop factored into shared `_ar_loop`. Verified runs: ref 39 frames + x-vector →
+    prefill 64 → 40 frames → 3.2 s wav (`onnx/base17_cpu_int4/`, model_type=base).
+  ✅ **speaker_encoder parity** (`eval_speaker.py`): ONNX x-vector vs PyTorch
+  `extract_speaker_embedding` = **cosine 1.000000, max|Δ| ~1e-6** at 2/3.5/6 s → the
+  reimplemented mel/STFT is exact. With tok_encoder (100%), residual_embed/talker/predictor
+  (100% TTS parity) all verified and the ICL prefill mirroring the reference, clone is correct
+  by construction. (A real-voice listen is still the only thing measuring perceptual fidelity.)
+- ✅ **KV-cache talker — SOLVED (was blocked).** Fix: feed the cache as a plain list of
+  [k,v] tensor pairs under input name **`past_kv`** (NOT `past_key_values`, so Olive's
+  DynamicCache-pytree auto-conversion — which torch.export rejected — doesn't fire), and
+  build the `DynamicCache` in-graph (`from_legacy_cache`/`to_legacy_cache`). `talker_cache.onnx`
+  exports (59 in / 58 out: flattened 28×2 K/V). Verified: empty-past output == no-cache talker
+  (logits/hidden cos 1.0, max|Δ| 0.0); end-to-end greedy **100% exact** vs no-cache. `inference.py`
+  auto-uses it (`_ar_loop_cached`) when present. Speedup grows with length (talker O(n) vs O(n²);
+  ~1.1x at 36 frames where the 15-call predictor dominates, more for long utterances).
+  Exported for `onnx/voicedesign/cpu_int4/` so far.
+- ⏳ **Export dirs:** the earlier `onnx/voicedesign/` & `onnx/customvoice/` were cleared in a
+  reorg; only `onnx/base17_cpu_int4/` survives. Re-exporting voicedesign cpu_int4 now.
+- TODO: validate generate on base export; export 0.6B; optional tok_decoder dynamic-frames.
+
+## Targets (3 models, auto-detected by config `model_type`) — all downloaded locally
+- `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign`  (qwen3_tts)  ← `voicedesign/`
+- `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice`  (qwen3_tts)  ← `customvoice/`
+- `Qwen/Qwen3-TTS-Tokenizer-12Hz`         (qwen3_tts_tokenizer_12hz) ← `tokenizer/` (also embedded in each TTS as `speech_tokenizer/`)
+
+## Architecture (verified from repo + load)
+TTS `Qwen3TTSForConditionalGeneration` → `talker`:
+- `talker.model` — Qwen3-style **28 L, hidden 2048, MROPE**, dual codec/text embedding
+- `talker.codec_head` — Linear 2048→3072 (first-codebook logits)
+- `talker.code_predictor` — 5 L, hidden 1024, **16 per-group heads** (residual codes)
+- `talker.text_projection`
+Tokenizer `Qwen3TTSTokenizerV2Model` → **encoder** (wav→codes) + **decoder** (codes→wav),
+24 kHz, 12.5/12 fps, 16 quantizers (decode), conv+transformer+RVQ.
+
+## Decomposition → `onnx/{device}_{precision}/` (flat)
+- TTS: `talker.onnx`, `code_predictor.onnx` (+ `tok_encoder/decoder.onnx` from embedded codec)
+- Tokenizer: `tok_encoder.onnx`, `tok_decoder.onnx`
+- codecs forced **fp32** (DAC int4/fp16 too lossy); LLM/predictor honor `--precision`.
+
+## create_model (ModelBuilder) — empirically checked, NOT viable for the talker
+onnxruntime-genai ModelBuilder supports these archs: Llama, Mistral, Qwen2, **Qwen3**,
+Gemma/2/3, Phi*, Granite, Nemotron, Olmo, SmolLM3, GptOss, InternLM2, Lfm2, HunYuan,
+VideoChatFlashQwen. The talker is **`qwen3_tts_talker`** (not listed) and, decisively,
+uses **MROPE** (`mrope_section=[24,20,20]`, interleaved) + dual codec/text embedding +
+`codec_head`. Even remapped to `Qwen3ForCausalLM`, ModelBuilder applies **standard RoPE
+(no MROPE)** → positionally wrong output. So create_model can't correctly convert it →
+**Olive export (with the real MROPE forward) is the correct path.** Confirmed via the
+supported-architecture list + the talker rope config.
+
+## Key decisions / findings
+- **Talker is NOT ModelBuilder-able** (custom MROPE + dual embedding ≠ stock Qwen3) →
+  exported via **Olive** (`OnnxConversion` + RtnQuant/Float16), like CSM backbone.
+  (ModelBuilder remap attempted only if a talker is detectably stock.)
+- **transformers version**: vendored `qwen_tts` (in `codes/`) needs **transformers==4.57.3**;
+  the shared venv is 5.10.2 (incompatible — check_model_inputs, config defaults,
+  ROPE_INIT_FUNCTIONS, …). Solved by making **optimize.py a PEP-723 uv script** that pins
+  4.57.3 in an isolated env (`uv run optimize.py ...`). Shared venv untouched.
+  One vendored patch: `@check_model_inputs()` → `@check_model_inputs` (tf5 API); 25Hz
+  tokenizer import made optional (needs `sox`).
+- **Model loads cleanly** under 4.57.3 (instantiates; submodules mapped).
+
+## Fixed — external-data relink on flattened models
+`optimize.py` flattens Olive's `model.onnx[.data]` → `{name}.onnx[.data]`, but the proto's
+`external_data.location` still pointed at `model.onnx.data` → onnxruntime failed to load any
+model with external data (`External data path does not exist`). Affected talker fp32/fp16
+(int4 talker is self-contained, no `.data`). Fixed: `relink_external_data()` rewrites each
+tensor's `location` (proto-only, no multi-GB RAM load) after the move; existing dirs relinked
+in place (cpu/cuda fp32=313 refs, fp16=255). Verified by `inference.py --selftest` loading talker.
+
+## inference.py (manifest-driven) — full generation implemented
+`Pipeline(onnx/{dev}_{prec}, tts_dir=...)` loads all parts with the manifest EP (CPU fallback).
+Building blocks verified in eval_*.py: `embed_text`, `embed_codec`, `talker_step`,
+`predict_residual`, `step_embed`, `encode`, `decode`/`decode_chunked`.
+
+`generate(text, language, instruct=, speaker=)` now implements the real text→speech path,
+faithful to `Qwen3TTSForConditionalGeneration.generate` with `non_streaming_mode=True`:
+  • **text-to-speech** — text + language
+  • **voice design**   — text + `instruct` (natural-language style)        [VoiceDesign]
+  • **custom voice**   — text + `speaker` name (+ optional instruct)       [CustomVoice]
+Prefill assembly mirrors modeling lines 2068-2234 (role + codec tags + pad/bos + text body +
+eos + codec_bos; instruct embeds prepended). AR talker loop is no-cache (re-runs the growing
+prefix; MROPE collapses to `arange` since `get_rope_index`=cumsum(mask)-1 with 3 identical rows
+and no padding). Each step: talker→first-codebook logits + last hidden → 15 residual codes via
+the **causal** teacher-forced predictor (fill known codes, read `group_logits[j-1]`) → next input
+= `residual_embed(codes16)` (=`codec_hiddens.sum(1)`) + `tts_pad`. suppress_tokens + repetition
+penalty + top-k/top-p sampling replicated in numpy. `--text` CLI bug fixed (was `if selftest or
+True:`; now proper branching).
+
+**Two new exported components** were required (the old 6 couldn't roll out generation):
+  • `talker.onnx` now emits **(logits, hidden_states)** — the predictor is conditioned on the
+    talker's last hidden state, not its logits. (re-export needed)
+  • `residual_embed.onnx` — `codec_ids[B,16] → codec_hiddens.sum(1)[B,2048]`, summing
+    `talker.model.codec_embedding(code0)` + `code_predictor.model.codec_embedding[i](code_{i+1})`.
+    Those residual per-group embeddings were buried in the predictor graph; needed for the
+    next-step talker input.
+
+**Audio-clone (ICL ref_audio/ref_text)** is a `base`-model feature: `create_voice_clone_prompt`
+raises for non-base, and `speaker_encoder` is None for VoiceDesign/CustomVoice. Not applicable to
+our two targets → intentionally not implemented (documented in inference.py).
+
+Re-export to refresh a dir:
+  `uv run optimize.py --model voicedesign --skip-download --device cpu --precision fp32 \
+       --components residual_embed talker`
+
+## Fixed — talker needed `dynamic_shapes` (dynamo ignores `dynamic_axes`)
+The first talker re-export locked seq to the dummy's 32 (`attention_mask Got 26 Expected 32`
+at gen step 0), breaking the AR loop. Root cause: with `use_dynamo_exporter`, Olive
+(`conversion.py` ~L368) passes **`dynamic_shapes`** to torch.export and **drops `dynamic_axes`**.
+Fix: added `dynamic_shapes` to the talker io_config (`{input: {axis:int → "dimname"}}`, shared
+`"seq"` ties the 3 inputs; torch 2.12 accepts string dim names). Re-exported → talker inputs are
+now `['batch','seq',2048]` etc. (Same root cause as the tok_decoder fixed-25-frame limitation —
+that decoder could be re-exported with `dynamic_shapes` too if a dynamic-frame decoder is wanted.)
+
+## ✅ Generation validated end-to-end (cpu_fp32, voicedesign)
+`uv run inference.py --model-path onnx/customvoice/cpu_fp32 --tts-dir voicedesign \
+   --text "Hello, this is a test." --instruct "A calm female voice." --out out.wav`
+→ AR loop runs (dynamic-seq talker → first code + hidden; 15 residuals via causal predictor;
+`residual_embed` next-step input), decodes real audio (RMS 0.10, peak 0.64, non-silent).
+
+**Greedy parity vs PyTorch (`eval_generate.py`, voicedesign, cpu_fp32):** tokenization identical;
+first-codebook 27/27 frames (100%); **all-16-codebooks 432/432 = 100.00%** over 27 frames — the
+ONNX generation reproduces `model.generate` exactly (prefill + MROPE + AR talker + causal residual
+predictor + residual_embed + EOS). Frame count matches the reference (sampling varies length).
+NOTE: `onnx/customvoice/cpu_fp32/` is hand-organized and actually holds **voicedesign** content
+(manifest `model_id: .\voicedesign\`). Only this dir has the upgraded talker + residual_embed so
+far; the other device×precision dirs still need the `residual_embed`/2-output-talker re-export.
+
+## Known operational note — memory pressure on full-run export
+A single `optimize.py ... --model voicedesign` (all components) can OOM-kill the process at
+`code_predictor`: the `talker` fp16 step holds ~2.8 GB external data + a ~150 s float16 pass, and
+stacking `code_predictor`'s float16 pass right after exhausts RAM. **No code defect** — every
+component converts cleanly in isolation (verified: code_predictor cuda/fp16 = 354 MB, ½ of fp32).
+Workaround: export heavy parts separately, e.g.
+`--components talker` then `--components code_predictor text_embed codec_embed tok_encoder tok_decoder`.
+`onnx/cuda_fp16/` now holds all 6 (text_embed 639 MB, codec_embed 12.5 MB, talker+.data,
+code_predictor 354 MB, tok_encoder/decoder fp32). TODO(optional): make optimize.py export each
+component in a subprocess so one full-run command can't OOM.
+
+## Build / run
+```
+uv run optimize.py --model Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign --device cpu --precision int4
+uv run optimize.py --model Qwen/Qwen3-TTS-Tokenizer-12Hz        --device cpu --precision fp32
+uv run optimize.py --model voicedesign --skip-download --components tok_decoder   # subset
+```
+
+## Export recipe (learned)
+- Load all sub-models with `attn_implementation="eager"`.
+- Codec transformer parts use `create_causal_mask` (not torch.onnx/TorchScript-traceable)
+  → set **`use_dynamo_exporter: true`** on the OnnxConversion pass (done for codec parts).
+- Bypass high-level wrappers with Python loops: `tok_decoder` calls `self.tok.decoder(...)`
+  directly (skips `model.decode`'s `chunked_decode` while-loop); transpose codes
+  `[B,T,16]→[B,16,T]` + clamp ≥0 first.
+- Codes layout: `model.decode` expects `[B, codes_length, num_quantizers]` (=[B,T,16]).
+
+## Status / next
+- ✅ Vendored `qwen_tts` (codes/), uv-script env (tf 4.57.3) validated, model loads.
+- ✅ `optimize.py` (uv script, model-name + device/precision dispatch) + `user_script.py`.
+- ✅ **6 sub-models export** (text_embed, codec_embed, talker MROPE LLM, code_predictor, tok_encoder, tok_decoder).
+- ✅ **Talker embedding primitives added + verified** (`text_embed`, `codec_embed`). The talker
+  prefill itself is control flow (variable text len, voice-clone/ICL branches, concat, MROPE
+  positions) → stays in Python (inference.py); only the *learned* lookups are in ONNX:
+  - `text_embed` = `text_projection(text_embedding(ids))` [B,T,2048] (also covers tts_bos/eos/pad
+    — specific ids); 1.28 GB.  `codec_embed` = `codec_embedding(ids)` [B,T,2048]; 25 MB.
+  - Parity (`eval_embed.py`, fp32): text_embed cosine 1.000000 / max|Δ| 1.2e-7;
+    codec_embed cosine 1.000000 / max|Δ| 0. **Every learned weight now runs through ONNX.**
+- ✅ **`code_predictor` corrected + verified** (was missing `small_to_mtp_projection` 2048→1024 +
+  codec-embedding assembly). New `CodePredictorWrapper` folds those in-graph (mirrors
+  `forward_sub_talker_finetune`); interface = `(talker_hidden[B,2048], codec_ids[B,16]) →
+  group_logits[B,15,vocab]`. Parity (`eval_predictor.py`, fp32, voicedesign):
+  wrapper-vs-native cosine 1.000000 / max|Δ| 0 / argmax 100%; ONNX-vs-wrapper cosine 1.000000 /
+  max|Δ| 5.7e-5 / argmax 100%. This is the **teacher-forced** variant (needs all 16 codes) — for
+  parity; the single-step AR variant for real generation is still TODO.
+- ✅ **tokenizer verified** (`eval_tokenizer.py`): encoder 100% exact index match; decoder cosine
+  1.00000 vs PyTorch. Note: `tok_decoder.onnx` is **fixed at 25 frames** (dynamic frames axis did
+  not survive dynamo export) — fine for parity, but inference must chunk/pad to 25.
+- ✅ **precision check** (`check_precision.py`): inspects weight dtypes + quant ops (size is
+  misleading — codecs are forced fp32, hence byte-identical across int4/fp16/fp32 dirs).
+- ✅ `tok_encoder` (wav→codes) **resolved** (225 MB, fp32, dynamo) → all 4 sub-models export.
+  Mirrors `MimiModel._encode_frame` (bypasses streaming `encode()`), **plus two static-shape
+  patches in `user_script.py`** so torch.export gets concrete conv lengths:
+  1. `_patch_mimi_static_padding` — rewrites `MimiConv1d._get_extra_padding_for_conv1d` to
+     pure-Python int math (stock builds padding as 0-dim tensors → `.item()` → unbacked
+     symints → RVQ `torch.cdist`'s `npoints>25` guard can't resolve).
+  2. `_intify_mimi_convs` — converts each `MimiConv1d`'s `stride`/`kernel_size`/`padding_total`
+     buffers (0-dim int64 tensors in this model) + derived `padding_left/right` to Python ints
+     (reading a buffer in forward, even via `int()`, is a `.item()` under export).
+  Fixed input length (24000); only batch dynamic. Unblocks CustomVoice voice-clone (clones via
+  reference codes from this encoder; no separate speaker encoder for custom_voice/voice_design).
+- Vendored patches applied for the trimmed-25Hz case: guarded `AutoConfig.register` in
+  `inference/qwen3_tts_tokenizer.py` (skip None 25Hz config; idempotent 12Hz).
+- ⏳ `inference.py` (text→talker→code_predictor→codes→tok_decoder→wav) + `eval.py`.
+- Other models: CustomVoice = same TTS arch; Tokenizer-12Hz = standalone codec
+  (tok_decoder works; tok_encoder shares the deferred Mimi issue).

check_precision.py

@@ -0,0 +1,73 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = ["onnx", "numpy"]
+# ///
+"""Verify the *actual* numeric precision of exported ONNX sub-models.
+
+File size is misleading: the codec parts (tok_encoder/tok_decoder) are forced
+fp32 in every precision dir, so they're byte-identical across cpu_fp16 / cpu_fp32
+/ cpu_int4. What really tells you the precision is the weight (initializer)
+dtype histogram + presence of quant ops (MatMulNBits/DequantizeLinear for int4).
+
+Usage:
+  uv run check_precision.py onnx/cpu_int4 onnx/cpu_fp16 onnx/cpu_fp32
+  uv run check_precision.py onnx/cpu_int4/talker.onnx        # single file
+"""
+import sys
+from collections import Counter
+from pathlib import Path
+
+import onnx
+from onnx import numpy_helper, TensorProto
+
+DT = {v: k for k, v in TensorProto.DataType.items()}
+QUANT_OPS = {"MatMulNBits", "DequantizeLinear", "QuantizeLinear", "MatMulInteger",
+             "DynamicQuantizeLinear", "ConvInteger"}
+
+
+def inspect(path: Path):
+    m = onnx.load(str(path), load_external_data=False)
+    g = m.graph
+    # initializer dtype histogram, weighted by element count
+    bytes_by_dt, count_by_dt = Counter(), Counter()
+    for init in g.initializer:
+        dt = DT.get(init.data_type, str(init.data_type))
+        n = 1
+        for d in init.dims:
+            n *= d
+        count_by_dt[dt] += n
+        bytes_by_dt[dt] += 1
+    ops = Counter(n.op_type for n in g.node)
+    quant = {op: ops[op] for op in QUANT_OPS if op in ops}
+
+    total = sum(count_by_dt.values()) or 1
+    dt_summary = ", ".join(
+        f"{dt}:{100*c/total:.1f}%" for dt, c in count_by_dt.most_common()
+    )
+    # infer label
+    if quant:
+        label = "INT4/quantized"
+    elif count_by_dt.get("FLOAT16", 0) > count_by_dt.get("FLOAT", 0):
+        label = "FP16"
+    elif count_by_dt.get("FLOAT", 0) > 0:
+        label = "FP32"
+    else:
+        label = "?"
+    print(f"  {path.name:<22} -> {label}")
+    print(f"      weight dtypes (by #elements): {dt_summary}")
+    if quant:
+        print(f"      quant ops: {quant}")
+
+
+def main():
+    targets = sys.argv[1:] or ["onnx/cpu_int4", "onnx/cpu_fp16", "onnx/cpu_fp32"]
+    for t in targets:
+        p = Path(t)
+        files = sorted(p.glob("*.onnx")) if p.is_dir() else [p]
+        print(f"\n=== {t} ===")
+        for f in files:
+            inspect(f)
+
+
+if __name__ == "__main__":
+    main()

cpu_fp16/code_predictor.onnx

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

cpu_fp16/codec_embed.onnx

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

cpu_fp16/manifest.json

@@ -0,0 +1,30 @@
+{
+  "model_id": ".\\voicedesign\\",
+  "model_kind": "tts",
+  "device": "cpu",
+  "precision": "fp16",
+  "execution_provider": "CPUExecutionProvider",
+  "sub_models": {
+    "code_predictor": {
+      "filename": "code_predictor.onnx"
+    },
+    "codec_embed": {
+      "filename": "codec_embed.onnx"
+    },
+    "residual_embed": {
+      "filename": "residual_embed.onnx"
+    },
+    "talker_cache": {
+      "filename": "talker_cache.onnx"
+    },
+    "text_embed": {
+      "filename": "text_embed.onnx"
+    },
+    "tok_decoder": {
+      "filename": "tok_decoder.onnx"
+    },
+    "tok_encoder": {
+      "filename": "tok_encoder.onnx"
+    }
+  }
+}

cpu_fp16/residual_embed.onnx

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

cpu_fp16/talker_cache.onnx

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

cpu_fp16/text_embed.onnx

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

cpu_fp16/tok_decoder.onnx

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

cpu_fp16/tok_encoder.onnx

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

cpu_fp32/code_predictor.onnx

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

cpu_fp32/codec_embed.onnx

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

cpu_fp32/manifest.json

@@ -0,0 +1,30 @@
+{
+  "model_id": ".\\voicedesign\\",
+  "model_kind": "tts",
+  "device": "cpu",
+  "precision": "fp32",
+  "execution_provider": "CPUExecutionProvider",
+  "sub_models": {
+    "code_predictor": {
+      "filename": "code_predictor.onnx"
+    },
+    "codec_embed": {
+      "filename": "codec_embed.onnx"
+    },
+    "residual_embed": {
+      "filename": "residual_embed.onnx"
+    },
+    "talker_cache": {
+      "filename": "talker_cache.onnx"
+    },
+    "text_embed": {
+      "filename": "text_embed.onnx"
+    },
+    "tok_decoder": {
+      "filename": "tok_decoder.onnx"
+    },
+    "tok_encoder": {
+      "filename": "tok_encoder.onnx"
+    }
+  }
+}

cpu_fp32/residual_embed.onnx

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

cpu_fp32/talker_cache.onnx

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

cpu_fp32/text_embed.onnx

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

cpu_fp32/tok_decoder.onnx

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

cpu_fp32/tok_encoder.onnx

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

cpu_int4/code_predictor.onnx

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

cpu_int4/codec_embed.onnx

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

cpu_int4/manifest.json

@@ -0,0 +1,30 @@
+{
+  "model_id": ".\\voicedesign\\",
+  "model_kind": "tts",
+  "device": "cpu",
+  "precision": "int4",
+  "execution_provider": "CPUExecutionProvider",
+  "sub_models": {
+    "code_predictor": {
+      "filename": "code_predictor.onnx"
+    },
+    "codec_embed": {
+      "filename": "codec_embed.onnx"
+    },
+    "residual_embed": {
+      "filename": "residual_embed.onnx"
+    },
+    "talker_cache": {
+      "filename": "talker_cache.onnx"
+    },
+    "text_embed": {
+      "filename": "text_embed.onnx"
+    },
+    "tok_decoder": {
+      "filename": "tok_decoder.onnx"
+    },
+    "tok_encoder": {
+      "filename": "tok_encoder.onnx"
+    }
+  }
+}

cpu_int4/residual_embed.onnx

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

cpu_int4/talker_cache.onnx

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

cpu_int4/text_embed.onnx

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

cpu_int4/tok_decoder.onnx

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

cpu_int4/tok_encoder.onnx

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

cuda_fp16/codec_embed.onnx

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

cuda_fp16/manifest.json

@@ -0,0 +1,30 @@
+{
+  "model_id": ".\\voicedesign\\",
+  "model_kind": "tts",
+  "device": "cuda",
+  "precision": "fp16",
+  "execution_provider": "CUDAExecutionProvider",
+  "sub_models": {
+    "code_predictor": {
+      "filename": "code_predictor.onnx"
+    },
+    "codec_embed": {
+      "filename": "codec_embed.onnx"
+    },
+    "residual_embed": {
+      "filename": "residual_embed.onnx"
+    },
+    "talker_cache": {
+      "filename": "talker_cache.onnx"
+    },
+    "text_embed": {
+      "filename": "text_embed.onnx"
+    },
+    "tok_decoder": {
+      "filename": "tok_decoder.onnx"
+    },
+    "tok_encoder": {
+      "filename": "tok_encoder.onnx"
+    }
+  }
+}

cuda_fp16/residual_embed.onnx

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

cuda_fp32/manifest.json

@@ -0,0 +1,30 @@
+{
+  "model_id": ".\\voicedesign\\",
+  "model_kind": "tts",
+  "device": "cuda",
+  "precision": "fp32",
+  "execution_provider": "CUDAExecutionProvider",
+  "sub_models": {
+    "code_predictor": {
+      "filename": "code_predictor.onnx"
+    },
+    "codec_embed": {
+      "filename": "codec_embed.onnx"
+    },
+    "residual_embed": {
+      "filename": "residual_embed.onnx"
+    },
+    "talker_cache": {
+      "filename": "talker_cache.onnx"
+    },
+    "text_embed": {
+      "filename": "text_embed.onnx"
+    },
+    "tok_decoder": {
+      "filename": "tok_decoder.onnx"
+    },
+    "tok_encoder": {
+      "filename": "tok_encoder.onnx"
+    }
+  }
+}

cuda_fp32/talker_cache.onnx.data

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

cuda_fp32/text_embed.onnx

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

cuda_int4/manifest.json

@@ -0,0 +1,30 @@
+{
+  "model_id": ".\\voicedesign\\",
+  "model_kind": "tts",
+  "device": "cuda",
+  "precision": "int4",
+  "execution_provider": "CUDAExecutionProvider",
+  "sub_models": {
+    "code_predictor": {
+      "filename": "code_predictor.onnx"
+    },
+    "codec_embed": {
+      "filename": "codec_embed.onnx"
+    },
+    "residual_embed": {
+      "filename": "residual_embed.onnx"
+    },
+    "talker_cache": {
+      "filename": "talker_cache.onnx"
+    },
+    "text_embed": {
+      "filename": "text_embed.onnx"
+    },
+    "tok_decoder": {
+      "filename": "tok_decoder.onnx"
+    },
+    "tok_encoder": {
+      "filename": "tok_encoder.onnx"
+    }
+  }
+}

cuda_int4/talker_cache.onnx

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

cuda_int4/text_embed.onnx

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

cuda_int4/tok_decoder.onnx

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

cuda_int4/tok_encoder.onnx

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

eval_cache.py

@@ -0,0 +1,49 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#   "onnxruntime>=1.20", "numpy", "soundfile", "librosa", "transformers",
+#   "numba>=0.60.0", "llvmlite>=0.43.0",
+# ]
+# ///
+"""Parity + speed: KV-cache talker loop vs no-cache loop. Greedy → must match exactly.
+
+Usage:
+  uv run eval_cache.py --model-path onnx/voicedesign/cpu_int4 --tts-dir voicedesign
+"""
+import argparse, sys, time
+from pathlib import Path
+import numpy as np
+
+HERE = Path(__file__).parent
+sys.path.insert(0, str(HERE))
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model-path", required=True)
+    ap.add_argument("--tts-dir", required=True)
+    ap.add_argument("--text", default="Hello, this is a KV cache parity test.")
+    ap.add_argument("--max-new-tokens", type=int, default=40)
+    args = ap.parse_args()
+
+    from inference import Pipeline
+    pipe = Pipeline(args.model_path, tts_dir=args.tts_dir)
+    if pipe.talker_cache is None:
+        print("no talker_cache.onnx in this dir — nothing to compare"); return
+
+    kw = dict(language="Auto", max_new_tokens=args.max_new_tokens, do_sample=False,
+              sub_do_sample=False, seed=0, verbose=False)
+    t0 = time.time(); codes_cached = pipe.generate(args.text, **kw); t_cached = time.time() - t0
+    pipe.talker_cache = None                         # force no-cache path
+    t0 = time.time(); codes_nocache = pipe.generate(args.text, **kw); t_nocache = time.time() - t0
+
+    n = min(len(codes_cached), len(codes_nocache))
+    match = float((codes_cached[:n] == codes_nocache[:n]).mean()) if n else -1
+    print(f"cached frames={len(codes_cached)} ({t_cached:.1f}s)  "
+          f"no-cache frames={len(codes_nocache)} ({t_nocache:.1f}s)  "
+          f"speedup={t_nocache/max(t_cached,1e-9):.2f}x")
+    print(f"all-16-codes exact match over {n} frames: {match*100:.2f}%")
+
+
+if __name__ == "__main__":
+    main()

eval_embed.py

@@ -0,0 +1,56 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#   "transformers==4.57.3", "torch", "torchvision", "torchaudio",
+#   "onnx", "onnxruntime>=1.20", "numpy", "safetensors",
+#   "huggingface_hub", "accelerate", "librosa", "soundfile",
+# ]
+# ///
+"""Parity for text_embed / codec_embed ONNX vs the PyTorch wrappers.
+
+Usage:
+  uv run eval_embed.py --model-path onnx/cpu_fp32 --tts-path voicedesign
+"""
+import argparse, json, sys
+from pathlib import Path
+import numpy as np
+
+HERE = Path(__file__).parent
+sys.path.insert(0, str(HERE)); sys.path.insert(0, str(HERE / "codes"))
+
+
+def cosine(a, b):
+    a, b = a.ravel().astype(np.float64), b.ravel().astype(np.float64)
+    return float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-9))
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model-path", required=True)
+    ap.add_argument("--tts-path", default="voicedesign")
+    args = ap.parse_args()
+
+    import torch, onnxruntime as ort
+    from user_script import (get_text_embed_model, get_codec_embed_model, _tts_dims, _load_tts)
+
+    mdir = Path(args.model_path)
+    sub = json.loads((mdir / "manifest.json").read_text())["sub_models"]
+    dims = _tts_dims(_load_tts(args.tts_path))
+
+    for name, get_model, lo, hi, in_name, out_name in [
+        ("text_embed",  get_text_embed_model,  0, 1000,            "text_ids",  "text_embeds"),
+        ("codec_embed", get_codec_embed_model, 0, dims["codec_vocab"], "codec_ids", "codec_embeds"),
+    ]:
+        sess = ort.InferenceSession(str(mdir / sub[name]["filename"]),
+                                    providers=["CPUExecutionProvider"])
+        wrap = get_model(args.tts_path)
+        ids = torch.randint(lo, hi, (1, 16), dtype=torch.int64)
+        with torch.no_grad():
+            ref = wrap(ids).numpy()
+        got = sess.run(None, {in_name: ids.numpy()})[0]
+        print(f"=== {name} ===  onnx={got.shape} ref={ref.shape}  "
+              f"cosine={cosine(got, ref):.6f}  max|Δ|={np.abs(got-ref).max():.3e}")
+
+
+if __name__ == "__main__":
+    main()

eval_generate.py

@@ -0,0 +1,98 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#   "transformers==4.57.3", "torch", "torchvision", "torchaudio",
+#   "onnx", "onnxruntime>=1.20", "numpy", "safetensors",
+#   "huggingface_hub", "accelerate", "librosa", "soundfile",
+# ]
+# ///
+"""Parity: ONNX inference.py generate() vs the PyTorch reference (greedy).
+
+Compares, on identical text/instruct, greedy (do_sample=False):
+  (1) tokenization  — reference processor ids  vs  inference.py AutoTokenizer ids
+  (2) first-codebook code sequence — reference talker codes vs ONNX codes
+      (frame count, first divergence step, % match up to min length)
+
+Greedy isolates correctness from sampling RNG. Small fp diffs can still flip an
+argmax and diverge late; report first-divergence step so a late split (good
+prefill) is distinguishable from step-0 (broken prefill).
+
+Usage:
+  uv run eval_generate.py --model-path onnx/customvoice/cpu_fp32 --tts-dir voicedesign \
+      --text "Hello, this is a test." --instruct "A calm female voice." --max-new-tokens 60
+"""
+import argparse, sys
+from pathlib import Path
+import numpy as np
+
+HERE = Path(__file__).parent
+sys.path.insert(0, str(HERE)); sys.path.insert(0, str(HERE / "codes"))
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model-path", required=True)
+    ap.add_argument("--tts-dir", default="voicedesign")
+    ap.add_argument("--text", default="Hello, this is a test.")
+    ap.add_argument("--instruct", default="A calm female voice.")
+    ap.add_argument("--language", default="Auto")
+    ap.add_argument("--max-new-tokens", type=int, default=60)
+    args = ap.parse_args()
+
+    import torch
+    from qwen_tts.inference.qwen3_tts_model import Qwen3TTSModel
+    from inference import Pipeline
+
+    # ── reference (PyTorch) ──────────────────────────────────────────────────
+    ref = Qwen3TTSModel.from_pretrained(args.tts_dir)
+    assistant = f"<|im_start|>assistant\n{args.text}<|im_end|>\n<|im_start|>assistant\n"
+    ref_ids = ref._tokenize_texts([assistant])[0]           # [1,1,L] or [1,L]
+    ref_ids_flat = ref_ids.reshape(-1).tolist()
+
+    input_ids = ref._tokenize_texts([ref._build_assistant_text(args.text)])
+    instruct_ids = [ref._tokenize_texts([ref._build_instruct_text(args.instruct)])[0]]
+    with torch.no_grad():
+        codes_list, _ = ref.model.generate(
+            input_ids=input_ids, instruct_ids=instruct_ids, languages=[args.language],
+            non_streaming_mode=True, do_sample=False, subtalker_dosample=False,
+            max_new_tokens=args.max_new_tokens, output_hidden_states=True,
+            return_dict_in_generate=True)
+    ref_codes = codes_list[0].cpu().numpy().astype(np.int64)   # [T,16]
+    ref_first = ref_codes[:, 0]
+
+    # ── ONNX (inference.py) ──────────────────────────────────────────────────
+    pipe = Pipeline(args.model_path, tts_dir=args.tts_dir)
+    onnx_ids = pipe._ids(assistant).reshape(-1).tolist()
+    onnx_codes = pipe.generate(
+        args.text, language=args.language, instruct=args.instruct,
+        max_new_tokens=args.max_new_tokens, do_sample=False, sub_do_sample=False,
+        verbose=False)
+    onnx_first = onnx_codes[:, 0]
+
+    # ── report ───────────────────────────────────────────────────────────────
+    print("\n=== tokenization ===")
+    print(f"  ref  ids ({len(ref_ids_flat)}): {ref_ids_flat}")
+    print(f"  onnx ids ({len(onnx_ids)}): {onnx_ids}")
+    print(f"  MATCH: {ref_ids_flat == onnx_ids}")
+
+    print("\n=== first-codebook codes (greedy) ===")
+    print(f"  ref  frames={len(ref_first)}  onnx frames={len(onnx_first)}")
+    n = min(len(ref_first), len(onnx_first))
+    if n:
+        eq = ref_first[:n] == onnx_first[:n]
+        div = int(np.argmax(~eq)) if not eq.all() else n
+        print(f"  match up to min-len: {int(eq.sum())}/{n} ({100*eq.mean():.1f}%)")
+        print(f"  first divergence step: {div}{' (none)' if div==n else ''}")
+        print(f"  ref [:12]: {ref_first[:12].tolist()}")
+        print(f"  onnx[:12]: {onnx_first[:12].tolist()}")
+
+    print("\n=== ALL 16 codebooks (greedy) ===")
+    m = min(len(ref_codes), len(onnx_codes))
+    if m:
+        full_eq = (ref_codes[:m] == onnx_codes[:m])
+        print(f"  full-grid match: {int(full_eq.sum())}/{full_eq.size} "
+              f"({100*full_eq.mean():.2f}%) over {m} frames × 16 groups")
+
+
+if __name__ == "__main__":
+    main()

eval_predictor.py

@@ -0,0 +1,74 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#   "transformers==4.57.3", "torch", "torchvision", "torchaudio",
+#   "onnx", "onnxruntime>=1.20", "numpy", "safetensors",
+#   "huggingface_hub", "accelerate", "librosa", "soundfile",
+# ]
+# ///
+"""Parity for the teacher-forced code_predictor ONNX.
+
+Two checks:
+  (1) wrapper vs model.forward_sub_talker_finetune  — is the in-graph wrapper
+      faithful to the original PyTorch composition?
+  (2) ONNX vs wrapper                               — did export preserve it?
+
+Usage:
+  uv run eval_predictor.py --model-path onnx/cpu_fp32 --tts-path voicedesign
+"""
+import argparse, json, sys
+from pathlib import Path
+import numpy as np
+
+HERE = Path(__file__).parent
+sys.path.insert(0, str(HERE)); sys.path.insert(0, str(HERE / "codes"))
+
+
+def cosine(a, b):
+    a, b = a.ravel().astype(np.float64), b.ravel().astype(np.float64)
+    return float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-9))
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model-path", required=True)
+    ap.add_argument("--tts-path", default="voicedesign")
+    args = ap.parse_args()
+
+    import torch, onnxruntime as ort
+    from user_script import get_code_predictor_model, _load_tts
+
+    mdir = Path(args.model_path)
+    fn = json.loads((mdir / "manifest.json").read_text())["sub_models"]["code_predictor"]["filename"]
+    sess = ort.InferenceSession(str(mdir / fn), providers=["CPUExecutionProvider"])
+
+    wrap = get_code_predictor_model(args.tts_path)
+    talker = _load_tts(args.tts_path).talker
+    ng = wrap.n_groups
+
+    torch.manual_seed(0)
+    hidden = torch.randn(1, 2048)
+    codes = torch.randint(0, 2048, (1, ng), dtype=torch.int64)
+
+    with torch.no_grad():
+        wrap_logits = wrap(hidden, codes).numpy()
+        native_logits, _ = talker.forward_sub_talker_finetune(codes, hidden)  # [B,15,vocab]
+        native_logits = native_logits.numpy()
+    onnx_logits = sess.run(None, {"talker_hidden": hidden.numpy(),
+                                  "codec_ids": codes.numpy()})[0]
+
+    print("=== (1) wrapper vs native forward_sub_talker_finetune ===")
+    print(f"  shapes wrap={wrap_logits.shape} native={native_logits.shape}")
+    print(f"  cosine={cosine(wrap_logits, native_logits):.6f}  "
+          f"max|Δ|={np.abs(wrap_logits-native_logits).max():.3e}  "
+          f"argmax agree={(wrap_logits.argmax(-1)==native_logits.argmax(-1)).mean():.3%}")
+
+    print("=== (2) ONNX vs wrapper ===")
+    print(f"  shapes onnx={onnx_logits.shape}")
+    print(f"  cosine={cosine(onnx_logits, wrap_logits):.6f}  "
+          f"max|Δ|={np.abs(onnx_logits-wrap_logits).max():.3e}  "
+          f"argmax agree={(onnx_logits.argmax(-1)==wrap_logits.argmax(-1)).mean():.3%}")
+
+
+if __name__ == "__main__":
+    main()

eval_speaker.py

@@ -0,0 +1,55 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#   "transformers==4.57.3", "torch", "torchvision", "torchaudio",
+#   "onnx", "onnxruntime>=1.20", "numpy", "safetensors",
+#   "huggingface_hub", "accelerate", "librosa", "soundfile",
+#   "numba>=0.60.0", "llvmlite>=0.43.0",
+# ]
+# ///
+"""Parity: speaker_encoder.onnx (x-vector, with reimplemented mel/STFT) vs PyTorch
+extract_speaker_embedding. Validates the one genuinely-new ONNX component for Base
+voice cloning across several audio lengths.
+
+Usage:
+  uv run eval_speaker.py --tts-dir base/1.7B --onnx onnx/base17_cpu_int4/speaker_encoder.onnx
+"""
+import argparse, sys
+from pathlib import Path
+import numpy as np
+
+HERE = Path(__file__).parent
+sys.path.insert(0, str(HERE)); sys.path.insert(0, str(HERE / "codes"))
+
+
+def cosine(a, b):
+    a, b = a.ravel().astype(np.float64), b.ravel().astype(np.float64)
+    return float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-9))
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--tts-dir", default="base/1.7B")
+    ap.add_argument("--onnx", default="onnx/base17_cpu_int4/speaker_encoder.onnx")
+    args = ap.parse_args()
+
+    import torch, onnxruntime as ort
+    from user_script import _load_tts
+
+    model = _load_tts(args.tts_dir)
+    assert model.speaker_encoder is not None, "not a base checkpoint"
+    sess = ort.InferenceSession(args.onnx, providers=["CPUExecutionProvider"])
+    rng = np.random.default_rng(0)
+
+    for secs in (2.0, 3.5, 6.0):
+        n = int(24000 * secs)
+        audio = (0.5 * rng.standard_normal(n)).astype(np.float32)
+        with torch.no_grad():
+            ref = model.extract_speaker_embedding(audio, 24000).cpu().numpy().ravel()
+        got = sess.run(None, {"audio": audio[None]})[0].ravel()
+        print(f"len={secs:>4}s  onnx{got.shape} ref{ref.shape}  "
+              f"cosine={cosine(got, ref):.6f}  max|Δ|={np.abs(got - ref).max():.3e}")
+
+
+if __name__ == "__main__":
+    main()

eval_tokenizer.py

@@ -0,0 +1,174 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#   "transformers==4.57.3",
+#   "torch",
+#   "torchvision",
+#   "torchaudio",
+#   "onnx",
+#   "onnxruntime>=1.20",
+#   "numpy",
+#   "safetensors",
+#   "huggingface_hub",
+#   "accelerate",
+#   "librosa",
+#   "soundfile",
+# ]
+# ///
+"""Evaluate the exported Qwen3-TTS tokenizer ONNX vs the original PyTorch.
+
+Reference = the exact `TokEncoderWrapper` / `TokDecoderWrapper` forwards from
+user_script.py (what optimize.py exported), loaded in fp32 from the tokenizer
+checkpoint. Pins transformers==4.57.3 via PEP-723, like optimize.py.
+
+Checks
+------
+(a) Encoder parity : audio[1,1,24000] → codes[1,frames,16]; exact index match %
+    (ONNX vs PyTorch) + per-codebook agreement.
+(b) Decoder parity : codes → waveform; cosine + max|Δ| (ONNX vs PyTorch).
+(c) Round-trip     : encode→decode entirely in ONNX; reconstruction vs input
+    (and vs the PyTorch round-trip) — cosine / max|Δ|.
+
+Usage:
+  uv run eval_tokenizer.py --model-path onnx/cpu_fp16
+  uv run eval_tokenizer.py --model-path onnx/cpu_fp16 --tok-path tokenizer --save-wav
+"""
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+
+HERE = Path(__file__).parent
+sys.path.insert(0, str(HERE))
+sys.path.insert(0, str(HERE / "codes"))
+
+SR = 24000
+N = 24000  # encoder is exported with a static 1-second input
+
+
+def cosine(a, b):
+    a, b = a.ravel().astype(np.float64), b.ravel().astype(np.float64)
+    return float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-9))
+
+
+def test_signal(seed=0):
+    """Deterministic 1 s @ 24 kHz mix of sweeps + harmonics (speech-band-ish)."""
+    rng = np.random.default_rng(seed)
+    t = np.arange(N) / SR
+    sweep = np.sin(2 * np.pi * (120 + 400 * t) * t)
+    harm = 0.4 * np.sin(2 * np.pi * 220 * t) + 0.2 * np.sin(2 * np.pi * 440 * t)
+    env = np.clip(np.sin(2 * np.pi * 2.5 * t), 0, 1)          # syllable-like envelope
+    sig = (sweep + harm) * (0.5 + 0.5 * env) + 0.01 * rng.standard_normal(N)
+    sig = sig / (np.abs(sig).max() + 1e-6) * 0.95
+    return sig.astype(np.float32)[None, None, :]               # [1, 1, N]
+
+
+def main():
+    ap = argparse.ArgumentParser(description="Qwen3-TTS tokenizer ONNX vs PyTorch")
+    ap.add_argument("--model-path", required=True, help="onnx/{device}_{precision} dir")
+    ap.add_argument("--tok-path", default="tokenizer", help="PyTorch tokenizer checkpoint dir")
+    ap.add_argument("--save-wav", action="store_true", help="dump input + reconstructions")
+    args = ap.parse_args()
+
+    import torch
+    import onnxruntime as ort
+    from user_script import get_tok_encoder_model, get_tok_decoder_model
+
+    mdir = Path(args.model_path)
+    manifest = json.loads((mdir / "manifest.json").read_text())
+    enc_onnx = mdir / manifest["sub_models"]["tok_encoder"]["filename"]
+    dec_onnx = mdir / manifest["sub_models"]["tok_decoder"]["filename"]
+    print(f"ONNX dir   : {mdir}")
+    print(f"  encoder  : {enc_onnx.name}")
+    print(f"  decoder  : {dec_onnx.name}")
+    print(f"PyTorch ref: {args.tok_path}\n")
+
+    so = ort.SessionOptions()
+    so.log_severity_level = 3
+    enc_sess = ort.InferenceSession(str(enc_onnx), so, providers=["CPUExecutionProvider"])
+    dec_sess = ort.InferenceSession(str(dec_onnx), so, providers=["CPUExecutionProvider"])
+
+    dec_shape = dec_sess.get_inputs()[0].shape          # [batch, frames, 16]
+    dec_frames = dec_shape[1] if isinstance(dec_shape[1], int) else None
+    print(f"  decoder input shape: {dec_shape}"
+          + ("  (frames FIXED — see note)" if dec_frames else "  (frames dynamic)"))
+
+    def fit_frames(codes, F):
+        """Tile/trim codes [1,T,16] to F frames so a fixed-length decoder accepts them."""
+        if F is None or codes.shape[1] == F:
+            return codes
+        T = codes.shape[1]
+        idx = np.arange(F) % T
+        return codes[:, idx, :]
+
+    print("Loading PyTorch reference wrappers (fp32) ...")
+    enc_pt = get_tok_encoder_model(args.tok_path)
+    dec_pt = get_tok_decoder_model(args.tok_path)
+
+    audio = test_signal()
+
+    # ── (a) Encoder parity ──────────────────────────────────────────────────
+    print("\n=== (a) Encoder parity  audio[1,1,24000] → codes[1,T,16] ===")
+    codes_onnx = enc_sess.run(None, {"audio": audio})[0]
+    with torch.no_grad():
+        codes_pt = enc_pt(torch.from_numpy(audio)).cpu().numpy()
+    codes_onnx = np.asarray(codes_onnx).astype(np.int64)
+    codes_pt = codes_pt.astype(np.int64)
+    print(f"  shapes   onnx={codes_onnx.shape}  pytorch={codes_pt.shape}")
+    if codes_onnx.shape == codes_pt.shape:
+        match = float((codes_onnx == codes_pt).mean())
+        print(f"  exact index match : {match:.4%}")
+        per_cb = (codes_onnx == codes_pt).mean(axis=(0, 1))   # [16]
+        worst = int(np.argmin(per_cb))
+        print(f"  per-codebook match: min={per_cb.min():.3f} (cb{worst})  "
+              f"mean={per_cb.mean():.3f}  max={per_cb.max():.3f}")
+        enc_ok = match > 0.99
+    else:
+        print("  SHAPE MISMATCH — cannot compare indices")
+        enc_ok = False
+
+    # ── (b) Decoder parity ──────────────────────────────────────────────────
+    print("\n=== (b) Decoder parity  codes → waveform ===")
+    codes_in = fit_frames(codes_pt, dec_frames)  # same codes into both decoders
+    if dec_frames and codes_in.shape[1] != codes_pt.shape[1]:
+        print(f"  (decoder is fixed at {dec_frames} frames; tiled {codes_pt.shape[1]}→{dec_frames})")
+    wav_onnx = dec_sess.run(None, {"audio_codes": codes_in})[0]
+    with torch.no_grad():
+        wav_pt = dec_pt(torch.from_numpy(codes_in)).cpu().numpy()
+    n = min(wav_onnx.shape[-1], wav_pt.shape[-1])
+    c = cosine(wav_onnx[..., :n], wav_pt[..., :n])
+    d = float(np.abs(wav_onnx[..., :n] - wav_pt[..., :n]).max())
+    print(f"  shapes   onnx={wav_onnx.shape}  pytorch={wav_pt.shape}")
+    print(f"  cosine={c:.5f}  max|Δ|={d:.4e}")
+    dec_ok = c > 0.999
+
+    # ── (c) Round-trip (ONNX encode→decode) ─────────────────────────────────
+    print("\n=== (c) Round-trip  ONNX encode→decode vs input / PyTorch ===")
+    codes_rt = fit_frames(codes_onnx, dec_frames)
+    wav_rt_onnx = dec_sess.run(None, {"audio_codes": codes_rt})[0]
+    with torch.no_grad():
+        wav_rt_pt = dec_pt(torch.from_numpy(codes_rt)).cpu().numpy()
+    m = min(wav_rt_onnx.shape[-1], audio.shape[-1])
+    print(f"  recon vs input    : cosine={cosine(wav_rt_onnx[..., :m], audio[..., :m]):.4f}")
+    k = min(wav_rt_onnx.shape[-1], wav_rt_pt.shape[-1])
+    print(f"  ONNX vs PyTorch RT: cosine={cosine(wav_rt_onnx[..., :k], wav_rt_pt[..., :k]):.5f}  "
+          f"max|Δ|={float(np.abs(wav_rt_onnx[..., :k] - wav_rt_pt[..., :k]).max()):.4e}")
+
+    if args.save_wav:
+        import soundfile as sf
+        sf.write(mdir / "eval_input.wav", audio[0, 0], SR)
+        sf.write(mdir / "eval_recon_onnx.wav", wav_rt_onnx.reshape(-1), SR)
+        sf.write(mdir / "eval_recon_pytorch.wav", wav_rt_pt.reshape(-1), SR)
+        print(f"\n  wrote eval_input.wav / eval_recon_onnx.wav / eval_recon_pytorch.wav → {mdir}")
+
+    print("\n=== verdict ===")
+    print(f"  encoder index parity : {'PASS' if enc_ok else 'CHECK'}")
+    print(f"  decoder waveform     : {'PASS' if dec_ok else 'CHECK'}")
+    sys.exit(0 if (enc_ok and dec_ok) else 1)
+
+
+if __name__ == "__main__":
+    main()

inference.py

@@ -0,0 +1,591 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#   "onnxruntime>=1.20", "numpy", "soundfile", "librosa", "transformers",
+#   "numba>=0.60.0", "llvmlite>=0.43.0",
+# ]
+# ///
+# numba/llvmlite pinned: librosa otherwise pulls numba 0.53.1 → llvmlite 0.36,
+# which fails to build on Python 3.12.
+"""Run the exported Qwen3-TTS ONNX sub-parts (build blocks + full text→speech).
+
+One --model-path points at  onnx/{device}_{precision}/  (flat, manifest-driven):
+  text_embed.onnx      text_ids[B,T]            → text_embeds[B,T,2048]   (text_projection∘text_embedding)
+  codec_embed.onnx     codec_ids[B,T]           → codec_embeds[B,T,2048]  (talker first-codebook embed)
+  talker.onnx          inputs_embeds[B,T,2048] + position_ids[3,B,T] + mask → (logits[B,T,V], hidden[B,T,2048])
+  code_predictor.onnx  talker_hidden[B,2048] + codec_ids[B,16] → group_logits[B,15,vocab]  (causal teacher-forced)
+  residual_embed.onnx  codec_ids[B,16]          → step_embed[B,2048]      (codec_hiddens.sum — next talker input)
+  tok_encoder.onnx     audio[B,1,24000]         → codes[B,frames,16]
+  tok_decoder.onnx     codes[B,25,16]           → waveform[B,1,L]   (FIXED 25 frames)
+
+GENERATION (`generate`) mirrors `Qwen3TTSForConditionalGeneration.generate` with
+`non_streaming_mode=True` for the two convertible checkpoints:
+  • text-to-speech  — text + language
+  • voice design    — text + instruct (natural-language style)            [VoiceDesign]
+  • custom voice    — text + speaker name (+ optional instruct)           [CustomVoice]
+ICL audio-clone (ref_audio/ref_text) is a `base`-model feature (needs the speaker
+encoder, which is absent here) and is intentionally not implemented.
+
+The talker is a no-cache forward, so the AR loop re-runs the growing prefix each
+step (correct, O(n²)). MROPE reduces to arange here (`get_rope_index` =
+cumsum(mask)-1, 3 identical rows for an unpadded single sequence).
+
+Usage:
+  uv run inference.py --model-path onnx/cpu_fp32 --selftest
+  uv run inference.py --model-path onnx/cpu_fp32 --tts-dir voicedesign \
+      --text "Hello there." --instruct "A calm, low female voice." --out out.wav
+  uv run inference.py --model-path onnx/cpu_fp32 --tts-dir customvoice \
+      --text "你好。" --speaker ethan --language chinese --out out.wav
+"""
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+
+SR = 24000
+DEC_FRAMES = 25            # tok_decoder is exported at a fixed 25-frame length
+N_GROUPS = 16
+
+
+def cosine(a, b):
+    a, b = a.ravel().astype(np.float64), b.ravel().astype(np.float64)
+    return float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-9))
+
+
+# ── sampling helpers (numpy; mirror HF generate logic) ──────────────────────────
+def _apply_repetition_penalty(logits, prev_ids, penalty):
+    if penalty == 1.0 or not prev_ids:
+        return logits
+    idx = np.array(sorted(set(int(i) for i in prev_ids)), dtype=np.int64)
+    sc = logits[idx]
+    logits[idx] = np.where(sc < 0, sc * penalty, sc / penalty)
+    return logits
+
+
+def _sample(logits, do_sample, top_k, top_p, temperature, rng):
+    logits = logits.astype(np.float64)
+    if not do_sample or temperature <= 0:
+        return int(np.argmax(logits))
+    logits = logits / max(temperature, 1e-6)
+    if top_k and top_k > 0:
+        k = min(top_k, logits.shape[-1])
+        kth = np.partition(logits, -k)[-k]
+        logits = np.where(logits < kth, -np.inf, logits)
+    logits -= logits.max()
+    probs = np.exp(logits)
+    probs /= probs.sum()
+    if top_p and top_p < 1.0:
+        order = np.argsort(probs)[::-1]
+        csum = np.cumsum(probs[order])
+        cut = np.searchsorted(csum, top_p) + 1
+        keep = order[:cut]
+        mask = np.zeros_like(probs)
+        mask[keep] = probs[keep]
+        probs = mask / mask.sum()
+    return int(rng.choice(len(probs), p=probs))
+
+
+class Pipeline:
+    """Manifest-driven loader for the exported Qwen3-TTS ONNX sub-parts."""
+
+    def __init__(self, model_path: str, tts_dir: str = None):
+        import onnxruntime as ort
+
+        self.root = Path(model_path)
+        self.manifest = json.loads((self.root / "manifest.json").read_text())
+        sm = self.manifest["sub_models"]
+        prov = self.manifest.get("execution_provider", "CPUExecutionProvider")
+        avail = ort.get_available_providers()
+        if prov not in avail:
+            print(f"  [warn] manifest EP {prov} unavailable; falling back to CPU", file=sys.stderr)
+            prov = "CPUExecutionProvider"
+        self.provider = prov
+
+        so = ort.SessionOptions()
+        so.log_severity_level = 3
+
+        def sess(name):
+            if name not in sm:
+                return None
+            return ort.InferenceSession(str(self.root / sm[name]["filename"]),
+                                        so, providers=[prov])
+
+        self.text_embed = sess("text_embed")
+        self.codec_embed = sess("codec_embed")
+        self.talker = sess("talker")
+        self.code_predictor = sess("code_predictor")
+        self.residual_embed = sess("residual_embed")
+        self.tok_encoder = sess("tok_encoder")
+        self.tok_decoder = sess("tok_decoder")
+        self.speaker_encoder = sess("speaker_encoder")   # Base only (voice-clone x-vector)
+        self.talker_cache = sess("talker_cache")         # optional O(n) KV-cache talker
+        if self.talker_cache is not None:                # ordered past-input names (flattened)
+            self._past_names = [i.name for i in self.talker_cache.get_inputs()][3:]
+
+        # config + tokenizer (only needed for full generation)
+        self.tts_dir = tts_dir
+        self._cfg = None
+        self._tok = None
+        if tts_dir is not None:
+            self._cfg = json.loads((Path(tts_dir) / "config.json").read_text())
+
+    # ── building blocks (each verified against PyTorch in eval_*.py) ──────────
+    def embed_text(self, text_ids):                 # [B,T] int64 → [B,T,2048]
+        return self.text_embed.run(None, {"text_ids": np.asarray(text_ids, np.int64)})[0]
+
+    def embed_codec(self, codec_ids):               # [B,T] int64 → [B,T,2048]
+        return self.codec_embed.run(None, {"codec_ids": np.asarray(codec_ids, np.int64)})[0]
+
+    def talker_step(self, inputs_embeds, position_ids, attention_mask):
+        """→ (logits[B,T,V], hidden[B,T,2048]).  Talker emits both since the
+        code_predictor is conditioned on the talker's last hidden state."""
+        return self.talker.run(None, {
+            "inputs_embeds": inputs_embeds.astype(np.float32),
+            "position_ids": np.asarray(position_ids, np.int64),
+            "attention_mask": np.asarray(attention_mask, np.int64)})
+
+    def talker_cache_step(self, inputs_embeds, position_ids, attention_mask, past):
+        """KV-cache talker: → (logits[B,cur,V], hidden[B,cur,2048], present[list of 56]).
+        `past`/`present` are ordered lists of the flattened K/V tensors (layer0_k, layer0_v,
+        layer1_k, …). Empty past = prefill; len-1 cur = decode."""
+        feed = {"inputs_embeds": inputs_embeds.astype(np.float32),
+                "position_ids": np.asarray(position_ids, np.int64),
+                "attention_mask": np.asarray(attention_mask, np.int64)}
+        for name, t in zip(self._past_names, past):
+            feed[name] = t.astype(np.float32)
+        out = self.talker_cache.run(None, feed)
+        return out[0], out[1], list(out[2:])      # logits, hidden, present
+
+    def predict_residual(self, talker_hidden, codec_ids):   # causal teacher-forced
+        return self.code_predictor.run(None, {
+            "talker_hidden": talker_hidden.astype(np.float32),
+            "codec_ids": np.asarray(codec_ids, np.int64)})[0]
+
+    def step_embed(self, codec_ids):                # [B,16] → [B,2048]  (sum of group embeds)
+        return self.residual_embed.run(None, {"codec_ids": np.asarray(codec_ids, np.int64)})[0]
+
+    @staticmethod
+    def _load_ref_wav(path):
+        """Load a reference wav → mono float32 @ 24 kHz."""
+        import soundfile as sf
+        wav, sr = sf.read(path, dtype="float32")
+        if wav.ndim > 1:
+            wav = wav.mean(axis=1)
+        if sr != SR:
+            import librosa
+            wav = librosa.resample(wav, orig_sr=sr, target_sr=SR)
+        return wav.astype(np.float32)
+
+    def encode_chunked(self, wav):                  # wav [L] 24k → ref codes [T,16]
+        """tok_encoder is fixed at 1 s (24000 samples); encode in 1 s windows + concat."""
+        out = []
+        for s in range(0, max(len(wav), 1), SR):
+            c = wav[s:s + SR]
+            if len(c) < SR:
+                c = np.pad(c, (0, SR - len(c)))
+            out.append(self.encode(c.reshape(1, 1, SR).astype(np.float32))[0])  # [frames,16]
+        return np.concatenate(out, axis=0).astype(np.int64)
+
+    def encode(self, audio):                        # [B,1,24000] → [B,frames,16]
+        return self.tok_encoder.run(None, {"audio": audio.astype(np.float32)})[0]
+
+    def decode(self, codes):                        # [B,F,16] → [B,1,L]
+        return self.tok_decoder.run(None, {"audio_codes": np.asarray(codes, np.int64)})[0]
+
+    def decode_chunked(self, codes):
+        """Decode arbitrary-length codes through the fixed-25-frame decoder by
+        tiling each 25-frame chunk; tail is padded by repetition then trimmed."""
+        F = codes.shape[1]
+        outs = []
+        for s in range(0, F, DEC_FRAMES):
+            chunk = codes[:, s:s + DEC_FRAMES]
+            if chunk.shape[1] < DEC_FRAMES:                     # pad tail by repeat
+                idx = np.arange(DEC_FRAMES) % chunk.shape[1]
+                chunk = chunk[:, idx]
+                wav = self.decode(chunk)
+                keep = int(round(wav.shape[-1] * (F - s) / DEC_FRAMES))
+                outs.append(wav[..., :keep]); break
+            outs.append(self.decode(chunk))
+        return np.concatenate(outs, axis=-1)
+
+    # ── config / tokenizer accessors ──────────────────────────────────────────
+    @property
+    def cfg(self):
+        if self._cfg is None:
+            raise RuntimeError("Pass --tts-dir (the HF model dir) for generation: "
+                               "config token ids + tokenizer live there.")
+        return self._cfg
+
+    @property
+    def model_type(self):
+        """tts_model_type from config: 'voice_design' | 'custom_voice' | 'base'.
+        Each model exposes different features (see _check_features)."""
+        c = self.cfg
+        return (c.get("tts_model_type")
+                or c.get("talker_config", {}).get("tts_model_type") or "unknown")
+
+    def _check_features(self, instruct=None, speaker=None, ref_audio=None):
+        """Gate features by model type so a flag that the loaded model can't honor
+        fails loudly instead of silently doing nothing:
+          voice_design → instruct (natural-language style); no speaker/ref
+          custom_voice → speaker (built-in voices) + optional instruct; no ref
+          base         → voice cloning (ref_audio/ref_text); no speaker/instruct
+        """
+        mt = self.model_type
+        if speaker and mt != "custom_voice":
+            raise ValueError(f"--speaker is a CustomVoice feature, but this model is '{mt}'. "
+                             "Use a customvoice checkpoint, or drop --speaker.")
+        if instruct and mt not in ("voice_design", "custom_voice"):
+            raise ValueError(f"--instruct is a VoiceDesign/CustomVoice feature, but this model "
+                             f"is '{mt}'. Drop --instruct (Base clones from --ref-audio instead).")
+        if ref_audio and mt != "base":
+            raise ValueError(f"voice cloning (--ref-audio) is a Base-model feature, but this "
+                             f"model is '{mt}'. Use a base checkpoint.")
+        if mt == "custom_voice" and not speaker:
+            print("  [note] CustomVoice with no --speaker → model's default voice.",
+                  file=sys.stderr)
+
+    @property
+    def tokenizer(self):
+        if self._tok is None:
+            from transformers import AutoTokenizer
+            self._tok = AutoTokenizer.from_pretrained(self.tts_dir, trust_remote_code=True)
+        return self._tok
+
+    def _ids(self, text):
+        enc = self.tokenizer(text, return_tensors="np")
+        ids = enc["input_ids"]
+        return ids if ids.ndim == 2 else ids[None]
+
+    # ── full text→speech generation ────────────────────────────────────────────
+    def generate(self, text, language="Auto", instruct=None, speaker=None,
+                 ref_audio=None, ref_text=None,
+                 max_new_tokens=2048, do_sample=True, top_k=50, top_p=1.0,
+                 temperature=0.9, repetition_penalty=1.05,
+                 sub_do_sample=True, sub_top_k=50, sub_top_p=1.0, sub_temperature=0.9,
+                 seed=0, verbose=True):
+        """Mirror Qwen3TTSForConditionalGeneration.generate (non_streaming_mode=True).
+        Features are gated by model type (voice_design=instruct, custom_voice=speaker,
+        base=clone). Returns codes [T,16] (int64). Decode with `decode_chunked(codes[None])`.
+        """
+        cfg = self.cfg
+        self._check_features(instruct=instruct, speaker=speaker, ref_audio=ref_audio)
+        if ref_audio:
+            return self._generate_clone(text, ref_audio, ref_text, language=language,
+                                        max_new_tokens=max_new_tokens, do_sample=do_sample,
+                                        top_k=top_k, top_p=top_p, temperature=temperature,
+                                        repetition_penalty=repetition_penalty,
+                                        sub_do_sample=sub_do_sample, sub_top_k=sub_top_k,
+                                        sub_top_p=sub_top_p, sub_temperature=sub_temperature,
+                                        seed=seed, verbose=verbose)
+        tc = cfg["talker_config"]
+        H = tc["hidden_size"]
+        rng = np.random.default_rng(seed)
+
+        # token ids
+        tts_bos, tts_eos, tts_pad = (cfg["tts_bos_token_id"], cfg["tts_eos_token_id"],
+                                     cfg["tts_pad_token_id"])
+        codec_eos = tc["codec_eos_token_id"]
+        codec_pad, codec_bos = tc["codec_pad_id"], tc["codec_bos_id"]
+        vocab = tc["vocab_size"]
+
+        # 0) text → ids (assistant template).  role = first 3, trailing = last 5.
+        assistant = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
+        input_id = self._ids(assistant)                       # [1, L]
+        if input_id.shape[1] < 9:
+            raise ValueError("text tokenized too short for the assistant template")
+
+        # 1) special text embeds (tts_bos/eos/pad)
+        spec = self.embed_text([[tts_bos, tts_eos, tts_pad]])   # [1,3,H]
+        bos_e, eos_e, pad_e = spec[:, 0:1], spec[:, 1:2], spec[:, 2:3]
+
+        # 2) language + codec prefill tags
+        lang = (language or "auto").lower()
+        if lang == "auto" or lang not in tc.get("codec_language_id", {}):
+            language_id = None
+        else:
+            language_id = tc["codec_language_id"][lang]
+        if language_id is None:
+            codec_prefill = [[tc["codec_nothink_id"], tc["codec_think_bos_id"],
+                              tc["codec_think_eos_id"]]]
+        else:
+            codec_prefill = [[tc["codec_think_id"], tc["codec_think_bos_id"],
+                              language_id, tc["codec_think_eos_id"]]]
+        codec0 = self.embed_codec(codec_prefill)                # [1,P,H]
+        codec1 = self.embed_codec([[codec_pad, codec_bos]])     # [1,2,H]
+
+        # speaker (custom voice): speaker name → spk_id → codec table embed
+        speaker_embed = None
+        if speaker:
+            spk_map = tc.get("spk_id", {})
+            if speaker.lower() not in spk_map:
+                raise ValueError(f"Speaker '{speaker}' not in spk_id {list(spk_map)[:8]}…")
+            speaker_embed = self.embed_codec([[spk_map[speaker.lower()]]])   # [1,1,H]
+
+        if speaker_embed is None:
+            codec_input = np.concatenate([codec0, codec1], axis=1)
+        else:
+            codec_input = np.concatenate([codec0, speaker_embed, codec1], axis=1)
+
+        # 3) instruct prefix (voice design) — prepended text_projection embeds
+        prefix = []
+        if instruct:
+            instruct_text = f"<|im_start|>user\n{instruct}<|im_end|>\n"
+            prefix.append(self.embed_text(self._ids(instruct_text)))
+
+        # 4) assemble talker prefill (non_streaming_mode=True)
+        role = self.embed_text(input_id[:, :3])                 # <|im_start|>assistant\n
+        pad_block = np.concatenate(
+            [np.repeat(pad_e, codec_input.shape[1] - 2, axis=1), bos_e], axis=1)
+        talker_in = np.concatenate([role, pad_block + codec_input[:, :-1]], axis=1)
+
+        body_ids = input_id[:, 3:-5]                            # pure text tokens
+        Ltext = body_ids.shape[1]
+        text_body = self.embed_text(body_ids)                   # [1,Ltext,H]
+        block1 = (np.concatenate([text_body, eos_e], axis=1)
+                  + self.embed_codec([[codec_pad] * (Ltext + 1)]))
+        block2 = pad_e + self.embed_codec([[codec_bos]])        # [1,1,H]
+        talker_in = np.concatenate([talker_in, block1, block2], axis=1)
+        if prefix:
+            talker_in = np.concatenate(prefix + [talker_in], axis=1)
+        # trailing_text_hidden is just tts_pad in non_streaming mode → add pad_e each step
+        trailing = pad_e[:, 0]                                  # [1,H]
+
+        return self._ar_loop(talker_in, trailing, vocab, codec_eos, max_new_tokens,
+                             do_sample, top_k, top_p, temperature, repetition_penalty,
+                             sub_do_sample, sub_top_k, sub_top_p, sub_temperature, seed, verbose)
+
+    def _ar_loop(self, talker_in, trailing, vocab, codec_eos, max_new_tokens, do_sample,
+                 top_k, top_p, temperature, repetition_penalty, sub_do_sample, sub_top_k,
+                 sub_top_p, sub_temperature, seed, verbose):
+        """AR talker loop (MROPE→arange). Uses the O(n) KV-cache talker if exported,
+        else the no-cache O(n²) talker. Shared by all generation paths. Returns codes [T,16]."""
+        if self.talker_cache is not None:
+            return self._ar_loop_cached(talker_in, trailing, vocab, codec_eos, max_new_tokens,
+                                        do_sample, top_k, top_p, temperature, repetition_penalty,
+                                        sub_do_sample, sub_top_k, sub_top_p, sub_temperature,
+                                        seed, verbose)
+        if self.talker is None:
+            raise RuntimeError("no talker model found: need talker_cache.onnx (preferred) or "
+                               "talker.onnx in the model dir.")
+        rng = np.random.default_rng(seed)
+        suppress = np.array([i for i in range(vocab - 1024, vocab) if i != codec_eos],
+                            dtype=np.int64)
+        all_codes, prev_first = [], []
+        for step in range(max_new_tokens):
+            T = talker_in.shape[1]
+            pos = np.broadcast_to(np.arange(T), (3, 1, T)).copy()
+            mask = np.ones((1, T), dtype=np.int64)
+            logits, hidden = self.talker_step(talker_in, pos, mask)
+            first = logits[0, -1].astype(np.float64).copy()
+            first[suppress] = -np.inf
+            first = _apply_repetition_penalty(first, prev_first, repetition_penalty)
+            code0 = _sample(first, do_sample, top_k, top_p, temperature, rng)
+            if code0 == codec_eos:
+                break
+            prev_first.append(code0)
+            th = hidden[0, -1][None].astype(np.float32)
+            codes16 = np.zeros((1, N_GROUPS), dtype=np.int64)
+            codes16[0, 0] = code0
+            for j in range(1, N_GROUPS):
+                gl = self.predict_residual(th, codes16)
+                codes16[0, j] = _sample(gl[0, j - 1], sub_do_sample, sub_top_k, sub_top_p,
+                                        sub_temperature, rng)
+            all_codes.append(codes16[0].copy())
+            nxt = self.step_embed(codes16)[:, None] + trailing[:, None]
+            talker_in = np.concatenate([talker_in, nxt], axis=1)
+            if verbose and (step + 1) % 25 == 0:
+                print(f"    …{step + 1} frames", file=sys.stderr)
+        codes = np.stack(all_codes, axis=0).astype(np.int64) if all_codes \
+            else np.zeros((0, N_GROUPS), np.int64)
+        if verbose:
+            print(f"  generated {codes.shape[0]} frames")
+        return codes
+
+    def _ar_loop_cached(self, talker_in, trailing, vocab, codec_eos, max_new_tokens, do_sample,
+                        top_k, top_p, temperature, repetition_penalty, sub_do_sample, sub_top_k,
+                        sub_top_p, sub_temperature, seed, verbose):
+        """O(n) KV-cache AR loop: prefill once, then decode one token/step feeding the cache.
+        Numerically identical to the no-cache loop (same positions, full causal attention)."""
+        rng = np.random.default_rng(seed)
+        suppress = np.array([i for i in range(vocab - 1024, vocab) if i != codec_eos], dtype=np.int64)
+        past = [np.zeros((1, 8, 0, 128), np.float32) for _ in self._past_names]
+        T0 = talker_in.shape[1]
+        pos = np.broadcast_to(np.arange(T0), (3, 1, T0)).copy()
+        logits, hidden, past = self.talker_cache_step(talker_in, pos, np.ones((1, T0), np.int64), past)
+        total = T0
+        all_codes, prev_first = [], []
+        for step in range(max_new_tokens):
+            first = logits[0, -1].astype(np.float64).copy()
+            first[suppress] = -np.inf
+            first = _apply_repetition_penalty(first, prev_first, repetition_penalty)
+            code0 = _sample(first, do_sample, top_k, top_p, temperature, rng)
+            if code0 == codec_eos:
+                break
+            prev_first.append(code0)
+            th = hidden[0, -1][None].astype(np.float32)
+            codes16 = np.zeros((1, N_GROUPS), dtype=np.int64)
+            codes16[0, 0] = code0
+            for j in range(1, N_GROUPS):
+                gl = self.predict_residual(th, codes16)
+                codes16[0, j] = _sample(gl[0, j - 1], sub_do_sample, sub_top_k, sub_top_p,
+                                        sub_temperature, rng)
+            all_codes.append(codes16[0].copy())
+            nxt = self.step_embed(codes16)[:, None] + trailing[:, None]    # [1,1,H]
+            pos = np.broadcast_to(np.array([total]), (3, 1, 1)).copy()
+            logits, hidden, past = self.talker_cache_step(
+                nxt, pos, np.ones((1, total + 1), np.int64), past)
+            total += 1
+            if verbose and (step + 1) % 25 == 0:
+                print(f"    …{step + 1} frames (cached)", file=sys.stderr)
+        codes = np.stack(all_codes, axis=0).astype(np.int64) if all_codes \
+            else np.zeros((0, N_GROUPS), np.int64)
+        if verbose:
+            print(f"  generated {codes.shape[0]} frames (KV-cache)")
+        return codes
+
+    def _generate_clone(self, text, ref_audio, ref_text, language="Auto", max_new_tokens=2048,
+                        do_sample=True, top_k=50, top_p=1.0, temperature=0.9,
+                        repetition_penalty=1.05, sub_do_sample=True, sub_top_k=50, sub_top_p=1.0,
+                        sub_temperature=0.9, seed=0, verbose=True):
+        """Base-model voice cloning (ICL), faithful to generate_icl_prompt (modeling L1968)
+        + the x-vector speaker prompt. Reference audio → codes (tok_encoder) + x-vector
+        (speaker_encoder); prompt = role + codec tags(+x-vector) + [ref_text+text+eos / codec_bos
+        + per-frame ref-code sum]. Returns generated codes [T,16]."""
+        if self.speaker_encoder is None:
+            raise RuntimeError("speaker_encoder.onnx missing — export it for the Base model: "
+                               "`uv run optimize.py --model base/1.7B --components speaker_encoder`.")
+        if not ref_text:
+            raise ValueError("voice clone requires --ref-text (transcript of --ref-audio).")
+        cfg = self.cfg; tc = cfg["talker_config"]; H = tc["hidden_size"]
+        tts_bos, tts_eos, tts_pad = (cfg["tts_bos_token_id"], cfg["tts_eos_token_id"],
+                                     cfg["tts_pad_token_id"])
+        codec_eos = tc["codec_eos_token_id"]
+        codec_pad, codec_bos = tc["codec_pad_id"], tc["codec_bos_id"]; vocab = tc["vocab_size"]
+
+        # reference audio → codes (tok_encoder, 1 s windows) + x-vector (speaker_encoder)
+        wav = self._load_ref_wav(ref_audio)
+        ref_code = self.encode_chunked(wav)                                  # [T_ref,16]
+        spk = self.speaker_encoder.run(None, {"audio": wav[None].astype(np.float32)})[0]
+        spk = spk.reshape(1, 1, H)                                           # x-vector [1,1,H]
+
+        assistant = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
+        input_id = self._ids(assistant)
+        ref_id = self._ids(f"<|im_start|>assistant\n{ref_text}<|im_end|>\n")[:, 3:-2]
+        text_id = input_id[:, 3:-5]
+
+        spec = self.embed_text([[tts_bos, tts_eos, tts_pad]])
+        bos_e, eos_e, pad_e = spec[:, 0:1], spec[:, 1:2], spec[:, 2:3]
+        lang = (language or "auto").lower()
+        language_id = (tc["codec_language_id"][lang]
+                       if lang != "auto" and lang in tc.get("codec_language_id", {}) else None)
+        codec_prefill = ([[tc["codec_nothink_id"], tc["codec_think_bos_id"], tc["codec_think_eos_id"]]]
+                         if language_id is None else
+                         [[tc["codec_think_id"], tc["codec_think_bos_id"], language_id,
+                           tc["codec_think_eos_id"]]])
+        codec0 = self.embed_codec(codec_prefill)
+        codec1 = self.embed_codec([[codec_pad, codec_bos]])
+        codec_input = np.concatenate([codec0, spk, codec1], axis=1)         # x-vector injected
+
+        role = self.embed_text(input_id[:, :3])
+        pad_block = np.concatenate([np.repeat(pad_e, codec_input.shape[1] - 2, axis=1), bos_e], axis=1)
+        base = np.concatenate([role, pad_block + codec_input[:, :-1]], axis=1)
+
+        # ICL block (generate_icl_prompt, non_streaming): the per-frame ref-code sum IS step_embed
+        text_embed = np.concatenate([self.embed_text(np.concatenate([ref_id, text_id], axis=1)),
+                                     eos_e], axis=1)                         # [1,T1,H]
+        T1 = text_embed.shape[1]
+        codec_embed = np.concatenate([self.embed_codec([[codec_bos]]),
+                                      self.step_embed(ref_code)[None]], axis=1)   # [1,1+T_ref,H]
+        icl = text_embed + self.embed_codec([[codec_pad] * T1])
+        icl = np.concatenate([icl, codec_embed + pad_e], axis=1)
+        talker_in = np.concatenate([base, icl], axis=1)
+        trailing = pad_e[:, 0]
+        if verbose:
+            print(f"  [clone] ref {ref_code.shape[0]} frames + ref_text {ref_id.shape[1]} toks "
+                  f"+ text {text_id.shape[1]} toks → prefill {talker_in.shape[1]}", file=sys.stderr)
+        return self._ar_loop(talker_in, trailing, vocab, codec_eos, max_new_tokens, do_sample,
+                             top_k, top_p, temperature, repetition_penalty, sub_do_sample,
+                             sub_top_k, sub_top_p, sub_temperature, seed, verbose)
+
+
+def selftest(pipe: Pipeline):
+    """Codec round-trip + building-block smoke test on the loaded EP."""
+    print(f"Provider: {pipe.provider}")
+    rng = np.random.default_rng(0)
+    t = np.arange(SR) / SR
+    audio = (0.6 * np.sin(2 * np.pi * (180 + 300 * t) * t)
+             + 0.01 * rng.standard_normal(SR)).astype(np.float32)[None, None, :]
+
+    codes = pipe.encode(audio)
+    print(f"  encode : audio{audio.shape} -> codes{codes.shape}")
+    wav = pipe.decode_chunked(codes)
+    print(f"  decode : codes{codes.shape} -> wav{wav.shape}")
+
+    te = pipe.embed_text(rng.integers(0, 1000, (1, 8)))
+    ce = pipe.embed_codec(rng.integers(0, 2048, (1, 8)))
+    print(f"  text_embed{te.shape}  codec_embed{ce.shape}")
+    th = rng.standard_normal((1, 2048)).astype(np.float32)
+    gl = pipe.predict_residual(th, rng.integers(0, 2048, (1, N_GROUPS)))
+    print(f"  code_predictor group_logits{gl.shape}")
+    if pipe.residual_embed is not None:
+        se = pipe.step_embed(rng.integers(0, 2048, (1, N_GROUPS)))
+        print(f"  residual_embed step_embed{se.shape}")
+    else:
+        print("  [warn] residual_embed.onnx missing — re-export for generate()")
+    print("OK — codec path + building blocks run on this EP.")
+
+
+def main():
+    ap = argparse.ArgumentParser(description="Qwen3-TTS ONNX inference")
+    ap.add_argument("--model-path", required=True, help="onnx/{device}_{precision} dir")
+    ap.add_argument("--tts-dir", help="HF model dir (config + tokenizer) — needed for --text")
+    ap.add_argument("--selftest", action="store_true", help="codec round-trip + block smoke test")
+    ap.add_argument("--text", help="text to synthesize (full generation)")
+    ap.add_argument("--instruct", help="voice-design style instruction (VoiceDesign only)")
+    ap.add_argument("--speaker", help="built-in speaker name (CustomVoice only)")
+    ap.add_argument("--ref-audio", help="reference wav for voice cloning (Base only)")
+    ap.add_argument("--ref-text", help="transcript of --ref-audio (Base clone)")
+    ap.add_argument("--language", default="Auto", help="language (default Auto)")
+    ap.add_argument("--out", default="out.wav", help="output wav path for --text")
+    ap.add_argument("--max-new-tokens", type=int, default=2048)
+    ap.add_argument("--greedy", action="store_true", help="disable sampling (argmax)")
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--save-wav", help="write the self-test reconstruction to this path")
+    args = ap.parse_args()
+
+    pipe = Pipeline(args.model_path, tts_dir=args.tts_dir)
+
+    if args.text is not None:
+        print(f"  model_type={pipe.model_type}", file=sys.stderr)
+        codes = pipe.generate(
+            args.text, language=args.language, instruct=args.instruct,
+            speaker=args.speaker, ref_audio=args.ref_audio, ref_text=args.ref_text,
+            max_new_tokens=args.max_new_tokens,
+            do_sample=not args.greedy, sub_do_sample=not args.greedy, seed=args.seed)
+        if codes.shape[0] == 0:
+            print("  [warn] no frames generated (immediate EOS)"); return
+        wav = pipe.decode_chunked(codes[None]).reshape(-1)
+        import soundfile as sf
+        sf.write(args.out, wav, SR)
+        print(f"  wrote {args.out}  ({wav.shape[0] / SR:.2f}s)")
+        return
+
+    if args.selftest:
+        selftest(pipe)
+        if args.save_wav:
+            import soundfile as sf
+            rng = np.random.default_rng(0)
+            t = np.arange(SR) / SR
+            audio = (0.6 * np.sin(2 * np.pi * (180 + 300 * t) * t)).astype(np.float32)[None, None, :]
+            sf.write(args.save_wav, pipe.decode_chunked(pipe.encode(audio)).reshape(-1), SR)
+            print(f"  wrote {args.save_wav}")
+        return
+
+    ap.error("nothing to do: pass --selftest or --text ...")
+
+
+if __name__ == "__main__":
+    main()

info.yml

@@ -0,0 +1,11 @@
+keywords:
+  - olive-ai
+  - tts
+  - speech-synthesis
+recipes:
+  - name: Qwen3-TTS-VoiceDesign
+    file: optimize.py
+    eps:
+      - CPUExecutionProvider
+    devices:
+      - cpu

optimize.py

@@ -0,0 +1,263 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#   "transformers==4.57.3",
+#   "torch",
+#   "torchvision",
+#   "torchaudio",
+#   "olive-ai",
+#   "onnx",
+#   "onnxruntime>=1.20",
+#   "onnxruntime-genai",
+#   "safetensors",
+#   "numpy",
+#   "tabulate",
+#   "huggingface_hub",
+#   "accelerate",
+#   "librosa",
+#   "numba>=0.60.0",
+#   "llvmlite>=0.43.0",
+#   "soundfile",
+# ]
+# ///
+# numba/llvmlite pinned >=0.60/0.43: librosa otherwise resolves to numba 0.53.1 →
+# llvmlite 0.36, which fails to build on Python 3.12 ("only versions <3.10 supported").
+"""Qwen3-TTS → ONNX sub-parts.  Run with:  uv run optimize.py --model <name|path> ...
+
+Self-contained uv script: pins transformers==4.57.3 (the version the vendored
+qwen_tts modeling needs) in an isolated env, leaving the shared 5.10.2 venv alone.
+
+Handles three model kinds (auto-detected from config.json `model_type`):
+  qwen3_tts                 → TTS  : talker (+ code_predictor) [+ embedded tokenizer]
+  qwen3_tts_tokenizer_12hz  → codec: tok_encoder + tok_decoder
+
+Outputs flat sub-parts under  onnx/{device}_{precision}/  + manifest.json.
+
+  --device {cpu,cuda}   --precision {int4,fp16,fp32}
+
+The talker is a custom Qwen3+MROPE LLM (dual codec/text embedding, codec_head) and is
+NOT onnxruntime-genai ModelBuilder-compatible (MROPE + dual embedding don't map to a
+stock Qwen3) — it is exported via Olive like the rest. (ModelBuilder remap is attempted
+only when the talker is detectably stock; otherwise Olive is used.)
+
+Usage:
+  uv run optimize.py --model Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign  --device cpu --precision int4
+  uv run optimize.py --model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice  --device cpu --precision fp16
+  uv run optimize.py --model Qwen/Qwen3-TTS-Tokenizer-12Hz         --device cpu --precision fp32
+  uv run optimize.py --model voicedesign --skip-download           # use local dir
+"""
+import argparse
+import json
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+
+HERE = Path(__file__).parent
+PROVIDER = {"cpu": "CPUExecutionProvider", "cuda": "CUDAExecutionProvider"}
+OLIVE_DEV = {"cpu": "cpu", "cuda": "gpu"}
+
+# component → (loader, io_config, dummy_inputs, model_subdir or "")
+TTS_COMPONENTS = {
+    "text_embed":     ("get_text_embed_model", "get_text_embed_io_config", "get_text_embed_dummy_inputs", ""),
+    "codec_embed":    ("get_codec_embed_model", "get_codec_embed_io_config", "get_codec_embed_dummy_inputs", ""),
+    "talker":         ("get_talker_model", "get_talker_io_config", "get_talker_dummy_inputs", ""),
+    "talker_cache":   ("get_talker_cache_model", "get_talker_cache_io_config", "get_talker_cache_dummy_inputs", ""),
+    "code_predictor": ("get_code_predictor_model", "get_code_predictor_io_config", "get_code_predictor_dummy_inputs", ""),
+    "residual_embed": ("get_residual_embed_model", "get_residual_embed_io_config", "get_residual_embed_dummy_inputs", ""),
+    "tok_encoder":    ("get_tok_encoder_model", "get_tok_encoder_io_config", "get_tok_encoder_dummy_inputs", "speech_tokenizer"),
+    "tok_decoder":    ("get_tok_decoder_model", "get_tok_decoder_io_config", "get_tok_decoder_dummy_inputs", "speech_tokenizer"),
+    "speaker_encoder": ("get_speaker_encoder_model", "get_speaker_encoder_io_config", "get_speaker_encoder_dummy_inputs", ""),
+}
+TOKENIZER_COMPONENTS = {
+    "tok_encoder": ("get_tok_encoder_model", "get_tok_encoder_io_config", "get_tok_encoder_dummy_inputs", ""),
+    "tok_decoder": ("get_tok_decoder_model", "get_tok_decoder_io_config", "get_tok_decoder_dummy_inputs", ""),
+}
+# audio codecs + speaker encoder stay fp32 (DAC/conv int4/fp16 too lossy; ECAPA small)
+FP32_ONLY = {"tok_encoder", "tok_decoder", "speaker_encoder"}
+
+
+def output_root(device, precision):
+    return HERE / "onnx" / f"{device}_{precision}"
+
+
+def resolve_model(model: str, skip_download: bool) -> Path:
+    """Local dir (relative or absolute) or HF id → local checkpoint dir."""
+    p = Path(model)
+    if p.is_dir():
+        return p
+    local = HERE / model.split("/")[-1].lower()
+    if local.is_dir():
+        return local
+    if skip_download:
+        raise FileNotFoundError(f"{model} not found locally and --skip-download set")
+    from huggingface_hub import snapshot_download
+    dst = HERE / model.split("/")[-1]
+    print(f"  Downloading {model} → {dst}")
+    snapshot_download(repo_id=model, local_dir=str(dst))
+    return dst
+
+
+def detect_type(model_dir: Path) -> str:
+    cfg = json.loads((model_dir / "config.json").read_text())
+    return cfg.get("model_type", "")
+
+
+def build_component(name, funcs, model_dir, out_dir, device, precision):
+    from olive import run
+    loader, io_cfg, dummy, subdir = funcs
+    src = model_dir / subdir if subdir else model_dir
+    prec = "fp32" if name in FP32_ONLY else precision
+    passes = {"c": {"type": "OnnxConversion", "target_opset": 20}}
+    # Transformers using create_causal_mask aren't torch.onnx(TorchScript)-traceable
+    # → dynamo exporter (codec parts + the MROPE talker).
+    if name in FP32_ONLY or name in ("talker", "talker_cache"):
+        passes["c"]["use_dynamo_exporter"] = True
+    if prec == "int4":
+        passes["q"] = {"type": "OnnxBlockwiseRtnQuantization", "bits": 4,
+                       "block_size": 32, "is_symmetric": False}
+    elif prec == "fp16":
+        passes["h"] = {"type": "OnnxFloatToFloat16"}
+
+    tmp = (out_dir / f"_{name}_tmp").resolve()
+    cfg = {
+        "input_model": {"type": "PyTorchModel", "model_path": str(src.resolve()),
+                        "model_loader": loader,
+                        "model_script": str((HERE / "user_script.py").resolve()),
+                        "io_config": io_cfg, "dummy_inputs_func": dummy},
+        "systems": {"local_system": {"type": "LocalSystem", "accelerators": [
+            {"device": OLIVE_DEV[device], "execution_providers": [PROVIDER[device]]}]}},
+        "passes": passes, "target": "local_system", "log_severity_level": 2,
+        "output_dir": str(tmp), "cache_dir": str((out_dir / "_ocache").resolve()),
+        "no_artifacts": True,
+    }
+    tmp_json = out_dir / f"_{name}.json"
+    tmp_json.write_text(json.dumps(cfg, indent=2))
+    print(f"  Olive export {name} (precision={prec}) → {out_dir}/{name}.onnx")
+    run(str(tmp_json))
+    tmp_json.unlink(missing_ok=True)
+    moved_data = False
+    for f in tmp.glob("model.onnx*"):
+        dst = out_dir / f.name.replace("model.onnx", f"{name}.onnx")
+        if dst.exists():
+            dst.unlink()
+        shutil.move(str(f), str(dst))
+        if dst.name.endswith(".onnx.data"):
+            moved_data = True
+    shutil.rmtree(tmp, ignore_errors=True)
+    shutil.rmtree(out_dir / "_ocache", ignore_errors=True)
+    if moved_data:
+        relink_external_data(out_dir / f"{name}.onnx", f"{name}.onnx.data")
+
+
+def relink_external_data(onnx_path, new_data_name):
+    """Point a renamed model's external_data references at the renamed .data file.
+
+    Olive writes `model.onnx` + `model.onnx.data`; we flatten to `{name}.onnx[.data]`
+    but the proto still references `model.onnx.data`, so onnxruntime fails to load
+    (`External data path does not exist`). Rewrite each tensor's `location` entry.
+    Proto-only (load_external_data=False) so we never pull the multi-GB blob into RAM.
+    """
+    import onnx
+    m = onnx.load(str(onnx_path), load_external_data=False)
+    n = 0
+    for t in m.graph.initializer:
+        if t.HasField("data_location") and t.data_location == onnx.TensorProto.EXTERNAL:
+            for kv in t.external_data:
+                if kv.key == "location":
+                    kv.value = new_data_name
+                    n += 1
+    onnx.save(m, str(onnx_path))      # proto only; .data already on disk
+    print(f"    relinked {n} external-data refs → {new_data_name}")
+
+
+def write_manifest(out_dir, model, device, precision, model_kind, built):
+    # include every sub-part present on disk (robust to incremental builds)
+    present = sorted({p.stem for p in out_dir.glob("*.onnx")})
+    manifest = {
+        "model_id": model, "model_kind": model_kind, "device": device,
+        "precision": precision, "execution_provider": PROVIDER[device],
+        "sub_models": {n: {"filename": f"{n}.onnx"} for n in present},
+    }
+    (out_dir / "manifest.json").write_text(json.dumps(manifest, indent=2))
+    print(f"  Wrote {out_dir / 'manifest.json'}")
+
+
+def main():
+    ap = argparse.ArgumentParser(description="Qwen3-TTS → ONNX sub-parts")
+    ap.add_argument("--model", required=True, help="HF id or local dir")
+    ap.add_argument("--device", choices=["cpu", "cuda"], default="cpu")
+    ap.add_argument("--precision", choices=["int4", "fp16", "fp32"], default="int4")
+    ap.add_argument("--components", nargs="*", default=None, help="subset to build")
+    ap.add_argument("--skip-download", action="store_true")
+    ap.add_argument("--no-isolate", action="store_true",
+                    help="build everything in ONE process (may OOM on the talker). Default: "
+                         "each component runs in its own subprocess for fresh memory.")
+    ap.add_argument("--_child", action="store_true", help=argparse.SUPPRESS)  # internal
+    args = ap.parse_args()
+
+    model_dir = resolve_model(args.model, args.skip_download)
+    mtype = detect_type(model_dir)
+    if mtype == "qwen3_tts":
+        kind, comp_map = "tts", TTS_COMPONENTS
+    elif mtype.startswith("qwen3_tts_tokenizer"):
+        kind, comp_map = "tokenizer", TOKENIZER_COMPONENTS
+    else:
+        raise ValueError(f"Unknown model_type '{mtype}' in {model_dir}")
+
+    if args.components:
+        components = [c for c in args.components if c in comp_map]
+    else:
+        components = list(comp_map)
+        # speaker_encoder exists only in `base` checkpoints — skip it by default elsewhere
+        # (it raises in the wrapper otherwise). Still buildable if requested explicitly.
+        tts_type = json.loads((model_dir / "config.json").read_text()).get("tts_model_type")
+        if "speaker_encoder" in components and tts_type != "base":
+            components.remove("speaker_encoder")
+        # talker and talker_cache hold the SAME transformer weights (no-cache vs KV-cache
+        # forward). Shipping both duplicates ~870 MB. Default to talker_cache only (it does
+        # prefill+decode, is faster O(n), and is what inference auto-uses); the plain no-cache
+        # `talker` is still buildable explicitly if a simpler graph is wanted.
+        if "talker" in components and "talker_cache" in components:
+            components.remove("talker")
+    out_dir = output_root(args.device, args.precision)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    print(f"Model {args.model} [{kind}] | {args.device}/{args.precision} → {out_dir}")
+    print(f"Components: {components}\n")
+
+    # ── memory-isolation: re-invoke this script once per component in a fresh subprocess.
+    # The talker's big dynamo/float16 pass commits a lot of virtual memory; stacking it on
+    # top of others in one process OOM/pagefile-kills on Windows. One process per component
+    # releases everything between builds. (Reuses sys.executable = the resolved uv env.)
+    if not args._child and not args.no_isolate and len(components) > 1:
+        print(f"Isolating {len(components)} components in subprocesses (avoids OOM)…\n")
+        ok, failed = [], []
+        for comp in components:
+            cmd = [sys.executable, str(Path(__file__).resolve()),
+                   "--model", args.model, "--device", args.device,
+                   "--precision", args.precision, "--components", comp, "--_child"]
+            if args.skip_download:
+                cmd.append("--skip-download")
+            print(f"=== [subprocess] {comp} ===")
+            rc = subprocess.run(cmd).returncode
+            (ok if rc == 0 else failed).append(comp)
+            if rc != 0:
+                print(f"  [warn] component '{comp}' failed (rc {rc}) — continuing")
+            print()
+        write_manifest(out_dir, args.model, args.device, args.precision, kind, ok)
+        print(f"\nDone → {out_dir}   built={ok}" + (f"  FAILED={failed}" if failed else ""))
+        return
+
+    # direct build (single component, --_child subprocess, or --no-isolate)
+    built = []
+    for name in components:
+        print(f"=== {name} ===")
+        build_component(name, comp_map[name], model_dir, out_dir, args.device, args.precision)
+        built.append(name); print()
+    if not args._child:                      # children skip manifest; parent writes it
+        write_manifest(out_dir, args.model, args.device, args.precision, kind, built)
+    print(f"\nDone → {out_dir}")
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,10 @@
+# Qwen3-TTS-VoiceDesign → ONNX (CPU INT4) build/run/eval dependencies
+olive-ai
+onnxruntime>=1.20
+onnxruntime-genai           # ModelBuilder (create_model) for the Qwen3 LLM sub-part
+torch
+transformers>=5.10
+safetensors
+numpy
+soundfile
+librosa