feat(nemotron_voicechat): load the VoiceChat checkpoint and run offline duplex inference (FastConformer+RNNT, Nemotron-H dual-head, EAR-TTS, codec)
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Summary
Port the four networks of NemotronLabs VoiceChat and the offline timeline loop that ties them together: a 16 kHz WAV goes in, the assistant's text, the function-channel text, the user transcript, and a 22.05 kHz response WAV come out. This sub-issue owns every module, config struct, weight-key rule, and the non-streaming (full-history) loop; the cache-aware online session and the WebSocket endpoint are the next two sub-issues and only add state management on top of the forward functions defined here.
## Current behavior
- `src/models/detection.rs` has no `nemotron_voicechat` arm (`"nemotron_h"` at line 300 and `"nemotron_h_nano_omni"` at 301 are the neighbours).
- `src/models/nemotron_h.rs:1491-1500` (`NemotronHModel { embeddings, layers, norm_f, lm_head, block_types, sequence_state }`) loads `backbone.embeddings`, `backbone.layers.N.*`, `backbone.norm_f.weight`, `lm_head` (`:2176-2500`) and has `forward_stage(inputs, layer_range, has_embedding, has_lm_head, caches)` (`:1583-1662`) which accepts pre-computed embeddings when `has_embedding == false` but returns the hidden state BEFORE `norm_f` when `has_lm_head == false` (`:1656-1661`). There is no second head and no public final-norm entry point.
- `src/models/gemma3.rs:499-560` (`TransformerBlock::from_weights(weights, args, layer_idx)` with the `model.layers.{idx}` prefix hardcoded at line 552) and `Gemma3Model::forward_with_caches_and_embeddings` (`:735-760`) always multiply injected embeddings by `sqrt(hidden_size)` (line 755-757); the TTS backbone must not scale and has no `embed_tokens` / `lm_head`.
- `src/audio/whisper_mel.rs:200` computes a Whisper-specific 80/128-mel log spectrogram (Slaney filters, `log10`, Whisper normalization); the VoiceChat frontend needs preemphasis, a 25 ms / 10 ms Hann STFT of size 512, natural `log(x + 2^-24)`, and no normalization.
- `src/models/kokoro/lstm.rs:43-58` (`load(w, prefix, hidden)`, `forward(x, t)`) and `src/models/kokoro/stft.rs:70-179` (`stft`, `istft`, `irdft_basis`) are the only LSTM and STFT/iSTFT implementations; both are `pub(crate)` and tied to Kokoro's key names and window sizes.
- The tiktoken/HF tokenizer loaders in `src/tokenizer/mod.rs` handle the LLM's `tokenizer.json`; nothing reads `rnnt_tokenizer/vocab.json` or the config's `rnnt_vocabulary` list.
## Expected behavior
```
./target/release/mlxcel generate -m models/NemotronLabs-VoiceChat-11B-4bit \
--audio question.wav --output-audio response.wav -p "Be concise and answer in one sentence." -n 0
```
prints `[user] ` and the assistant text, and writes `response.wav` (22050 Hz mono PCM16) containing the spoken answer in the built-in `Aria` voice. `-p` is the system prompt (empty means none); `-n` is ignored because the output length equals the input timeline length (plus `--extra-decoding-seconds`, default 3, of appended silence so the model can finish answering).
### config.json
Top level: `model_type "nemotron_voicechat"`, `bos_token_id 1`, `eos_token_id 2`, `pad_token_id 12`, `silence_token_id 11`, `rnnt_blank_id 1024`, `input_sample_rate 16000`, `output_sample_rate 22050`, `frame_duration 0.08`, `function_channel_weight 2.0`, `speaker "Aria"`, `rnnt_vocabulary` (list of 1024 SentencePiece pieces), `quantization` / `quantization_config` (per-module map: `group_size 64, bits 4` with one entry per quantized linear; only `stt_model.embed_tokens`, `stt_model.lm_head`, `stt_model.function_head` and the Nemotron-H mixer projections are quantized; the speech encoder, RNNT, TTS and codec stay bf16).
`text_config`: `model_type "nemotron_h"`, `hidden_size 4480`, `intermediate_size 15680`, `num_hidden_layers 56`, `num_attention_heads 40`, `num_key_value_heads 8`, `head_dim 128`, `mamba_num_heads 128`, `mamba_head_dim 80`, `ssm_state_size 128`, `conv_kernel 4`, `n_groups 8`, `use_conv_bias true`, `mamba_proj_bias false`, `mlp_bias false`, `attention_bias false`, `layer_norm_epsilon 1e-5`, `vocab_size 131072`, `max_position_embeddings 131072`, `hybrid_override_pattern "M-M-M-MM-M-M-M*-M-M-M*-M-M-M-M*-M-M-M-M*-M-MM-M-M-M-M-M-"` (attention at layers 14, 21, 30, 39; `-` layers are MLPs).
`audio_config.preprocessor`: `sample_rate 16000`, `features 128`, `n_fft 512`, `window_size 0.025` (400 samples), `window_stride 0.01` (160), `window "hann"`, `preemph 0.97`, `log true`, `normalize "NA"`, `dither 1e-5` (training only; not applied), `log_zero_guard_value 2^-24`, `pad_to 0`.
`audio_config.encoder`: `feat_in 128`, `n_layers 24`, `d_model 1024`, `n_heads 8`, `ff_expansion_factor 4`, `subsampling_factor 8`, `subsampling_conv_channels 256`, `conv_kernel_size 9`, `causal_downsampling true`, `conv_context_size "causal"`, `conv_norm_type "layer_norm"`, `self_attention_model "rel_pos"`, `att_context_style "chunked_limited"`, `att_context_size [[70, 0]]`, `pos_emb_max_len 5000`, `use_bias false`, `xscaling false`.
`audio_config.decoder` (RNNT prediction net): `pred_hidden 640`, `pred_rnn_layers 2`, `vocab_size 1024`, `blank_as_pad true`. `audio_config.joint`: `joint_hidden 640`, `activation "relu"`, `encoder_hidden 1024`, `pred_hidden 640`, `num_classes 1024`. `audio_config.output_dim 4480`, `audio_config.max_symbols 10`.
`tts_config`: `hidden_size 1152`, `intermediate_size 4608`, `num_hidden_layers 28`, `num_attention_heads 16`, `num_key_value_heads 16`, `head_dim 72`, `sliding_window 7500`, `sliding_window_pattern 6`, `rms_norm_eps 1e-6`, `query_pre_attn_scalar 256`, `rope_global_base_freq 1e6`, `rope_local_base_freq 1e4`, `latent_size 512`, `num_quantizers 31`, `codebook_size 1024`, `num_delay_speech_tokens 2`, `num_iterations 8`, `guidance_scale 0.2`, `top_p 0.95`, `noise_scale 0.001`, `exponent 3.0`, `use_gated_fusion_for_text_audio true`, `use_subword_flag_emb true`, `use_bos_eos_emb true`, `use_audio_prompt_frozen_projection true`, `audio_prompt_duration 3.0`; `character_encoder {hidden_size 1152, intermediate_size 4608, num_hidden_layers 1, num_attention_heads 16, num_key_value_heads 16, head_dim 72, rms_norm_eps 1e-6, query_pre_attn_scalar 256, attn_logit_softcapping 50, rope_base 1e4, char_vocab_size 257}`; `mog_head {intermediate_size 4608, low_rank 64, min_log_std -4, num_layers 3, num_predictions 1024, eps 1e-6}`.
`codec_config`: `sample_rate 22050`, `base_channels 384`, `channel_multipliers [1, 2, 4]`, `downsample_rates [7, 7, 9]`, `blocks_per_stage 3`, `block_kernel_size 7`, `latent_dim 512`, `n_fft 16`, `hop_length 4`, `num_quantizers 31`, `codebook_size 1024`. Derived: `waveform_to_token_ratio = 4 * 7 * 7 * 9 = 1764` samples per code frame (80 ms at 22050 Hz), `stft_channels = 2 * (16/2 + 1) = 18`.
### Network 1: log-mel frontend and FastConformer (`stt_model.perception.*`)
Log-mel for a waveform `x` (f32, 16 kHz): `x' = [x0, x1 - 0.97 x0, x2 - 0.97 x1, ...]`; STFT with `n_fft 512`, `hop 160`, a symmetric 400-point Hann window zero-padded to 512 (56 zeros left, 56 right), reflect padding of 256 samples each side (`center = true`); power spectrum `|X|^2` over 257 bins; `mel = filters[128, 257] @ power` with Slaney-normalized Slaney-scale filters (`hz_to_mel_slaney` exists in `src/audio/whisper_mel.rs:60`); `log(mel + 2^-24)`; no normalization. Output `[1, T, 128]` with `T = 1 + floor(len / 160)`.
Subsampling (`pre_encode`): input `[B, T, 128, 1]` channels-last. Three stride-2 stages; every 3x3 stride-2 conv is preceded by asymmetric padding `(left 2, right 1)` on both the time and frequency axes: `conv.0` = Conv2d(1 -> 256, k3, s2) + ReLU; `conv.2` = depthwise Conv2d(256, k3, s2, groups 256), `conv.3` = Conv2d(256 -> 256, k1) + ReLU; `conv.5` / `conv.6` repeat. Output lengths per stage `floor((n + 3 - 3) / 2) + 1`; frequency 128 -> 64 -> 32 -> 16. Flatten `[B, T', 16, 256]` channel-major as `transpose(0, 1, 3, 2).reshape(B, T', 4096)` then `out: Linear(4096 -> 1024, bias)`. Ten ms mel frames become 80 ms encoder frames.
Relative positional encoding: the sinusoid table `pe[2*max_len - 1, 1024]` over positions `max_len-1 .. -(max_len-1)`; for a length-L input take the centered slice of `2L - 1` rows; `xscaling false` so the input is not scaled.
Each of 24 blocks (`layers.N`):
```
x = x + 0.5 * ff1(LN(norm_feed_forward1)(x)) # ff: linear1 1024->4096 (no bias), SiLU, linear2
x = x + attn(LN(norm_self_att)(x), pos_emb, mask)
x = x + conv(LN(norm_conv)(x))
x = x + 0.5 * ff2(LN(norm_feed_forward2)(x))
x = LN(norm_out)(x)
```
Attention (`self_attn`, no bias on `linear_q/k/v/out`, `linear_pos` without bias, per-layer `pos_bias_u` / `pos_bias_v` `[8, 128]`): `q_u = (q + pos_bias_u)`, `q_v = (q + pos_bias_v)`, `p = linear_pos(pos_emb)`; `bd = rel_shift(q_v @ p^T)[..., :L] * scale` where `rel_shift` left-pads the last axis by one, views as `[B, H, pos_len + 1, L]`, drops the first row and views back; scores = `(q_u @ k^T) * scale + bd + mask`; softmax; `@ v`; `linear_out`. Mask for `att_context_size [70, 0]` (`chunked_limited`, chunk size `right + 1 = 1`): frame i may attend to frames `i - 70 ..= i`, everything else is `-1e30`.
Convolution module: `pointwise_conv1` (1024 -> 2048, k1, no bias) -> GLU on the channel axis -> left-pad 8 frames -> depthwise conv k9 (groups 1024, no bias) -> LayerNorm (weight key `conv.batch_norm`, it is a LayerNorm despite the name) -> SiLU -> `pointwise_conv2` (k1, no bias).
`perception.proj`: Linear(1024 -> 4480, bias). The perception returns both the projected `[B, T', 4480]` (LLM audio channel) and the raw `[B, T', 1024]` encoder output (RNNT input).
### Network 1b: RNNT transcript branch (`stt_model.rnnt_decoder.*`, `stt_model.rnnt_joint.*`)
Prediction net: `prediction.embed` Embedding(1025, 640) (index 1024 is blank, never embedded), `prediction.dec_rnn.lstm` two stacked LSTM(640) layers. Joint: `enc` Linear(1024 -> 640, bias), `pred` Linear(640 -> 640, bias), ReLU, `joint_net.2` Linear(640 -> 1025, bias). Greedy decode per encoder frame `e[1, 1, 1024]`: loop up to `max_symbols = 10` times: `pred_in = embed(last_token)` if `last_token != blank` else zeros `[1, 1, 640]`; `(out, h') = lstm(pred_in, h)`; `logits = joint_net(relu(enc(e) + pred(out)))`; `tok = argmax`; if `tok == 1024` break; else `last_token = tok`, `h = h'`, and append `tok` unless it is a special piece (``, ``, ``, ``, or a `` language tag). Text = join of pieces with `▁` -> space, stripped. Pieces come from `rnnt_vocabulary` (config) or `rnnt_tokenizer/vocab.json`.
LSTM weight layout: `weight_ih_l{n}` `[2560, in]`, `weight_hh_l{n}` `[2560, 640]`, gate order `i, f, g, o`; bias = `bias_ih_l{n} + bias_hh_l{n}`. `src/models/kokoro/lstm.rs` implements the same cell; generalize its loader to take explicit keys or add a small `Lstm` in `src/audio/` (shared by both).
### Network 2: Nemotron-H with a fused input and two heads (`stt_model.embed_tokens`, `stt_model.llm.*`, `stt_model.lm_head`, `stt_model.function_head`)
Per timeline position t the LLM input is one embedding row:
```
fused[t] = E(prev_text[t]) + audio[t] + 2.0 * E(prev_function[t])
```
where `E = stt_model.embed_tokens` (quantized), `audio[t]` is the projected perception frame (or a system-prompt token embedding during the prefix), and `prev_*[0] = pad (12)`. Run the 56 Nemotron-H layers on `fused` (no embedding lookup inside the LLM), apply `norm_f`, then `text_logits = lm_head(h)`, `function_logits = function_head(h)`; both heads are `[4480 -> 131072]` quantized, no bias. Text and function ids are the argmax of the last position (greedy; no sampling).
Key mapping for the existing loader: `stt_model.llm.layers.N.X -> backbone.layers.N.X`, `stt_model.llm.norm_f.weight -> backbone.norm_f.weight`, `stt_model.embed_tokens.* -> backbone.embeddings.*`, `stt_model.lm_head.* -> lm_head.*`; `stt_model.function_head.*` is loaded separately as a `UnifiedLinear`. Mamba `conv1d.weight` is `[channels, 4, 1]` (MLX layout) in the converted checkpoints; keep the `[0, 2, 1]` transpose gate for torch `[channels, 1, 4]` exports.
### Network 3: EAR-TTS (`tts_model.tts_model.*`, `tts_model.audio_prompt_latents.Aria`, `tts_model.codec_silence_tokens`, `tts_model._control_codes`)
Backbone: 28 Gemma-3-style blocks (`backbone.layers.N.{input_layernorm, self_attn.{q,k,v,o}_proj, self_attn.{q,k}_norm, post_attention_layernorm, pre_feedforward_layernorm, mlp.{gate,up,down}_proj, post_feedforward_layernorm}`, `backbone.norm`), hidden 1152, 16 heads of 72 (no GQA), attention scale `256^-0.5`, gelu-tanh gated MLP 4608, offset RMSNorm (`weight + 1`), layer `i` global when `(i + 1) % 6 == 0` (RoPE base 1e6, full `KVCache`) else sliding (base 1e4, rotating cache of 7500). Inputs are injected embeddings with NO `sqrt(hidden)` scaling and there is no token embedding or LM head.
Code embedding: `depthsum(code[.., 31]) = sum_q rvq_embs[q][code_q]` with `rvq_embs [31, 1024, 512]` extended by one zero row per codebook so index 1024 (mask) contributes nothing; `embed_code: Linear(512 -> 1152, no bias)`.
Text conditioning (`embed_subword.*`): for each conditioned subword id build its character sequence through the dense character vocabulary (the tokenizer's single-character tokens sorted by id, 256 entries + 1 padding = 257; the constructor must fail if the count differs), embed with `embed_tokens` Embedding(257, 1152), scale by `sqrt(1152)`, run one T5Gemma encoder layer (pre/post self-attn offset norms, q/k/v/o without bias, RoPE base 1e4, scores `tanh(s / 50) * 50` softcap, mask over padded chars, pre/post feed-forward offset norms, gated gelu-tanh MLP 4608), final offset norm, masked mean-pool over characters, `proj_embedding: Linear(1152 -> 1152, no bias)`; then add `subword_flag_emb.cont_emb[is_continuation[id]]` (Embedding(2, 1152) indexed by the int buffer `is_continuation[131073]`, OOV -> `pad_tensor`) and `bos_eos_emb.special_emb[special_flags[id]]` (Embedding(3, 1152) indexed by `special_flags[131072]`, OOV -> `pad_tensor`).
Classifier-free guidance: every TTS forward runs batch 2: row 0 conditioned on the text embedding, row 1 on `null_emb` broadcast; code embeddings are duplicated.
Gated fusion (`gated_fusion_audio_text`): `final_norm(sigmoid(residual_scale) * (sigmoid(gate) * audio_proj(code_embed / 31) + (1 - sigmoid(gate)) * text_proj(text_embed)))`, `audio_proj` / `text_proj` Linear(1152, 1152) with bias, `gate [1152]`, `residual_scale` scalar, `final_norm` offset RMSNorm.
Warmup (once per session, the Aria prompt): let `frames = 37` (`audio_prompt_latents.Aria` is `[1, 37, 1152]`). Encode `(37 + 1) * 1764` zero samples with the codec to get 38 code frames; set frame 0 and frame 36 to all-mask (1024); keep frames `0..37` (drop the last) as `code [1, 37, 31]`; `subwords = [12] * 37`, `subword_mask = false` except the last two, `audio_mask = false` except the last one. Then `shifted = concat(zeros, code[:, :-1])`, `code_embed = embed_code(depthsum(shifted))`, `bos_mask = audio_mask & !prev(audio_mask)`, `pre_bos = cumsum(bos_mask) == 0`, `code_embed = where(pre_bos, Aria_latent, code_embed) + bos_mask * bos_emb`, duplicate for CFG, `cond = condition(subwords, subword_mask)`, `inputs = fusion(code_embed, cond)`, run the backbone to fill the caches. `previous_code = code[:, -1:]`.
Per step (timeline position t >= 1, current text id `cur`): if `cur == 2` (EOS) set `previous_code = codec_silence_tokens`; `code_embed = embed_code(depthsum(previous_code))` duplicated, `cond = condition([[cur]], mask true)`, `hidden = backbone(fusion(code_embed, cond), caches)` `[2, 1, 1152]`; then `generate_codes(hidden)`:
```
masked_i = ceil((1 - (i/8)^3)^(1/3) * 31) for i in 0..8 -> counts_i = masked_i - masked_{i+1} (last: masked_7)
code = [1024] * 31; done = 0
for count in counts (skip zeros):
emb = embed_code(depthsum(code))
x = concat(emb + hidden[0], emb + hidden[1]) # cond / uncond rows
mu, logs = mog_head.infer(x, guidance 0.2, top_p 0.95)
residual = mu + exp(logs) * N(0, 1) * 0.001
for q in done..done+count: code[q] = argmin_k ||residual - rvq_embs[q][k]||^2 ; residual -= rvq_embs[q][code[q]]
done += count
```
`mog_head.infer`: 3 x (`pre_norm`, gated gelu-tanh MLP 1152 -> 4608 -> 1152, `post_norm`, residual) then the final offset norm (`mlp_stack.3.weight`); CFG `x = cond + 0.2 * (cond - uncond)`; `logits = proj_logits(x)` `[.., 1024]` with top-p 0.95 nucleus masking; pick the component by Gumbel-max on `log softmax(logits)`; `mu = low_mat[c] @ (proj_mus.weight.reshape(1024, 64, 1152)[c] @ x)` `[512]`; `logs = max(proj_logs(x), -4)`; return `mu * exp(logs) + proj_else(x)` and `logs`. Randomness: seed the generator from the session seed (default 0) so runs are reproducible.
Before codec decoding, replace any of the three `_control_codes` values in the generated codes by `codec_silence_tokens` (the per-codebook silence code).
### Network 4: codec (`tts_model.audio_codec.*`)
Encoder (`encoder.layers.{0..15}`): `0` Conv1d(18 -> 384, k1, no bias); stages `(384, 7)`, `(768, 7)`, `(1536, 9)`: each stage has 3 ConvNeXt blocks then a strided Conv1d to the next width (`384 -> 768, k7 s7`; `768 -> 1536, k7 s7`; `1536 -> 512, k9 s9`, no bias). ConvNeXt block (`dwconv` depthwise k7 with bias, causal left pad 6; `norm` channel LayerNorm eps 1e-6; `pwconv1` k1 C -> 4C; GELU; `pwconv2` 4C -> C; residual). Spectrogram input: pad the waveform by 6 samples each side (zeros), periodic Hann(16), `center = false`, hop 4 -> `[B, 9 complex bins, T]`, features = `concat(real, imag)` -> 18 channels, transpose to `[B, T, 18]`.
PRVQ: 31 codebooks `prvq.mus_list.{q}` `[1024, 512]`; encode = iterative nearest-mean residual quantization; decode = sum of selected means. `prvq.variance_list.{q}.variance` scalars are loaded but unused.
Decoder (`decoder.layers.{0..15}`): three stages in reverse, each `ConvTranspose1d(src -> C, k = rate, stride = rate, no bias)` followed by 3 ConvNeXt blocks, then Conv1d(384 -> 18, k1). Output `[B, T, 18]` -> transpose -> `magnitude = 100 * exp(-softplus(-m + ln 100))` over the first 9 channels, `phase` the last 9, `real = mag cos(phase)`, `imag = mag sin(phase)` with imag of bin 0 and bin 8 forced to zero; iSTFT (n_fft 16, hop 4, periodic Hann, `center = false`) then drop 6 samples at each end. One code frame becomes 1764 samples.
Weight layout in the converted checkpoints: Conv1d kernels are `[out, k, in]` (MLX) and ConvTranspose1d kernels are `[out, k, in]` after the `(1, 2, 0)` transpose of torch's `[in, out, k]`; `prvq._variance_list` is renamed to `prvq.variance_list`. Gate the transposes on `shape[-1] == in_channels` so a re-sanitize is a no-op.
### Offline timeline loop
1. Load WAV, resample to 16 kHz mono (`src/audio/whisper_mel.rs:176 resample_to_16k` or the `rubato` path in `src/audio/preprocessing_resample.rs`), append `extra_decoding_seconds * 16000` zeros.
2. `mel = log_mel(x)`; `(audio_embeds [1, F, 4480], lengths, asr_embeds [1, F, 1024]) = perception(mel)`; `F = lengths[0]`.
3. System prompt (when non-empty): `ids = [1] + encode(prompt) + [2]`; `audio_embeds = concat(E(ids), audio_embeds)` (prompt tokens enter on the audio channel); `prompt_frames = len(ids)`.
4. `text[t] = function[t] = 12` for all `t < prompt_frames + F`; TTS warmup; `previous_code = prompt_code[:, -1:]`.
5. For `t in 0..timeline`: `fused = E(text[t-1]) + audio_embeds[t] + 2 E(function[t-1])` (pad at t = 0); `h = llm(fused, caches)`; if `t >= prompt_frames`: `text[t] = argmax(text_logits)`, `function[t] = argmax(function_logits)`; if `t == 0` continue (no TTS at position 0); TTS step with `text[t]` -> `codes[t]`.
6. Drop the prompt prefix from `text`, `function`, `codes`; replace control codes; `audio = codec.decode(codes)` (full-history decode, no cache); transcript = RNNT greedy over `asr_embeds[:, :F]`.
7. Assistant text = decode of `text` ids with `{12, 11, 1, 2}` removed (`skip_special_tokens = false`); function text likewise.
## Implementation plan
1. **Detection and registry.** `src/models/detection.rs`: `"nemotron_voicechat" => Ok(ModelType::NemotronVoiceChat)` placed before the `nemotron_h` arms. `src/models/mod.rs`: enum, supported list, description `("Nemotron VoiceChat (FastConformer + Nemotron-H + EAR-TTS + codec, full duplex)", "Speech")`. `LoadedModel::NemotronVoiceChat(models::NemotronVoiceChatModel)` in `src/loaded_model.rs` with a `LanguageModel` impl that refuses decoder-only calls (pattern: `Florence2VlmModel`, `src/models/florence2/runtime.rs:228`). Detection test.
2. **Config (`src/models/nemotron_voicechat/config.rs`).** `VoiceChatConfig { text: NemotronHConfig, audio: AudioConfig { preprocessor: MelArgs, encoder: ConformerArgs, decoder: PredictArgs, joint: JointArgs, output_dim, max_symbols }, tts: TtsConfig { .., character_encoder: CharEncoderConfig, mog_head: MogConfig }, codec: CodecConfig, bos/eos/pad/silence ids, rnnt_blank_id, input/output sample rates, frame_duration, function_channel_weight, speaker, rnnt_vocabulary: Vec, quantization }` with serde defaults equal to the values above; `att_context_size` accepts `[70, 0]` or `[[70, 0]]`. Per-module quantization: read the `quantization` map and hand `(group_size, bits)` to `UnifiedLinear::from_weights` only for keys present in it; everything else loads dense (`UnifiedLinear` already resolves dense vs quantized from the presence of `.scales`, so the map is a cross-check and the dense loaders must not require it).
3. **Frontend (`src/audio/nemotron_mel.rs`).** `fn log_mel_spectrogram(x: &[f32], args: &MelArgs) -> (Vec, usize /*frames*/)` (preemphasis, centered 400-in-512 Hann, reflect pad 256, power, Slaney mel 128, `ln(x + 2^-24)`). Reuse `periodic_hann` / `hz_to_mel_slaney` / `reflect_pad` from `src/audio/whisper_mel.rs` by making them `pub(crate)`. Tests: `frame_count_is_one_plus_floor_len_over_hop`, `silence_gives_log_guard_floor` (every value equals `ln(2^-24)`), `preemphasis_first_sample_unchanged`.
4. **Speech encoder (`src/audio/fastconformer.rs`).** `FastConformerEncoder::from_weights(weights, "stt_model.perception.encoder", &ConformerArgs)` with `CausalDwStridingSubsampling`, `RelPositionalEncoding` (table built once for `max_len = 5000`, extended on demand), `RelPosMultiHeadAttention` (`rel_shift` as above), `ConformerConvolution`, `ConformerBlock`; `forward(mel [1, T, 128], lengths) -> ([1, T', 1024], lengths')` and `fn chunked_limited_mask(len, left, right)`. Conv2d kernels are `[out, kh, kw, in]` (converted) with the `shape[-1] == in_channels` gate for torch `[out, in, kh, kw]`; depthwise kernels carry `groups = channels` (`mlxcel_core::conv2d` takes a `groups` argument; copy the existing depthwise call in `deepseekocr_sam.rs`). `Perception { encoder, proj: Linear }` returns `(projected, lengths, encoded)`. Tests on random weights: `subsampling_lengths_128_mel_frames_give_16`, `chunked_limited_mask_70_0_allows_self_and_70_left`, `rel_shift_matches_hand_indexing` (a 1x1x3x5 hand-checked case), `block_output_shape`.
5. **RNNT (`src/audio/rnnt.rs`).** `PredictNetwork` (Embedding + 2-layer LSTM; reuse/generalize `src/models/kokoro/lstm.rs`), `JointNetwork`, `fn greedy_decode_frame(enc [1,1,1024], state) -> Vec`, `fn decode_pieces(ids, vocabulary) -> String` with the special-piece filter. Tests: `blank_only_frame_emits_nothing`, `max_symbols_bounds_loop`, `piece_decode_maps_underscore_to_space_and_drops_lang_tags`.
6. **LLM glue (`src/models/nemotron_h.rs`).** Add `pub fn forward_embeds_to_hidden(&self, inputs_embeds, caches) -> UniquePtr` = `forward_stage(.., has_embedding false, has_lm_head false)` followed by `norm_f`; keep `lm_head` inside the model and expose `pub fn lm_head(&self, h)`; VoiceChat holds `function_head: UnifiedLinear` beside it. Loader renames keys as listed in the Network 2 section before `NemotronHModel::from_weights(config, weights, block_types)`. Test: `forward_embeds_to_hidden_matches_forward_on_embedded_ids` on a tiny random config.
7. **TTS (`src/models/nemotron_voicechat/tts.rs`).** `OffsetRmsNorm`, `GatedMlp`, `MogHead`, `CharAwareSubwordEncoder` (with `set_vocabulary(&HashMap)` building the 257-entry char table and the subword -> chars map), `GatedFusion`, `RvqEarTtsModel { backbone: Gemma3Backbone, bos_emb, null_emb, embed_code, embed_subword, fusion, audio_prompt_projection_W, mog_head, rvq_embs }`, `SpeechDecoder { codec, tts, control_codes: [i32; 3], aria_latent [1, 37, 1152], codec_silence_tokens [31] }`. Add `Gemma3Backbone` to `src/models/gemma3.rs`: `from_weights(weights, prefix, &ModelArgs)` building `TransformerBlock`s from `{prefix}.layers.{i}` (parameterize the prefix at line 552) and `{prefix}.norm`, plus `forward_embeds(inputs, caches, scale_inputs: bool)`; `make_caches()` returning `KVCache` for `(i+1) % 6 == 0` and `RotatingKVCache::new(7500, keep 0)` otherwise (`src/lib/mlxcel-core/src/cache.rs:4094`). Tests: `offset_rms_norm_equals_rms_norm_with_weight_plus_one`, `top_p_keeps_at_least_the_largest_logit`, `mask_schedule_counts_sum_to_31_for_exponent_3`, `depthsum_ignores_mask_index`, `char_encoder_scatters_only_masked_positions`, `mog_infer_shapes_and_finite` (random weights, batch 2).
8. **Codec (`src/audio/nemotron_codec.rs`).** `ConvNeXtBlock1d`, `AudioEncoder`, `AudioDecoder`, `Prvq`, `NemotronCodec { encode(wave [B, N]) -> codes [B, 31, T], decode(codes) -> wave [B, 1, N] }`, `CausalConv1dCache` (per-layer overlap buffers, used by the streaming sub-issue), STFT/iSTFT at n_fft 16 hop 4 (generalize `src/models/kokoro/stft.rs` to take `n_fft`, `hop`, `window`, `center`, or add a small `src/audio/stft.rs`). `mlxcel_core::conv1d` (bridge line `src/lib/mlxcel-core/src/lib.rs:2246`) with `groups = channels` for the depthwise conv and `mlxcel_core::conv_transpose1d` (`:2402`) for the decoder stage starts. Tests: `waveform_to_token_ratio_is_1764`, `stft_channels_is_18`, `encode_decode_round_trip_on_silence_is_near_zero` (random-init codebooks make this a shape test; the numeric check uses the real checkpoint), `decoder_transpose_gate_is_idempotent`.
9. **Model and offline runtime (`src/models/nemotron_voicechat/{mod.rs, model.rs, session.rs}`).** `NemotronVoiceChatModel { config, embed_tokens: UnifiedEmbedding, llm: NemotronHModel, function_head, perception, rnnt: (PredictNetwork, JointNetwork), speech_decoder: SpeechDecoder, tokenizer: MlxcelTokenizer }` with `load(model_path)`; `VoiceChatResult { text, function_text, user_transcript, audio: Vec, sample_rate: 22050, text_tokens, function_tokens, audio_codes }`; `fn generate_offline(&self, wav: &[f32], system_prompt: Option<&str>, extra_decoding_seconds: f32, seed: u64) -> Result` implementing the timeline loop. Model-owned caches: Nemotron-H `NemotronLayerCache` vector from `make_caches`, TTS caches from `Gemma3Backbone::make_caches`.
10. **CLI.** In `src/commands/generate.rs`, branch on `LoadedModel::NemotronVoiceChat` before the text path (pattern: the Florence-2 branch at line 2431): require `--audio`, treat `-p` as the system prompt, add `--extra-decoding-seconds` (f32, default 3.0) and `--seed` reuse, write `--output-audio` with `src/audio/wav_writer.rs::encode_wav_pcm16(samples, 22050, 1)`, print `[user] ...`, the assistant text, and `[function] ...` when non-empty. Reject `--image` / `--video`.
11. **Docs.** `docs/supported-models.md` speech section; `docs/nemotron-voicechat.md` with the architecture summary above and the CLI usage.
## Validation
(a) Unit tests listed per step, in `src/audio/nemotron_mel_tests.rs`, `src/audio/fastconformer_tests.rs`, `src/audio/rnnt_tests.rs`, `src/audio/nemotron_codec_tests.rs`, `src/models/nemotron_voicechat/tts_tests.rs`, `src/models/nemotron_h_tests.rs`, `src/models/detection_tests.rs`.
(b) Real checkpoint: `mlx-community/NemotronLabs-VoiceChat-11B-4bit` (9.2 GB; `rnnt_tokenizer/`, `tokenizer.json`, two safetensors shards) and `mlx-community/NemotronLabs-VoiceChat-11B-8bit` for the quality check.
```
# Make a 16 kHz mono question WAV: macOS `say -o /tmp/q.aiff "What is the capital of France?"` then
# `ffmpeg -i /tmp/q.aiff -ar 16000 -ac 1 /tmp/question.wav`, or Kokoro through the server's
# POST /v1/audio/speech (24 kHz output; resample to 16 kHz the same way).
./target/release/mlxcel generate -m models/NemotronLabs-VoiceChat-11B-4bit --audio /tmp/question.wav \
--output-audio /tmp/response.wav -p "Be concise and answer in one sentence." -n 0
```
Acceptance: the printed `[user]` transcript contains "capital of France"; the assistant text mentions Paris; `/tmp/response.wav` is 22050 Hz mono, non-silent (RMS > 0.01 over the answer region), and its length equals `(timeline_frames) * 1764` samples; every intermediate (mel, encoder, logits, codes) is finite; with `--seed 0` two runs are byte-identical. Parity: the codec alone round-trips a 1 s 440 Hz tone with SNR > 20 dB through `encode` then `decode` on the real weights; the speech encoder output for a 1 s WAV matches the model publisher's own inference code run once on CPU within RMS 5e-3 at bf16; the first 20 assistant text ids match exactly for the same WAV and system prompt.
## Acceptance criteria
- [ ] `nemotron_voicechat` resolves to `ModelType::NemotronVoiceChat` and the 4-bit checkpoint loads with per-module quantization (only the listed modules quantized)
- [ ] Log-mel + FastConformer + RNNT transcribe a spoken sentence correctly
- [ ] Nemotron-H runs on fused embeddings with `norm_f` and both heads; text/function ids are greedy argmax
- [ ] EAR-TTS warmup with the Aria prompt and per-step code generation produce 31 codes per frame; control codes are replaced
- [ ] Codec round-trip SNR > 20 dB on a tone; decoded answer audio is non-silent and 1764 samples per frame
- [ ] `mlxcel generate --audio --output-audio` prints transcript + answer and writes a playable WAV; seeded runs are deterministic
- [ ] docs/supported-models.md updated
- [ ] detection table in src/models/detection.rs updated with a test
- [ ] cargo test --workspace --profile test-fast --features metal,accelerate passes
- [ ] cargo clippy --workspace --all-targets -- -D warnings and cargo fmt --all -- --check pass
## Out of scope
- Cache-aware online inference, event streaming, WebSocket (next sub-issues).
- Other speakers; the `speaker` field is validated to equal `"Aria"`.
## Dependencies
- Part of epic #1372
Contributor guide
Research direction
Start by reading src/models/detection.rs and the Nemotron-H loading and forward paths in src/models/nemotron_h.rs, then compare the existing frontend and reusable audio implementations in src/audio/whisper_mel.rs and src/models/kokoro/. Check the tokenizer loader in src/tokenizer/mod.rs as well. Done means the documented generate command accepts a 16 kHz WAV, prints the transcript and assistant output, and writes a 22.05 kHz mono PCM16 response WAV.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- ai, audio-video-rtc, backend
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 25/100