# 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).