diff --git a/examples/offline_inference/nemotron_voicechat/end2end.py b/examples/offline_inference/nemotron_voicechat/end2end.py
new file mode 100644
index 000000000..b020e5bad
--- /dev/null
+++ b/examples/offline_inference/nemotron_voicechat/end2end.py
@@ -0,0 +1,229 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline speech-to-speech with NVIDIA-NemotronLabs-VoiceChat-11B.
+
+Feeds a 16 kHz user utterance (plus a spoken-style system prompt) through the
+3-stage nemotron_voicechat pipeline and writes the agent's reply as text and a
+22.05 kHz WAV.
+
+Usage:
+ python end2end.py \
+ --checkpoint /path/to/NVIDIA-NemotronLabs-VoiceChat-11B \
+ --wav /path/to/user_question_16k.wav \
+ --output-dir results/nemotron_voicechat
+
+The system prompt defaults to the NeMo reference default. The text tokenizer
+resolves from the ``nvidia/NVIDIA-Nemotron-Nano-9B-v2`` HF id (or the
+``NEMOTRON_VOICECHAT_LLM_PATH`` env var for air-gapped runs).
+
+Sizing the input WAV: the timeline is frame-locked, so the reply budget IS the
+input duration — the acoustic channel trails the text channel, and a reply
+that has not finished when the frames run out is truncated SILENTLY (no error;
+ASR of the WAV will just disagree with the emitted text). Leave generous
+trailing silence after the question (a question ending at ~4.5 s truncated in
+an 8 s WAV but completed cleanly in a 16 s one). The script warns when the
+text channel is still emitting non-PAD tokens near the last frame.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from pathlib import Path
+
+import numpy as np
+import soundfile as sf
+
+# Warn if the text channel is still speaking within this many final frames
+# (12 frames = ~1 s) — the frame-locked reply likely got cut off.
+TRUNCATION_GUARD_FRAMES = 12
+
+# The NeMo reference offline script's default system prompt.
+DEFAULT_SYSTEM_PROMPT = (
+ "You are an AI voice assistant developed by NVIDIA. "
+ "Your name is NVIDIA Voice Chat. "
+ "Answer in a spoken, conversational style rather than a written one. "
+ "Do not repeat the same sentence over and over again. "
+ "Start the conversation by greeting the user."
+)
+
+
+def load_wav_16k_mono(path: str) -> np.ndarray:
+ wav, sr = sf.read(path, dtype="float32")
+ if wav.ndim > 1:
+ wav = wav.mean(axis=1)
+ if sr != 16000:
+ from vllm.multimodal.audio import resample_audio_scipy
+
+ wav = resample_audio_scipy(wav, orig_sr=sr, target_sr=16000)
+ return wav.astype(np.float32)
+
+
+def main() -> None:
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument("--checkpoint", required=True, help="NVIDIA-NemotronLabs-VoiceChat-11B directory")
+ parser.add_argument("--wav", required=True, help="user speech (any sr; resampled to 16 kHz mono)")
+ parser.add_argument("--output-dir", default="results/nemotron_voicechat")
+ parser.add_argument("--system-prompt", default=DEFAULT_SYSTEM_PROMPT)
+ parser.add_argument(
+ "--stage-configs-path",
+ default=None,
+ help=(
+ "Deploy yaml override. Pass vllm_omni/deploy/nemotron_labs_voicechat_streaming.yaml "
+ "for the async-chunk streaming pipeline (stages overlap; audio starts decoding "
+ "before the thinker finishes). Default: the bundled non-streaming yaml (bit-exact "
+ "parity path)."
+ ),
+ )
+ args = parser.parse_args()
+
+ if not Path(args.wav).is_file():
+ raise FileNotFoundError(f"Input WAV not found: {args.wav}")
+ ckpt = Path(args.checkpoint)
+ raw_cfg = json.loads((ckpt / "config.json").read_text())
+ stt_cfg = raw_cfg["model"]["stt"]["model"]
+
+ from transformers import AutoTokenizer
+ from vllm.sampling_params import SamplingParams
+
+ from vllm_omni.entrypoints.omni import Omni
+ from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_thinker import (
+ compute_acoustic_frame_count,
+ )
+
+ tok_ref = os.environ.get("NEMOTRON_VOICECHAT_LLM_PATH") or stt_cfg.get(
+ "pretrained_llm", "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
+ )
+ tokenizer = AutoTokenizer.from_pretrained(tok_ref, trust_remote_code=False)
+ bos_id = tokenizer.convert_tokens_to_ids(stt_cfg.get("bos_token", ""))
+ eos_id = tokenizer.convert_tokens_to_ids(stt_cfg.get("eos_token", ""))
+ pad_id = tokenizer.convert_tokens_to_ids(stt_cfg.get("pad_token", ""))
+ # NeMo encode_system_prompt: [bos] + tokens + [eos].
+ prompt_ids = [bos_id] + tokenizer.encode(args.system_prompt, add_special_tokens=False) + [eos_id]
+
+ wav = load_wav_16k_mono(args.wav)
+ n_frames = compute_acoustic_frame_count(stt_cfg, int(wav.shape[0]))
+ print(f"input: {wav.shape[0]} samples @16kHz -> {n_frames} acoustic frames; prompt {len(prompt_ids)} tokens")
+
+ inputs = {
+ # vLLM prompt = system prompt + one placeholder position (acoustic
+ # frame 0); the thinker replaces every position with fused embeddings.
+ "prompt_token_ids": prompt_ids + [pad_id],
+ "additional_information": {
+ # Plain float list: raw ndarrays do not survive the engine-core
+ # message serialization.
+ "nvc_audio": wav.astype(np.float32).tolist(),
+ "nvc_sr": 16000,
+ "nvc_prompt_token_ids": prompt_ids,
+ "nvc_expected_frames": n_frames,
+ },
+ }
+ sampling_params_list = [
+ # thinker: greedy, frame-locked generation length. NO min_tokens: the
+ # tokenizer's eos_token is — the frame-locked PAD/silence
+ # token — and vLLM's min_tokens masks "EOS" to -inf, which would forbid
+ # silence and force the agent to babble through the whole utterance.
+ SamplingParams(temperature=0.0, max_tokens=n_frames, ignore_eos=True, seed=0),
+ # talker: placeholder loop, stops on the stage's stop token. 16383:
+ # the 1-token placeholder prompt + max_tokens must fit the stage's
+ # max_model_len (16384).
+ SamplingParams(temperature=0.0, max_tokens=16383, stop_token_ids=[1], detokenize=False, seed=0),
+ # code2wav: single decode step.
+ SamplingParams(temperature=0.0, max_tokens=1, detokenize=False, seed=0),
+ ]
+
+ omni_kwargs: dict = {}
+ if args.stage_configs_path:
+ omni_kwargs["stage_configs_path"] = args.stage_configs_path
+ omni = Omni(model=str(ckpt), trust_remote_code=True, **omni_kwargs)
+ outputs = omni.generate(inputs, sampling_params_list)
+
+ out_dir = Path(args.output_dir)
+ out_dir.mkdir(parents=True, exist_ok=True)
+ stem = Path(args.wav).stem
+
+ agent_text = None
+ audio = None
+ sr = 22050
+ for stage_outputs in outputs:
+ # One entry per final stage (text from the thinker, audio from code2wav);
+ # mirrors the qwen2_5_omni offline example's collection loop.
+ output = getattr(stage_outputs, "request_output", stage_outputs)
+ final_type = getattr(stage_outputs, "final_output_type", None)
+ completions = getattr(output, "outputs", None) or []
+ if not completions:
+ continue
+ completion = completions[0]
+ if final_type in (None, "text"):
+ text = getattr(completion, "text", None)
+ if text:
+ agent_text = text
+ token_ids = getattr(completion, "token_ids", None)
+ if token_ids:
+ (out_dir / f"{stem}_text_tokens.json").write_text(json.dumps(list(map(int, token_ids))))
+ # Frame lock: the reply budget IS the input length. If the text
+ # channel is still mid-sentence when the frames run out, the
+ # spoken answer is silently truncated — warn instead.
+ tail = list(map(int, token_ids))[-TRUNCATION_GUARD_FRAMES:]
+ if any(t != pad_id for t in tail):
+ print(
+ f"WARNING: the text channel emits non-PAD tokens within the last "
+ f"{TRUNCATION_GUARD_FRAMES} frames ({TRUNCATION_GUARD_FRAMES * 80} ms) — the reply "
+ "likely ran out of acoustic frames and is cut off mid-sentence. "
+ "Pad the input WAV with more trailing silence and rerun."
+ )
+ # Best-effort: dump the function-channel debug timeline if the
+ # engine surfaced the request info on this output.
+ info = getattr(output, "additional_information", None)
+ func_timeline = info.get("nvc_function_tokens") if isinstance(info, dict) else None
+ if func_timeline is not None:
+ values = func_timeline.tolist() if hasattr(func_timeline, "tolist") else list(func_timeline)
+ (out_dir / f"{stem}_function_tokens.json").write_text(json.dumps([int(v) for v in values]))
+ print(f"function-channel timeline: {len(values)} entries")
+ if final_type in (None, "audio"):
+ # multimodal_output is a MultimodalPayload (Mapping) — duck-type it.
+ mm = getattr(completion, "multimodal_output", None)
+ get = getattr(mm, "get", None)
+ if not callable(get):
+ continue
+ candidate = None
+ for key in ("audio", "model_outputs", "wav", "waveform"):
+ value = get(key)
+ if value is not None:
+ candidate = value
+ break
+ if isinstance(candidate, list):
+ candidate = candidate[0] if candidate else None
+ if candidate is not None:
+ arr = np.asarray(
+ candidate.float().cpu() if hasattr(candidate, "cpu") else candidate, dtype=np.float32
+ ).reshape(-1)
+ if arr.size > 0:
+ audio = arr
+ sr_val = get("sr")
+ if sr_val is not None:
+ sr_item = sr_val[0] if isinstance(sr_val, list) else sr_val
+ sr = int(sr_item.item() if hasattr(sr_item, "item") else sr_item)
+
+ if agent_text is not None:
+ # The frame-locked text channel emits PAD tokens on silent frames;
+ # strip them (NeMo's post-inference detokenizer does the same).
+ pad_token = stt_cfg.get("pad_token", "")
+ agent_text = " ".join(agent_text.replace(pad_token, " ").split())
+ text_path = out_dir / f"{stem}_output.txt"
+ text_path.write_text(agent_text + "\n")
+ print(f"agent text -> {text_path}\n{agent_text}")
+ else:
+ print("WARNING: no agent text in outputs")
+
+ if audio is not None and audio.size > 0:
+ wav_path = out_dir / f"{stem}_output.wav"
+ sf.write(wav_path, audio, sr)
+ print(f"agent audio -> {wav_path} ({audio.size / sr:.2f}s @ {sr} Hz)")
+ else:
+ raise RuntimeError("No audio produced by the code2wav stage.")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/mkdocs.yml b/mkdocs.yml
index fd97b9b37..242eb33ae 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -111,6 +111,7 @@ plugins:
- "vllm_omni.entrypoints.openai" # avoid importing vllm in mkdocs building
- "vllm_omni.entrypoints.openai.protocol" # avoid importing vllm in mkdocs building
- "vllm_omni.experimental" # avoid importing optional serving deps (openai, fastapi, uvicorn, httpx) in mkdocs building
+ - "vllm_omni.model_executor.models.nemotron_voicechat.nemo_vendored" # vendored NeMo code; upstream docstrings trip griffe strict warnings
- mkdocstrings:
handlers:
python:
diff --git a/pyproject.toml b/pyproject.toml
index f9fbdfb2a..0e91bfae6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -206,6 +206,12 @@ ignore = [
# These files keep upstream conventions (star imports, long lines, math-symbol
# argument names) so they stay easy to diff against upstream releases.
"vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py" = ["E402", "E501", "F403", "F405", "N803"]
+# NemotronVoiceChat inference modules vendored from NVIDIA-NeMo/Speech (branch
+# nemotron-labs-voicechat). Numerics-bearing code is kept byte-close to
+# upstream to stay diffable and parity-safe, so upstream style (long lines,
+# legacy typing, whitespace, unused-import shims) is exempted. TID251 and the
+# F-correctness rules still apply.
+"vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/**" = ["E501", "E721", "E741", "F401", "F841", "I001", "N803", "N814", "UP006", "UP007", "UP008", "UP011", "UP032", "UP035", "W291", "W293"]
"vllm_omni/experimental/fullduplex/joyvl/decision/prompts.py" = ["E501"] # System/summarizer prompts are long prose lines
[tool.ruff.lint.flake8-tidy-imports.banned-api]
@@ -314,6 +320,10 @@ extend-ignore-identifiers-re = [
".*[Nn][Dd]3.*",
".*updator.*",
"byt5.*",
+ # NemotronVoiceChat vendored triton kernel: dout/DOUT abbreviate d_out
+ # (output feature dim), not "doubt".
+ ".*dout.*",
+ ".*DOUT.*",
]
extend-ignore-re = [
"\\.\\.\\.n,d->\\.\\.\\.nd",
@@ -334,6 +344,9 @@ byt5 = "byt5"
# "seeked" is the standard HTML media element DOM event name (fires when a
# seek completes), used by the word-timestamps Gradio demo. Not a typo.
seeked = "seeked"
+# Ryan Leary: author credited in the NemotronVoiceChat vendored NeMo
+# features.py copyright header. A surname, not a typo of "leery".
+Leary = "Leary"
# IndexTTS 2.5 language codes and speech-emotion-recognition event token.
fo = "fo"
ba = "ba"
diff --git a/recipes/NVIDIA/NemotronLabs-VoiceChat.md b/recipes/NVIDIA/NemotronLabs-VoiceChat.md
new file mode 100644
index 000000000..8faa268d2
--- /dev/null
+++ b/recipes/NVIDIA/NemotronLabs-VoiceChat.md
@@ -0,0 +1,113 @@
+# NemotronLabs VoiceChat 11B
+
+> Offline speech-to-speech: a 16 kHz user utterance in, the agent's spoken reply
+> (text + 22.05 kHz WAV) out, on a 3-stage vLLM-Omni pipeline.
+
+## Summary
+
+- Vendor: NVIDIA
+- Model: [`nvidia/NVIDIA-NemotronLabs-VoiceChat-11B`](https://huggingface.co/nvidia/NVIDIA-NemotronLabs-VoiceChat-11B)
+- Task: Speech-to-speech voice chat. The model runs a frame-locked 12.5 Hz
+ timeline: a Conformer + NemotronH hybrid-Mamba thinker emits one text token
+ per acoustic frame, a Gemma3-1B EAR-TTS talker turns the text timeline into
+ 31-quantizer RVQ code stacks, and an RVQ-VAE codec decodes them to audio.
+- Mode: Offline single-turn inference (batch=1). Online/duplex serving,
+ function calling, and batch>1 are documented follow-ups.
+- Maintainer: [`@yuekaizhang`](https://github.com/yuekaizhang)
+
+## When to use this recipe
+
+Use this recipe to run a single-turn voice-chat exchange with
+`NVIDIA-NemotronLabs-VoiceChat-11B` on one GPU: you provide a user utterance as
+a WAV file (any sample rate; it is resampled to 16 kHz mono) plus an optional
+spoken-style system prompt, and get back the agent's reply as text and a
+22.05 kHz WAV. The integration is NeMo-free at runtime — the perception
+Conformer, EAR-TTS talker, and RVQ-VAE codec are vendored, dependency-stripped
+NeMo modules (`nemo_vendored/`), so no `nemo_toolkit` install is needed.
+
+## References
+
+- Offline example:
+ [`examples/offline_inference/nemotron_voicechat/end2end.py`](../../examples/offline_inference/nemotron_voicechat/end2end.py)
+- Model modules (thinker / talker / code2wav / vendored NeMo):
+ [`vllm_omni/model_executor/models/nemotron_voicechat/`](../../vllm_omni/model_executor/models/nemotron_voicechat/)
+- Staged pipeline config:
+ [`vllm_omni/deploy/nemotron_labs_voicechat.yaml`](../../vllm_omni/deploy/nemotron_labs_voicechat.yaml)
+- Upstream model card:
+ [`nvidia/NVIDIA-NemotronLabs-VoiceChat-11B`](https://huggingface.co/nvidia/NVIDIA-NemotronLabs-VoiceChat-11B)
+- Reference implementation: NVIDIA-NeMo/Speech, branch `nemotron-labs-voicechat`
+
+## Pipeline
+
+| stage | arch | dtype | role |
+|---|---|---|---|
+| 0 thinker | `NemotronVoiceChatThinkerForConditionalGeneration` (LLM_AR) | fp32 | WAV + system prompt -> frame-locked text-token timeline (+ function channel) |
+| 1 talker | `NemotronVoiceChatTalker` (LLM_AR) | fp32 | text timeline -> 31-quantizer RVQ code stacks (one per 80 ms frame) |
+| 2 code2wav | `NemotronVoiceChatCode2Wav` (LLM_GENERATION) | fp32 | RVQ-VAE decode -> 22.05 kHz PCM |
+
+Stages 0/1 default to fp32 for exact parity with the NeMo reference
+implementation (greedy decoding matches it token for token on the acceptance
+fixture). The deploy yaml documents a ~2x-faster bf16 thinker option whose
+output stayed within one word of the reference in testing.
+
+## Hardware Support
+
+## GPU
+
+### 1x H100 80GB
+
+#### Environment
+
+- OS: Linux
+- Python: 3.12
+- vLLM version: 0.26.0
+- vLLM-Omni version or commit: this PR / current `main`
+
+#### Command
+
+```bash
+# Tokenizer: the checkpoint ships no HF tokenizer; it resolves from the
+# nvidia/NVIDIA-Nemotron-Nano-9B-v2 HF id automatically. For air-gapped runs,
+# point NEMOTRON_VOICECHAT_LLM_PATH at a local snapshot of that repo instead.
+python examples/offline_inference/nemotron_voicechat/end2end.py \
+ --checkpoint /path/to/NVIDIA-NemotronLabs-VoiceChat-11B \
+ --wav /path/to/user_question.wav \
+ --output-dir results/nemotron_voicechat
+```
+
+#### Verification
+
+```bash
+ls results/nemotron_voicechat
+# _output.txt the agent reply as text
+# _output.wav the agent reply as 22.05 kHz audio
+# _text_tokens.json the frame-locked text-token timeline
+```
+
+The reply text should read as a coherent spoken-style answer to the question in
+the input WAV, and the WAV should transcribe to (approximately) the same text
+with any ASR model.
+
+#### Notes
+
+- Memory usage: the shipped yaml runs all three stages on one GPU
+ (`gpu_memory_utilization` 0.62 / 0.12 / 0.06); peak usage is dominated by the
+ fp32 thinker. The fp32 default has a hard floor of roughly 43 GB of thinker
+ weights alone (9B backbone + 587M `embed_tokens` + 587M `function_head` +
+ 0.6B Conformer), so 48 GB cards cannot run it — use the bf16 thinker option
+ documented in the deploy yaml on anything smaller than an 80 GB part.
+- Input sizing: the timeline is frame-locked, so the reply budget IS the input
+ duration. The acoustic channel trails the text channel; if the WAV does not
+ carry enough trailing silence for the reply to finish, the spoken answer is
+ truncated silently. Leave generous trailing silence (a question ending at
+ ~4.5 s truncated in an 8 s WAV but completed cleanly in 16 s); the offline
+ example warns when the text channel is still speaking near the last frame.
+- Key flags: sampling is greedy end to end. The thinker is frame-locked —
+ `max_tokens` equals the acoustic frame count with `ignore_eos=True`. Do NOT
+ set `min_tokens` on the thinker: the tokenizer's EOS token is also the
+ frame-locked PAD/silence token, so masking it forces the model to babble
+ instead of pausing.
+- The talker's `max_tokens` is 16383 (its stage prompt is one placeholder
+ token, and the stage context is 16384).
+- Known limitations: batch=1 only; offline single turn (no online/duplex
+ serving yet); the function-call channel is decoded but not yet acted on.
diff --git a/recipes/README.md b/recipes/README.md
index ca287bf70..e010ade67 100644
--- a/recipes/README.md
+++ b/recipes/README.md
@@ -66,6 +66,7 @@ recipes/
| [`OpenBMB/VoxCPM2.md`](./OpenBMB/VoxCPM2.md) | Online + offline TTS with native AR pipeline (48 kHz, 30+ languages) | 1x RTX 4090 24GB |
| [`OpenMOSS/MOSS-TTS.md`](./OpenMOSS/MOSS-TTS.md) | Online + offline multilingual TTS (MOSS-TTS family, 8B) | 1x H100 80GB |
| [`NVIDIA/Nemotron-Labs-Audex.md`](./NVIDIA/Nemotron-Labs-Audex.md) | TTS / text-to-audio / audio understanding / cascaded S2S (Audex 2B + 30B-A3B) | 1x H100 80GB |
+| [`NVIDIA/NemotronLabs-VoiceChat.md`](./NVIDIA/NemotronLabs-VoiceChat.md) | Offline speech-to-speech voice chat (11B, frame-locked 12.5 Hz, 3-stage thinker/talker/code2wav) | 1x H100 80GB |
| [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |
| [`Qwen/Qwen-Image-2512.md`](./Qwen/Qwen-Image-2512.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt FP8 / mixed FP8/NVFP4 | 1x A800 80GB / 2x B200 |
| [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |
diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py
index 5676a17a6..1fb32239d 100644
--- a/tests/config/test_omni_config.py
+++ b/tests/config/test_omni_config.py
@@ -513,6 +513,7 @@ def test_sub_config_fields_match_structured_scopes():
"enable_prefix_caching",
"disable_hybrid_kv_cache_manager",
"mm_processor_cache_gb",
+ "mamba_ssm_cache_dtype",
}
assert {f.name for f in fields(OmniStageSchedulerConfig)} == {
"max_num_seqs",
diff --git a/tests/core/sched/test_omni_scheduling_coordinator.py b/tests/core/sched/test_omni_scheduling_coordinator.py
index a93929306..444d21c34 100644
--- a/tests/core/sched/test_omni_scheduling_coordinator.py
+++ b/tests/core/sched/test_omni_scheduling_coordinator.py
@@ -120,6 +120,9 @@ class TestFullPayloadCoordinatorSelection(unittest.TestCase):
("MiMoAudioModel", "code2wav"),
("Qwen3TTSCode2Wav", "code2wav"),
("MiniCPMO45Code2Wav", "code2wav"),
+ # nemotron_voicechat: sync-path RVQ-VAE decode waits for the
+ # talker's full [frames, 31] code-stack payload.
+ ("NemotronVoiceChatCode2Wav", "code2wav"),
("CosyVoice3Model", "cosyvoice3_code2wav"),
# audex: sync-path TTS decoder (async_chunk: false deploys) and
# the TTA XCodec1 stage (always sync full-payload).
diff --git a/tests/model_executor/models/test_nemotron_voicechat_registration.py b/tests/model_executor/models/test_nemotron_voicechat_registration.py
new file mode 100644
index 000000000..09d85b271
--- /dev/null
+++ b/tests/model_executor/models/test_nemotron_voicechat_registration.py
@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""L1 registration + config-shim tests for NemotronVoiceChat (CPU, no weights)."""
+
+import subprocess
+from pathlib import Path
+
+import pytest
+
+from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
+from vllm_omni.config.stage_config import StageExecutionType
+from vllm_omni.model_executor.models.nemotron_voicechat.configuration_nemotron_voicechat import (
+ NemotronVoiceChatConfig,
+)
+from vllm_omni.model_executor.models.registry import _OMNI_MODELS
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_ARCHS = (
+ "NemotronVoiceChatThinkerForConditionalGeneration",
+ "NemotronVoiceChatTalkerForConditionalGeneration",
+ "NemotronVoiceChatCode2Wav",
+)
+
+
+def _minimal_nemo_dict() -> dict:
+ return {
+ "data": {"frame_length": 0.08, "source_sample_rate": 16000, "target_sample_rate": 22050},
+ "model": {
+ "inference_speaker_name": "Aria",
+ "stt": {
+ "model": {
+ "perception": {"encoder": {}, "preprocessor": {}, "modality_adapter": {}},
+ "pretrained_llm": "nvidia/NVIDIA-Nemotron-Nano-9B-v2",
+ "use_function_head": True,
+ }
+ },
+ "speech_generation": {
+ "model": {
+ "tts_config": {"backbone_config": {"hidden_size": 1152, "sliding_window": 7500}},
+ "codec_config": {"num_quantizers": 31, "codebook_size": 1024},
+ }
+ },
+ },
+ }
+
+
+def test_pipeline_registered_with_three_stages() -> None:
+ pipeline = OMNI_PIPELINES["nemotron_voicechat"]
+ assert OMNI_PIPELINES["nemotron_labs_voicechat"] is pipeline
+ assert len(pipeline.stages) == 3
+ thinker, talker, code2wav = pipeline.stages
+ assert thinker.execution_type == StageExecutionType.LLM_AR
+ assert thinker.final_output and thinker.final_output_type == "text"
+ assert thinker.owns_tokenizer
+ assert talker.execution_type == StageExecutionType.LLM_AR
+ assert talker.model_arch == "NemotronVoiceChatTalkerForConditionalGeneration"
+ assert talker.input_sources == (0,)
+ assert code2wav.execution_type == StageExecutionType.LLM_GENERATION
+ assert code2wav.final_output and code2wav.final_output_type == "audio"
+ assert code2wav.input_sources == (1,)
+ assert pipeline.validate() == []
+
+
+def test_architectures_registered() -> None:
+ for arch in _ARCHS:
+ assert arch in _OMNI_MODELS, f"{arch} missing from _OMNI_MODELS"
+
+
+def test_full_payload_whitelists_cover_code2wav_only() -> None:
+ from vllm_omni.core.sched.omni_scheduling_coordinator import _FULL_PAYLOAD_INPUT_STAGES
+
+ assert ("NemotronVoiceChatCode2Wav", "code2wav") in _FULL_PAYLOAD_INPUT_STAGES
+ # The thinker->talker hop is token-path only: the talker must NOT wait on a
+ # connector payload or stage 1 hangs.
+ assert ("NemotronVoiceChatTalkerForConditionalGeneration", "talker") not in _FULL_PAYLOAD_INPUT_STAGES
+
+
+def test_config_parses_nemo_layout_without_hub_access(monkeypatch) -> None:
+ # Provide the thinker text config explicitly so the test never touches HF.
+ cfg = NemotronVoiceChatConfig(
+ thinker_text_config={"model_type": "llama", "hidden_size": 4480, "num_hidden_layers": 56},
+ **_minimal_nemo_dict(),
+ )
+ assert cfg.get_text_config().hidden_size == 4480
+ assert cfg.talker_config.hidden_size == 1152
+ assert cfg.talker_config.sliding_window == 7500
+ assert cfg.code2wav_config.num_quantizers == 31
+ assert cfg.inference_speaker_name == "Aria"
+
+
+@pytest.mark.parametrize(
+ "mutation",
+ ["drop_stt", "drop_speech_generation", "drop_perception", "drop_codec"],
+)
+def test_config_rejects_non_voicechat_checkpoints(mutation: str) -> None:
+ raw = _minimal_nemo_dict()
+ if mutation == "drop_stt":
+ del raw["model"]["stt"]
+ elif mutation == "drop_speech_generation":
+ del raw["model"]["speech_generation"]
+ elif mutation == "drop_perception":
+ del raw["model"]["stt"]["model"]["perception"]
+ else:
+ del raw["model"]["speech_generation"]["model"]["codec_config"]
+ with pytest.raises(ValueError, match="NemotronVoiceChat|missing"):
+ NemotronVoiceChatConfig(
+ thinker_text_config={"model_type": "llama", "hidden_size": 4480},
+ **raw,
+ )
+
+
+def test_no_nemo_imports_in_nemotron_voicechat_tree() -> None:
+ """Dependency guard: the vendored tree must not import the nemo package.
+
+ Scoped to the nemotron_voicechat model tree (other models may have their
+ own optional nemo integrations).
+ """
+ repo_root = Path(__file__).resolve().parents[3]
+ targets = [
+ str(repo_root / "vllm_omni/model_executor/models/nemotron_voicechat"),
+ str(repo_root / "vllm_omni/model_executor/stage_input_processors/nemotron_voicechat.py"),
+ ]
+ result = subprocess.run(
+ ["grep", "-rnE", r"^\s*(import nemo\b|from nemo(\.|\s))", *targets, "--include=*.py"],
+ capture_output=True,
+ text=True,
+ )
+ assert result.returncode != 0, f"Found nemo imports in the nemotron_voicechat tree:\n{result.stdout}"
diff --git a/tests/model_executor/models/test_nemotron_voicechat_shapes.py b/tests/model_executor/models/test_nemotron_voicechat_shapes.py
new file mode 100644
index 000000000..8bd279bb8
--- /dev/null
+++ b/tests/model_executor/models/test_nemotron_voicechat_shapes.py
@@ -0,0 +1,227 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Golden shape test for the NemotronVoiceChat frame-locked timeline contract.
+
+Pins the exact prefill/off-by-one arithmetic against values verified with the
+NeMo reference run on sample_general.wav (fp32, greedy): perception produced
+196 acoustic frames for the 249,734-sample 16 kHz fixture (the NeMo dump's
+250,760 was the loader's padded batch width; perception consumed the true
+length); the system prompt tokenized to 56 ids ([bos] + 54 + [eos]); NeMo's
+timeline T = 252 = 56 + 196; and our fp32 pipeline matched the reference text
+channel 196/196 tokens.
+"""
+
+import pytest
+
+from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_thinker import (
+ compute_acoustic_frame_count,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# Perception config subset matching the checkpoint (mel 10 ms hop, n_fft 512,
+# dw-striding 8x CAUSAL subsampling, conv kernel 3). causal_downsampling is
+# what the real checkpoint sets; omitting it here is exactly what let the
+# non-causal padding formula pass review round 0 (helper vs itself).
+_STT_CFG = {
+ "perception": {
+ "preprocessor": {"sample_rate": 16000, "window_stride": 0.01, "window_size": 0.025, "n_fft": 512},
+ "encoder": {
+ "subsampling": "dw_striding",
+ "subsampling_factor": 8,
+ "subsampling_conv_kernel_size": 3,
+ "causal_downsampling": True,
+ },
+ }
+}
+
+# (num_16k_samples, expected 12.5 Hz frames), all verified against the REAL
+# vendored AudioPerceptionModule built from the checkpoint config. The first
+# row is the sample_general.wav acceptance fixture (196 frames -> the NeMo
+# T=252 timeline). 71271 is the reviewer's 4.45 s repro that the non-causal
+# formula undersized by one; 250760 is the NeMo loader's padded width, which
+# genuinely yields 197 (not the fixture's 196).
+_GOLDEN = [
+ (249734, 196),
+ (250760, 197),
+ (71271, 57),
+ (16000, 14),
+ (16000 * 2, 26),
+ (16000 * 4, 51),
+ (16000 * 8, 101),
+]
+
+
+@pytest.mark.parametrize(("num_samples", "expected_frames"), _GOLDEN)
+def test_acoustic_frame_count_golden(num_samples: int, expected_frames: int) -> None:
+ assert compute_acoustic_frame_count(_STT_CFG, num_samples) == expected_frames
+
+
+def test_frame_count_matches_vendored_subsampling_geometry() -> None:
+ """Pin the helper against the real ConvSubsampling, not against itself.
+
+ Builds the vendored subsampling module with the checkpoint geometry (both
+ causal and non-causal) and checks the helper's padding arithmetic against
+ the module's actual ``_left_padding + _right_padding`` and its
+ ``calc_length`` output over a spread of mel lengths.
+ """
+ import torch
+
+ from vllm_omni.model_executor.models.nemotron_voicechat.nemo_vendored.asr.subsampling import (
+ ConvSubsampling,
+ calc_length,
+ )
+
+ for causal in (True, False):
+ sub = ConvSubsampling(
+ subsampling="dw_striding",
+ subsampling_factor=8,
+ feat_in=128,
+ feat_out=8,
+ conv_channels=8,
+ is_causal=causal,
+ )
+ expected_paddings = sub._left_padding + sub._right_padding
+ cfg = {
+ "perception": {
+ "preprocessor": _STT_CFG["perception"]["preprocessor"],
+ "encoder": {
+ "subsampling": "dw_striding",
+ "subsampling_factor": 8,
+ "subsampling_conv_kernel_size": 3,
+ "causal_downsampling": causal,
+ },
+ }
+ }
+ for num_samples in (16000, 71271, 249734, 250760):
+ # Recompute the helper's mel length, then ask the module's own
+ # calc_length what the subsampled length should be.
+ mel_len = (num_samples + 512 - 512) // 160 + 1
+ module_frames = int(
+ calc_length(
+ torch.tensor([mel_len]),
+ all_paddings=expected_paddings,
+ kernel_size=3,
+ stride=2,
+ ceil_mode=False,
+ repeat_num=3,
+ )[0]
+ )
+ assert compute_acoustic_frame_count(cfg, num_samples) == module_frames, (
+ f"helper diverged from vendored geometry (causal={causal}, samples={num_samples})"
+ )
+
+
+def test_timeline_max_model_len_validation() -> None:
+ from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_thinker import (
+ validate_timeline_fits_max_model_len,
+ )
+
+ # The acceptance fixture (56 + 1 + 196 = 253) fits the shipped 8192 limit.
+ validate_timeline_fits_max_model_len(56, 196, 8192)
+ # Exactly at the limit: the vLLM request occupies prompt + 1 (frame-0
+ # placeholder) + frames tokens, so 57 + 8135 == 8192 is allowed...
+ validate_timeline_fits_max_model_len(56, 8135, 8192)
+ # ...and one more frame must be rejected, naming every offending length.
+ with pytest.raises(ValueError) as excinfo:
+ validate_timeline_fits_max_model_len(56, 8136, 8192)
+ message = str(excinfo.value)
+ for fragment in (
+ "vllm_prefill_len=57",
+ "acoustic_frame_count=8136",
+ "vllm_total_len=8193",
+ "max_model_len=8192",
+ "logical_prompt_token_len=56",
+ "logical_total_len=8192",
+ ):
+ assert fragment in message
+
+
+def _bare_thinker():
+ import torch
+ from torch import nn
+
+ from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_thinker import (
+ NemotronVoiceChatThinkerForConditionalGeneration,
+ )
+
+ thinker = NemotronVoiceChatThinkerForConditionalGeneration.__new__(NemotronVoiceChatThinkerForConditionalGeneration)
+ nn.Module.__init__(thinker)
+ thinker._use_function_head = True
+ thinker._dtype = torch.float32
+ thinker.function_head = nn.Linear(8, 16, bias=False)
+ return thinker
+
+
+def test_postprocess_appends_function_token_to_timeline(monkeypatch: pytest.MonkeyPatch) -> None:
+ import torch
+
+ thinker = _bare_thinker()
+ hidden = torch.randn(3, 8)
+ expected = int(thinker.function_head(hidden[-1:]).argmax(dim=-1))
+
+ # Default: only the load-bearing prev-token; the cumulative debug timeline
+ # (a per-step host sync + growing concat) stays off.
+ update = thinker.postprocess(hidden)
+ assert int(update["nvc_prev_function_token"]) == expected
+ assert "nvc_function_tokens" not in update
+
+ monkeypatch.setenv("NEMOTRON_VOICECHAT_DEBUG_FUNCTION_TIMELINE", "1")
+ # No existing timeline: starts one.
+ update = thinker.postprocess(hidden)
+ assert update["nvc_function_tokens"].tolist() == [expected]
+ # Existing timeline: appended, not replaced.
+ existing = torch.tensor([7, 8, 9], dtype=torch.long)
+ update = thinker.postprocess(hidden, nvc_function_tokens=existing)
+ assert update["nvc_function_tokens"].tolist() == [7, 8, 9, expected]
+
+
+def test_postprocess_skips_intermediate_prefill_chunks(monkeypatch: pytest.MonkeyPatch) -> None:
+ import torch
+
+ monkeypatch.setenv("NEMOTRON_VOICECHAT_DEBUG_FUNCTION_TIMELINE", "1")
+ thinker = _bare_thinker()
+ hidden = torch.randn(4, 8)
+ # Intermediate chunk: 0 computed + 4 rows < prompt_len 10 -> no token.
+ assert thinker.postprocess(hidden, _omni_is_prefill=True, _omni_num_computed_tokens=0, _omni_prompt_len=10) == {}
+ # Final chunk reaches the last prefill position (acoustic frame 0).
+ update = thinker.postprocess(hidden, _omni_is_prefill=True, _omni_num_computed_tokens=6, _omni_prompt_len=10)
+ assert "nvc_prev_function_token" in update
+ assert update["nvc_function_tokens"].numel() == 1
+
+
+def test_prefill_contract_arithmetic() -> None:
+ # vllm_prefill_len = logical_prompt_token_len + 1 (the +1 is acoustic
+ # frame 0); decode steps == acoustic_frame_count; NeMo timeline
+ # T == prompt + frames. Values from the verified sample_general.wav run.
+ logical_prompt_len = 56
+ frames = compute_acoustic_frame_count(_STT_CFG, 249734)
+ vllm_prefill_len = logical_prompt_len + 1
+ assert vllm_prefill_len == 57
+ assert logical_prompt_len + frames == 252 # NeMo reference T
+ # The talker consumes the full timeline from t=1 and the producer trims
+ # the prompt region: rows P-1 onward of the (T-1)-row code stack.
+ talker_rows = (logical_prompt_len + frames) - 1
+ trimmed = talker_rows - (logical_prompt_len - 1)
+ assert trimmed == frames
+
+
+def test_torchaudio_mel_filterbank_matches_librosa() -> None:
+ """The perception frontend builds its mel filterbank with torchaudio.
+
+ torchaudio's melscale_fbanks(norm="slaney", mel_scale="slaney") implements
+ the same slaney-convention math as librosa.filters.mel (the NeMo reference
+ dependency) up to float32 rounding. The 1e-6 tolerance is ~4x the observed
+ 2.6e-7 max diff; e2e this stays token- and WAV-identical to the reference
+ on the parity fixture (per-feature normalization absorbs it).
+ """
+ librosa = pytest.importorskip("librosa")
+ torchaudio = pytest.importorskip("torchaudio")
+
+ # The checkpoint's perception preprocessor config (16 kHz, n_fft=512, 128 mels).
+ ta = torchaudio.functional.melscale_fbanks(
+ n_freqs=257, f_min=0.0, f_max=8000.0, n_mels=128, sample_rate=16000, norm="slaney", mel_scale="slaney"
+ ).T.numpy()
+ reference = librosa.filters.mel(sr=16000, n_fft=512, n_mels=128, fmin=0.0, fmax=8000.0, norm="slaney")
+ assert ta.shape == reference.shape
+ assert abs(ta - reference).max() < 1e-6
diff --git a/tests/model_executor/stage_input_processors/test_nemotron_voicechat.py b/tests/model_executor/stage_input_processors/test_nemotron_voicechat.py
new file mode 100644
index 000000000..dcdff5bbb
--- /dev/null
+++ b/tests/model_executor/stage_input_processors/test_nemotron_voicechat.py
@@ -0,0 +1,422 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.stage_input_processors.nemotron_voicechat import (
+ _FULL_PAYLOAD_REPLACE_KEYS,
+ talker2code2wav_full_payload,
+ talker2code2wav_token_only,
+ thinker2talker_token_only,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_PAD = 12
+
+
+def _thinker_output(prompt_len: int, generated: list[int]):
+ return SimpleNamespace(
+ finished=True,
+ prompt_token_ids=list(range(100, 100 + prompt_len)) + [_PAD],
+ additional_information=None,
+ outputs=[SimpleNamespace(token_ids=generated, cumulative_token_ids=None)],
+ )
+
+
+def test_thinker2talker_builds_pad_prompt_timeline() -> None:
+ generated = [7, 8, 9, 10]
+ inputs = thinker2talker_token_only([_thinker_output(prompt_len=5, generated=generated)])
+ assert len(inputs) == 1
+ prompt = inputs[0]
+ # 1-token placeholder prompt; the timeline rides additional_information.
+ assert prompt["prompt_token_ids"] == [0]
+ info = prompt["additional_information"]
+ assert info["nvc_logical_prompt_len"] == 5
+ assert info["nvc_frame_count"] == 4
+ assert info["nvc_text_timeline"] == [_PAD] * 5 + generated
+
+
+def test_thinker2talker_skips_unfinished() -> None:
+ out = _thinker_output(3, [1])
+ out.finished = False
+ assert thinker2talker_token_only([out]) == []
+
+
+def test_full_payload_trims_prompt_region_rows() -> None:
+ # Rows correspond to timeline steps t=1..T-1; with prompt_len P the NeMo
+ # trim keeps rows for t >= P, i.e. row index P-1 onward.
+ prompt_len = 4
+ frames, q = 10, 31
+ codes = torch.arange(frames * q, dtype=torch.long).reshape(frames, q) % 1024
+ request = SimpleNamespace(
+ additional_information={
+ "codes": {"audio": codes},
+ "nvc_logical_prompt_len": prompt_len,
+ },
+ additional_information_cpu=None,
+ )
+ payload = talker2code2wav_full_payload(request=request, is_finished=True)
+ assert payload.codes.audio.shape == (frames - (prompt_len - 1), q)
+ assert torch.equal(payload.codes.audio, codes[prompt_len - 1 :])
+ assert bool(payload.meta.finished)
+
+
+def test_full_payload_without_codes_is_empty_finished() -> None:
+ # The sync producer must never emit a meta-only "pending" struct (it would
+ # be shipped as the real payload); with no codes it ships an empty finished
+ # payload regardless of the request's is_finished flag.
+ payload = talker2code2wav_full_payload(request=SimpleNamespace(), is_finished=False)
+ assert payload is not None
+ assert bool(payload.meta.finished)
+ assert payload.codes.audio.numel() == 0
+
+
+def test_full_payload_replace_semantics_declared() -> None:
+ # Full-sequence tensors must REPLACE in the accumulator, not concatenate.
+ # The AR runner flattens nested payloads before accumulation, so the
+ # flattened "codes.audio" key is the one that matters.
+ assert "codes.audio" in _FULL_PAYLOAD_REPLACE_KEYS
+ assert "codes" in _FULL_PAYLOAD_REPLACE_KEYS
+
+
+def test_full_payload_missing_prompt_len_fails_explicitly() -> None:
+ codes = torch.zeros((5, 31), dtype=torch.long)
+ request = SimpleNamespace(
+ additional_information={"codes": {"audio": codes}},
+ additional_information_cpu=None,
+ )
+ with pytest.raises(ValueError, match="nvc_logical_prompt_len"):
+ talker2code2wav_full_payload(request=request, is_finished=True)
+
+
+def test_code2wav_distinguishes_empty_payload_from_missing_info() -> None:
+ from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_code2wav import (
+ NemotronVoiceChatCode2Wav,
+ )
+
+ extract = NemotronVoiceChatCode2Wav._codes_from_runtime_info
+ # Explicit empty payload (e.g. zero frames after the prompt trim): the
+ # empty tensor comes back as-is -> empty audio, never a placeholder-id
+ # decode.
+ empty = extract(None, {"codes": {"audio": torch.empty(0, dtype=torch.long)}})
+ assert isinstance(empty, torch.Tensor) and empty.numel() == 0
+ # Missing codes info entirely -> None (forward raises: placeholder
+ # input_ids are never decoded).
+ assert extract(None, {"meta": {}}) is None
+ assert extract(None, None) is None
+ # Normal payload passes through (nested and flattened key forms).
+ codes = torch.zeros((3, 31), dtype=torch.long)
+ assert extract(None, {"codes": {"audio": codes}}) is codes
+ assert extract(None, {"codes.audio": codes}) is codes
+
+
+def test_validate_code_stack_negatives() -> None:
+ from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_code2wav import (
+ validate_code_stack,
+ )
+
+ good = torch.zeros((4, 31), dtype=torch.long)
+ assert validate_code_stack(good, 31, 1024).shape == (4, 31)
+ with pytest.raises(ValueError, match="not divisible"):
+ validate_code_stack(torch.zeros(100, dtype=torch.long), 31, 1024)
+ with pytest.raises(ValueError, match="shaped"):
+ validate_code_stack(torch.zeros((4, 30), dtype=torch.long), 31, 1024)
+ bad = good.clone()
+ bad[1, 5] = 1024
+ with pytest.raises(ValueError, match="out-of-range"):
+ validate_code_stack(bad, 31, 1024)
+
+
+def test_tokenizer_resolution_failure_is_actionable(monkeypatch) -> None:
+ from transformers import AutoConfig
+
+ from vllm_omni.model_executor.models.nemotron_voicechat.configuration_nemotron_voicechat import (
+ NemotronVoiceChatConfig,
+ )
+
+ def _boom(*args, **kwargs):
+ raise OSError("offline")
+
+ monkeypatch.setattr(AutoConfig, "from_pretrained", _boom)
+ with pytest.raises(RuntimeError, match="NEMOTRON_VOICECHAT_LLM_PATH"):
+ NemotronVoiceChatConfig(
+ data={"frame_length": 0.08, "source_sample_rate": 16000, "target_sample_rate": 22050},
+ model={
+ "inference_speaker_name": "Aria",
+ "stt": {"model": {"perception": {"encoder": {}}, "pretrained_llm": "nvidia/nope"}},
+ "speech_generation": {"model": {"tts_config": {}, "codec_config": {}}},
+ },
+ )
+
+
+def test_sanitize_tts_model_cfg_disables_pretrained_loads(monkeypatch) -> None:
+ from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_talker import (
+ sanitize_tts_model_cfg,
+ )
+
+ original = {
+ "pretrained_lm_name": "nvidia/NVIDIA-Nemotron-Nano-9B-v2",
+ "pretrained_codec_model": "/some/external/codec.ckpt",
+ "tts_config": {
+ "pretrained_text_name": "some/hub-llm",
+ "backbone_type": "gemma3_text",
+ "cas_config": {"pretrained_tokenizer_name": "some/hub-tokenizer", "keep": 1},
+ },
+ }
+ monkeypatch.delenv("NEMOTRON_VOICECHAT_LLM_PATH", raising=False)
+ cfg = sanitize_tts_model_cfg(original)
+ # External weight-load triggers are gone; construction is config-only.
+ assert "pretrained_text_name" not in cfg["tts_config"]
+ assert "pretrained_codec_model" not in cfg
+ assert "pretrained_tokenizer_name" not in cfg["tts_config"]["cas_config"]
+ assert cfg["tts_config"]["cas_config"]["keep"] == 1
+ # Tokenizer reference survives.
+ assert cfg["pretrained_lm_name"] == "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
+ # The shared stage config is never mutated (deep copy).
+ assert original["tts_config"]["pretrained_text_name"] == "some/hub-llm"
+
+ # Air-gapped override reroutes the tokenizer reference.
+ monkeypatch.setenv("NEMOTRON_VOICECHAT_LLM_PATH", "/local/nemotron")
+ cfg = sanitize_tts_model_cfg(original)
+ assert cfg["pretrained_lm_name"] == "/local/nemotron"
+
+
+def test_code2wav_token_only_sizes_placeholder() -> None:
+ talker_output = SimpleNamespace(
+ finished=True,
+ outputs=[SimpleNamespace(token_ids=[0] * 9, cumulative_token_ids=None)],
+ )
+ inputs = talker2code2wav_token_only([talker_output])
+ assert len(inputs) == 1
+ assert len(inputs[0]["prompt_token_ids"]) == 9
+
+
+# =========================
+# Async-chunk (streaming) producers
+# =========================
+
+
+def _fake_transfer_manager(codec_chunk_frames: int | None = None):
+ extra = {}
+ if codec_chunk_frames is not None:
+ extra["codec_chunk_frames"] = codec_chunk_frames
+ return SimpleNamespace(
+ request_payload={},
+ connector=SimpleNamespace(config={"extra": extra}),
+ put_req_chunk={},
+ )
+
+
+def _streaming_request(prompt_len: int, generated: list[int], info: dict | None = None):
+ return SimpleNamespace(
+ external_req_id="req-0",
+ prompt_token_ids=list(range(100, 100 + prompt_len)) + [_PAD],
+ output_token_ids=generated,
+ additional_information=info,
+ )
+
+
+def test_thinker2talker_async_chunk_cumulative_timeline() -> None:
+ from vllm_omni.model_executor.stage_input_processors.nemotron_voicechat import (
+ thinker2talker_async_chunk,
+ )
+
+ tm = _fake_transfer_manager()
+ req = _streaming_request(prompt_len=3, generated=[7])
+ payload = thinker2talker_async_chunk(tm, {}, req, is_finished=False)
+ assert payload is not None
+ assert payload.ids.all == [_PAD] * 3 + [7]
+ assert payload.meta.num_processed_tokens == 3
+ assert not bool(payload.meta.finished)
+
+ # No new tokens since the last chunk -> skip.
+ assert thinker2talker_async_chunk(tm, {}, req, is_finished=False) is None
+
+ # More tokens -> longer cumulative timeline; final chunk sets finished.
+ req.output_token_ids = [7, 8, 9]
+ payload = thinker2talker_async_chunk(tm, {}, req, is_finished=True)
+ assert payload.ids.all == [_PAD] * 3 + [7, 8, 9]
+ assert bool(payload.meta.finished)
+
+
+def test_talker2code2wav_async_chunk_cadence_and_left_context() -> None:
+ from vllm_omni.model_executor.stage_input_processors.nemotron_voicechat import (
+ talker2code2wav_async_chunk,
+ )
+
+ tm = _fake_transfer_manager(codec_chunk_frames=2)
+ prompt_len = 3
+ info = {"nvc_logical_prompt_len": prompt_len}
+ req = _streaming_request(prompt_len=prompt_len, generated=[], info=info)
+
+ def cumulative(rows: int) -> dict:
+ # Talker rows are timeline steps t=1..; prompt region occupies the
+ # first prompt_len-1 rows before acoustic frames start.
+ return {"codes": {"audio": torch.arange((prompt_len - 1 + rows) * 31).reshape(-1, 31)}}
+
+ # Only prompt-region rows so far -> nothing to ship.
+ assert talker2code2wav_async_chunk(tm, cumulative(0), req, is_finished=False) is None
+ # One new frame < chunk size 2 -> hold.
+ assert talker2code2wav_async_chunk(tm, cumulative(1), req, is_finished=False) is None
+ # Two frames -> first chunk, no left context.
+ payload = talker2code2wav_async_chunk(tm, cumulative(2), req, is_finished=False)
+ assert payload is not None
+ assert payload.meta.left_context_size == 0
+ assert payload.codes.audio.shape == (2, 31)
+ # Third frame + finished -> cumulative history with left context 2.
+ payload = talker2code2wav_async_chunk(tm, cumulative(3), req, is_finished=True)
+ assert payload.meta.left_context_size == 2
+ assert payload.codes.audio.shape == (3, 31)
+ assert bool(payload.meta.finished)
+
+
+def test_talker2code2wav_async_chunk_requires_prompt_len() -> None:
+ from vllm_omni.model_executor.stage_input_processors.nemotron_voicechat import (
+ talker2code2wav_async_chunk,
+ )
+
+ tm = _fake_transfer_manager()
+ req = _streaming_request(prompt_len=3, generated=[], info=None)
+ with pytest.raises(ValueError, match="logical prompt length"):
+ talker2code2wav_async_chunk(tm, {"codes": {"audio": torch.zeros(5, 31, dtype=torch.long)}}, req)
+
+
+# =========================
+# Talker drain-per-wake (async scheduling review P1/P2)
+# =========================
+
+
+class _FakeTTSModelOutput:
+ past_key_values = None
+
+
+class _FakeTTS:
+ """Stub DuplexEARTTS: infer_codes_one_step returns a constant code row."""
+
+ def __init__(self, q: int = 31):
+ self.q = q
+ self.steps = 0
+
+ def set_init_inputs(self, speaker_name: str) -> None:
+ del speaker_name
+
+ def get_init_inputs(self, B: int) -> dict: # noqa: N803 — mirrors the vendored DuplexEARTTS signature
+ return {
+ "code": torch.zeros((B, 2, self.q), dtype=torch.long),
+ "subword_ids": torch.zeros((B, 3), dtype=torch.long),
+ }
+
+ def _get_generation_config(self, guidance_enabled: bool):
+ del guidance_enabled
+ return {}
+
+ def tts_model(self, **kwargs):
+ del kwargs
+ return _FakeTTSModelOutput()
+
+ def infer_codes_one_step(self, **kwargs):
+ del kwargs
+ self.steps += 1
+ return torch.full((1, 1, self.q), self.steps, dtype=torch.long), None
+
+
+def _bare_talker():
+ from torch import nn
+
+ from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_talker import (
+ NemotronVoiceChatTalkerForConditionalGeneration,
+ )
+
+ talker = NemotronVoiceChatTalkerForConditionalGeneration.__new__(NemotronVoiceChatTalkerForConditionalGeneration)
+ nn.Module.__init__(talker)
+ talker.tts = _FakeTTS()
+ talker._sessions = {}
+ talker._hidden = 8
+ talker._dtype = torch.float32
+ talker._speaker_name = "Aria"
+ talker.config = SimpleNamespace(tts_cfg={"inference_guidance_enabled": False})
+ return talker
+
+
+def _async_info(timeline: list[int], prompt_len: int, finished: bool, request_id: str = "req-0") -> dict:
+ return {
+ "request_id": request_id,
+ "additional_information": {
+ "ids": {"all": list(timeline)},
+ "meta": {
+ "num_processed_tokens": prompt_len,
+ "finished": torch.tensor(finished),
+ "request_id": request_id,
+ },
+ },
+ }
+
+
+def test_talker_drains_all_received_positions_per_wake() -> None:
+ """P2: chunk 0's PAD prompt region must not cost one upstream chunk per step,
+ and P1: a coalesced chunk (2+ new positions, 1 wake) must be fully drained."""
+ talker = _bare_talker()
+ pad, prompt_len = 12, 4
+ ids = torch.zeros(1, dtype=torch.long)
+
+ # Prefill wake: chunk 0 carries only the PAD prompt region.
+ info = _async_info([pad] * prompt_len, prompt_len, finished=False)
+ talker.preprocess(ids, None, _omni_is_prefill=True, **info)
+ session = talker._sessions["req-0"]
+ assert session["step"] == 1 and not session["sync_mode"]
+
+ # First decode wake: drains ALL prompt-region steps (t=1..P-1) at once.
+ _, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info)
+ assert session["step"] == prompt_len # t advanced 1 -> P in one wake
+ assert update["codes"]["audio"].shape == (prompt_len - 1, 31)
+ assert float(embeds[0, 0]) == 0.0 # not finished
+
+ # Coalesced chunk: TWO new frame tokens arrive in ONE wake (delayed save
+ # thread merged two thinker steps). Both must be consumed by this wake.
+ info2 = _async_info([pad] * prompt_len + [7, 8], prompt_len, finished=False)
+ _, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info2)
+ assert session["step"] == prompt_len + 2
+ assert update["codes"]["audio"].shape == (prompt_len + 1, 31)
+ assert float(embeds[0, 0]) == 0.0
+
+ # Zero-progress wake (no new positions, not finished): plain CONTINUE, no
+ # fabricated frames, no codes update.
+ _, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info2)
+ assert session["step"] == prompt_len + 2
+ assert update == {}
+ assert float(embeds[0, 0]) == 0.0
+
+ # Final chunk: one more token + finished marker -> drain + stop flag.
+ info3 = _async_info([pad] * prompt_len + [7, 8, 9], prompt_len, finished=True)
+ _, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info3)
+ assert session["step"] == prompt_len + 3
+ assert update["codes"]["audio"].shape == (prompt_len + 2, 31)
+ assert float(embeds[0, 0]) == 1.0 # stop
+
+
+def test_talker_sync_mode_steps_once_per_wake() -> None:
+ """The verified sync parity path keeps its one-TTS-step-per-engine-step cadence."""
+ talker = _bare_talker()
+ pad, prompt_len = 12, 3
+ ids = torch.zeros(1, dtype=torch.long)
+ info = {
+ "request_id": "req-1",
+ "additional_information": {
+ "nvc_text_timeline": [pad] * prompt_len + [7, 8],
+ "nvc_logical_prompt_len": prompt_len,
+ },
+ }
+ talker.preprocess(ids, None, _omni_is_prefill=True, **info)
+ session = talker._sessions["req-1"]
+ assert session["sync_mode"] and session["upstream_finished"]
+ for expected_step in (2, 3, 4):
+ _, embeds, _ = talker.preprocess(ids, None, _omni_is_prefill=False, **info)
+ assert session["step"] == expected_step
+ # Last position -> stop flag on the final step.
+ _, embeds, _ = talker.preprocess(ids, None, _omni_is_prefill=False, **info)
+ assert session["step"] == 5
+ assert float(embeds[0, 0]) == 1.0
diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py
index 5ae83d1b9..4b5c73a39 100644
--- a/vllm_omni/config/omni_config.py
+++ b/vllm_omni/config/omni_config.py
@@ -128,6 +128,7 @@ class _CacheEngineOverrides(TypedDict, total=False):
enable_prefix_caching: bool
disable_hybrid_kv_cache_manager: bool
mm_processor_cache_gb: float
+ mamba_ssm_cache_dtype: str
class _SchedulerEngineOverrides(TypedDict, total=False):
@@ -377,6 +378,8 @@ class OmniStageCacheConfig:
enable_prefix_caching: bool | None = None
disable_hybrid_kv_cache_manager: bool | None = None
mm_processor_cache_gb: float | None = Field(default=None, ge=0.0)
+ # Hybrid-mamba SSM state dtype ("auto"/"float32"); vLLM CacheConfig field.
+ mamba_ssm_cache_dtype: str | None = None
@config
diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py
index ac2e9d38d..f02a3578e 100644
--- a/vllm_omni/config/pipeline_registry.py
+++ b/vllm_omni/config/pipeline_registry.py
@@ -88,6 +88,9 @@ from vllm_omni.model_executor.models.moss_tts.pipeline import (
MOSS_TTS_REALTIME_PIPELINE,
)
from vllm_omni.model_executor.models.moss_tts_nano.pipeline import MOSS_TTS_NANO_PIPELINE
+from vllm_omni.model_executor.models.nemotron_voicechat.pipeline import (
+ NEMOTRON_VOICECHAT_PIPELINE,
+)
from vllm_omni.model_executor.models.omnivoice.pipeline import OMNIVOICE_PIPELINE
from vllm_omni.model_executor.models.personaplex.pipeline import PERSONAPLEX_PIPELINE
from vllm_omni.model_executor.models.qwen2_5_omni.pipeline import (
@@ -118,6 +121,11 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
"qwen2_5_omni": QWEN2_5_OMNI_PIPELINE,
"qwen2_5_omni_thinker_only": QWEN2_5_OMNI_THINKER_ONLY_PIPELINE,
"personaplex": PERSONAPLEX_PIPELINE,
+ "nemotron_voicechat": NEMOTRON_VOICECHAT_PIPELINE,
+ # Alias: lets bare `vllm-omni serve `
+ # auto-detect through the path-basename fallback (the checkpoint config.json
+ # has no model_type key).
+ "nemotron_labs_voicechat": NEMOTRON_VOICECHAT_PIPELINE,
"qwen3_omni_moe": resolve_qwen3_omni_pipeline,
"qwen3_tts": QWEN3_TTS_PIPELINE,
"step_audio_2": STEP_AUDIO2_PIPELINE,
diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py
index fa8e45295..496aea208 100644
--- a/vllm_omni/config/stage_config.py
+++ b/vllm_omni/config/stage_config.py
@@ -357,6 +357,9 @@ class StageDeployConfig:
async_scheduling: bool | None = None
disable_hybrid_kv_cache_manager: bool | None = None
mm_processor_cache_gb: float | None = None
+ # Hybrid-mamba stages (e.g. the NemotronVoiceChat thinker's NemotronH
+ # backbone) pin the SSM state dtype; projected onto vLLM CacheConfig.
+ mamba_ssm_cache_dtype: str | None = None
# Generic compilation, profiling, tokenizer/config parsing, and model
# loading controls.
diff --git a/vllm_omni/core/sched/omni_scheduling_coordinator.py b/vllm_omni/core/sched/omni_scheduling_coordinator.py
index 424fed2d9..75f234d8c 100644
--- a/vllm_omni/core/sched/omni_scheduling_coordinator.py
+++ b/vllm_omni/core/sched/omni_scheduling_coordinator.py
@@ -51,6 +51,10 @@ _FULL_PAYLOAD_INPUT_STAGES: frozenset[tuple[str, str]] = frozenset(
("MiniCPMO45Code2Wav", "code2wav"),
# cosyvoice3: cosyvoice3_talker (Stage 0) -> cosyvoice3_code2wav (Stage 1).
("CosyVoice3Model", "cosyvoice3_code2wav"),
+ # nemotron_voicechat: talker (Stage 1) -> code2wav (Stage 2). Stage 2
+ # waits for the talker's full-payload code stacks; the thinker (Stage 0)
+ # -> talker (Stage 1) hop is token-path only and must NOT be listed.
+ ("NemotronVoiceChatCode2Wav", "code2wav"),
# audex TTS sync path: thinker (Stage 0) -> streaming decoder (Stage 1).
# The default deploy is async_chunk; this covers async_chunk: false.
("AudexCode2Wav", "audex_code2wav"),
diff --git a/vllm_omni/deploy/nemotron_labs_voicechat.yaml b/vllm_omni/deploy/nemotron_labs_voicechat.yaml
new file mode 100644
index 000000000..fe872caea
--- /dev/null
+++ b/vllm_omni/deploy/nemotron_labs_voicechat.yaml
@@ -0,0 +1,81 @@
+# NemotronVoiceChat-11B offline speech-to-speech deployment.
+# 3 stages on one GPU: thinker (NemotronH AR) -> talker (EAR-TTS AR) -> code2wav
+# (RVQ-VAE decode). thinker->talker is token-path only; talker->code2wav ships
+# the [frames, 31] code stack via the shared-memory connector.
+pipeline: nemotron_voicechat
+async_chunk: false
+
+connectors:
+ connector_of_shared_memory:
+ name: SharedMemoryConnector
+ extra:
+ shm_threshold_bytes: 65536
+
+stages:
+ - stage_id: 0 # thinker: speech -> frame-locked text (final text output)
+ max_num_seqs: 1
+ gpu_memory_utilization: 0.62
+ enforce_eager: true
+ enable_prefix_caching: false
+ async_scheduling: false
+ max_num_batched_tokens: 8192
+ max_model_len: 8192
+ # fp32 matches the NeMo reference exactly (196/196 token parity on the
+ # acceptance fixture). bf16 is ~2x faster and measured within one word of
+ # the reference on the same fixture — flip to bfloat16 if you prefer speed
+ # over exact parity.
+ dtype: float32
+ # NemotronH mamba SSM state in fp32 (matches NVIDIA's own vLLM deployment
+ # of this model; bf16 SSM state accumulates drift over long sequences).
+ mamba_ssm_cache_dtype: float32
+ devices: "0"
+ default_sampling_params:
+ # Greedy, frame-locked: the example sets max_tokens to the
+ # exact acoustic frame count per request; text EOS never stops generation.
+ temperature: 0.0
+ ignore_eos: true
+ seed: 0
+
+ - stage_id: 1 # talker: text timeline -> 31-quantizer code stacks
+ max_num_seqs: 1
+ gpu_memory_utilization: 0.12
+ enforce_eager: true
+ enable_prefix_caching: false
+ async_scheduling: false
+ skip_tokenizer_init: true
+ max_num_batched_tokens: 4096
+ max_model_len: 16384
+ # fp32 like the NeMo reference: the EAR-TTS MoG/codec-embedding path mixes
+ # fp32 buffers into the forward, and 1B params in fp32 are cheap.
+ dtype: float32
+ devices: "0"
+ output_connectors:
+ to_stage_2: connector_of_shared_memory
+ default_sampling_params:
+ # Placeholder-token loop; the stage emits stop token 1 after the last
+ # timeline position. 16383 (not 16384): the request also carries one
+ # placeholder prompt token, and prompt + max_tokens must fit
+ # max_model_len.
+ temperature: 0.0
+ max_tokens: 16383
+ stop_token_ids: [1]
+ detokenize: false
+ seed: 0
+
+ - stage_id: 2 # code2wav: codes -> 22.05 kHz PCM (fp32 codec)
+ max_num_seqs: 1
+ gpu_memory_utilization: 0.06
+ enforce_eager: true
+ enable_prefix_caching: false
+ async_scheduling: false
+ skip_tokenizer_init: true
+ max_num_batched_tokens: 65536
+ max_model_len: 65536
+ devices: "0"
+ input_connectors:
+ from_stage_1: connector_of_shared_memory
+ default_sampling_params:
+ temperature: 0.0
+ max_tokens: 1
+ detokenize: false
+ seed: 0
diff --git a/vllm_omni/deploy/nemotron_labs_voicechat_streaming.yaml b/vllm_omni/deploy/nemotron_labs_voicechat_streaming.yaml
new file mode 100644
index 000000000..133bbee34
--- /dev/null
+++ b/vllm_omni/deploy/nemotron_labs_voicechat_streaming.yaml
@@ -0,0 +1,84 @@
+# NemotronVoiceChat-11B STREAMING speech-to-speech deployment (async-chunk).
+#
+# Same 3 stages as nemotron_labs_voicechat.yaml, but the stages overlap:
+# the thinker streams its frame-locked text timeline to the talker per decode
+# step (cumulative ids.all payloads), and the talker streams prompt-trimmed
+# cumulative code stacks to code2wav every `codec_chunk_frames` frames. The
+# codec re-decodes the full prefix per chunk and emits only the new samples
+# (NeMo decode_one_audio_step semantics), so the concatenated stream is
+# seam-free. First audio arrives after ~codec_chunk_frames thinker steps
+# instead of after the whole pipeline completes.
+#
+# The non-streaming yaml remains the parity-acceptance path (bit-exact WAV);
+# streamed audio is validated by ASR/WER instead (prefix decodes are not
+# bit-identical to the one-shot decode of the final stack).
+pipeline: nemotron_voicechat
+async_chunk: true
+
+connectors:
+ connector_of_shared_memory:
+ name: SharedMemoryConnector
+ extra:
+ shm_threshold_bytes: 65536
+ # Streaming cadence for talker -> code2wav: emit a chunk once this many
+ # NEW frames (80 ms each) have accumulated; ~1 s of audio per chunk.
+ codec_chunk_frames: 13
+
+stages:
+ - stage_id: 0 # thinker: speech -> frame-locked text (final text output)
+ max_num_seqs: 1
+ gpu_memory_utilization: 0.62
+ enforce_eager: true
+ enable_prefix_caching: false
+ async_scheduling: false
+ max_num_batched_tokens: 8192
+ max_model_len: 8192
+ dtype: float32
+ mamba_ssm_cache_dtype: float32
+ devices: "0"
+ output_connectors:
+ to_stage_1: connector_of_shared_memory
+ default_sampling_params:
+ temperature: 0.0
+ ignore_eos: true
+ seed: 0
+
+ - stage_id: 1 # talker: text timeline -> 31-quantizer code stacks
+ max_num_seqs: 1
+ gpu_memory_utilization: 0.12
+ enforce_eager: true
+ enable_prefix_caching: false
+ async_scheduling: false
+ skip_tokenizer_init: true
+ max_num_batched_tokens: 4096
+ max_model_len: 16384
+ dtype: float32
+ devices: "0"
+ input_connectors:
+ from_stage_0: connector_of_shared_memory
+ output_connectors:
+ to_stage_2: connector_of_shared_memory
+ default_sampling_params:
+ temperature: 0.0
+ max_tokens: 16383
+ stop_token_ids: [1]
+ detokenize: false
+ seed: 0
+
+ - stage_id: 2 # code2wav: codes -> 22.05 kHz PCM (fp32 codec)
+ max_num_seqs: 1
+ gpu_memory_utilization: 0.06
+ enforce_eager: true
+ enable_prefix_caching: false
+ async_scheduling: false
+ skip_tokenizer_init: true
+ max_num_batched_tokens: 65536
+ max_model_len: 65536
+ devices: "0"
+ input_connectors:
+ from_stage_1: connector_of_shared_memory
+ default_sampling_params:
+ temperature: 0.0
+ max_tokens: 1
+ detokenize: false
+ seed: 0
diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py
index 1d9eb803b..a3aa96f60 100644
--- a/vllm_omni/engine/arg_utils.py
+++ b/vllm_omni/engine/arg_utils.py
@@ -30,6 +30,11 @@ _ARCH_TO_MODEL_TYPE: dict[str, str] = {
# the HF config without manual overrides.
"PersonaPlexTalkerForConditionalGeneration": "personaplex",
"PersonaPlexCode2Wav": "personaplex",
+ # NemotronVoiceChat ships a NeMo-style config.json without model_type;
+ # create_model_config() patches model_type=nemotron_voicechat from the arch.
+ "NemotronVoiceChatThinkerForConditionalGeneration": "nemotron_voicechat",
+ "NemotronVoiceChatTalkerForConditionalGeneration": "nemotron_voicechat",
+ "NemotronVoiceChatCode2Wav": "nemotron_voicechat",
"VoxCPM2TalkerForConditionalGeneration": "voxcpm2",
}
@@ -56,6 +61,9 @@ def _register_omni_hf_configs() -> None:
MossTTSLocalConfig,
MossTTSRealtimeConfig,
)
+ from vllm_omni.model_executor.models.nemotron_voicechat.configuration_nemotron_voicechat import (
+ NemotronVoiceChatConfig,
+ )
from vllm_omni.model_executor.models.personaplex.configuration_personaplex import (
PersonaPlexConfig,
)
@@ -87,6 +95,7 @@ def _register_omni_hf_configs() -> None:
("moss_tts_realtime", MossTTSRealtimeConfig),
("qwen3_tts", Qwen3TTSConfig),
("personaplex", PersonaPlexConfig),
+ ("nemotron_voicechat", NemotronVoiceChatConfig),
("cosyvoice3", CosyVoice3Config),
("glm_tts", GLMTTSConfig),
("omnivoice", OmniVoiceConfig),
@@ -324,6 +333,19 @@ class OmniEngineArgs(EngineArgs):
if model_type is not None:
self._patch_empty_hf_config(model_type)
+ # NemotronVoiceChat: the checkpoint's only tokenizer subfolder is the
+ # auxiliary rnnt_tokenizer/ (ASR loss vocab), which the generic
+ # subfolder auto-detect below would wrongly pick up. The text tokenizer
+ # is the Nemotron backbone's, resolved from a local override or its
+ # HF id (cache-friendly).
+ if not self.tokenizer and self.model_arch in (
+ "NemotronVoiceChatThinkerForConditionalGeneration",
+ "NemotronVoiceChatTalkerForConditionalGeneration",
+ "NemotronVoiceChatCode2Wav",
+ ):
+ self.tokenizer = os.environ.get("NEMOTRON_VOICECHAT_LLM_PATH") or "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
+ logger.info("NemotronVoiceChat: using text tokenizer from %s", self.tokenizer)
+
# Auto-detect tokenizer for models that store it in a subdirectory
# rather than the root (e.g. CosyVoice3 uses CosyVoice-BlankEN/).
if not self.tokenizer and self.model:
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/__init__.py b/vllm_omni/model_executor/models/nemotron_voicechat/__init__.py
new file mode 100644
index 000000000..f8046bfc5
--- /dev/null
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/__init__.py
@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""NVIDIA NemotronLabs VoiceChat-11B: offline speech-to-speech 3-stage pipeline."""
+
+from vllm_omni.model_executor.models.nemotron_voicechat.configuration_nemotron_voicechat import (
+ NemotronVoiceChatCode2WavConfig,
+ NemotronVoiceChatConfig,
+ NemotronVoiceChatTalkerConfig,
+)
+from vllm_omni.model_executor.models.nemotron_voicechat.pipeline import (
+ NEMOTRON_VOICECHAT_PIPELINE,
+)
+
+__all__ = [
+ "NemotronVoiceChatConfig",
+ "NemotronVoiceChatTalkerConfig",
+ "NemotronVoiceChatCode2WavConfig",
+ "NEMOTRON_VOICECHAT_PIPELINE",
+]
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/configuration_nemotron_voicechat.py b/vllm_omni/model_executor/models/nemotron_voicechat/configuration_nemotron_voicechat.py
new file mode 100644
index 000000000..1469f72df
--- /dev/null
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/configuration_nemotron_voicechat.py
@@ -0,0 +1,211 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""HF-compatible configs synthesized from the NeMo-style NemotronVoiceChat checkpoint.
+
+The ``nvidia/NVIDIA-NemotronLabs-VoiceChat-11B`` checkpoint ships a NeMo training
+config (top-level ``data`` / ``model`` keys) instead of a transformers config:
+there is no ``model_type`` or ``architectures``. vLLM-Omni's empty-config patching
+(:meth:`OmniEngineArgs._patch_empty_hf_config`) injects
+``model_type="nemotron_voicechat"`` and routes the raw dict here, where it is
+parsed into per-stage sub-configs:
+
+* ``get_text_config()`` -> the NemotronH backbone config (thinker stage; the AR
+ runner reads ``hidden_size`` etc. from it). Resolved from the checkpoint's
+ ``model.stt.model.pretrained_llm`` HF id via the local HF cache, or from the
+ ``NEMOTRON_VOICECHAT_LLM_PATH`` env var pointing at a local snapshot.
+* ``talker_config`` -> Gemma3-based EAR-TTS backbone dims (talker stage).
+* ``code2wav_config`` -> RVQ-VAE codec dims (code2wav stage).
+
+The raw NeMo sections are preserved on the config (``stt_cfg`` / ``tts_cfg`` /
+``nemo_data``) because the stage models consume them verbatim (perception layout,
+fusion weights, TTS generation parameters, codec architecture).
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+from transformers import AutoConfig, PretrainedConfig
+
+_EXPECTED_STRUCTURE = (
+ "Expected a NemotronVoiceChat checkpoint config.json with the NeMo layout: "
+ "top-level 'model' containing both 'stt' (with 'model.stt.model.perception' "
+ "and 'model.stt.model.pretrained_llm') and 'speech_generation' (with "
+ "'model.speech_generation.model.tts_config' and '...codec_config')."
+)
+
+# Default HF id of the thinker text backbone; the checkpoint records the same id
+# under model.stt.model.pretrained_llm.
+_DEFAULT_LLM_ID = "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
+
+# Env override for air-gapped runs: local directory containing the Nemotron
+# backbone's config.json (and optionally tokenizer files).
+_LLM_PATH_ENV = "NEMOTRON_VOICECHAT_LLM_PATH"
+
+
+class NemotronVoiceChatTalkerConfig(PretrainedConfig):
+ """Stage config for the EAR-TTS talker (28-layer Gemma3-style backbone)."""
+
+ model_type = "nemotron_voicechat_talker"
+
+ def __init__(self, **kwargs: Any) -> None:
+ backbone: dict[str, Any] = kwargs.pop("backbone_config", None) or {}
+ super().__init__(**kwargs)
+ self.backbone_config = backbone
+ self.hidden_size = int(backbone.get("hidden_size", 1152))
+ self.num_hidden_layers = int(backbone.get("num_hidden_layers", 28))
+ self.num_attention_heads = int(backbone.get("num_attention_heads", 16))
+ self.head_dim = int(backbone.get("head_dim", 72))
+ self.intermediate_size = int(backbone.get("intermediate_size", 4608))
+ self.sliding_window = int(backbone.get("sliding_window", 7500))
+ # The talker consumes/produces per-frame code stacks; vocab is nominal
+ # (codes are sampled by the vendored MOG head, not a vLLM sampler).
+ self.vocab_size = int(kwargs.get("vocab_size", 1024))
+ self.max_position_embeddings = int(backbone.get("max_position_embeddings", 32768))
+
+ def get_text_config(self, **_: Any) -> NemotronVoiceChatTalkerConfig:
+ return self
+
+
+class NemotronVoiceChatCode2WavConfig(PretrainedConfig):
+ """Stage config for the RVQ-VAE codec decoder."""
+
+ model_type = "nemotron_voicechat_code2wav"
+
+ def __init__(self, **kwargs: Any) -> None:
+ codec: dict[str, Any] = kwargs.pop("codec_config", None) or {}
+ super().__init__(**kwargs)
+ self.codec_config = codec
+ self.num_quantizers = int(codec.get("num_quantizers", 31))
+ self.codebook_size = int(codec.get("codebook_size", 1024))
+ self.latent_size = int(codec.get("latent_size", 512))
+ self.wav_to_token_ratio = int(codec.get("wav_to_token_ratio", 1764))
+ self.sample_rate = int(kwargs.get("sample_rate", 22050))
+ # The generation runner sizes buffers from hidden_size.
+ self.hidden_size = self.latent_size
+ self.num_attention_heads = 1
+ self.num_hidden_layers = 1
+ self.vocab_size = self.codebook_size
+ self.max_position_embeddings = int(kwargs.get("max_position_embeddings", 65536))
+
+ def get_text_config(self, **_: Any) -> NemotronVoiceChatCode2WavConfig:
+ return self
+
+
+class NemotronVoiceChatConfig(PretrainedConfig):
+ """Unified config: parses the raw NeMo dict into per-stage sub-configs."""
+
+ model_type = "nemotron_voicechat"
+
+ def __init__(self, **kwargs: Any) -> None:
+ nemo_model = kwargs.pop("model", None)
+ nemo_data = kwargs.pop("data", None)
+ thinker_text_config = kwargs.pop("thinker_text_config", None)
+ talker_config = kwargs.pop("talker_config", None)
+ code2wav_config = kwargs.pop("code2wav_config", None)
+ stt_cfg = kwargs.pop("stt_cfg", None)
+ tts_cfg = kwargs.pop("tts_cfg", None)
+ # Drop remaining NeMo training-only top-level sections.
+ for training_key in ("exp_manager", "trainer", "hf_export_dir", "name"):
+ kwargs.pop(training_key, None)
+ # transformers may call get_text_config() during super().__init__;
+ # give it a safe placeholder until the real sub-config is resolved.
+ self.thinker_text_config = None
+ super().__init__(**kwargs)
+
+ if nemo_model is not None:
+ stt = nemo_model.get("stt") or {}
+ speech_generation = nemo_model.get("speech_generation") or {}
+ if not (isinstance(stt, dict) and isinstance(stt.get("model"), dict)) or not (
+ isinstance(speech_generation, dict) and isinstance(speech_generation.get("model"), dict)
+ ):
+ raise ValueError(
+ "config.json is missing the NemotronVoiceChat NeMo signature keys "
+ "'model.stt.model' and/or 'model.speech_generation.model'. " + _EXPECTED_STRUCTURE
+ )
+ stt_cfg = stt["model"]
+ tts_cfg = speech_generation["model"]
+ self.inference_speaker_name = nemo_model.get("inference_speaker_name") or "Aria"
+ self.tts_data = speech_generation.get("data") or {}
+ elif stt_cfg is None or tts_cfg is None:
+ # Bare construction (e.g. transformers plumbing); leave a skeleton
+ # that fails loudly if a stage model tries to use it.
+ self.inference_speaker_name = kwargs.get("inference_speaker_name", "Aria")
+ self.tts_data = {}
+ if stt_cfg is None or tts_cfg is None:
+ self.stt_cfg: dict[str, Any] = {}
+ self.tts_cfg: dict[str, Any] = {}
+ else:
+ if "perception" not in stt_cfg:
+ raise ValueError("'model.stt.model.perception' is missing from config.json. " + _EXPECTED_STRUCTURE)
+ if "tts_config" not in tts_cfg or "codec_config" not in tts_cfg:
+ raise ValueError(
+ "'model.speech_generation.model.tts_config'/'codec_config' are missing "
+ "from config.json. " + _EXPECTED_STRUCTURE
+ )
+ self.stt_cfg = dict(stt_cfg)
+ self.tts_cfg = dict(tts_cfg)
+ self.nemo_data = dict(nemo_data) if isinstance(nemo_data, dict) else {}
+ self.frame_length = float(self.nemo_data.get("frame_length", 0.08))
+ self.source_sample_rate = int(self.nemo_data.get("source_sample_rate", 16000))
+ self.target_sample_rate = int(self.nemo_data.get("target_sample_rate", 22050))
+
+ self.thinker_text_config = self._resolve_thinker_text_config(thinker_text_config)
+ self.talker_config = self._coerce(
+ talker_config,
+ NemotronVoiceChatTalkerConfig,
+ defaults={"backbone_config": (self.tts_cfg.get("tts_config") or {}).get("backbone_config") or {}},
+ )
+ self.code2wav_config = self._coerce(
+ code2wav_config,
+ NemotronVoiceChatCode2WavConfig,
+ defaults={
+ "codec_config": self.tts_cfg.get("codec_config") or {},
+ "sample_rate": self.target_sample_rate,
+ },
+ )
+
+ if not getattr(self, "architectures", None):
+ self.architectures = ["NemotronVoiceChatThinkerForConditionalGeneration"]
+
+ @staticmethod
+ def _coerce(value: Any, cls: type[PretrainedConfig], defaults: dict[str, Any]) -> PretrainedConfig:
+ if isinstance(value, cls):
+ return value
+ if isinstance(value, PretrainedConfig):
+ value = value.to_dict()
+ if isinstance(value, dict):
+ merged = dict(defaults)
+ merged.update(value)
+ return cls(**merged)
+ return cls(**defaults)
+
+ def _resolve_thinker_text_config(self, provided: Any) -> PretrainedConfig:
+ """NemotronH backbone config: explicit > env-local > HF cache/hub."""
+ if isinstance(provided, PretrainedConfig):
+ return provided
+ if isinstance(provided, dict) and provided:
+ try:
+ return AutoConfig.for_model(**dict(provided))
+ except (KeyError, ValueError):
+ return PretrainedConfig(**provided)
+ if not self.stt_cfg:
+ # Skeleton config (bare construction); resolved lazily by the stage.
+ return PretrainedConfig()
+ llm_ref = os.environ.get(_LLM_PATH_ENV) or self.stt_cfg.get("pretrained_llm") or _DEFAULT_LLM_ID
+ try:
+ # nemotron_h is a native transformers model type; trust_remote_code=True
+ # would chase the auto_map .py modules and break air-gapped dirs that
+ # ship only config.json.
+ return AutoConfig.from_pretrained(llm_ref, trust_remote_code=False)
+ except Exception as exc:
+ raise RuntimeError(
+ f"Failed to load the NemotronVoiceChat thinker backbone config from {llm_ref!r}. "
+ f"Make sure the HF cache contains 'nvidia/NVIDIA-Nemotron-Nano-9B-v2' or set "
+ f"{_LLM_PATH_ENV} to a local directory with its config.json."
+ ) from exc
+
+ def get_text_config(self, **_: Any) -> PretrainedConfig:
+ text_config = self.__dict__.get("thinker_text_config")
+ return text_config if text_config is not None else self
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/__init__.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/__init__.py
new file mode 100644
index 000000000..d30a94542
--- /dev/null
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/__init__.py
@@ -0,0 +1,34 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat); dependency-stripped
+# copies of the NeMo inference modules needed to run NVIDIA-NemotronLabs-VoiceChat-11B.
+
+from .duplex_ear_tts import DuplexEARTTS
+from .ear_tts_commons import PreTrainedModel
+from .ear_tts_vae_codec import RVQVAEModel
+from .fusion import create_fusion_module
+from .perception import AudioPerceptionModule, IdentityConnector
+from .tokenizer import AutoTokenizer
+
+__all__ = [
+ "AudioPerceptionModule",
+ "AutoTokenizer",
+ "DuplexEARTTS",
+ "IdentityConnector",
+ "PreTrainedModel",
+ "RVQVAEModel",
+ "create_fusion_module",
+]
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/__init__.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/__init__.py
new file mode 100644
index 000000000..d1dda45b6
--- /dev/null
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/__init__.py
@@ -0,0 +1,25 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat); dependency-stripped
+# copies of the NeMo ASR conformer modules needed by the perception encoder.
+
+from .audio_preprocessing import AudioToMelSpectrogramPreprocessor
+from .conformer_encoder import ConformerEncoder
+
+__all__ = [
+ "AudioToMelSpectrogramPreprocessor",
+ "ConformerEncoder",
+]
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/audio_preprocessing.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/audio_preprocessing.py
new file mode 100644
index 000000000..e4343ef44
--- /dev/null
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/audio_preprocessing.py
@@ -0,0 +1,259 @@
+# SPDX-FileCopyrightText: Copyright (c) 2020 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/asr/modules/audio_preprocessing.py; training-only code removed.
+# Kept: AudioPreprocessor, AudioToMelSpectrogramPreprocessor (+ its config dataclass).
+# Removed: AudioToMFCCPreprocessor, SpectrogramAugmentation, MaskedPatchAugmentation,
+# CropOrPadSpectrogramAugmentation and numba/spec-augment imports (training-only).
+
+import math
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+from packaging import version
+
+from ..compat import (
+ Exportable,
+ NeuralModule,
+ logging,
+ logging_mode,
+ typecheck,
+)
+from .features import FilterbankFeatures, FilterbankFeaturesTA
+
+try:
+ import torchaudio
+ import torchaudio.functional
+ import torchaudio.transforms
+
+ TORCHAUDIO_VERSION = version.parse(torchaudio.__version__)
+ TORCHAUDIO_VERSION_MIN = version.parse("0.5")
+
+ HAVE_TORCHAUDIO = True
+except ModuleNotFoundError:
+ HAVE_TORCHAUDIO = False
+
+__all__ = [
+ "AudioToMelSpectrogramPreprocessor",
+]
+
+
+class AudioPreprocessor(NeuralModule, ABC):
+ """
+ An interface for Neural Modules that performs audio pre-processing,
+ transforming the wav files to features.
+ """
+
+ def __init__(self, win_length, hop_length):
+ super().__init__()
+
+ self.win_length = win_length
+ self.hop_length = hop_length
+
+ self.torch_windows = {
+ "hann": torch.hann_window,
+ "hamming": torch.hamming_window,
+ "blackman": torch.blackman_window,
+ "bartlett": torch.bartlett_window,
+ "ones": torch.ones,
+ None: torch.ones,
+ }
+
+ # Normally, when you call to(dtype) on a torch.nn.Module, all
+ # floating point parameters and buffers will change to that
+ # dtype, rather than being float32. The AudioPreprocessor
+ # classes, uniquely, don't actually have any parameters or
+ # buffers from what I see. In addition, we want the input to
+ # the preprocessor to be float32, but need to create the
+ # output in appropriate precision. We have this empty tensor
+ # here just to detect which dtype tensor this module should
+ # output at the end of execution.
+ self.register_buffer("dtype_sentinel_tensor", torch.tensor((), dtype=torch.float32), persistent=False)
+
+ @typecheck()
+ @torch.no_grad()
+ def forward(self, input_signal, length):
+ if input_signal.dtype != torch.float32:
+ logging.warning(
+ f"AudioPreprocessor received an input signal of dtype {input_signal.dtype}, rather than torch.float32. In sweeps across multiple datasets, we have found that the preprocessor is not robust to low precision mathematics. As such, it runs in float32. Your input will be cast to float32, but this is not necessarily enough to recovery full accuracy. For example, simply casting input_signal from torch.float32 to torch.bfloat16, then back to torch.float32 before running AudioPreprocessor causes drops in absolute WER of up to 0.1%. torch.bfloat16 simply does not have enough mantissa bits to represent enough values in the range [-1.0,+1.0] correctly.",
+ mode=logging_mode.ONCE,
+ )
+ processed_signal, processed_length = self.get_features(input_signal.to(torch.float32), length)
+ processed_signal = processed_signal.to(self.dtype_sentinel_tensor.dtype)
+ return processed_signal, processed_length
+
+ @abstractmethod
+ def get_features(self, input_signal, length):
+ # Called by forward(). Subclasses should implement this.
+ pass
+
+
+class AudioToMelSpectrogramPreprocessor(AudioPreprocessor, Exportable):
+ """Featurizer module that converts wavs to mel spectrograms.
+
+ Args:
+ sample_rate (int): Sample rate of the input audio data.
+ Defaults to 16000
+ window_size (float): Size of window for fft in seconds
+ Defaults to 0.02
+ window_stride (float): Stride of window for fft in seconds
+ Defaults to 0.01
+ n_window_size (int): Size of window for fft in samples
+ Defaults to None. Use one of window_size or n_window_size.
+ n_window_stride (int): Stride of window for fft in samples
+ Defaults to None. Use one of window_stride or n_window_stride.
+ window (str): Windowing function for fft. can be one of ['hann',
+ 'hamming', 'blackman', 'bartlett']
+ Defaults to "hann"
+ normalize (str): Can be one of ['per_feature', 'all_features']; all
+ other options disable feature normalization. 'all_features'
+ normalizes the entire spectrogram to be mean 0 with std 1.
+ 'pre_features' normalizes per channel / freq instead.
+ Defaults to "per_feature"
+ n_fft (int): Length of FT window. If None, it uses the smallest power
+ of 2 that is larger than n_window_size.
+ Defaults to None
+ preemph (float): Amount of pre emphasis to add to audio. Can be
+ disabled by passing None.
+ Defaults to 0.97
+ features (int): Number of mel spectrogram freq bins to output.
+ Defaults to 64
+ lowfreq (int): Lower bound on mel basis in Hz.
+ Defaults to 0
+ highfreq (int): Lower bound on mel basis in Hz.
+ Defaults to None
+ log (bool): Log features.
+ Defaults to True
+ log_zero_guard_type(str): Need to avoid taking the log of zero. There
+ are two options: "add" or "clamp".
+ Defaults to "add".
+ log_zero_guard_value(float, or str): Add or clamp requires the number
+ to add with or clamp to. log_zero_guard_value can either be a float
+ or "tiny" or "eps". torch.finfo is used if "tiny" or "eps" is
+ passed.
+ Defaults to 2**-24.
+ dither (float): Amount of white-noise dithering.
+ Defaults to 1e-5
+ pad_to (int): Ensures that the output size of the time dimension is
+ a multiple of pad_to.
+ Defaults to 16
+ frame_splicing (int): Defaults to 1
+ exact_pad (bool): If True, sets stft center to False and adds padding, such that num_frames = audio_length
+ // hop_length. Defaults to False.
+ pad_value (float): The value that shorter mels are padded with.
+ Defaults to 0
+ mag_power (float): The power that the linear spectrogram is raised to
+ prior to multiplication with mel basis.
+ Defaults to 2 for a power spec
+ rng : Random number generator
+ nb_augmentation_prob (float) : Probability with which narrowband augmentation would be applied to
+ samples in the batch.
+ Defaults to 0.0
+ nb_max_freq (int) : Frequency above which all frequencies will be masked for narrowband augmentation.
+ Defaults to 4000
+ use_torchaudio: Whether to use the `torchaudio` implementation.
+ mel_norm: Normalization used for mel filterbank weights.
+ Defaults to 'slaney' (area normalization)
+ stft_exact_pad: Deprecated argument, kept for compatibility with older checkpoints.
+ stft_conv: Deprecated argument, kept for compatibility with older checkpoints.
+ """
+
+ def save_to(self, save_path: str):
+ pass
+
+ @classmethod
+ def restore_from(cls, restore_path: str):
+ pass
+
+ def __init__(
+ self,
+ sample_rate=16000,
+ window_size=0.02,
+ window_stride=0.01,
+ n_window_size=None,
+ n_window_stride=None,
+ window="hann",
+ normalize="per_feature",
+ n_fft=None,
+ preemph=0.97,
+ features=64,
+ lowfreq=0,
+ highfreq=None,
+ log=True,
+ log_zero_guard_type="add",
+ log_zero_guard_value=2**-24,
+ dither=1e-5,
+ pad_to=16,
+ frame_splicing=1,
+ exact_pad=False,
+ pad_value=0,
+ mag_power=2.0,
+ rng=None,
+ nb_augmentation_prob=0.0,
+ nb_max_freq=4000,
+ use_torchaudio: bool = False,
+ mel_norm="slaney",
+ stft_exact_pad=False, # Deprecated arguments; kept for config compatibility
+ stft_conv=False, # Deprecated arguments; kept for config compatibility
+ ):
+ super().__init__(n_window_size, n_window_stride)
+
+ self._sample_rate = sample_rate
+ if window_size and n_window_size:
+ raise ValueError(f"{self} received both window_size and n_window_size. Only one should be specified.")
+ if window_stride and n_window_stride:
+ raise ValueError(f"{self} received both window_stride and n_window_stride. Only one should be specified.")
+ if window_size:
+ n_window_size = int(window_size * self._sample_rate)
+ if window_stride:
+ n_window_stride = int(window_stride * self._sample_rate)
+
+ # Given the long and similar argument list, point to the class and instantiate it by reference
+ if not use_torchaudio:
+ featurizer_class = FilterbankFeatures
+ else:
+ featurizer_class = FilterbankFeaturesTA
+ self.featurizer = featurizer_class(
+ sample_rate=self._sample_rate,
+ n_window_size=n_window_size,
+ n_window_stride=n_window_stride,
+ window=window,
+ normalize=normalize,
+ n_fft=n_fft,
+ preemph=preemph,
+ nfilt=features,
+ lowfreq=lowfreq,
+ highfreq=highfreq,
+ log=log,
+ log_zero_guard_type=log_zero_guard_type,
+ log_zero_guard_value=log_zero_guard_value,
+ dither=dither,
+ pad_to=pad_to,
+ frame_splicing=frame_splicing,
+ exact_pad=exact_pad,
+ pad_value=pad_value,
+ mag_power=mag_power,
+ rng=rng,
+ nb_augmentation_prob=nb_augmentation_prob,
+ nb_max_freq=nb_max_freq,
+ mel_norm=mel_norm,
+ stft_exact_pad=stft_exact_pad, # Deprecated arguments; kept for config compatibility
+ stft_conv=stft_conv, # Deprecated arguments; kept for config compatibility
+ )
+
+ def get_features(self, input_signal, length):
+ return self.featurizer(input_signal, length)
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/causal_convs.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/causal_convs.py
new file mode 100644
index 000000000..4d593895e
--- /dev/null
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/causal_convs.py
@@ -0,0 +1,154 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/asr/parts/submodules/causal_convs.py; training-only code removed.
+
+from typing import Union
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+__all__ = ["CausalConv2D", "CausalConv1D"]
+
+
+class CausalConv2D(nn.Conv2d):
+ """
+ A causal version of nn.Conv2d where each location in the 2D matrix would have no access to locations on its right or down
+ All arguments are the same as nn.Conv2d except padding which should be set as None
+ """
+
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ kernel_size: int,
+ stride: int = 1,
+ padding: Union[str, int] = 0,
+ dilation: int = 1,
+ groups: int = 1,
+ bias: bool = True,
+ padding_mode: str = "zeros",
+ device=None,
+ dtype=None,
+ ) -> None:
+ if padding is not None:
+ raise ValueError("Argument padding should be set to None for CausalConv2D.")
+ self._left_padding = kernel_size - 1
+ self._right_padding = stride - 1
+
+ padding = 0
+ super(CausalConv2D, self).__init__(
+ in_channels,
+ out_channels,
+ kernel_size,
+ stride,
+ padding,
+ dilation,
+ groups,
+ bias,
+ padding_mode,
+ device,
+ dtype,
+ )
+
+ def forward(
+ self,
+ x,
+ ):
+ x = F.pad(x, pad=(self._left_padding, self._right_padding, self._left_padding, self._right_padding))
+ x = super().forward(x)
+ return x
+
+
+class CausalConv1D(nn.Conv1d):
+ """
+ A causal version of nn.Conv1d where each step would have limited access to locations on its right or left
+ All arguments are the same as nn.Conv1d except padding.
+
+ If padding is set None, then paddings are set automatically to make it a causal convolution where each location would not see any steps on its right.
+
+ If padding is set as a list (size of 2), then padding[0] would be used as left padding and padding[1] as right padding.
+ It would make it possible to control the number of steps to be accessible on the right and left.
+ This mode is not supported when stride > 1. padding[0]+padding[1] should be equal to (kernel_size - 1).
+ """
+
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ kernel_size: int,
+ stride: int = 1,
+ padding: Union[str, int] = 0,
+ dilation: int = 1,
+ groups: int = 1,
+ bias: bool = True,
+ padding_mode: str = "zeros",
+ device=None,
+ dtype=None,
+ ) -> None:
+ self.cache_drop_size = None
+ if padding is None:
+ self._left_padding = kernel_size - 1
+ self._right_padding = stride - 1
+ else:
+ if stride != 1 and padding != kernel_size - 1:
+ raise ValueError("No striding allowed for non-symmetric convolutions!")
+ if isinstance(padding, int):
+ self._left_padding = padding
+ self._right_padding = padding
+ elif isinstance(padding, list) and len(padding) == 2 and padding[0] + padding[1] == kernel_size - 1:
+ self._left_padding = padding[0]
+ self._right_padding = padding[1]
+ else:
+ raise ValueError(f"Invalid padding param: {padding}!")
+
+ self._max_cache_len = self._left_padding
+
+ super(CausalConv1D, self).__init__(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=kernel_size,
+ stride=stride,
+ padding=0,
+ dilation=dilation,
+ groups=groups,
+ bias=bias,
+ padding_mode=padding_mode,
+ device=device,
+ dtype=dtype,
+ )
+
+ def update_cache(self, x, cache=None):
+ if cache is None:
+ new_x = F.pad(x, pad=(self._left_padding, self._right_padding))
+ next_cache = cache
+ else:
+ new_x = F.pad(x, pad=(0, self._right_padding))
+ new_x = torch.cat([cache, new_x], dim=-1)
+ if self.cache_drop_size > 0:
+ next_cache = new_x[:, :, : -self.cache_drop_size]
+ else:
+ next_cache = new_x
+ next_cache = next_cache[:, :, -cache.size(-1) :]
+ return new_x, next_cache
+
+ def forward(self, x, cache=None):
+ x, cache = self.update_cache(x, cache=cache)
+ x = super().forward(x)
+ if cache is None:
+ return x
+ else:
+ return x, cache
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/conformer_encoder.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/conformer_encoder.py
new file mode 100644
index 000000000..b6f8531a8
--- /dev/null
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/conformer_encoder.py
@@ -0,0 +1,850 @@
+# SPDX-FileCopyrightText: Copyright (c) 2020 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/asr/modules/conformer_encoder.py; training-only code removed.
+
+import math
+import random
+from collections import OrderedDict
+from dataclasses import dataclass
+from typing import List, Optional, Set, Tuple
+
+import torch
+import torch.distributed
+from omegaconf import DictConfig, ListConfig, open_dict
+from torch import nn
+
+from ..compat import (
+ AccessMixin,
+ CacheAwareStreamingConfig,
+ Exportable,
+ NeuralModule,
+ StreamingEncoder,
+ logging,
+ typecheck,
+)
+from .causal_convs import CausalConv1D
+from .conformer_modules import ConformerLayer
+from .multi_head_attention import (
+ LocalAttRelPositionalEncoding,
+ MultiHeadAttention,
+ PositionalEncoding,
+ RelPositionalEncoding,
+ RelPositionMultiHeadAttention,
+ RelPositionMultiHeadAttentionLongformer,
+)
+from .regularization_utils import compute_stochastic_depth_drop_probs
+from .subsampling import (
+ ConvSubsampling,
+ StackingSubsampling,
+ SubsamplingReductionModule,
+)
+
+__all__ = ["ConformerEncoder"]
+
+
+class ConformerEncoder(NeuralModule, StreamingEncoder, Exportable, AccessMixin):
+ """
+ The encoder for ASR model of Conformer.
+ Based on this paper:
+ 'Conformer: Convolution-augmented Transformer for Speech Recognition' by Anmol Gulati et al.
+ https://arxiv.org/abs/2005.08100
+
+ Args:
+ feat_in (int): the size of feature channels
+ n_layers (int): number of layers of ConformerBlock
+ d_model (int): the hidden size of the model
+ feat_out (int): the size of the output features
+ Defaults to -1 (means feat_out is d_model)
+ subsampling (str): the method of subsampling:
+ choices = ['vggnet', 'striding', 'dw-striding', 'stacking', 'stacking_norm']
+ Defaults to striding.
+ subsampling_factor (int): the subsampling factor which should be power of 2
+ Defaults to 4.
+ subsampling_conv_chunking_factor(int): optionally, force chunk inputs (helpful for large inputs)
+ Should be power of 2, 1 (auto-chunking, default), or -1 (no chunking)
+ subsampling_conv_channels (int): the size of the convolutions in the subsampling module
+ Defaults to -1 which would set it to d_model.
+ reduction (str, Optional): the method of reduction, choices=['pooling', 'striding']. If no value
+ is passed, then no reduction is performed and the models runs with the original 4x subsampling.
+ reduction_position (int, Optional): the index of the layer to apply reduction. If -1, apply reduction
+ at the end.
+ reduction_factor (int): the reduction factor which should be either 1 or a power of 2
+ Defaults to 1.
+ ff_expansion_factor (int): the expansion factor in feed forward layers
+ Defaults to 4.
+ self_attention_model (str): the type of the attention layer and positional encoding.
+
+ 'rel_pos':
+ relative positional embedding and Transformer-XL
+ 'rel_pos_local_attn':
+ relative positional embedding and Transformer-XL with local attention using
+ overlapping chunks. Attention context is determined by att_context_size parameter.
+ 'abs_pos':
+ absolute positional embedding and Transformer
+
+ Default is rel_pos.
+ pos_emb_max_len (int): the maximum length of positional embeddings
+ Defaults to 5000
+ n_heads (int): number of heads in multi-headed attention layers
+ Defaults to 4.
+ att_context_size (List[Union[List[int],int]]): specifies the context sizes on each side.
+ Each context size should be a list of two integers like `[100, 100]`.
+ A list of context sizes like `[[100,100]`, `[100,50]]` can also be passed. -1 means unlimited context.
+ Defaults to `[-1, -1]`
+ att_context_probs (List[float]): a list of probabilities of each one of the att_context_size
+ when a list of them is passed. If not specified, uniform distribution is being used.
+ Defaults to None
+ att_context_style (str): 'regular' or 'chunked_limited'.
+ Defaults to 'regular'
+ xscaling (bool): enables scaling the inputs to the multi-headed attention layers by `sqrt(d_model)`.
+ Defaults to True.
+ untie_biases (bool): whether to not share (untie) the bias weights between layers of Transformer-XL
+ Defaults to True.
+ conv_kernel_size (int): the size of the convolutions in the convolutional modules
+ Defaults to 31.
+ conv_norm_type (str): the type of the normalization in the convolutional modules
+ Defaults to 'batch_norm'.
+ conv_context_size (list): it can be"causal" or a list of two integers
+ while `conv_context_size[0]+conv_context_size[1]+1==conv_kernel_size`.
+ `None` means `[(conv_kernel_size-1)//2`, `(conv_kernel_size-1)//2]`, and 'causal' means
+ `[(conv_kernel_size-1), 0]`.
+ Defaults to None.
+ conv_dual_mode (bool): specifies if convolution should be dual mode when dual_offline mode is being used.
+ When enables, the left half of the convolution kernel would get masked in streaming cases.
+ Defaults to False.
+ use_bias (bool): Use bias in all Linear and Conv1d layers from each ConformerLayer to improve
+ activation flow and stabilize training of huge models.
+ Defaults to True.
+ dropout (float): the dropout rate used in all layers except the attention layers
+ Defaults to 0.1.
+ dropout_pre_encoder (float): the dropout rate used before the encoder
+ Defaults to 0.1.
+ dropout_emb (float): the dropout rate used for the positional embeddings
+ Defaults to 0.1.
+ dropout_att (float): the dropout rate used for the attention layer
+ Defaults to 0.0.
+ stochastic_depth_drop_prob (float): if non-zero, will randomly drop
+ layers during training. The higher this value, the more often layers
+ are dropped. Defaults to 0.0.
+ stochastic_depth_mode (str): can be either "linear" or "uniform". If
+ set to "uniform", all layers have the same probability of drop. If
+ set to "linear", the drop probability grows linearly from 0 for the
+ first layer to the desired value for the final layer. Defaults to
+ "linear".
+ stochastic_depth_start_layer (int): starting layer for stochastic depth.
+ All layers before this will never be dropped. Note that drop
+ probability will be adjusted accordingly if mode is "linear" when
+ start layer is > 1. Defaults to 1.
+ global_tokens (int): number of tokens to be used for global attention.
+ Only relevant if self_attention_model is 'rel_pos_local_attn'.
+ Defaults to 0.
+ global_tokens_spacing (int): how far apart the global tokens are
+ Defaults to 1.
+ global_attn_separate (bool): whether the q, k, v layers used for global tokens should be separate.
+ Defaults to False.
+ use_pytorch_sdpa (bool): use torch sdpa instead of manual attention.
+ Defaults to False.
+ use_pytorch_sdpa_backends (list[str]): list of backend names to use in sdpa.
+ None or empty list means all backends. e.g. ["MATH"]
+ Defaults to None.
+ bypass_pre_encode: if True, skip the pre-encoder module and the `audio_signal` should be pre-encoded
+ embeddings. The `audio_signal` input supports two formats depending on the `bypass_pre_encode`
+ boolean flag. This determines the required format of the input variable `audio_signal`.
+ Defaults to `bypass_pre_encode=False`. `bypass_pre_encode=True` is used for the cases
+ where frame-level, context-independent embeddings are needed to be saved or reused.
+ (e.g., speaker cache in streaming speaker diarization)
+ sync_max_audio_length (bool): when true, performs NCCL all_reduce to allocate the same amount of memory for
+ positional encoding buffers on all GPUs. Disabling this setting may help with deadlocks in certain
+ scenarios such as model parallelism, or generally when this module is not being ran on some GPUs
+ as a part of the training step.
+ """
+
+ def __init__(
+ self,
+ feat_in,
+ n_layers,
+ d_model,
+ feat_out=-1,
+ causal_downsampling=False,
+ subsampling="striding",
+ subsampling_factor=4,
+ subsampling_conv_chunking_factor=1,
+ subsampling_conv_channels=-1,
+ reduction=None,
+ reduction_position=None,
+ reduction_factor=1,
+ ff_expansion_factor=4,
+ self_attention_model="rel_pos",
+ n_heads=4,
+ att_context_size=None,
+ att_context_probs=None,
+ att_context_style="regular",
+ xscaling=True,
+ untie_biases=True,
+ pos_emb_max_len=5000,
+ conv_kernel_size=31,
+ conv_norm_type="batch_norm",
+ conv_context_size=None,
+ use_bias=True,
+ dropout=0.1,
+ dropout_pre_encoder=0.1,
+ dropout_emb=0.1,
+ dropout_att=0.0,
+ stochastic_depth_drop_prob: float = 0.0,
+ stochastic_depth_mode: str = "linear",
+ stochastic_depth_start_layer: int = 1,
+ global_tokens: int = 0,
+ global_tokens_spacing: int = 1,
+ global_attn_separate: bool = False,
+ use_pytorch_sdpa: bool = False,
+ use_pytorch_sdpa_backends=None,
+ sync_max_audio_length: bool = True,
+ ):
+ super().__init__()
+ d_ff = d_model * ff_expansion_factor
+ self.d_model = d_model
+ self.n_layers = n_layers
+ self._feat_in = feat_in
+ self.att_context_style = att_context_style
+ self.subsampling_factor = subsampling_factor
+ self.subsampling_conv_chunking_factor = subsampling_conv_chunking_factor
+
+ self.self_attention_model = self_attention_model
+ self.global_tokens = global_tokens
+ self.global_attn_separate = global_attn_separate
+ self.global_tokens_spacing = global_tokens_spacing
+ self.use_pytorch_sdpa = use_pytorch_sdpa
+ if use_pytorch_sdpa_backends is None:
+ use_pytorch_sdpa_backends = []
+ self.use_pytorch_sdpa_backends = use_pytorch_sdpa_backends
+ self.sync_max_audio_length = sync_max_audio_length
+
+ # Setting up the att_context_size
+ (
+ self.att_context_size_all,
+ self.att_context_size,
+ self.att_context_probs,
+ self.conv_context_size,
+ ) = self._calc_context_sizes(
+ att_context_style=att_context_style,
+ att_context_size=att_context_size,
+ att_context_probs=att_context_probs,
+ conv_context_size=conv_context_size,
+ conv_kernel_size=conv_kernel_size,
+ )
+
+ if xscaling:
+ self.xscale = math.sqrt(d_model)
+ else:
+ self.xscale = None
+
+ # Subsampling
+ if subsampling_conv_channels == -1:
+ subsampling_conv_channels = d_model
+ if subsampling and subsampling_factor > 1:
+ if subsampling in ["stacking", "stacking_norm"]:
+ # stacking_norm has an extra layer norm after stacking comparing to stacking
+ self.pre_encode = StackingSubsampling(
+ subsampling_factor=subsampling_factor,
+ feat_in=feat_in,
+ feat_out=d_model,
+ norm=True if subsampling == "stacking_norm" else False,
+ )
+ else:
+ self.pre_encode = ConvSubsampling(
+ subsampling=subsampling,
+ subsampling_factor=subsampling_factor,
+ feat_in=feat_in,
+ feat_out=d_model,
+ conv_channels=subsampling_conv_channels,
+ subsampling_conv_chunking_factor=subsampling_conv_chunking_factor,
+ activation=nn.ReLU(True),
+ is_causal=causal_downsampling,
+ )
+ else:
+ self.pre_encode = nn.Linear(feat_in, d_model)
+
+ # Reduction
+ if reduction and reduction_factor > 1:
+ assert reduction_position >= -1 and reduction_position < n_layers
+ self.reduction_subsampling = SubsamplingReductionModule(
+ reduction=reduction,
+ d_model=d_model,
+ reduction_factor=reduction_factor,
+ )
+ self.reduction_position = reduction_position
+ else:
+ self.reduction_subsampling = None
+ self.reduction_position = None
+
+ self._feat_out = d_model
+
+ # Biases for relative positional encoding
+ if not untie_biases and self_attention_model == "rel_pos":
+ d_head = d_model // n_heads
+ pos_bias_u = nn.Parameter(torch.Tensor(n_heads, d_head))
+ pos_bias_v = nn.Parameter(torch.Tensor(n_heads, d_head))
+ nn.init.zeros_(pos_bias_u)
+ nn.init.zeros_(pos_bias_v)
+ else:
+ pos_bias_u = None
+ pos_bias_v = None
+
+ # Positional encodings
+ self.pos_emb_max_len = pos_emb_max_len
+ if self_attention_model == "rel_pos":
+ self.pos_enc = RelPositionalEncoding(
+ d_model=d_model,
+ dropout_rate=dropout_pre_encoder,
+ max_len=pos_emb_max_len,
+ xscale=self.xscale,
+ dropout_rate_emb=dropout_emb,
+ )
+ elif self_attention_model == "rel_pos_local_attn":
+ if max(att_context_size) <= 0:
+ raise ValueError("When using local attention, context size must be set > 0")
+ self.pos_enc = LocalAttRelPositionalEncoding(
+ att_context_size=att_context_size,
+ d_model=d_model,
+ dropout_rate=dropout,
+ max_len=pos_emb_max_len,
+ xscale=self.xscale,
+ dropout_rate_emb=dropout_emb,
+ )
+ elif self_attention_model == "abs_pos":
+ pos_bias_u = None
+ pos_bias_v = None
+ self.pos_enc = PositionalEncoding(
+ d_model=d_model, dropout_rate=dropout_pre_encoder, max_len=pos_emb_max_len, xscale=self.xscale
+ )
+ else:
+ raise ValueError(f"Not valid self_attention_model: '{self_attention_model}'!")
+
+ self.layers = nn.ModuleList()
+ for i in range(n_layers):
+ layer = ConformerLayer(
+ d_model=d_model,
+ d_ff=d_ff,
+ self_attention_model=self_attention_model,
+ global_tokens=global_tokens,
+ global_tokens_spacing=global_tokens_spacing,
+ global_attn_separate=global_attn_separate,
+ n_heads=n_heads,
+ conv_kernel_size=conv_kernel_size,
+ conv_norm_type=conv_norm_type,
+ conv_context_size=self.conv_context_size,
+ dropout=dropout,
+ dropout_att=dropout_att,
+ pos_bias_u=pos_bias_u,
+ pos_bias_v=pos_bias_v,
+ att_context_size=self.att_context_size,
+ use_bias=use_bias,
+ use_pytorch_sdpa=self.use_pytorch_sdpa,
+ use_pytorch_sdpa_backends=self.use_pytorch_sdpa_backends,
+ )
+ self.layers.append(layer)
+
+ if feat_out > 0 and feat_out != self._feat_out:
+ self.out_proj = nn.Linear(self._feat_out, feat_out)
+ self._feat_out = feat_out
+ else:
+ self.out_proj = None
+ self._feat_out = d_model
+ self.set_max_audio_length(self.pos_emb_max_len)
+ self.use_pad_mask = True
+
+ self.setup_streaming_params()
+ self.export_cache_support = False
+
+ self.layer_drop_probs = compute_stochastic_depth_drop_probs(
+ len(self.layers), stochastic_depth_drop_prob, stochastic_depth_mode, stochastic_depth_start_layer
+ )
+ # will be set in self.forward() if defined in AccessMixin config
+ self.interctc_capture_at_layers = None
+
+ @typecheck()
+ def forward(
+ self,
+ audio_signal,
+ length,
+ cache_last_channel=None,
+ cache_last_time=None,
+ cache_last_channel_len=None,
+ bypass_pre_encode=False,
+ ):
+ """
+ Forward function for the ConformerEncoder accepting an audio signal and its corresponding length.
+ The `audio_signal` input supports two formats depending on the `bypass_pre_encode` boolean flag.
+ This determines the required format of the input variable `audio_signal`:
+ (1) bypass_pre_encode = False (default):
+ `audio_signal` must be a tensor containing audio features.
+ Shape: (batch, self._feat_in, n_frames)
+ (2) bypass_pre_encode = True:
+ `audio_signal` must be a tensor containing pre-encoded embeddings.
+ Shape: (batch, n_frame, self.d_model)
+ """
+ if not bypass_pre_encode and audio_signal.shape[-2] != self._feat_in:
+ raise ValueError(
+ f"If bypass_pre_encode is False, audio_signal should have shape "
+ f"(batch, {self._feat_in}, n_frame) but got last dimension {audio_signal.shape[-2]}."
+ )
+ if bypass_pre_encode and audio_signal.shape[-1] != self.d_model:
+ raise ValueError(
+ f"If bypass_pre_encode is True, audio_signal should have shape "
+ f"(batch, n_frame, {self.d_model}) but got last dimension {audio_signal.shape[-1]}."
+ )
+
+ if bypass_pre_encode:
+ self.update_max_seq_length(
+ seq_length=audio_signal.size(2) * self.subsampling_factor, device=audio_signal.device
+ )
+ else:
+ self.update_max_seq_length(seq_length=audio_signal.size(2), device=audio_signal.device)
+ return self.forward_internal(
+ audio_signal,
+ length,
+ cache_last_channel=cache_last_channel,
+ cache_last_time=cache_last_time,
+ cache_last_channel_len=cache_last_channel_len,
+ bypass_pre_encode=bypass_pre_encode,
+ )
+
+ def forward_internal(
+ self,
+ audio_signal,
+ length,
+ cache_last_channel=None,
+ cache_last_time=None,
+ cache_last_channel_len=None,
+ bypass_pre_encode=False,
+ ):
+ """
+ The `audio_signal` input supports two formats depending on the `bypass_pre_encode` boolean flag.
+ This determines the required format of the input variable `audio_signal`:
+ (1) bypass_pre_encode = False (default):
+ `audio_signal` must be a tensor containing audio features.
+ Shape: (batch, self._feat_in, n_frames)
+ (2) bypass_pre_encode = True:
+ `audio_signal` must be a tensor containing pre-encoded embeddings.
+ Shape: (batch, n_frame, self.d_model)
+
+ `bypass_pre_encode=True` is used in cases where frame-level, context-independent embeddings are
+ needed to be saved or reused (e.g., speaker cache in streaming speaker diarization).
+ """
+ if length is None:
+ length = audio_signal.new_full(
+ (audio_signal.size(0),), audio_signal.size(-1), dtype=torch.int64, device=audio_signal.device
+ )
+
+ # select a random att_context_size with the distribution specified by att_context_probs during training
+ # for non-validation cases like test, validation or inference, it uses the first mode in self.att_context_size
+ if self.training and len(self.att_context_size_all) > 1:
+ cur_att_context_size = random.choices(self.att_context_size_all, weights=self.att_context_probs)[0]
+ else:
+ cur_att_context_size = self.att_context_size
+
+ if not bypass_pre_encode:
+ audio_signal = torch.transpose(audio_signal, 1, 2)
+
+ if isinstance(self.pre_encode, nn.Linear):
+ audio_signal = self.pre_encode(audio_signal)
+ else:
+ audio_signal, length = self.pre_encode(x=audio_signal, lengths=length)
+ length = length.to(torch.int64)
+ # `self.streaming_cfg` is set by setup_streaming_cfg(), called in the init
+ if self.streaming_cfg.drop_extra_pre_encoded > 0 and cache_last_channel is not None:
+ audio_signal = audio_signal[:, self.streaming_cfg.drop_extra_pre_encoded :, :]
+ length = (length - self.streaming_cfg.drop_extra_pre_encoded).clamp(min=0)
+
+ if self.reduction_position is not None and cache_last_channel is not None:
+ raise ValueError("Caching with reduction feature is not supported yet!")
+
+ max_audio_length = audio_signal.size(1)
+ if cache_last_channel is not None:
+ cache_len = self.streaming_cfg.last_channel_cache_size
+ cache_keep_size = max_audio_length - self.streaming_cfg.cache_drop_size
+ max_audio_length = max_audio_length + cache_len
+ padding_length = length + cache_len
+ offset = torch.neg(cache_last_channel_len) + cache_len
+ else:
+ padding_length = length
+ cache_last_channel_next = None
+ cache_len = 0
+ offset = None
+
+ audio_signal, pos_emb = self.pos_enc(x=audio_signal, cache_len=cache_len)
+
+ # Create the self-attention and padding masks
+ pad_mask, att_mask = self._create_masks(
+ att_context_size=cur_att_context_size,
+ padding_length=padding_length,
+ max_audio_length=max_audio_length,
+ offset=offset,
+ device=audio_signal.device,
+ )
+
+ if cache_last_channel is not None:
+ pad_mask = pad_mask[:, cache_len:]
+ if att_mask is not None:
+ att_mask = att_mask[:, cache_len:]
+ # Convert caches from the tensor to list
+ cache_last_time_next = []
+ cache_last_channel_next = []
+
+ for lth, (drop_prob, layer) in enumerate(zip(self.layer_drop_probs, self.layers)):
+ original_signal = audio_signal
+ if cache_last_channel is not None:
+ cache_last_channel_cur = cache_last_channel[lth]
+ cache_last_time_cur = cache_last_time[lth]
+ else:
+ cache_last_channel_cur = None
+ cache_last_time_cur = None
+ audio_signal = layer(
+ x=audio_signal,
+ att_mask=att_mask,
+ pos_emb=pos_emb,
+ pad_mask=pad_mask,
+ cache_last_channel=cache_last_channel_cur,
+ cache_last_time=cache_last_time_cur,
+ )
+
+ if cache_last_channel_cur is not None:
+ (audio_signal, cache_last_channel_cur, cache_last_time_cur) = audio_signal
+ cache_last_channel_next.append(cache_last_channel_cur)
+ cache_last_time_next.append(cache_last_time_cur)
+
+ # applying stochastic depth logic from https://arxiv.org/abs/2102.03216
+ if self.training and drop_prob > 0.0:
+ should_drop = torch.rand(1) < drop_prob
+ # adjusting to match expectation
+ if should_drop:
+ # that's not efficient, but it's hard to implement distributed
+ # version of dropping layers without deadlock or random seed meddling
+ # so multiplying the signal by 0 to ensure all weights get gradients
+ audio_signal = audio_signal * 0.0 + original_signal
+ else:
+ # not doing this operation if drop prob is 0 as it's identity in that case
+ audio_signal = (audio_signal - original_signal) / (1.0 - drop_prob) + original_signal
+
+ if self.reduction_position == lth:
+ audio_signal, length = self.reduction_subsampling(x=audio_signal, lengths=length)
+ max_audio_length = audio_signal.size(1)
+ # Don't update the audio_signal here because then it will again scale the audio_signal
+ # and cause an increase in the WER
+ _, pos_emb = self.pos_enc(x=audio_signal, cache_len=cache_len)
+ pad_mask, att_mask = self._create_masks(
+ att_context_size=cur_att_context_size,
+ padding_length=length,
+ max_audio_length=max_audio_length,
+ offset=offset,
+ device=audio_signal.device,
+ )
+
+ # saving tensors if required for interctc loss
+ if self.is_access_enabled(getattr(self, "model_guid", None)):
+ if self.interctc_capture_at_layers is None:
+ self.interctc_capture_at_layers = self.access_cfg.get("interctc", {}).get("capture_layers", [])
+ if lth in self.interctc_capture_at_layers:
+ lth_audio_signal = audio_signal
+ if self.out_proj is not None:
+ lth_audio_signal = self.out_proj(audio_signal)
+ # shape is the same as the shape of audio_signal output, i.e. [B, D, T]
+ self.register_accessible_tensor(
+ name=f"interctc/layer_output_{lth}", tensor=torch.transpose(lth_audio_signal, 1, 2)
+ )
+ self.register_accessible_tensor(name=f"interctc/layer_length_{lth}", tensor=length)
+
+ if self.out_proj is not None:
+ audio_signal = self.out_proj(audio_signal)
+
+ # Reduction
+ if self.reduction_position == -1:
+ audio_signal, length = self.reduction_subsampling(x=audio_signal, lengths=length)
+
+ audio_signal = torch.transpose(audio_signal, 1, 2)
+ length = length.to(dtype=torch.int64)
+
+ if cache_last_channel is not None:
+ cache_last_channel_next = torch.stack(cache_last_channel_next, dim=0)
+ cache_last_time_next = torch.stack(cache_last_time_next, dim=0)
+ return (
+ audio_signal,
+ length,
+ cache_last_channel_next,
+ cache_last_time_next,
+ torch.clamp(cache_last_channel_len + cache_keep_size, max=cache_len),
+ )
+ else:
+ return audio_signal, length
+
+ def update_max_seq_length(self, seq_length: int, device):
+ """
+ Updates the maximum sequence length for the model.
+
+ Args:
+ seq_length (int): New maximum sequence length.
+ device (torch.device): Device to use for computations.
+ """
+ # Find global max audio length across all nodes
+ if self.sync_max_audio_length and torch.distributed.is_initialized():
+ global_max_len = torch.tensor([seq_length], dtype=torch.float32, device=device)
+
+ # Update across all ranks in the distributed system
+ torch.distributed.all_reduce(global_max_len, op=torch.distributed.ReduceOp.MAX)
+
+ seq_length = global_max_len.int().item()
+
+ if seq_length > self.max_audio_length:
+ self.set_max_audio_length(seq_length)
+
+ def set_max_audio_length(self, max_audio_length):
+ """
+ Sets maximum input length.
+ Pre-calculates internal seq_range mask.
+
+ Args:
+ max_audio_length (int): New maximum sequence length.
+ """
+ self.max_audio_length = max_audio_length
+ device = next(self.parameters()).device
+ dtype = next(self.parameters()).dtype
+ self.pos_enc.extend_pe(max_audio_length, device, dtype)
+
+ def _create_masks(self, att_context_size, padding_length, max_audio_length, offset, device):
+ if self.self_attention_model != "rel_pos_local_attn":
+ att_mask = torch.ones(1, max_audio_length, max_audio_length, dtype=torch.bool, device=device)
+
+ if self.att_context_style == "regular":
+ if att_context_size[0] >= 0:
+ att_mask = att_mask.triu(diagonal=-att_context_size[0])
+ if att_context_size[1] >= 0:
+ att_mask = att_mask.tril(diagonal=att_context_size[1])
+ elif self.att_context_style == "chunked_limited":
+ # When right context is unlimited, just the left side of the masking need to get updated
+ if att_context_size[1] == -1:
+ if att_context_size[0] >= 0:
+ att_mask = att_mask.triu(diagonal=-att_context_size[0])
+ else:
+ chunk_size = att_context_size[1] + 1
+ # left_chunks_num specifies the number of chunks to be visible by each chunk on the left side
+ if att_context_size[0] >= 0:
+ left_chunks_num = att_context_size[0] // chunk_size
+ else:
+ left_chunks_num = 10000
+
+ chunk_idx = torch.arange(0, max_audio_length, dtype=torch.int, device=att_mask.device)
+ chunk_idx = torch.div(chunk_idx, chunk_size, rounding_mode="trunc")
+ diff_chunks = chunk_idx.unsqueeze(1) - chunk_idx.unsqueeze(0)
+ chunked_limited_mask = torch.logical_and(
+ torch.le(diff_chunks, left_chunks_num), torch.ge(diff_chunks, 0)
+ )
+ att_mask = torch.logical_and(att_mask, chunked_limited_mask.unsqueeze(0))
+ else:
+ att_mask = None
+
+ # pad_mask is the masking to be used to ignore paddings
+ pad_mask = torch.arange(0, max_audio_length, device=device).expand(
+ padding_length.size(0), -1
+ ) < padding_length.unsqueeze(-1)
+
+ if offset is not None:
+ pad_mask_off = torch.arange(0, max_audio_length, device=device).expand(
+ padding_length.size(0), -1
+ ) >= offset.unsqueeze(-1)
+ pad_mask = pad_mask_off.logical_and(pad_mask)
+
+ if att_mask is not None:
+ # pad_mask_for_att_mask is the mask which helps to ignore paddings
+ pad_mask_for_att_mask = pad_mask.unsqueeze(1).repeat([1, max_audio_length, 1])
+ pad_mask_for_att_mask = torch.logical_and(pad_mask_for_att_mask, pad_mask_for_att_mask.transpose(1, 2))
+ # att_mask is the masking to be used by the MHA layers to ignore the tokens not supposed to be visible
+ att_mask = att_mask[:, :max_audio_length, :max_audio_length]
+ # paddings should also get ignored, so pad_mask_for_att_mask is used to ignore their corresponding scores
+ att_mask = torch.logical_and(pad_mask_for_att_mask, att_mask.to(pad_mask_for_att_mask.device))
+ att_mask = ~att_mask
+
+ pad_mask = ~pad_mask
+ return pad_mask, att_mask
+
+ def _calc_context_sizes(
+ self, att_context_size, att_context_probs, att_context_style, conv_context_size, conv_kernel_size
+ ):
+ # convert att_context_size to a standard list of lists
+ if att_context_size:
+ att_context_size_all = list(att_context_size)
+ if isinstance(att_context_size_all[0], int):
+ att_context_size_all = [att_context_size_all]
+ for i, att_cs in enumerate(att_context_size_all):
+ if isinstance(att_cs, ListConfig):
+ att_context_size_all[i] = list(att_cs)
+ if att_context_style == "chunked_limited":
+ if att_cs[0] > 0 and att_cs[0] % (att_cs[1] + 1) > 0:
+ raise ValueError(f"att_context_size[{i}][0] % (att_context_size[{i}][1] + 1) should be zero!")
+ if att_cs[1] < 0 and len(att_context_size_all) <= 1:
+ raise ValueError(
+ f"Right context (att_context_size[{i}][1]) can not be unlimited for chunked_limited style!"
+ )
+ else:
+ att_context_size_all = [[-1, -1]]
+
+ if att_context_probs:
+ if len(att_context_probs) != len(att_context_size_all):
+ raise ValueError("The size of the att_context_probs should be the same as att_context_size.")
+ att_context_probs = list(att_context_probs)
+ if sum(att_context_probs) != 1:
+ raise ValueError("The sum of numbers in att_context_probs should be equal to one to be a distribution.")
+ else:
+ att_context_probs = [1.0 / len(att_context_size_all)] * len(att_context_size_all)
+
+ if conv_context_size is not None:
+ if isinstance(conv_context_size, ListConfig):
+ conv_context_size = list(conv_context_size)
+ if not isinstance(conv_context_size, list) and not isinstance(conv_context_size, str):
+ raise ValueError(
+ "Invalid conv_context_size! It should be the string 'causal' or a list of two integers."
+ )
+ if conv_context_size == "causal":
+ conv_context_size = [conv_kernel_size - 1, 0]
+ else:
+ if conv_context_size[0] + conv_context_size[1] + 1 != conv_kernel_size:
+ raise ValueError(f"Invalid conv_context_size: {self.conv_context_size}!")
+ else:
+ conv_context_size = [(conv_kernel_size - 1) // 2, (conv_kernel_size - 1) // 2]
+ return att_context_size_all, att_context_size_all[0], att_context_probs, conv_context_size
+
+ def setup_streaming_params(
+ self,
+ chunk_size: int = None,
+ shift_size: int = None,
+ left_chunks: int = None,
+ att_context_size: list = None,
+ max_context: int = 10000,
+ ):
+ """
+ This function sets the needed values and parameters to perform streaming.
+ The configuration would be stored in self.streaming_cfg.
+ The streaming configuration is needed to simulate streaming inference.
+
+ Args:
+ chunk_size (int): overrides the chunk size
+ shift_size (int): overrides the shift size for chunks
+ left_chunks (int): overrides the number of left chunks visible to each chunk
+ max_context (int): the value used for the cache size of last_channel layers
+ if left context is set to infinity (-1)
+ Defaults to -1 (means feat_out is d_model)
+ """
+ streaming_cfg = CacheAwareStreamingConfig()
+
+ # When att_context_size is not specified, it uses the default_att_context_size
+ if att_context_size is None:
+ att_context_size = self.att_context_size
+
+ if chunk_size is not None:
+ if chunk_size < 1:
+ raise ValueError("chunk_size needs to be a number larger or equal to one.")
+ lookahead_steps = chunk_size - 1
+ streaming_cfg.cache_drop_size = chunk_size - shift_size
+ elif self.att_context_style == "chunked_limited":
+ lookahead_steps = att_context_size[1]
+ streaming_cfg.cache_drop_size = 0
+ elif self.att_context_style == "regular":
+ lookahead_steps = att_context_size[1] * self.n_layers + self.conv_context_size[1] * self.n_layers
+ streaming_cfg.cache_drop_size = lookahead_steps
+ else:
+ streaming_cfg.cache_drop_size = 0
+ lookahead_steps = None
+
+ if chunk_size is None:
+ streaming_cfg.last_channel_cache_size = att_context_size[0] if att_context_size[0] >= 0 else max_context
+ else:
+ if left_chunks is None:
+ streaming_cfg.last_channel_cache_size = att_context_size[0] if att_context_size[0] >= 0 else max_context
+ logging.warning(
+ f"left_chunks is not set. Setting it to default: {streaming_cfg.last_channel_cache_size}."
+ )
+ else:
+ streaming_cfg.last_channel_cache_size = left_chunks * chunk_size
+
+ if hasattr(self.pre_encode, "get_sampling_frames"):
+ sampling_frames = self.pre_encode.get_sampling_frames()
+ else:
+ sampling_frames = 0
+
+ if isinstance(sampling_frames, list):
+ streaming_cfg.chunk_size = [
+ sampling_frames[0] + self.subsampling_factor * lookahead_steps,
+ sampling_frames[1] + self.subsampling_factor * lookahead_steps,
+ ]
+ else:
+ streaming_cfg.chunk_size = sampling_frames * (1 + lookahead_steps)
+
+ if isinstance(sampling_frames, list):
+ streaming_cfg.shift_size = [
+ sampling_frames[0] + sampling_frames[1] * (lookahead_steps - streaming_cfg.cache_drop_size),
+ sampling_frames[1] + sampling_frames[1] * (lookahead_steps - streaming_cfg.cache_drop_size),
+ ]
+ else:
+ streaming_cfg.shift_size = sampling_frames * (1 + lookahead_steps - streaming_cfg.cache_drop_size)
+
+ if isinstance(streaming_cfg.shift_size, list):
+ streaming_cfg.valid_out_len = (
+ streaming_cfg.shift_size[1] - sampling_frames[1]
+ ) // self.subsampling_factor + 1
+ else:
+ streaming_cfg.valid_out_len = streaming_cfg.shift_size // self.subsampling_factor
+
+ if hasattr(self.pre_encode, "get_streaming_cache_size"):
+ streaming_cfg.pre_encode_cache_size = self.pre_encode.get_streaming_cache_size()
+ else:
+ streaming_cfg.pre_encode_cache_size = 0
+
+ if isinstance(streaming_cfg.pre_encode_cache_size, list):
+ if streaming_cfg.pre_encode_cache_size[1] >= 1:
+ streaming_cfg.drop_extra_pre_encoded = (
+ 1 + (streaming_cfg.pre_encode_cache_size[1] - 1) // self.subsampling_factor
+ )
+ else:
+ streaming_cfg.drop_extra_pre_encoded = 0
+ else:
+ streaming_cfg.drop_extra_pre_encoded = streaming_cfg.pre_encode_cache_size // self.subsampling_factor
+
+ for m in self.layers.modules():
+ if hasattr(m, "_max_cache_len"):
+ if isinstance(m, MultiHeadAttention):
+ m.cache_drop_size = streaming_cfg.cache_drop_size
+ if isinstance(m, CausalConv1D):
+ m.cache_drop_size = streaming_cfg.cache_drop_size
+
+ self.streaming_cfg = streaming_cfg
+
+ def change_subsampling_conv_chunking_factor(self, subsampling_conv_chunking_factor: int):
+ """
+ Update the conv_chunking_factor (int)
+ Default is 1 (auto)
+ Set it to -1 (disabled) or to a specific value (power of 2) if you OOM in the conv subsampling layers
+
+
+ Args:
+ subsampling_conv_chunking_factor (int)
+ """
+
+ if not hasattr(self.pre_encode, "change_subsampling_conv_chunking_factor"):
+ logging.info("Model pre_encoder doesn't have a change_subsampling_conv_chunking_factor method ")
+ return
+
+ self.pre_encode.change_subsampling_conv_chunking_factor(
+ subsampling_conv_chunking_factor=subsampling_conv_chunking_factor
+ )
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/conformer_modules.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/conformer_modules.py
new file mode 100644
index 000000000..4875ff734
--- /dev/null
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/conformer_modules.py
@@ -0,0 +1,372 @@
+# SPDX-FileCopyrightText: Copyright (c) 2020 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/asr/parts/submodules/conformer_modules.py; training-only code removed.
+
+import torch
+from torch import nn as nn
+from torch.nn import LayerNorm
+
+from ..compat import AccessMixin, AttentionAdapterModuleMixin, FusedBatchNorm1d, Swish, activation_registry
+from .causal_convs import CausalConv1D
+from .multi_head_attention import (
+ MultiHeadAttention,
+ RelPositionMultiHeadAttention,
+ RelPositionMultiHeadAttentionLongformer,
+)
+
+__all__ = ["ConformerConvolution", "ConformerFeedForward", "ConformerLayer"]
+
+
+class ConformerLayer(torch.nn.Module, AttentionAdapterModuleMixin, AccessMixin):
+ """A single block of the Conformer encoder.
+
+ Args:
+ d_model (int): input dimension of MultiheadAttentionMechanism and PositionwiseFeedForward
+ d_ff (int): hidden dimension of PositionwiseFeedForward
+ self_attention_model (str): type of the attention layer and positional encoding
+ 'rel_pos': relative positional embedding and Transformer-XL
+ 'rel_pos_local_attn': relative positional embedding and Transformer-XL with local attention using
+ overlapping chunks. Attention context is determined by att_context_size parameter.
+ 'abs_pos': absolute positional embedding and Transformer
+ Default is rel_pos.
+ global_tokens (int): number of tokens to be used for global attention.
+ Only relevant if self_attention_model is 'rel_pos_local_attn'.
+ Defaults to 0.
+ global_tokens_spacing (int): how far apart the global tokens are
+ Defaults to 1.
+ global_attn_separate (bool): whether the q, k, v layers used for global tokens should be separate.
+ Defaults to False.
+ n_heads (int): number of heads for multi-head attention
+ conv_kernel_size (int): kernel size for depthwise convolution in convolution module
+ dropout (float): dropout probabilities for linear layers
+ dropout_att (float): dropout probabilities for attention distributions
+ use_bias (bool): Apply bias to all Linear and Conv1d layers from each ConformerLayer to improve activation flow and stabilize training of huge models.
+ Defaults to True.
+ """
+
+ def __init__(
+ self,
+ d_model,
+ d_ff,
+ self_attention_model="rel_pos",
+ global_tokens=0,
+ global_tokens_spacing=1,
+ global_attn_separate=False,
+ n_heads=4,
+ conv_kernel_size=31,
+ conv_norm_type="batch_norm",
+ conv_context_size=None,
+ dropout=0.1,
+ dropout_att=0.1,
+ pos_bias_u=None,
+ pos_bias_v=None,
+ att_context_size=[-1, -1],
+ use_bias=True,
+ use_pytorch_sdpa=False,
+ use_pytorch_sdpa_backends=None,
+ ):
+ super(ConformerLayer, self).__init__()
+
+ self.use_pytorch_sdpa = use_pytorch_sdpa
+ if use_pytorch_sdpa_backends is None:
+ use_pytorch_sdpa_backends = []
+ self.use_pytorch_sdpa_backends = use_pytorch_sdpa_backends
+ self.self_attention_model = self_attention_model
+ self.n_heads = n_heads
+ self.fc_factor = 0.5
+
+ # first feed forward module
+ self.norm_feed_forward1 = LayerNorm(d_model)
+ self.feed_forward1 = ConformerFeedForward(d_model=d_model, d_ff=d_ff, dropout=dropout, use_bias=use_bias)
+
+ # convolution module
+ self.norm_conv = LayerNorm(d_model)
+ self.conv = ConformerConvolution(
+ d_model=d_model,
+ kernel_size=conv_kernel_size,
+ norm_type=conv_norm_type,
+ conv_context_size=conv_context_size,
+ use_bias=use_bias,
+ )
+
+ # multi-headed self-attention module
+ self.norm_self_att = LayerNorm(d_model)
+ MHA_max_cache_len = att_context_size[0]
+
+ if self_attention_model == "rel_pos":
+ self.self_attn = RelPositionMultiHeadAttention(
+ n_head=n_heads,
+ n_feat=d_model,
+ dropout_rate=dropout_att,
+ pos_bias_u=pos_bias_u,
+ pos_bias_v=pos_bias_v,
+ max_cache_len=MHA_max_cache_len,
+ use_bias=use_bias,
+ use_pytorch_sdpa=self.use_pytorch_sdpa,
+ use_pytorch_sdpa_backends=self.use_pytorch_sdpa_backends,
+ )
+ elif self_attention_model == "rel_pos_local_attn":
+ self.self_attn = RelPositionMultiHeadAttentionLongformer(
+ n_head=n_heads,
+ n_feat=d_model,
+ dropout_rate=dropout_att,
+ pos_bias_u=pos_bias_u,
+ pos_bias_v=pos_bias_v,
+ max_cache_len=MHA_max_cache_len,
+ att_context_size=att_context_size,
+ global_tokens=global_tokens,
+ global_tokens_spacing=global_tokens_spacing,
+ global_attn_separate=global_attn_separate,
+ use_bias=use_bias,
+ )
+ elif self_attention_model == "abs_pos":
+ self.self_attn = MultiHeadAttention(
+ n_head=n_heads,
+ n_feat=d_model,
+ dropout_rate=dropout_att,
+ max_cache_len=MHA_max_cache_len,
+ use_bias=use_bias,
+ use_pytorch_sdpa=self.use_pytorch_sdpa,
+ use_pytorch_sdpa_backends=self.use_pytorch_sdpa_backends,
+ )
+ else:
+ raise ValueError(
+ f"'{self_attention_model}' is not not a valid value for 'self_attention_model', "
+ f"valid values can be from ['rel_pos', 'rel_pos_local_attn', 'abs_pos']"
+ )
+
+ # second feed forward module
+ self.norm_feed_forward2 = LayerNorm(d_model)
+ self.feed_forward2 = ConformerFeedForward(d_model=d_model, d_ff=d_ff, dropout=dropout, use_bias=use_bias)
+
+ self.dropout = nn.Dropout(dropout)
+ self.norm_out = LayerNorm(d_model)
+
+ def forward(self, x, att_mask=None, pos_emb=None, pad_mask=None, cache_last_channel=None, cache_last_time=None):
+ """
+ Args:
+ x (torch.Tensor): input signals (B, T, d_model)
+ att_mask (torch.Tensor): attention masks(B, T, T)
+ pos_emb (torch.Tensor): (L, 1, d_model)
+ pad_mask (torch.tensor): padding mask
+ cache_last_channel (torch.tensor) : cache for MHA layers (B, T_cache, d_model)
+ cache_last_time (torch.tensor) : cache for convolutional layers (B, d_model, T_cache)
+ Returns:
+ x (torch.Tensor): (B, T, d_model)
+ cache_last_channel (torch.tensor) : next cache for MHA layers (B, T_cache, d_model)
+ cache_last_time (torch.tensor) : next cache for convolutional layers (B, d_model, T_cache)
+ """
+ residual = x
+ x = self.norm_feed_forward1(x)
+ x = self.feed_forward1(x)
+ residual = residual + self.dropout(x) * self.fc_factor
+
+ x = self.norm_self_att(residual)
+ if self.self_attention_model == "rel_pos":
+ x = self.self_attn(query=x, key=x, value=x, mask=att_mask, pos_emb=pos_emb, cache=cache_last_channel)
+ elif self.self_attention_model == "rel_pos_local_attn":
+ x = self.self_attn(query=x, key=x, value=x, pad_mask=pad_mask, pos_emb=pos_emb, cache=cache_last_channel)
+ elif self.self_attention_model == "abs_pos":
+ x = self.self_attn(query=x, key=x, value=x, mask=att_mask, cache=cache_last_channel)
+ else:
+ x = None
+
+ if x is not None and cache_last_channel is not None:
+ (x, cache_last_channel) = x
+
+ residual = residual + self.dropout(x)
+
+ if self.is_adapter_available():
+ # Call the MHA adapters
+ pack_input = {
+ "x": residual,
+ "loc": "mha",
+ "att_mask": att_mask,
+ "pos_emb": pos_emb,
+ }
+ pack_input = self.forward_enabled_adapters(pack_input)
+ residual = pack_input["x"]
+
+ x = self.norm_conv(residual)
+ x = self.conv(x, pad_mask=pad_mask, cache=cache_last_time)
+ if cache_last_time is not None:
+ (x, cache_last_time) = x
+ residual = residual + self.dropout(x)
+
+ x = self.norm_feed_forward2(residual)
+ x = self.feed_forward2(x)
+ residual = residual + self.dropout(x) * self.fc_factor
+
+ x = self.norm_out(residual)
+
+ if self.is_adapter_available():
+ # Call the adapters
+ pack_input = {
+ "x": x,
+ "loc": "post",
+ }
+ pack_input = self.forward_enabled_adapters(pack_input)
+ x = pack_input["x"]
+
+ if self.is_access_enabled(getattr(self, "model_guid", None)) and self.access_cfg.get(
+ "save_encoder_tensors", False
+ ):
+ self.register_accessible_tensor(name="encoder", tensor=x)
+ if cache_last_channel is None:
+ return x
+ else:
+ return x, cache_last_channel, cache_last_time
+
+
+class ConformerConvolution(nn.Module):
+ """The convolution module for the Conformer model.
+ Args:
+ d_model (int): hidden dimension
+ kernel_size (int): kernel size for depthwise convolution
+ pointwise_activation (str): name of the activation function to be used for the pointwise conv.
+ Note that Conformer uses a special key `glu_` which is treated as the original default from
+ the paper.
+ use_bias (bool): Use bias in all Linear and Conv1d layers improve activation flow and stabilize training of huge models.
+ Defaults to True
+ """
+
+ def __init__(
+ self,
+ d_model,
+ kernel_size,
+ norm_type="batch_norm",
+ conv_context_size=None,
+ pointwise_activation="glu_",
+ use_bias=True,
+ ):
+ super(ConformerConvolution, self).__init__()
+ assert (kernel_size - 1) % 2 == 0
+ self.d_model = d_model
+ self.kernel_size = kernel_size
+ self.norm_type = norm_type
+ self.use_bias = use_bias
+
+ if conv_context_size is None:
+ conv_context_size = (kernel_size - 1) // 2
+
+ if pointwise_activation in activation_registry:
+ self.pointwise_activation = activation_registry[pointwise_activation]()
+ dw_conv_input_dim = d_model * 2
+
+ if hasattr(self.pointwise_activation, "inplace"):
+ self.pointwise_activation.inplace = True
+ else:
+ self.pointwise_activation = pointwise_activation
+ dw_conv_input_dim = d_model
+
+ self.pointwise_conv1 = nn.Conv1d(
+ in_channels=d_model,
+ out_channels=d_model * 2,
+ kernel_size=1,
+ stride=1,
+ padding=0,
+ bias=self.use_bias,
+ )
+
+ self.depthwise_conv = CausalConv1D(
+ in_channels=dw_conv_input_dim,
+ out_channels=dw_conv_input_dim,
+ kernel_size=kernel_size,
+ stride=1,
+ padding=conv_context_size,
+ groups=dw_conv_input_dim,
+ bias=self.use_bias,
+ )
+
+ if norm_type == "batch_norm":
+ self.batch_norm = nn.BatchNorm1d(dw_conv_input_dim)
+ elif norm_type == "instance_norm":
+ self.batch_norm = nn.InstanceNorm1d(dw_conv_input_dim)
+ elif norm_type == "layer_norm":
+ self.batch_norm = nn.LayerNorm(dw_conv_input_dim)
+ elif norm_type == "fused_batch_norm":
+ self.batch_norm = FusedBatchNorm1d(dw_conv_input_dim)
+ elif norm_type.startswith("group_norm"):
+ num_groups = int(norm_type.replace("group_norm", ""))
+ self.batch_norm = nn.GroupNorm(num_groups=num_groups, num_channels=d_model)
+ else:
+ raise ValueError(f"conv_norm_type={norm_type} is not valid!")
+
+ self.activation = Swish()
+ self.pointwise_conv2 = nn.Conv1d(
+ in_channels=dw_conv_input_dim,
+ out_channels=d_model,
+ kernel_size=1,
+ stride=1,
+ padding=0,
+ bias=self.use_bias,
+ )
+
+ def forward(self, x, pad_mask=None, cache=None):
+ x = x.transpose(1, 2)
+ x = self.pointwise_conv1(x)
+
+ # Compute the activation function or use GLU for original Conformer
+ if self.pointwise_activation == "glu_":
+ x = nn.functional.glu(x, dim=1)
+ else:
+ x = self.pointwise_activation(x)
+
+ if pad_mask is not None:
+ x = x.masked_fill(pad_mask.unsqueeze(1), 0.0)
+
+ x = self.depthwise_conv(x, cache=cache)
+ if cache is not None:
+ x, cache = x
+
+ if self.norm_type == "layer_norm":
+ x = x.transpose(1, 2)
+ x = self.batch_norm(x)
+ x = x.transpose(1, 2)
+ else:
+ x = self.batch_norm(x)
+
+ x = self.activation(x)
+ x = self.pointwise_conv2(x)
+ x = x.transpose(1, 2)
+ if cache is None:
+ return x
+ else:
+ return x, cache
+
+
+class ConformerFeedForward(nn.Module):
+ """
+ feed-forward module of Conformer model.
+ use_bias (bool): Apply bias to all Linear and Conv1d layers improve activation flow and stabilize training of huge models.
+ """
+
+ def __init__(self, d_model, d_ff, dropout, activation=Swish(), use_bias=True):
+ super(ConformerFeedForward, self).__init__()
+ self.d_model = d_model
+ self.d_ff = d_ff
+ self.use_bias = use_bias
+ self.linear1 = nn.Linear(d_model, d_ff, bias=self.use_bias)
+ self.activation = activation
+ self.dropout = nn.Dropout(p=dropout)
+ self.linear2 = nn.Linear(d_ff, d_model, bias=self.use_bias)
+
+ def forward(self, x):
+ x = self.linear1(x)
+ x = self.activation(x)
+ x = self.dropout(x)
+ x = self.linear2(x)
+ return x
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/features.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/features.py
new file mode 100644
index 000000000..2261d739e
--- /dev/null
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/features.py
@@ -0,0 +1,633 @@
+# Copyright (c) 2018 Ryan Leary
+# SPDX-FileCopyrightText: Copyright (c) 2020 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+# This file contains code artifacts adapted from https://github.com/ryanleary/patter
+# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/asr/parts/preprocessing/features.py; training-only code removed.
+# Removed: WaveformFeaturizer / FeaturizerFactory (data-loading only, required perturb/segment).
+
+import math
+import random
+from typing import Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from ..compat import logging
+
+try:
+ import torchaudio
+
+ HAVE_TORCHAUDIO = True
+except ModuleNotFoundError:
+ HAVE_TORCHAUDIO = False
+
+
+CONSTANT = 1e-5
+
+
+def normalize_batch(x, seq_len, normalize_type):
+ x_mean = None
+ x_std = None
+ if normalize_type == "per_feature":
+ batch_size = x.shape[0]
+ max_time = x.shape[2]
+
+ # When doing stream capture to a graph, item() is not allowed
+ # because it calls cudaStreamSynchronize(). Therefore, we are
+ # sacrificing some error checking when running with cuda graphs.
+ if (
+ torch.cuda.is_available()
+ and not torch.cuda.is_current_stream_capturing()
+ and torch.any(seq_len == 1).item()
+ ):
+ raise ValueError(
+ "normalize_batch with `per_feature` normalize_type received a tensor of length 1. This will result "
+ "in torch.std() returning nan. Make sure your audio length has enough samples for a single "
+ "feature (ex. at least `hop_length` for Mel Spectrograms)."
+ )
+ time_steps = torch.arange(max_time, device=x.device).unsqueeze(0).expand(batch_size, max_time)
+ valid_mask = time_steps < seq_len.unsqueeze(1)
+ x_mean_numerator = torch.where(valid_mask.unsqueeze(1), x, 0.0).sum(axis=2)
+ x_mean_denominator = valid_mask.sum(axis=1)
+ x_mean = x_mean_numerator / x_mean_denominator.unsqueeze(1)
+
+ # Subtract 1 in the denominator to correct for the bias.
+ x_std = torch.sqrt(
+ torch.sum(torch.where(valid_mask.unsqueeze(1), x - x_mean.unsqueeze(2), 0.0) ** 2, axis=2)
+ / (x_mean_denominator.unsqueeze(1) - 1.0)
+ )
+ # make sure x_std is not zero
+ x_std += CONSTANT
+ return (x - x_mean.unsqueeze(2)) / x_std.unsqueeze(2), x_mean, x_std
+ elif normalize_type == "all_features":
+ x_mean = torch.zeros(seq_len.shape, dtype=x.dtype, device=x.device)
+ x_std = torch.zeros(seq_len.shape, dtype=x.dtype, device=x.device)
+ for i in range(x.shape[0]):
+ x_mean[i] = x[i, :, : seq_len[i].item()].mean()
+ x_std[i] = x[i, :, : seq_len[i].item()].std()
+ # make sure x_std is not zero
+ x_std += CONSTANT
+ return (x - x_mean.view(-1, 1, 1)) / x_std.view(-1, 1, 1), x_mean, x_std
+ elif "fixed_mean" in normalize_type and "fixed_std" in normalize_type:
+ x_mean = torch.tensor(normalize_type["fixed_mean"], device=x.device)
+ x_std = torch.tensor(normalize_type["fixed_std"], device=x.device)
+ return (
+ (x - x_mean.view(x.shape[0], x.shape[1]).unsqueeze(2)) / x_std.view(x.shape[0], x.shape[1]).unsqueeze(2),
+ x_mean,
+ x_std,
+ )
+ else:
+ return x, x_mean, x_std
+
+
+def splice_frames(x, frame_splicing):
+ """Stacks frames together across feature dim
+
+ input is batch_size, feature_dim, num_frames
+ output is batch_size, feature_dim*frame_splicing, num_frames
+
+ """
+ seq = [x]
+ for n in range(1, frame_splicing):
+ seq.append(torch.cat([x[:, :, :n], x[:, :, n:]], dim=2))
+ return torch.cat(seq, dim=1)
+
+
+@torch.jit.script_if_tracing
+def make_seq_mask_like(
+ lengths: torch.Tensor, like: torch.Tensor, time_dim: int = -1, valid_ones: bool = True
+) -> torch.Tensor:
+ """
+
+ Args:
+ lengths: Tensor with shape [B] containing the sequence length of each batch element
+ like: The mask will contain the same number of dimensions as this Tensor, and will have the same max
+ length in the time dimension of this Tensor.
+ time_dim: Time dimension of the `shape_tensor` and the resulting mask. Zero-based.
+ valid_ones: If True, valid tokens will contain value `1` and padding will be `0`. Else, invert.
+
+ Returns:
+ A :class:`torch.Tensor` containing 1's and 0's for valid and invalid tokens, respectively, if `valid_ones`, else
+ vice-versa. Mask will have the same number of dimensions as `like`. Batch and time dimensions will match
+ the `like`. All other dimensions will be singletons. E.g., if `like.shape == [3, 4, 5]` and
+ `time_dim == -1', mask will have shape `[3, 1, 5]`.
+ """
+ # Mask with shape [B, T]
+ mask = torch.arange(like.shape[time_dim], device=like.device).repeat(lengths.shape[0], 1).lt(lengths.view(-1, 1))
+ # [B, T] -> [B, *, T] where * is any number of singleton dimensions to expand to like tensor
+ for _ in range(like.dim() - mask.dim()):
+ mask = mask.unsqueeze(1)
+ # If needed, transpose time dim
+ if time_dim != -1 and time_dim != mask.dim() - 1:
+ mask = mask.transpose(-1, time_dim)
+ # Maybe invert the padded vs. valid token values
+ if not valid_ones:
+ mask = ~mask
+ return mask
+
+
+class FilterbankFeatures(nn.Module):
+ """Featurizer that converts wavs to Mel Spectrograms.
+ See AudioToMelSpectrogramPreprocessor for args.
+ """
+
+ def __init__(
+ self,
+ sample_rate=16000,
+ n_window_size=320,
+ n_window_stride=160,
+ window="hann",
+ normalize="per_feature",
+ n_fft=None,
+ preemph=0.97,
+ nfilt=64,
+ lowfreq=0,
+ highfreq=None,
+ log=True,
+ log_zero_guard_type="add",
+ log_zero_guard_value=2**-24,
+ dither=CONSTANT,
+ pad_to=16,
+ max_duration=16.7,
+ frame_splicing=1,
+ exact_pad=False,
+ pad_value=0,
+ mag_power=2.0,
+ use_grads=False,
+ rng=None,
+ nb_augmentation_prob=0.0,
+ nb_max_freq=4000,
+ mel_norm="slaney",
+ stft_exact_pad=False, # Deprecated arguments; kept for config compatibility
+ stft_conv=False, # Deprecated arguments; kept for config compatibility
+ ):
+ super().__init__()
+ if stft_conv or stft_exact_pad:
+ logging.warning(
+ "Using torch_stft is deprecated and has been removed. The values have been forcibly set to False "
+ "for FilterbankFeatures and AudioToMelSpectrogramPreprocessor. Please set exact_pad to True "
+ "as needed."
+ )
+ if exact_pad and n_window_stride % 2 == 1:
+ raise NotImplementedError(
+ f"{self} received exact_pad == True, but hop_size was odd. If audio_length % hop_size == 0. Then the "
+ "returned spectrogram would not be of length audio_length // hop_size. Please use an even hop_size."
+ )
+ self.log_zero_guard_value = log_zero_guard_value
+ if (
+ n_window_size is None
+ or n_window_stride is None
+ or not isinstance(n_window_size, int)
+ or not isinstance(n_window_stride, int)
+ or n_window_size <= 0
+ or n_window_stride <= 0
+ ):
+ raise ValueError(
+ f"{self} got an invalid value for either n_window_size or n_window_stride. Both must be positive ints."
+ )
+ logging.info(f"PADDING: {pad_to}")
+
+ self.win_length = n_window_size
+ self.hop_length = n_window_stride
+ self.n_fft = n_fft or 2 ** math.ceil(math.log2(self.win_length))
+ self.stft_pad_amount = (self.n_fft - self.hop_length) // 2 if exact_pad else None
+ self.exact_pad = exact_pad
+
+ if exact_pad:
+ logging.info("STFT using exact pad")
+ torch_windows = {
+ "hann": torch.hann_window,
+ "hamming": torch.hamming_window,
+ "blackman": torch.blackman_window,
+ "bartlett": torch.bartlett_window,
+ "none": None,
+ }
+ window_fn = torch_windows.get(window, None)
+ window_tensor = window_fn(self.win_length, periodic=False) if window_fn else None
+ self.register_buffer("window", window_tensor)
+
+ self.normalize = normalize
+ self.log = log
+ self.dither = dither
+ self.frame_splicing = frame_splicing
+ self.nfilt = nfilt
+ self.preemph = preemph
+ self.pad_to = pad_to
+ highfreq = highfreq or sample_rate / 2
+
+ # librosa.filters.mel replaced by the torchaudio equivalent (norm="slaney" +
+ # mel_scale="slaney" matches librosa's default HTK=False convention; filterbank
+ # max abs diff 2.6e-7, verified token- and WAV-identical on the parity fixture).
+ if not HAVE_TORCHAUDIO:
+ raise ValueError(f"Need to install torchaudio to instantiate a {self.__class__.__name__}")
+ filterbanks = (
+ torchaudio.functional.melscale_fbanks(
+ n_freqs=self.n_fft // 2 + 1,
+ f_min=lowfreq,
+ f_max=highfreq,
+ n_mels=nfilt,
+ sample_rate=sample_rate,
+ norm=mel_norm,
+ mel_scale="slaney",
+ )
+ .T.to(torch.float)
+ .unsqueeze(0)
+ )
+ self.register_buffer("fb", filterbanks)
+
+ # Calculate maximum sequence length
+ max_length = self.get_seq_len(torch.tensor(max_duration * sample_rate, dtype=torch.float))
+ max_pad = pad_to - (max_length % pad_to) if pad_to > 0 else 0
+ self.max_length = max_length + max_pad
+ self.pad_value = pad_value
+ self.mag_power = mag_power
+
+ # We want to avoid taking the log of zero
+ # There are two options: either adding or clamping to a small value
+ if log_zero_guard_type not in ["add", "clamp"]:
+ raise ValueError(
+ f"{self} received {log_zero_guard_type} for the "
+ f"log_zero_guard_type parameter. It must be either 'add' or "
+ f"'clamp'."
+ )
+
+ self.use_grads = use_grads
+ if not use_grads:
+ self.forward = torch.no_grad()(self.forward)
+ self._rng = random.Random() if rng is None else rng
+ self.nb_augmentation_prob = nb_augmentation_prob
+ if self.nb_augmentation_prob > 0.0:
+ if nb_max_freq >= sample_rate / 2:
+ self.nb_augmentation_prob = 0.0
+ else:
+ self._nb_max_fft_bin = int((nb_max_freq / sample_rate) * n_fft)
+
+ # log_zero_guard_value is the the small we want to use, we support
+ # an actual number, or "tiny", or "eps"
+ self.log_zero_guard_type = log_zero_guard_type
+ logging.debug(f"sr: {sample_rate}")
+ logging.debug(f"n_fft: {self.n_fft}")
+ logging.debug(f"win_length: {self.win_length}")
+ logging.debug(f"hop_length: {self.hop_length}")
+ logging.debug(f"n_mels: {nfilt}")
+ logging.debug(f"fmin: {lowfreq}")
+ logging.debug(f"fmax: {highfreq}")
+ logging.debug(f"using grads: {use_grads}")
+ logging.debug(f"nb_augmentation_prob: {nb_augmentation_prob}")
+
+ def stft(self, x):
+ return torch.stft(
+ x,
+ n_fft=self.n_fft,
+ hop_length=self.hop_length,
+ win_length=self.win_length,
+ center=False if self.exact_pad else True,
+ window=self.window.to(dtype=torch.float, device=x.device),
+ return_complex=True,
+ )
+
+ def log_zero_guard_value_fn(self, x):
+ if isinstance(self.log_zero_guard_value, str):
+ if self.log_zero_guard_value == "tiny":
+ return torch.finfo(x.dtype).tiny
+ elif self.log_zero_guard_value == "eps":
+ return torch.finfo(x.dtype).eps
+ else:
+ raise ValueError(
+ f"{self} received {self.log_zero_guard_value} for the "
+ f"log_zero_guard_type parameter. It must be either a "
+ f"number, 'tiny', or 'eps'"
+ )
+ else:
+ return self.log_zero_guard_value
+
+ def get_seq_len(self, seq_len):
+ # Assuming that center is True is stft_pad_amount = 0
+ pad_amount = self.stft_pad_amount * 2 if self.stft_pad_amount is not None else self.n_fft // 2 * 2
+ seq_len = torch.floor_divide((seq_len + pad_amount - self.n_fft), self.hop_length) + 1
+ return seq_len.to(dtype=torch.long)
+
+ def forward(self, x, seq_len, linear_spec=False):
+ # Initialize debug dictionary
+ # debug_snapshots = {}
+ # debug_snapshots['step0_input'] = x[0, :1280].cpu() if x.shape[1] >= 1280 else x[0].cpu()
+ # debug_snapshots['input_len'] = seq_len[0].item()
+
+ seq_len = self.get_seq_len(seq_len)
+ # debug_snapshots['output_seq_len'] = seq_len[0].item()
+
+ if self.stft_pad_amount is not None:
+ x = torch.nn.functional.pad(
+ x.unsqueeze(1), (self.stft_pad_amount, self.stft_pad_amount), "reflect"
+ ).squeeze(1)
+ # debug_snapshots['step1_after_reflect_pad'] = x[0, :1280].cpu() if x.shape[1] >= 1280 else x[0].cpu()
+
+ # dither (only in training mode for eval determinism)
+ if self.training and self.dither > 0:
+ x += self.dither * torch.randn_like(x)
+ # debug_snapshots['step2_after_dither'] = x[0, :1280].cpu() if x.shape[1] >= 1280 else x[0].cpu()
+
+ # do preemphasis
+ if self.preemph is not None:
+ x = torch.cat((x[:, 0].unsqueeze(1), x[:, 1:] - self.preemph * x[:, :-1]), dim=1)
+ # debug_snapshots['step3_after_preemph'] = x[0, :1280].cpu() if x.shape[1] >= 1280 else x[0].cpu()
+
+ # disable autocast to get full range of stft values
+ with torch.amp.autocast(x.device.type, enabled=False):
+ x = self.stft(x)
+
+ # torch stft returns complex tensor (of shape [B,N,T]); so convert to magnitude
+ # guard is needed for sqrt if grads are passed through
+ guard = 0 if not self.use_grads else CONSTANT
+ x = torch.view_as_real(x)
+ x = torch.sqrt(x.pow(2).sum(-1) + guard)
+ # debug_snapshots['step4_after_stft_magnitude'] = x[0, :, :10].cpu() # First 10 frames, all freq bins
+
+ if self.training and self.nb_augmentation_prob > 0.0:
+ for idx in range(x.shape[0]):
+ if self._rng.random() < self.nb_augmentation_prob:
+ x[idx, self._nb_max_fft_bin :, :] = 0.0
+
+ # get power spectrum
+ if self.mag_power != 1.0:
+ x = x.pow(self.mag_power)
+ # debug_snapshots['step5_after_power'] = x[0, :, :10].cpu()
+
+ # return plain spectrogram if required
+ if linear_spec:
+ return x, seq_len
+
+ # disable autocast, otherwise it might be automatically casted to fp16
+ # on fp16 compatible GPUs and get NaN values for input value of 65520
+ with torch.amp.autocast(x.device.type, enabled=False):
+ # dot with filterbank energies
+ x = torch.matmul(self.fb.to(x.dtype), x)
+ # debug_snapshots['step6_after_mel_filterbank'] = x[0, :, :10].cpu()
+
+ # log features if required
+ if self.log:
+ if self.log_zero_guard_type == "add":
+ x = torch.log(x + self.log_zero_guard_value_fn(x))
+ elif self.log_zero_guard_type == "clamp":
+ x = torch.log(torch.clamp(x, min=self.log_zero_guard_value_fn(x)))
+ else:
+ raise ValueError("log_zero_guard_type was not understood")
+ # debug_snapshots['step7_after_log'] = x[0, :, :10].cpu()
+
+ # frame splicing if required
+ if self.frame_splicing > 1:
+ x = splice_frames(x, self.frame_splicing)
+ # debug_snapshots['step8_after_frame_splicing'] = x[0, :, :10].cpu()
+
+ # normalize if required
+ if self.normalize:
+ x, _, _ = normalize_batch(x, seq_len, normalize_type=self.normalize)
+ # debug_snapshots['step9_after_normalize'] = x[0, :, :10].cpu()
+
+ # mask to zero any values beyond seq_len in batch, pad to multiple of `pad_to` (for efficiency)
+ max_len = x.size(-1)
+ mask = torch.arange(max_len, device=x.device)
+ mask = mask.repeat(x.size(0), 1) >= seq_len.unsqueeze(1)
+ x = x.masked_fill(mask.unsqueeze(1).type(torch.bool).to(device=x.device), self.pad_value)
+ # debug_snapshots['step10_after_masking'] = x[0, :, :10].cpu()
+ del mask
+
+ pad_to = self.pad_to
+ if pad_to == "max":
+ x = nn.functional.pad(x, (0, self.max_length - x.size(-1)), value=self.pad_value)
+ elif pad_to > 0:
+ pad_amt = x.size(-1) % pad_to
+ if pad_amt != 0:
+ x = nn.functional.pad(x, (0, pad_to - pad_amt), value=self.pad_value)
+ # debug_snapshots['step11_after_padding'] = x[0, :, :10].cpu()
+
+ # Debug: Save all intermediate steps
+ # import os
+
+ # # Add final output and config to debug snapshots
+ # debug_snapshots['final_output'] = x[0].cpu()
+ # debug_snapshots['final_shape'] = list(x.shape)
+ # debug_snapshots['config'] = {
+ # 'pad_to': pad_to,
+ # 'normalize': self.normalize,
+ # 'exact_pad': self.exact_pad,
+ # 'n_fft': self.n_fft,
+ # 'hop_length': self.hop_length,
+ # 'preemph': self.preemph,
+ # 'mag_power': self.mag_power,
+ # }
+
+ # torch.save(debug_snapshots, debug_path)
+ # print(f"[Preprocessor Debug] Saved {len(debug_snapshots)} steps to {debug_path}")
+ # print(f" Input length: {debug_snapshots['input_len']}, Output seq_len: {seq_len[0].item()}, Final shape: {x.shape}")
+
+ return x, seq_len
+
+
+class FilterbankFeaturesTA(nn.Module):
+ """
+ Exportable, `torchaudio`-based implementation of Mel Spectrogram extraction.
+
+ See `AudioToMelSpectrogramPreprocessor` for args.
+
+ """
+
+ def __init__(
+ self,
+ sample_rate: int = 16000,
+ n_window_size: int = 320,
+ n_window_stride: int = 160,
+ normalize: str | None = "per_feature",
+ nfilt: int = 64,
+ n_fft: int | None = None,
+ preemph: float = 0.97,
+ lowfreq: float = 0,
+ highfreq: float | None = None,
+ log: bool = True,
+ log_zero_guard_type: str = "add",
+ log_zero_guard_value: Union[float, str] = 2**-24,
+ dither: float = 1e-5,
+ window: str = "hann",
+ pad_to: int = 0,
+ pad_value: float = 0.0,
+ mel_norm="slaney",
+ # Seems like no one uses these options anymore. Don't convolute the code by supporting them.
+ use_grads: bool = False, # Deprecated arguments; kept for config compatibility
+ max_duration: float = 16.7, # Deprecated arguments; kept for config compatibility
+ frame_splicing: int = 1, # Deprecated arguments; kept for config compatibility
+ exact_pad: bool = False, # Deprecated arguments; kept for config compatibility
+ nb_augmentation_prob: float = 0.0, # Deprecated arguments; kept for config compatibility
+ nb_max_freq: int = 4000, # Deprecated arguments; kept for config compatibility
+ mag_power: float = 2.0, # Deprecated arguments; kept for config compatibility
+ rng: random.Random | None = None, # Deprecated arguments; kept for config compatibility
+ stft_exact_pad: bool = False, # Deprecated arguments; kept for config compatibility
+ stft_conv: bool = False, # Deprecated arguments; kept for config compatibility
+ ):
+ super().__init__()
+ if not HAVE_TORCHAUDIO:
+ raise ValueError(f"Need to install torchaudio to instantiate a {self.__class__.__name__}")
+
+ # Make sure log zero guard is supported, if given as a string
+ supported_log_zero_guard_strings = {"eps", "tiny"}
+ if isinstance(log_zero_guard_value, str) and log_zero_guard_value not in supported_log_zero_guard_strings:
+ raise ValueError(
+ f"Log zero guard value must either be a float or a member of {supported_log_zero_guard_strings}"
+ )
+
+ # Copied from `AudioPreprocessor` due to the ad-hoc structuring of the Mel Spec extractor class
+ self.torch_windows = {
+ "hann": torch.hann_window,
+ "hamming": torch.hamming_window,
+ "blackman": torch.blackman_window,
+ "bartlett": torch.bartlett_window,
+ "ones": torch.ones,
+ None: torch.ones,
+ }
+
+ # Ensure we can look up the window function
+ if window not in self.torch_windows:
+ raise ValueError(f"Got window value '{window}' but expected a member of {self.torch_windows.keys()}")
+
+ self.win_length = n_window_size
+ self.hop_length = n_window_stride
+ self._sample_rate = sample_rate
+ self._normalize_strategy = normalize
+ self._use_log = log
+ self._preemphasis_value = preemph
+ self.log_zero_guard_type = log_zero_guard_type
+ self.log_zero_guard_value: Union[str, float] = log_zero_guard_value
+ self.dither = dither
+ self.pad_to = pad_to
+ self.pad_value = pad_value
+ self.n_fft = n_fft
+ self._mel_spec_extractor: torchaudio.transforms.MelSpectrogram = torchaudio.transforms.MelSpectrogram(
+ sample_rate=self._sample_rate,
+ win_length=self.win_length,
+ hop_length=self.hop_length,
+ n_mels=nfilt,
+ window_fn=self.torch_windows[window],
+ mel_scale="slaney",
+ norm=mel_norm,
+ n_fft=n_fft,
+ f_max=highfreq,
+ f_min=lowfreq,
+ wkwargs={"periodic": False},
+ )
+
+ def _resolve_log_zero_guard_value(self, dtype: torch.dtype) -> float:
+ if isinstance(self.log_zero_guard_value, float):
+ return self.log_zero_guard_value
+ return getattr(torch.finfo(dtype), self.log_zero_guard_value)
+
+ def _apply_dithering(self, signals: torch.Tensor) -> torch.Tensor:
+ if self.training and self.dither > 0.0:
+ noise = torch.randn_like(signals) * self.dither
+ signals = signals + noise
+ return signals
+
+ def _apply_preemphasis(self, signals: torch.Tensor) -> torch.Tensor:
+ if self._preemphasis_value is not None:
+ padded = torch.nn.functional.pad(signals, (1, 0))
+ signals = signals - self._preemphasis_value * padded[:, :-1]
+ return signals
+
+ def _compute_output_lengths(self, input_lengths: torch.Tensor) -> torch.Tensor:
+ out_lengths = input_lengths.div(self.hop_length, rounding_mode="floor").add(1).long()
+ return out_lengths
+
+ def _apply_pad_to(self, features: torch.Tensor) -> torch.Tensor:
+ # Only apply during training; else need to capture dynamic shape for exported models
+ if not self.training or self.pad_to == 0 or features.shape[-1] % self.pad_to == 0:
+ return features
+ pad_length = self.pad_to - (features.shape[-1] % self.pad_to)
+ return torch.nn.functional.pad(features, pad=(0, pad_length), value=self.pad_value)
+
+ def _apply_log(self, features: torch.Tensor) -> torch.Tensor:
+ if self._use_log:
+ zero_guard = self._resolve_log_zero_guard_value(features.dtype)
+ if self.log_zero_guard_type == "add":
+ features = features + zero_guard
+ elif self.log_zero_guard_type == "clamp":
+ features = features.clamp(min=zero_guard)
+ else:
+ raise ValueError(f"Unsupported log zero guard type: '{self.log_zero_guard_type}'")
+ features = features.log()
+ return features
+
+ def _extract_spectrograms(self, signals: torch.Tensor) -> torch.Tensor:
+ # Complex FFT needs to be done in single precision
+ with torch.amp.autocast("cuda", enabled=False):
+ features = self._mel_spec_extractor(waveform=signals)
+ return features
+
+ def _apply_normalization(self, features: torch.Tensor, lengths: torch.Tensor, eps: float = 1e-5) -> torch.Tensor:
+ # For consistency, this function always does a masked fill even if not normalizing.
+ mask: torch.Tensor = make_seq_mask_like(lengths=lengths, like=features, time_dim=-1, valid_ones=False)
+ features = features.masked_fill(mask, 0.0)
+ # Maybe don't normalize
+ if self._normalize_strategy is None:
+ return features
+ # Use the log zero guard for the sqrt zero guard
+ guard_value = self._resolve_log_zero_guard_value(features.dtype)
+ if self._normalize_strategy == "per_feature" or self._normalize_strategy == "all_features":
+ # 'all_features' reduces over each sample; 'per_feature' reduces over each channel
+ reduce_dim = 2
+ if self._normalize_strategy == "all_features":
+ reduce_dim = [1, 2]
+ # [B, D, T] -> [B, D, 1] or [B, 1, 1]
+ means = features.sum(dim=reduce_dim, keepdim=True).div(lengths.view(-1, 1, 1))
+ stds = (
+ features.sub(means)
+ .masked_fill(mask, 0.0)
+ .pow(2.0)
+ .sum(dim=reduce_dim, keepdim=True) # [B, D, T] -> [B, D, 1] or [B, 1, 1]
+ .div(lengths.view(-1, 1, 1) - 1) # assume biased estimator
+ .clamp(min=guard_value) # avoid sqrt(0)
+ .sqrt()
+ )
+ features = (features - means) / (stds + eps)
+ else:
+ # Deprecating constant std/mean
+ raise ValueError(f"Unsupported norm type: '{self._normalize_strategy}")
+ features = features.masked_fill(mask, 0.0)
+ return features
+
+ def forward(self, input_signal: torch.Tensor, length: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+ feature_lengths = self._compute_output_lengths(input_lengths=length)
+ signals = self._apply_dithering(signals=input_signal)
+ signals = self._apply_preemphasis(signals=signals)
+ features = self._extract_spectrograms(signals=signals)
+ features = self._apply_log(features=features)
+ features = self._apply_normalization(features=features, lengths=feature_lengths)
+ features = self._apply_pad_to(features=features)
+ return features, feature_lengths
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/multi_head_attention.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/multi_head_attention.py
new file mode 100644
index 000000000..fcfc36f4b
--- /dev/null
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/multi_head_attention.py
@@ -0,0 +1,1133 @@
+# Copyright 2017 Johns Hopkins University (Shinji Watanabe)
+# SPDX-FileCopyrightText: Copyright (c) 2020 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/asr/parts/submodules/multi_head_attention.py; training-only code removed.
+
+"""
+Part of this code is adopted from https://github.com/espnet/espnet
+"""
+
+import math
+from functools import lru_cache
+from typing import List, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.attention
+import torch.nn.functional as F
+
+from ..compat import avoid_float16_autocast_context
+
+__all__ = [
+ "RelPositionMultiHeadAttention",
+ "RelPositionalEncoding",
+ "PositionalEncoding",
+]
+
+INF_VAL = 10000.0
+
+
+class MultiHeadAttention(nn.Module):
+ """Multi-Head Attention layer of Transformer.
+ Args:
+ n_head (int): number of heads
+ n_feat (int): size of the features
+ dropout_rate (float): dropout rate
+ use_bias (bool): whether to remove bias in linear and conv layers
+ use_pytorch_sdpa (bool): use torch sdpa instead of manual attention
+ use_pytorch_sdpa_backends list[str]: list of backend names to use in sdpa. None or empty list means all backends. e.g. ["MATH"]
+ """
+
+ def __init__(
+ self,
+ n_head,
+ n_feat,
+ dropout_rate,
+ max_cache_len=0,
+ use_bias=True,
+ use_pytorch_sdpa=False,
+ use_pytorch_sdpa_backends=None,
+ ):
+ """Construct an MultiHeadedAttention object."""
+ super(MultiHeadAttention, self).__init__()
+ self.use_pytorch_sdpa = use_pytorch_sdpa
+ if self.use_pytorch_sdpa and use_pytorch_sdpa_backends:
+ use_pytorch_sdpa_backends = list(
+ map(
+ lambda backend_name: getattr(torch.nn.attention.SDPBackend, backend_name),
+ use_pytorch_sdpa_backends,
+ )
+ )
+ self.use_pytorch_sdpa_backends = use_pytorch_sdpa_backends
+
+ self.cache_drop_size = None
+ self.use_bias = use_bias
+ self.dropout_rate = dropout_rate
+ assert n_feat % n_head == 0
+ # We assume d_v always equals d_k
+ self.d_k = n_feat // n_head
+ self.s_d_k = math.sqrt(self.d_k)
+ self.h = n_head
+ self.linear_q = nn.Linear(n_feat, n_feat, bias=use_bias)
+ self.linear_k = nn.Linear(n_feat, n_feat, bias=use_bias)
+ self.linear_v = nn.Linear(n_feat, n_feat, bias=use_bias)
+ self.linear_out = nn.Linear(n_feat, n_feat, bias=use_bias)
+ self.dropout = nn.Dropout(p=dropout_rate)
+
+ self._max_cache_len = max_cache_len
+
+ def forward_qkv(self, query, key, value):
+ """Transforms query, key and value.
+ Args:
+ query (torch.Tensor): (batch, time1, size)
+ key (torch.Tensor): (batch, time2, size)
+ value (torch.Tensor): (batch, time2, size)
+ returns:
+ q (torch.Tensor): (batch, head, time1, size)
+ k (torch.Tensor): (batch, head, time2, size)
+ v (torch.Tensor): (batch, head, time2, size)
+ """
+ n_batch = query.size(0)
+ q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k)
+ k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k)
+ v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k)
+ q = q.transpose(1, 2)
+ k = k.transpose(1, 2)
+ v = v.transpose(1, 2)
+
+ return q, k, v
+
+ def forward_attention(self, value, scores, mask):
+ """Compute attention context vector.
+ Args:
+ value (torch.Tensor): (batch, time2, size)
+ scores(torch.Tensor): (batch, time1, time2)
+ mask(torch.Tensor): (batch, time1, time2)
+ returns:
+ value (torch.Tensor): transformed `value` (batch, time2, d_model) weighted by the attention scores
+ """
+ n_batch = value.size(0)
+ if mask is not None:
+ mask = mask.unsqueeze(1) # (batch, 1, time1, time2)
+ scores = scores.masked_fill(mask, -INF_VAL)
+ attn = torch.softmax(scores, dim=-1).masked_fill(mask, 0.0) # (batch, head, time1, time2)
+ else:
+ attn = torch.softmax(scores, dim=-1) # (batch, head, time1, time2)
+
+ p_attn = self.dropout(attn)
+ x = torch.matmul(p_attn, value) # (batch, head, time1, d_k)
+ x = x.transpose(1, 2).reshape(n_batch, -1, self.h * self.d_k) # (batch, time1, d_model)
+
+ return self.linear_out(x) # (batch, time1, d_model)
+
+ def forward(self, query, key, value, mask, pos_emb=None, cache=None):
+ """Compute 'Scaled Dot Product Attention'.
+ Args:
+ query (torch.Tensor): (batch, time1, size)
+ key (torch.Tensor): (batch, time2, size)
+ value(torch.Tensor): (batch, time2, size)
+ mask (torch.Tensor): (batch, time1, time2)
+ cache (torch.Tensor) : (batch, time_cache, size)
+
+ returns:
+ output (torch.Tensor): transformed `value` (batch, time1, d_model) weighted by the query dot key attention
+ cache (torch.Tensor) : (batch, time_cache_next, size)
+ """
+ key, value, query, cache = self.update_cache(key=key, value=value, query=query, cache=cache)
+
+ if torch.is_autocast_enabled():
+ query, key, value = query.to(torch.float32), key.to(torch.float32), value.to(torch.float32)
+
+ # temporary until we solve this more gracefully
+ with avoid_float16_autocast_context():
+ q, k, v = self.forward_qkv(query, key, value)
+
+ if self.use_pytorch_sdpa:
+ n_batch = value.size(0)
+
+ if mask is not None:
+ mask = ~mask.unsqueeze(1)
+
+ dropout_rate = self.dropout_rate if self.training else 0
+ if self.use_pytorch_sdpa_backends:
+ with torch.nn.attention.sdpa_kernel(self.use_pytorch_sdpa_backends):
+ out = torch.nn.functional.scaled_dot_product_attention(
+ q, k, v, attn_mask=mask, dropout_p=dropout_rate
+ )
+ else:
+ out = torch.nn.functional.scaled_dot_product_attention(
+ q, k, v, attn_mask=mask, dropout_p=dropout_rate
+ )
+
+ # this IF block can be deleted when https://github.com/pytorch/pytorch/pull/131863 is in the stable version
+ if mask is not None:
+ all_masked_rows = torch.all(~mask, dim=-1)
+ all_masked_rows.unsqueeze_(-1)
+ out = out.masked_fill(all_masked_rows, 0.0)
+
+ out = out.transpose(1, 2).reshape(n_batch, -1, self.h * self.d_k) # (batch, time1, d_model)
+ out = self.linear_out(out) # (batch, time1, d_model)
+ else:
+ scores = torch.matmul(q, k.transpose(-2, -1)) / self.s_d_k
+ out = self.forward_attention(v, scores, mask)
+
+ if cache is None:
+ return out
+ else:
+ return out, cache
+
+ def update_cache(self, key, value, query, cache):
+ if cache is not None:
+ key = value = torch.cat([cache, key], dim=1)
+ q_keep_size = query.shape[1] - self.cache_drop_size
+ cache = torch.cat([cache[:, q_keep_size:, :], query[:, :q_keep_size, :]], dim=1)
+ return key, value, query, cache
+
+
+class RelPositionMultiHeadAttention(MultiHeadAttention):
+ """Multi-Head Attention layer of Transformer-XL with support of relative positional encoding.
+ Paper: https://arxiv.org/abs/1901.02860
+ Args:
+ n_head (int): number of heads
+ n_feat (int): size of the features
+ dropout_rate (float): dropout rate
+ use_bias (bool): whether to apply bias in linear and conv layers of MultiHeadAttention
+ """
+
+ def __init__(
+ self,
+ n_head,
+ n_feat,
+ dropout_rate,
+ pos_bias_u,
+ pos_bias_v,
+ max_cache_len=0,
+ use_bias=True,
+ use_pytorch_sdpa=False,
+ use_pytorch_sdpa_backends=None,
+ ):
+ """Construct an RelPositionMultiHeadedAttention object."""
+ super().__init__(
+ n_head=n_head,
+ n_feat=n_feat,
+ dropout_rate=dropout_rate,
+ max_cache_len=max_cache_len,
+ use_bias=use_bias,
+ use_pytorch_sdpa=use_pytorch_sdpa,
+ use_pytorch_sdpa_backends=use_pytorch_sdpa_backends,
+ )
+ # linear transformation for positional encoding
+ self.linear_pos = nn.Linear(n_feat, n_feat, bias=False)
+ # these two learnable biases are used in matrix c and matrix d
+ # as described in https://arxiv.org/abs/1901.02860 Section 3.3
+ if pos_bias_u is None or pos_bias_v is None:
+ self.pos_bias_u = nn.Parameter(torch.FloatTensor(self.h, self.d_k))
+ self.pos_bias_v = nn.Parameter(torch.FloatTensor(self.h, self.d_k))
+ # nn.init.normal_(self.pos_bias_u, 0.0, 0.02)
+ # nn.init.normal_(self.pos_bias_v, 0.0, 0.02)
+ nn.init.zeros_(self.pos_bias_u)
+ nn.init.zeros_(self.pos_bias_v)
+ else:
+ self.pos_bias_u = pos_bias_u
+ self.pos_bias_v = pos_bias_v
+
+ def rel_shift(self, x):
+ """Compute relative positional encoding.
+ Args:
+ x (torch.Tensor): (batch, nheads, time, 2*time-1)
+ """
+ b, h, qlen, pos_len = x.size() # (b, h, t1, t2)
+ # need to add a column of zeros on the left side of last dimension to perform the relative shifting
+ x = torch.nn.functional.pad(x, pad=(1, 0)) # (b, h, t1, t2+1)
+ x = x.view(b, h, -1, qlen) # (b, h, t2+1, t1)
+ # need to drop the first row
+ x = x[:, :, 1:].view(b, h, qlen, pos_len) # (b, h, t1, t2)
+ return x
+
+ def forward(self, query, key, value, mask, pos_emb, cache=None):
+ """Compute 'Scaled Dot Product Attention' with rel. positional encoding.
+ Args:
+ query (torch.Tensor): (batch, time1, size)
+ key (torch.Tensor): (batch, time2, size)
+ value(torch.Tensor): (batch, time2, size)
+ mask (torch.Tensor): (batch, time1, time2)
+ pos_emb (torch.Tensor) : (batch, time1, size)
+ cache (torch.Tensor) : (batch, time_cache, size)
+
+ Returns:
+ output (torch.Tensor): transformed `value` (batch, time1, d_model) weighted by the query dot key attention
+ cache (torch.Tensor) : (batch, time_cache_next, size)
+ """
+ key, value, query, cache = self.update_cache(key=key, value=value, query=query, cache=cache)
+
+ if torch.is_autocast_enabled():
+ query, key, value = query.to(torch.float32), key.to(torch.float32), value.to(torch.float32)
+
+ # temporary until we solve this more gracefully
+ with avoid_float16_autocast_context():
+ q, k, v = self.forward_qkv(query, key, value)
+ q = q.transpose(1, 2) # (batch, time1, head, d_k)
+
+ n_batch_pos = pos_emb.size(0)
+ n_batch = value.size(0)
+ p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k)
+ p = p.transpose(1, 2) # (batch, head, time1, d_k)
+
+ # (batch, head, time1, d_k)
+ q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2)
+ # (batch, head, time1, d_k)
+ q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2)
+
+ # compute attention score
+ # first compute matrix a and matrix c
+ # as described in https://arxiv.org/abs/1901.02860 Section 3.3
+ # (batch, head, time1, time2)
+
+ # compute matrix b and matrix d
+ # (batch, head, time1, time2)
+ matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1))
+ matrix_bd = self.rel_shift(matrix_bd)
+
+ if self.use_pytorch_sdpa:
+ scale_factor = 1 / math.sqrt(q_with_bias_u.size(-1))
+ matrix_bd = matrix_bd[:, :, :, : k.size(-2)] * scale_factor
+
+ if mask is not None:
+ mask = mask.unsqueeze(1)
+ matrix_bd.masked_fill_(mask, -INF_VAL)
+
+ dropout_rate = self.dropout_rate if self.training else 0
+ if self.use_pytorch_sdpa_backends:
+ with torch.nn.attention.sdpa_kernel(self.use_pytorch_sdpa_backends):
+ out = torch.nn.functional.scaled_dot_product_attention(
+ q_with_bias_u, k, v, attn_mask=matrix_bd, dropout_p=dropout_rate
+ )
+ else:
+ out = torch.nn.functional.scaled_dot_product_attention(
+ q_with_bias_u, k, v, attn_mask=matrix_bd, dropout_p=dropout_rate
+ )
+
+ # this IF block can be deleted when https://github.com/pytorch/pytorch/pull/131863 is in the stable version
+ if mask is not None:
+ all_masked_rows = torch.all(mask, dim=-1)
+ all_masked_rows.unsqueeze_(-1)
+ all_masked_rows = all_masked_rows.expand(-1, out.size(1), -1, out.size(-1))
+ out = out.masked_fill(all_masked_rows, 0.0)
+
+ out = out.transpose(1, 2).reshape(n_batch, -1, self.h * self.d_k) # (batch, time1, d_model)
+ out = self.linear_out(out) # (batch, time1, d_model)
+ else:
+ # drops extra elements in the matrix_bd to match the matrix_ac's size
+ matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1))
+ matrix_bd = matrix_bd[:, :, :, : matrix_ac.size(-1)]
+ scores = (matrix_ac + matrix_bd) / self.s_d_k # (batch, head, time1, time2)
+ out = self.forward_attention(v, scores, mask)
+
+ if cache is None:
+ return out
+ else:
+ return out, cache
+
+
+class RelPositionMultiHeadAttentionLongformer(RelPositionMultiHeadAttention):
+ """Multi-Head Attention layer of Transformer-XL with sliding window local+global attention from Longformer.
+ Partially adapted from allenai (https://github.com/allenai/longformer/blob/master/longformer/sliding_chunks.py)
+ and huggingface (https://github.com/huggingface/transformers/blob/main/src/transformers/models/longformer/modeling_longformer.py)
+ Paper: https://arxiv.org/abs/1901.02860 (Transformer-XL),
+ https://arxiv.org/abs/2004.05150 (Longformer)
+ Args:
+ n_head (int): number of heads
+ n_feat (int): size of the features
+ dropout_rate (float): dropout rate
+ pos_bias_u (Tensor): the positional bias matrix U
+ pos_bias_v (Tensor): the positional bias matrix V
+ att_context_size (List[int]): List of 2 ints corresponding to left and right attention context sizes.
+ max_cache_len (int): the maximum size of cache
+ global_tokens (int): number of tokens to be used for global attention
+ global_tokens_spacing (int): how far apart the global tokens are
+ global_attn_separate (bool): whether the q, k, v layers used for global tokens should be separate
+ use_bias (bool): whether to apply bias in linear and conv layers of MultiHeadAttention
+ """
+
+ def __init__(
+ self,
+ n_head,
+ n_feat,
+ dropout_rate,
+ pos_bias_u,
+ pos_bias_v,
+ att_context_size,
+ max_cache_len=0,
+ global_tokens=0,
+ global_tokens_spacing=1,
+ global_attn_separate=False,
+ use_bias=True,
+ use_pytorch_sdpa=False,
+ use_pytorch_sdpa_backends=None,
+ ):
+ """Construct an RelPositionMultiHeadAttentionLongformer object."""
+ super().__init__(
+ n_head=n_head,
+ n_feat=n_feat,
+ dropout_rate=dropout_rate,
+ pos_bias_u=pos_bias_u,
+ pos_bias_v=pos_bias_v,
+ max_cache_len=max_cache_len,
+ use_bias=use_bias,
+ use_pytorch_sdpa=use_pytorch_sdpa,
+ use_pytorch_sdpa_backends=use_pytorch_sdpa_backends,
+ )
+
+ if use_pytorch_sdpa:
+ raise NotImplementedError("Not implemented for Longformer yet")
+
+ self.att_context_size = att_context_size
+ self.global_tokens = global_tokens
+ self.global_tokens_spacing = global_tokens_spacing
+ self.global_attn_separate = global_attn_separate
+
+ if self.global_attn_separate:
+ self.global_q = nn.Linear(n_feat, n_feat, bias=use_bias)
+ self.global_k = nn.Linear(n_feat, n_feat, bias=use_bias)
+ self.global_v = nn.Linear(n_feat, n_feat, bias=use_bias)
+
+ def forward(self, query, key, value, pad_mask, pos_emb, cache=None):
+ """Compute Scaled Dot Product Local Attention with rel. positional encoding. using overlapping chunks
+ Args:
+ query (torch.Tensor): (batch, time, size)
+ key (torch.Tensor): (batch, time, size)
+ value(torch.Tensor): (batch, time, size)
+ pad_mask (torch.Tensor): (batch, time)
+ pos_emb (torch.Tensor) : (batch, 2w + 1, size)
+ cache (torch.Tensor) : (batch, time_cache, size)
+ Returns:
+ output (torch.Tensor): transformed `value` (batch, time1, d_model) weighted by the query dot key attention
+ cache (torch.Tensor) : (batch, time_cache_next, size)
+ """
+
+ key, value, query, cache = self.update_cache(key=key, value=value, query=query, cache=cache)
+
+ if torch.is_autocast_enabled():
+ query, key, value = query.to(torch.float32), key.to(torch.float32), value.to(torch.float32)
+
+ # temporary until we solve this more gracefully
+ with avoid_float16_autocast_context():
+ q, k, v = self.forward_qkv(query, key, value)
+ n_batch, _, T, _ = q.size()
+
+ w = max(self.att_context_size[0], self.att_context_size[1])
+ if w <= 0:
+ raise ValueError("When using local attention, context size must be set > 0")
+ pad_len = (2 * w - T % (2 * w)) % (2 * w) # pad time to 2w
+ q = F.pad(q, (0, 0, 0, pad_len)) # (batch, head, time, size)
+ k = F.pad(k, (0, 0, 0, pad_len)) # (batch, head, time, size)
+ v = F.pad(v, (0, 0, 0, pad_len)) # (batch, head, time, size)
+ mask = F.pad(pad_mask, (0, pad_len), value=1.0)
+
+ q_with_bias_u = q + self.pos_bias_u.unsqueeze(1) # (batch, head, time, size)
+ q_with_bias_v = q + self.pos_bias_v.unsqueeze(1) # (batch, head, time, size)
+
+ diagonal_matrix_ac = self.sliding_chunks_matmul_qk(
+ q_with_bias_u, k, w, padding_value=0.0
+ ) # (batch, head, time, 2w + 1)
+
+ # add relative positional embedding
+
+ n_batch_pos = pos_emb.size(0)
+ p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k).transpose(1, 2)
+ # (batch, head, 2w, size)
+ diagonal_matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1))
+ # (batch, head, time, 2w + 1)
+
+ start_pos = w - self.att_context_size[0]
+ end_pos = w + self.att_context_size[1]
+
+ diagonal_matrix_ac[:, :, :, : self.att_context_size[0]] += diagonal_matrix_bd[
+ :, :, :, : self.att_context_size[0]
+ ]
+ diagonal_matrix_ac[:, :, :, -(self.att_context_size[1] + 1) :] += diagonal_matrix_bd[
+ :, :, :, self.att_context_size[0] :
+ ]
+ scores = diagonal_matrix_ac / self.s_d_k
+ # (batch, head, time, 2w + 1)
+
+ # mask invalid positions
+ scores[:, :, :, :start_pos] = -INF_VAL
+ scores[:, :, :, end_pos + 1 :] = -INF_VAL
+
+ # This implementation is fast and takes very little memory because num_heads x hidden_size = 1
+ # from (bsz x seq_len) to (bsz x num_heads x seqlen x hidden_size)
+ mask = mask.unsqueeze(dim=1).unsqueeze(dim=-1)
+ # cast to float/half then replace 1's with -inf
+ float_mask = mask.type_as(scores).masked_fill(mask, -INF_VAL)
+ ones = float_mask.new_ones(size=float_mask.size()) # tensor of ones
+ # diagonal mask with zeros everywhere and -inf inplace of padding
+ d_mask = self.sliding_chunks_matmul_qk(ones, float_mask, w, padding_value=0.0)
+ # (batch, head, time, 2w + 1)
+
+ scores += d_mask
+
+ if self.global_tokens > 0:
+ # create q, k, v for global attn
+ if self.global_attn_separate:
+ global_q = self.global_q(query).view(n_batch, -1, self.h, self.d_k)
+ global_k = self.global_k(key).view(n_batch, -1, self.h, self.d_k)
+ global_v = self.global_v(value).view(n_batch, -1, self.h, self.d_k)
+ global_q = global_q.transpose(1, 2)
+ global_k = global_k.transpose(1, 2)
+ global_v = global_v.transpose(1, 2)
+ global_q = F.pad(global_q, (0, 0, 0, pad_len)) # (batch, head, time, size)
+ global_k = F.pad(global_k, (0, 0, 0, pad_len)) # (batch, head, time, size)
+ global_v = F.pad(global_v, (0, 0, 0, pad_len)) # (batch, head, time, size)
+ else:
+ global_q, global_k, global_v = q, k, v
+
+ global_q /= self.s_d_k
+
+ # assign which tokens are global
+ is_index_global_attn = torch.zeros_like(pad_mask)
+ is_index_global_attn[
+ :, : self.global_tokens * self.global_tokens_spacing : self.global_tokens_spacing
+ ] = 1.0
+
+ # compute global attn indices
+ (
+ max_num_global_attn_indices,
+ is_index_global_attn_nonzero,
+ is_local_index_global_attn_nonzero,
+ is_local_index_no_global_attn_nonzero,
+ ) = self._get_global_attn_indices(is_index_global_attn=is_index_global_attn)
+
+ # calculate global attn probs with global keys
+ # (batch, time, head, max_num_global_attn_indices)
+ global_key_attn = self._compute_global_key_attn(
+ query=global_q.transpose(1, 2),
+ key=global_k.transpose(1, 2),
+ max_num_global_attn_indices=max_num_global_attn_indices,
+ is_index_global_attn_nonzero=is_index_global_attn_nonzero,
+ is_local_index_global_attn_nonzero=is_local_index_global_attn_nonzero,
+ is_local_index_no_global_attn_nonzero=is_local_index_no_global_attn_nonzero,
+ ).transpose(1, 2)
+
+ # concat to local_attn_probs
+ # (batch, time, head, max_num_global_attn_indices + 2*w)
+ scores = torch.cat((global_key_attn, scores), dim=-1)
+
+ # free memory
+ del global_key_attn
+
+ attn = torch.softmax(scores, dim=-1).masked_fill(mask, 0.0)
+ p_attn = self.dropout(attn)
+ # (batch, head, time, 2w + 1)
+
+ if self.global_tokens > 0:
+ # compute sum of global and local attn
+ out = self._compute_attn_output_with_global_indices(
+ value=v,
+ attn_probs=p_attn,
+ max_num_global_attn_indices=max_num_global_attn_indices,
+ is_index_global_attn_nonzero=is_index_global_attn_nonzero,
+ is_local_index_global_attn_nonzero=is_local_index_global_attn_nonzero,
+ w=w,
+ )
+ else:
+ # compute local attn only
+ out = self.sliding_chunks_matmul_pv(p_attn, v, w)
+
+ out = out.reshape(n_batch, -1, self.h * self.d_k)[:, :T]
+
+ if self.global_tokens > 0:
+ out_global_to_all = self._compute_out_global_to_all(
+ query=global_q,
+ key=global_k,
+ value=global_v,
+ max_num_global_attn_indices=max_num_global_attn_indices,
+ is_local_index_global_attn_nonzero=is_local_index_global_attn_nonzero,
+ is_index_global_attn_nonzero=is_index_global_attn_nonzero,
+ is_local_index_no_global_attn_nonzero=is_local_index_no_global_attn_nonzero,
+ is_index_masked=mask,
+ )
+
+ # overwrite values with global attention
+ out[is_index_global_attn_nonzero] = out_global_to_all
+
+ ret = self.linear_out(out)
+
+ if cache is None:
+ return ret
+ else:
+ return ret, cache
+
+ def _get_global_attn_indices(self, is_index_global_attn: torch.Tensor) -> Tuple:
+ """
+ Compute global attention indices.
+
+ Args:
+ is_index_global_attn (torch.Tensor): (batch, time) A boolean tensor indicating if an index is a global attention index.
+
+ Returns:
+ max_num_global_attn_indices (int): Maximum number of global attention indices in the batch.
+ is_index_global_attn_nonzero (tuple): Indices of global attention (non-zero elements).
+ is_local_index_global_attn_nonzero (tuple): Indices of non-padding values within global attention indices.
+ is_local_index_no_global_attn_nonzero (tuple): Indices of padding values within global attention indices.
+ """
+ # Calculate the number of global attention indices in the batch
+ num_global_attn_indices = is_index_global_attn.long().sum(dim=1)
+
+ # Find the maximum number of global attention indices in the batch
+ max_num_global_attn_indices = num_global_attn_indices.max()
+
+ # Get the indices of global attention (non-zero elements)
+ is_index_global_attn_nonzero = is_index_global_attn.nonzero(as_tuple=True)
+
+ # Create a helper tensor to find the local indices of global attention
+ is_local_index_global_attn = torch.arange(
+ max_num_global_attn_indices, device=is_index_global_attn.device
+ ) < num_global_attn_indices.unsqueeze(dim=-1)
+
+ # Find the non-padding values within global attention indices
+ is_local_index_global_attn_nonzero = is_local_index_global_attn.nonzero(as_tuple=True)
+
+ # Find the padding values within global attention indices
+ is_local_index_no_global_attn_nonzero = (is_local_index_global_attn == 0).nonzero(as_tuple=True)
+
+ return (
+ max_num_global_attn_indices,
+ is_index_global_attn_nonzero,
+ is_local_index_global_attn_nonzero,
+ is_local_index_no_global_attn_nonzero,
+ )
+
+ def _compute_global_key_attn(
+ self,
+ key: torch.Tensor,
+ query: torch.Tensor,
+ max_num_global_attn_indices: int,
+ is_index_global_attn_nonzero: tuple,
+ is_local_index_global_attn_nonzero: tuple,
+ is_local_index_no_global_attn_nonzero: tuple,
+ ) -> torch.Tensor:
+ """
+ Compute the attention probabilities using only global key vectors.
+
+ Args:
+ key (torch.Tensor): (batch, time, head, head_dim) The key vectors.
+ query (torch.Tensor): (batch, time, head, head_dim) The query vectors.
+ max_num_global_attn_indices (int): Maximum number of global attention indices in the batch.
+ is_index_global_attn_nonzero (tuple): Indices of global attention (non-zero elements).
+ is_local_index_global_attn_nonzero (tuple): Non-padding values within global attention indices.
+ is_local_index_no_global_attn_nonzero (tuple): Padding values within global attention indices.
+
+ Returns:
+ attn_probs_from_global_key (torch.Tensor): (batch, time, head, max_num_global_attn_indices) The computed attention probabilities using only global key vectors.
+ """
+ batch_size = key.shape[0]
+
+ # create only global key vectors
+ key_only_global = key.new_zeros(batch_size, max_num_global_attn_indices, self.h, self.d_k)
+
+ key_only_global[is_local_index_global_attn_nonzero] = key[is_index_global_attn_nonzero]
+
+ # (batch_size, seq_len, head, max_num_global_attn_indices)
+ attn_probs_from_global_key = torch.einsum("blhd,bshd->blhs", (query, key_only_global))
+
+ # need to transpose since ONNX export only supports consecutive indexing: https://pytorch.org/docs/stable/onnx.html#writes-sets
+ attn_probs_from_global_key = attn_probs_from_global_key.transpose(1, 3)
+ attn_probs_from_global_key[
+ is_local_index_no_global_attn_nonzero[0], is_local_index_no_global_attn_nonzero[1], :, :
+ ] = torch.finfo(attn_probs_from_global_key.dtype).min
+ attn_probs_from_global_key = attn_probs_from_global_key.transpose(1, 3)
+
+ return attn_probs_from_global_key
+
+ def _compute_attn_output_with_global_indices(
+ self,
+ value: torch.Tensor,
+ attn_probs: torch.Tensor,
+ max_num_global_attn_indices: int,
+ is_index_global_attn_nonzero: tuple,
+ is_local_index_global_attn_nonzero: tuple,
+ w: int,
+ ) -> torch.Tensor:
+ """
+ Compute the attention output with global indices.
+
+ Args:
+ value (torch.Tensor): (batch, head, time, head_dim) The value vectors for global attention.
+ attn_probs (torch.Tensor): (batch, time, head, 2w) The attention probabilities.
+ max_num_global_attn_indices (int): Maximum number of global attention indices in the batch.
+ is_index_global_attn_nonzero (tuple): Indices of global attention (non-zero elements).
+ is_local_index_global_attn_nonzero (tuple): Non-padding values within global attention indices.
+ w (int): Local context size
+ Returns:
+ torch.Tensor: (batch, time, head x head_dim) The attention output of all tokens attending to global.
+ """
+ batch_size, time = attn_probs.shape[0], attn_probs.shape[2]
+
+ value = value.transpose(1, 2)
+
+ # get value vectors for global only
+ value_vectors_only_global = value.new_zeros(batch_size, max_num_global_attn_indices, self.h, self.d_k)
+ value_vectors_only_global[is_local_index_global_attn_nonzero] = value[is_index_global_attn_nonzero]
+
+ # cut local attn probs to global only
+ attn_probs_only_global = attn_probs.narrow(-1, 0, max_num_global_attn_indices)
+ # compute attn output only global
+ attn_output_only_global = torch.matmul(
+ attn_probs_only_global.clone(), value_vectors_only_global.transpose(1, 2).clone()
+ ).transpose(1, 2)
+
+ # reshape attn probs
+ attn_probs_without_global = attn_probs.narrow(
+ -1, max_num_global_attn_indices, attn_probs.size(-1) - max_num_global_attn_indices
+ ).contiguous()
+
+ # compute attn output with global
+ attn_output_without_global = self.sliding_chunks_matmul_pv(attn_probs_without_global, value.transpose(1, 2), w)
+
+ return attn_output_only_global + attn_output_without_global
+
+ def _compute_out_global_to_all(
+ self,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ max_num_global_attn_indices: int,
+ is_local_index_global_attn_nonzero: tuple,
+ is_index_global_attn_nonzero: tuple,
+ is_local_index_no_global_attn_nonzero: tuple,
+ is_index_masked: torch.Tensor,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """
+ Compute the attention output of global tokens attending to all.
+
+ Args:
+ query (torch.Tensor): (batch, head, time, head_dim) The queries for global attention.
+ key (torch.Tensor): (batch, head, time, head_dim) The keys for global attention.
+ value (torch.Tensor): (batch, head, time, head_dim) The values for global attention.
+ max_num_global_attn_indices (int): Maximum number of global attention indices in the batch.
+ is_local_index_global_attn_nonzero (tuple): Non-padding values within global attention indices.
+ is_index_global_attn_nonzero (tuple): Indices of global attention (non-zero elements).
+ is_local_index_no_global_attn_nonzero (tuple): Padding values within global attention indices.
+ is_index_masked (torch.Tensor): (batch, time) A boolean tensor indicating if an index is masked.
+
+ Returns:
+ global_attn_output (torch.Tensor): (batch, max_num_global_attn_indices, head x head_dim)
+ The attention output of global tokens attending to all.
+ """
+
+ batch_size = key.shape[0]
+ seq_len = key.shape[2]
+
+ global_k = key.reshape(batch_size * self.h, -1, self.d_k)
+ global_v = value.reshape(batch_size * self.h, -1, self.d_k)
+
+ global_q = query.transpose(1, 2)
+ global_q_from_global = global_q.new_zeros(batch_size, max_num_global_attn_indices, self.h, self.d_k)
+ global_q_from_global[is_local_index_global_attn_nonzero] = global_q[is_index_global_attn_nonzero]
+ global_q_from_global = global_q_from_global.transpose(0, 1).reshape(batch_size * self.h, -1, self.d_k)
+
+ # compute attn scores
+ global_attn_scores = torch.bmm(global_q_from_global, global_k.transpose(1, 2))
+ global_attn_scores = global_attn_scores.view(batch_size, self.h, max_num_global_attn_indices, seq_len)
+
+ # need to transpose since ONNX export only supports consecutive indexing: https://pytorch.org/docs/stable/onnx.html#writes-sets
+ global_attn_scores = global_attn_scores.transpose(1, 2)
+ global_attn_scores[is_local_index_no_global_attn_nonzero[0], is_local_index_no_global_attn_nonzero[1], :, :] = (
+ torch.finfo(global_attn_scores.dtype).min
+ )
+ global_attn_scores = global_attn_scores.transpose(1, 2)
+
+ global_attn_scores = global_attn_scores.masked_fill(
+ is_index_masked.transpose(2, 3),
+ torch.finfo(global_attn_scores.dtype).min,
+ )
+
+ global_attn_scores = global_attn_scores.view(batch_size * self.h, max_num_global_attn_indices, seq_len)
+
+ # compute global attn probs
+ if self.training:
+ global_attn_probs_float = nn.functional.softmax(global_attn_scores, dim=-1, dtype=torch.float32)
+ else:
+ global_attn_probs_float = nn.functional.softmax(global_attn_scores, dim=-1)
+
+ global_attn_probs = self.dropout(global_attn_probs_float)
+
+ # global attn output
+ global_attn_output = torch.bmm(global_attn_probs, global_v)
+ global_attn_output = global_attn_output.view(batch_size, self.h, max_num_global_attn_indices, self.d_k)
+
+ global_attn_output = global_attn_output[
+ is_local_index_global_attn_nonzero[0], :, is_local_index_global_attn_nonzero[1]
+ ]
+
+ global_attn_output = global_attn_output.reshape(global_attn_output.shape[0], -1)
+
+ return global_attn_output
+
+ # Longformer implementation for overlap case
+ #
+ def _skew(self, x: torch.Tensor, direction: List[int], padding_value: float) -> torch.Tensor:
+ """Convert diagonals into columns (or columns into diagonals depending on `direction`
+
+ Args:
+ x (torch.Tensor): (batch x head, chunk_count, 2w, 2w)
+ direction (List[int]): padding directions
+ padding_value (float): value to pad with
+
+ Returns:
+ output (torch.Tensor): (batch x head, chunk_count, 2w, 2w + 1)
+
+ """
+ x_padded = F.pad(x, direction, value=padding_value)
+ x_padded = x_padded.view(*x_padded.size()[:-2], x_padded.size(-1), x_padded.size(-2))
+ return x_padded
+
+ def _skew2(self, x: torch.Tensor, padding_value: float) -> torch.Tensor:
+ """Shift every row 1 step to right converting columns into diagonals
+
+ Args:
+ x (torch.Tensor): (batch x head, chunks_count + 1, w, 2w + 1)
+ padding_value (float): value to pad with
+
+ Returns:
+ output (torch.Tensor): (batch x head, chunks_count + 1, w, 3w)
+ """
+ # X = B x C x M x L
+ B, C, M, L = x.size()
+ x = F.pad(x, (0, M + 1), value=padding_value) # B x C x M x (L+M+1)
+ x = x.view(B, C, -1) # B x C x ML+MM+M
+ x = x[:, :, :-M] # B x C x ML+MM
+ x = x.view(B, C, M, M + L) # B x C, M x L+M
+ x = x[:, :, :, :-1]
+ return x
+
+ def _chunk_overlap(self, x: torch.Tensor, w: int) -> torch.Tensor:
+ """Convert into overlapping chunks.
+
+ Args:
+ x (torch.Tensor): # (batch x head, time, size)
+ w (int): Chunk overlap size
+
+ Returns:
+ output (torch.Tensor): # (batch x head, chunk_count, 2w, size)
+ """
+
+ # non-overlapping chunks of size = 2w
+ x = x.view(x.size(0), x.size(1) // (w * 2), w * 2, x.size(2))
+
+ # use `as_strided` to make the chunks overlap with an overlap size = w
+ chunk_size = list(x.size())
+ chunk_size[1] = chunk_size[1] * 2 - 1
+
+ chunk_stride = list(x.stride())
+ chunk_stride[1] = chunk_stride[1] // 2
+ return x.as_strided(size=chunk_size, stride=chunk_stride)
+
+ @lru_cache()
+ def _get_invalid_locations_mask(self, w: int, device: str):
+ diagonals_list = []
+ for j in range(-w, 1):
+ diagonal_mask = torch.zeros(w, device="cpu", dtype=torch.uint8)
+ diagonal_mask[:-j] = 1
+ diagonals_list.append(diagonal_mask)
+
+ mask = torch.stack(diagonals_list, dim=-1)
+ mask = mask[None, None, :, :]
+
+ ending_mask = mask.flip(dims=(2, 3)).bool().to(device)
+ return mask.bool().to(device), ending_mask
+
+ def mask_invalid_locations(
+ self,
+ input_tensor: torch.Tensor,
+ w: int,
+ ):
+ """
+ Mask locations invalid for the sliding window attention
+
+ Args:
+ input_tensor (torch.Tensor): # (batch x head, time, size)
+ w (int): Chunk overlap size
+ """
+ beginning_mask, ending_mask = self._get_invalid_locations_mask(w, input_tensor.device)
+ seq_len = input_tensor.size(2)
+ beginning_input = input_tensor[:, :, :w, : w + 1]
+ beginning_mask = beginning_mask[:, :, :seq_len].expand(beginning_input.size())
+ beginning_input.masked_fill_(beginning_mask, -float("inf"))
+
+ ending_input = input_tensor[:, :, -w:, -(w + 1) :]
+ ending_mask = ending_mask[:, :, -seq_len:].expand(ending_input.size())
+ ending_input.masked_fill_(ending_mask, -float("inf"))
+
+ def sliding_chunks_matmul_qk(self, q: torch.Tensor, k: torch.Tensor, w: int, padding_value: float) -> torch.Tensor:
+ """Matrix multiplication of query x key tensors using with a sliding window attention pattern.
+ This implementation splits the input into overlapping chunks of size 2w
+ with an overlap of size w
+
+ Args:
+ q (torch.Tensor): (batch, head, time, size)
+ k (torch.Tensor): (batch, head, time, size)
+ w (int): Chunk overlap size
+ padding_value (float): Value to pad with
+
+ Returns:
+ output (torch.Tensor): (batch, head, time, 2w + 1)
+ """
+ bsz, num_heads, seqlen, head_dim = q.size()
+ assert seqlen % (w * 2) == 0
+ assert q.size() == k.size()
+
+ chunks_count = seqlen // w - 1
+
+ # group bsz and num_heads dimensions into one, then chunk seqlen into chunks of size w * 2
+ q = q.reshape(bsz * num_heads, seqlen, head_dim)
+ k = k.reshape(bsz * num_heads, seqlen, head_dim)
+
+ chunk_q = self._chunk_overlap(q, w) # (batch x head, chunk_count, 2w, size)
+ chunk_k = self._chunk_overlap(k, w) # (batch x head, chunk_count, 2w, size)
+
+ # matrix multiplication
+ # bcxd: bsz*num_heads x chunks x 2w x head_dim
+ # bcyd: bsz*num_heads x chunks x 2w x head_dim
+ # bcxy: bsz*num_heads x chunks x 2w x 2w
+ chunk_attn = torch.einsum("bcxd,bcyd->bcxy", (chunk_q, chunk_k)) # multiply
+ # (batch x head, chunk_count, 2w, 2w)
+
+ # convert diagonals into columns
+ diagonal_chunk_attn = self._skew(chunk_attn, direction=(0, 0, 0, 1), padding_value=padding_value)
+ # (batch x head, chunk_count, 2w, 2w + 1)
+
+ # allocate space for the overall attention matrix where the chunks are combined. The last dimension
+ # has (w * 2 + 1) columns. The first (w) columns are the w lower triangles (attention from a word to
+ # w previous words). The following column is attention score from each word to itself, then
+ # followed by w columns for the upper triangle.
+
+ diagonal_attn = diagonal_chunk_attn.new_empty((bsz * num_heads, chunks_count + 1, w, w * 2 + 1))
+ # (batch x head, chunk_count + 1, w, 2w + 1)
+
+ # copy parts from diagonal_chunk_attn into the compined matrix of attentions
+ # - copying the main diagonal and the upper triangle
+ diagonal_attn[:, :-1, :, w:] = diagonal_chunk_attn[:, :, :w, : w + 1]
+ diagonal_attn[:, -1, :, w:] = diagonal_chunk_attn[:, -1, w:, : w + 1]
+ # - copying the lower triangle
+ diagonal_attn[:, 1:, :, :w] = diagonal_chunk_attn[:, :, -(w + 1) : -1, w + 1 :]
+ diagonal_attn[:, 0, 1:w, 1:w] = diagonal_chunk_attn[:, 0, : w - 1, 1 - w :]
+
+ # separate bsz and num_heads dimensions again
+ diagonal_attn = diagonal_attn.view(bsz, num_heads, seqlen, 2 * w + 1)
+ # (batch, head, time, 2w + 1)
+
+ self.mask_invalid_locations(diagonal_attn, w)
+
+ return diagonal_attn
+
+ def sliding_chunks_matmul_pv(self, prob: torch.Tensor, v: torch.Tensor, w: int):
+ """Same as sliding_chunks_matmul_qk but for prob and value tensors.
+
+ Args:
+ prob (torch.Tensor): (batch, head, time, size)
+ v (torch.Tensor): (batch, head, time, size)
+ w (int): Chunk overlap size
+
+ Returns:
+ output (torch.Tensor): (batch, time, head, size)
+ """
+ bsz, num_heads, seqlen, head_dim = v.size()
+ chunks_count = seqlen // w - 1
+ # group bsz and num_heads dimensions into one, then chunk seqlen into chunks of size 2w
+ chunk_prob = prob.reshape(bsz * num_heads, seqlen // w, w, 2 * w + 1)
+ # (batch x head, chunks_count + 1, w, 2w + 1)
+
+ # group bsz and num_heads dimensions into one
+ v = v.reshape(bsz * num_heads, seqlen, head_dim)
+ # (batch x head, time, size)
+
+ # pad seqlen with w at the beginning of the sequence and another w at the end
+ padded_v = F.pad(v, (0, 0, w, w), value=-1)
+ # (batch x head, time + 2w, size)
+
+ # chunk padded_v into chunks of size 3w and an overlap of size w
+ chunk_v_size = (bsz * num_heads, chunks_count + 1, 3 * w, head_dim)
+ chunk_v_stride = padded_v.stride()
+ chunk_v_stride = chunk_v_stride[0], w * chunk_v_stride[1], chunk_v_stride[1], chunk_v_stride[2]
+ chunk_v = padded_v.as_strided(size=chunk_v_size, stride=chunk_v_stride)
+ # (batch x head, chunks_count + 1, 3w, size)
+
+ skewed_prob = self._skew2(chunk_prob, padding_value=0)
+ # (batch x head, chunks_count + 1, w, 3w)
+
+ context = torch.einsum("bcwd,bcdh->bcwh", (skewed_prob, chunk_v))
+ # (batch x head, chunks_count + 1, w, size)
+
+ return context.view(bsz, num_heads, seqlen, head_dim).transpose(1, 2)
+
+
+class PositionalEncoding(torch.nn.Module):
+ """Fixed sinusoidal positional encoding.
+ Args:
+ d_model (int): embedding dim
+ dropout_rate (float): dropout rate
+ max_len (int): maximum input length
+ xscale (bool): whether to scale the input by sqrt(d_model)
+ dropout_rate_emb (float): dropout rate for the positional embeddings
+ """
+
+ def __init__(self, d_model, dropout_rate, max_len=5000, xscale=None, dropout_rate_emb=0.0):
+ """Construct an PositionalEncoding object."""
+ super(PositionalEncoding, self).__init__()
+ self.d_model = d_model
+ self.xscale = xscale
+ self.dropout = torch.nn.Dropout(p=dropout_rate)
+ self.max_len = max_len
+ if dropout_rate_emb > 0:
+ self.dropout_emb = nn.Dropout(dropout_rate_emb)
+ else:
+ self.dropout_emb = None
+
+ def create_pe(self, positions, dtype):
+ pos_length = positions.size(0)
+ pe = torch.zeros(pos_length, self.d_model, device=positions.device)
+ div_term = torch.exp(
+ torch.arange(0, self.d_model, 2, dtype=torch.float32, device=positions.device)
+ * -(math.log(INF_VAL) / self.d_model)
+ )
+ pe[:, 0::2] = torch.sin(positions * div_term)
+ pe[:, 1::2] = torch.cos(positions * div_term)
+ pe = pe.unsqueeze(0).to(dtype)
+ if hasattr(self, "pe"):
+ self.pe = pe
+ else:
+ self.register_buffer("pe", pe, persistent=False)
+
+ def extend_pe(self, length, device, dtype):
+ """Reset and extend the positional encodings if needed."""
+ if hasattr(self, "pe") and self.pe.size(1) >= length:
+ return
+ positions = torch.arange(0, length, dtype=torch.float32, device=device).unsqueeze(1)
+ self.create_pe(positions=positions, dtype=dtype)
+
+ def forward(self, x: torch.Tensor, cache_len=0):
+ """Adds positional encoding.
+ Args:
+ x (torch.Tensor): Input. Its shape is (batch, time, feature_size)
+ cache_len (int): the size of the cache which is used to shift positions
+ Returns:
+ x+pos_emb (torch.Tensor): Its shape is (batch, time, feature_size)
+ pos_emb (torch.Tensor): Its shape is (1, time, feature_size)
+ """
+ input_len = x.size(1) + cache_len
+ if self.xscale:
+ x = x * self.xscale
+ pos_emb = self.pe[:, :input_len]
+ if self.dropout_emb:
+ pos_emb = self.dropout_emb(pos_emb)
+ x = x + pos_emb
+ return self.dropout(x), pos_emb
+
+
+class RelPositionalEncoding(PositionalEncoding):
+ """Relative positional encoding for TransformerXL's layers
+ See : Appendix B in https://arxiv.org/abs/1901.02860
+ Args:
+ d_model (int): embedding dim
+ dropout_rate (float): dropout rate
+ max_len (int): maximum input length
+ xscale (bool): whether to scale the input by sqrt(d_model)
+ dropout_rate_emb (float): dropout rate for the positional embeddings
+ """
+
+ def extend_pe(self, length, device, dtype):
+ """Reset and extend the positional encodings if needed."""
+ needed_size = 2 * length - 1
+ if hasattr(self, "pe") and self.pe.size(1) >= needed_size:
+ return
+ # positions would be from negative numbers to positive
+ # positive positions would be used for left positions and negative for right positions
+ positions = torch.arange(length - 1, -length, -1, dtype=torch.float32, device=device).unsqueeze(1)
+ self.create_pe(positions=positions, dtype=dtype)
+
+ def forward(self, x, cache_len=0):
+ """Compute positional encoding.
+ Args:
+ x (torch.Tensor): Input. Its shape is (batch, time, feature_size)
+ cache_len (int): the size of the cache which is used to shift positions
+ Returns:
+ x (torch.Tensor): Its shape is (batch, time, feature_size)
+ pos_emb (torch.Tensor): Its shape is (1, time, feature_size)
+ """
+
+ if self.xscale:
+ x = x * self.xscale
+
+ # center_pos would be the index of position 0
+ # negative positions would be used for right and positive for left tokens
+ # for input of length L, 2*L-1 positions are needed, positions from (L-1) to -(L-1)
+ input_len = x.size(1) + cache_len
+ center_pos = self.pe.size(1) // 2 + 1
+ start_pos = center_pos - input_len
+ end_pos = center_pos + input_len - 1
+ pos_emb = self.pe[:, start_pos:end_pos]
+ if self.dropout_emb:
+ pos_emb = self.dropout_emb(pos_emb)
+ return self.dropout(x), pos_emb
+
+
+class LocalAttRelPositionalEncoding(PositionalEncoding):
+ """Relative positional encoding for sliding window attention or chunked attention.
+ See above for relative positional encoding based on Transformer-XL paper
+ Args:
+ left_chunk_size (int): number of frames to in past chunks
+ chunk size (int): number of frames (max frames if using multimode) in current chunk
+ d_model (int): embedding dim
+ dropout_rate (float): dropout rate
+ max_len (int): maximum input length
+ xscale (bool): whether to scale the input by sqrt(d_model)
+ dropout_rate_emb (float): dropout rate for the positional embeddings
+ """
+
+ def __init__(self, att_context_size, **kwargs):
+ super(LocalAttRelPositionalEncoding, self).__init__(**kwargs)
+ self.left_context = att_context_size[0]
+ self.right_context = att_context_size[1]
+
+ def extend_pe(self, length, device, dtype):
+ """Reset and extend the positional encodings only at the beginning"""
+ if hasattr(self, "pe"):
+ return
+
+ positions = torch.arange(
+ self.left_context, -self.right_context - 1, -1, dtype=torch.float32, device=device
+ ).unsqueeze(1)
+ self.create_pe(positions=positions, dtype=dtype)
+
+ def forward(self, x, cache_len=0):
+ """Compute positional encoding.
+ Args:
+ x (torch.Tensor): Input. Its shape is (batch, time, feature_size)
+ Returns:
+ x (torch.Tensor): Its shape is (batch, time, feature_size)
+ pos_emb (torch.Tensor): Its shape is (1, time, feature_size)
+ """
+
+ if self.xscale:
+ x = x * self.xscale
+
+ end_pos = self.left_context + self.right_context + 1
+ pos_emb = self.pe[:, :end_pos]
+ if self.dropout_emb:
+ pos_emb = self.dropout_emb(pos_emb)
+ return self.dropout(x), pos_emb
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/regularization_utils.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/regularization_utils.py
new file mode 100644
index 000000000..a4ccc14ae
--- /dev/null
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/regularization_utils.py
@@ -0,0 +1,67 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/asr/parts/utils/regularization_utils.py; training-only code removed.
+
+from typing import List
+
+
+def compute_stochastic_depth_drop_probs(
+ num_layers: int,
+ stochastic_depth_drop_prob: float = 0.0,
+ stochastic_depth_mode: str = "linear",
+ stochastic_depth_start_layer: int = 1,
+) -> List[float]:
+ """Computes drop probabilities for stochastic depth regularization technique.
+ The first layer is never dropped and the starting layer needs to be greater
+ or equal to 1.
+
+ Args:
+ num_layers (int): number of layers in the network.
+ stochastic_depth_drop_prob (float): if non-zero, will randomly drop
+ layers during training. The higher this value, the more often layers
+ are dropped. Defaults to 0.0.
+ stochastic_depth_mode (str): can be either "linear" or "uniform". If
+ set to "uniform", all layers have the same probability of drop. If
+ set to "linear", the drop probability grows linearly from 0 for the
+ first layer to the desired value for the final layer. Defaults to
+ "linear".
+ stochastic_depth_start_layer (int): starting layer for stochastic depth.
+ All layers before this will never be dropped. Note that drop
+ probability will be adjusted accordingly if mode is "linear" when
+ start layer is > 1. Defaults to 1.
+ Returns:
+ List[float]: list of drop probabilities for all layers
+ """
+ if not (0 <= stochastic_depth_drop_prob < 1.0):
+ raise ValueError("stochastic_depth_drop_prob has to be in [0, 1).")
+ if not (1 <= stochastic_depth_start_layer <= num_layers):
+ raise ValueError("stochastic_depth_start_layer has to be in [1, num layers].")
+
+ # Layers before `stochastic_depth_start_layer` are never dropped
+ layer_drop_probs = [0.0] * stochastic_depth_start_layer
+
+ # Layers starting with `stochastic_depth_start_layer` may be dropped
+ if (L := num_layers - stochastic_depth_start_layer) > 0:
+ if stochastic_depth_mode == "linear":
+ # we start with 1/L * drop_prob and and end with the desired drop probability.
+ layer_drop_probs += [l / L * stochastic_depth_drop_prob for l in range(1, L + 1)]
+ elif stochastic_depth_mode == "uniform":
+ layer_drop_probs += [stochastic_depth_drop_prob] * L
+ else:
+ raise ValueError(
+ f'stochastic_depth_mode has to be one of ["linear", "uniform"]. Current value: {stochastic_depth_mode}'
+ )
+ return layer_drop_probs
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/subsampling.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/subsampling.py
new file mode 100644
index 000000000..1faffee7d
--- /dev/null
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/subsampling.py
@@ -0,0 +1,628 @@
+# SPDX-FileCopyrightText: Copyright (c) 2020 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/asr/parts/submodules/subsampling.py; training-only code removed.
+
+import math
+
+import torch
+import torch.nn as nn
+from torch.nn import LayerNorm
+
+from .causal_convs import CausalConv1D, CausalConv2D
+from ..compat import logging
+
+
+class StackingSubsampling(torch.nn.Module):
+ """Stacking subsampling which simply stacks consecutive frames to reduce the sampling rate
+ Args:
+ subsampling_factor (int): The subsampling factor
+ feat_in (int): size of the input features
+ feat_out (int): size of the output features
+ norm (bool): whether to use an MLP layer after the stacking along with normalization. default is False.
+ """
+
+ def __init__(self, subsampling_factor, feat_in, feat_out, norm=False):
+ super(StackingSubsampling, self).__init__()
+ self.subsampling_factor = subsampling_factor
+ self.proj_out = torch.nn.Linear(subsampling_factor * feat_in, feat_out)
+ if norm:
+ self.pre_norm = LayerNorm(feat_in)
+ else:
+ self.pre_norm = None
+
+ def get_sampling_frames(self):
+ return self.subsampling_factor
+
+ def get_streaming_cache_size(self):
+ return 0
+
+ def forward(self, x, lengths):
+ b, t, h = x.size()
+ pad_size = (self.subsampling_factor - (t % self.subsampling_factor)) % self.subsampling_factor
+ x = torch.nn.functional.pad(x, (0, 0, 0, pad_size))
+ if self.pre_norm is not None:
+ x = self.pre_norm(x)
+ _, t, _ = x.size()
+ x = torch.reshape(x, (b, t // self.subsampling_factor, h * self.subsampling_factor))
+ x = self.proj_out(x)
+ lengths = torch.div(lengths + pad_size, self.subsampling_factor, rounding_mode="floor")
+ return x, lengths
+
+
+class ConvSubsampling(torch.nn.Module):
+ """Convolutional subsampling which supports VGGNet and striding approach introduced in:
+ VGGNet Subsampling: Transformer-transducer: end-to-end speech recognition with self-attention (https://arxiv.org/pdf/1910.12977.pdf)
+ Striding Subsampling: "Speech-Transformer: A No-Recurrence Sequence-to-Sequence Model for Speech Recognition" by Linhao Dong et al. (https://ieeexplore.ieee.org/document/8462506)
+ Args:
+ subsampling (str): The subsampling technique from {"vggnet", "striding", "dw-striding"}
+ subsampling_factor (int): The subsampling factor which should be a power of 2
+ subsampling_conv_chunking_factor (int): Input chunking factor which can be -1 (no chunking)
+ 1 (auto) or a power of 2. Default is 1
+ feat_in (int): size of the input features
+ feat_out (int): size of the output features
+ conv_channels (int): Number of channels for the convolution layers.
+ activation (Module): activation function, default is nn.ReLU()
+ """
+
+ def __init__(
+ self,
+ subsampling,
+ subsampling_factor,
+ feat_in,
+ feat_out,
+ conv_channels,
+ subsampling_conv_chunking_factor=1,
+ activation=nn.ReLU(),
+ is_causal=False,
+ ):
+ super(ConvSubsampling, self).__init__()
+ self._subsampling = subsampling
+ self._conv_channels = conv_channels
+ self._feat_in = feat_in
+ self._feat_out = feat_out
+
+ if subsampling_factor % 2 != 0:
+ raise ValueError("Sampling factor should be a multiply of 2!")
+ self._sampling_num = int(math.log(subsampling_factor, 2))
+ self.subsampling_factor = subsampling_factor
+ self.is_causal = is_causal
+
+ if (
+ subsampling_conv_chunking_factor != -1
+ and subsampling_conv_chunking_factor != 1
+ and subsampling_conv_chunking_factor % 2 != 0
+ ):
+ raise ValueError("subsampling_conv_chunking_factor should be -1, 1, or a power of 2")
+ self.subsampling_conv_chunking_factor = subsampling_conv_chunking_factor
+
+ in_channels = 1
+ layers = []
+
+ if subsampling == "vggnet":
+ self._stride = 2
+ self._kernel_size = 2
+ self._ceil_mode = True
+
+ self._left_padding = 0
+ self._right_padding = 0
+
+ for i in range(self._sampling_num):
+ layers.append(
+ torch.nn.Conv2d(
+ in_channels=in_channels, out_channels=conv_channels, kernel_size=3, stride=1, padding=1
+ )
+ )
+ layers.append(activation)
+ layers.append(
+ torch.nn.Conv2d(
+ in_channels=conv_channels, out_channels=conv_channels, kernel_size=3, stride=1, padding=1
+ )
+ )
+ layers.append(activation)
+ layers.append(
+ torch.nn.MaxPool2d(
+ kernel_size=self._kernel_size,
+ stride=self._stride,
+ padding=self._left_padding,
+ ceil_mode=self._ceil_mode,
+ )
+ )
+ in_channels = conv_channels
+
+ elif subsampling == "dw_striding":
+ self._stride = 2
+ self._kernel_size = 3
+ self._ceil_mode = False
+
+ if self.is_causal:
+ self._left_padding = self._kernel_size - 1
+ self._right_padding = self._stride - 1
+ self._max_cache_len = subsampling_factor + 1
+ else:
+ self._left_padding = (self._kernel_size - 1) // 2
+ self._right_padding = (self._kernel_size - 1) // 2
+ self._max_cache_len = 0
+
+ # Layer 1
+ if self.is_causal:
+ layers.append(
+ CausalConv2D(
+ in_channels=in_channels,
+ out_channels=conv_channels,
+ kernel_size=self._kernel_size,
+ stride=self._stride,
+ padding=None,
+ )
+ )
+ else:
+ layers.append(
+ torch.nn.Conv2d(
+ in_channels=in_channels,
+ out_channels=conv_channels,
+ kernel_size=self._kernel_size,
+ stride=self._stride,
+ padding=self._left_padding,
+ )
+ )
+ in_channels = conv_channels
+ layers.append(activation)
+
+ for i in range(self._sampling_num - 1):
+ if self.is_causal:
+ layers.append(
+ CausalConv2D(
+ in_channels=in_channels,
+ out_channels=in_channels,
+ kernel_size=self._kernel_size,
+ stride=self._stride,
+ padding=None,
+ groups=in_channels,
+ )
+ )
+ else:
+ layers.append(
+ torch.nn.Conv2d(
+ in_channels=in_channels,
+ out_channels=in_channels,
+ kernel_size=self._kernel_size,
+ stride=self._stride,
+ padding=self._left_padding,
+ groups=in_channels,
+ )
+ )
+
+ layers.append(
+ torch.nn.Conv2d(
+ in_channels=in_channels,
+ out_channels=conv_channels,
+ kernel_size=1,
+ stride=1,
+ padding=0,
+ groups=1,
+ )
+ )
+ layers.append(activation)
+ in_channels = conv_channels
+
+ elif subsampling == "striding":
+ self._stride = 2
+ self._kernel_size = 3
+ self._ceil_mode = False
+
+ if self.is_causal:
+ self._left_padding = self._kernel_size - 1
+ self._right_padding = self._stride - 1
+ self._max_cache_len = subsampling_factor + 1
+ else:
+ self._left_padding = (self._kernel_size - 1) // 2
+ self._right_padding = (self._kernel_size - 1) // 2
+ self._max_cache_len = 0
+
+ for i in range(self._sampling_num):
+ if self.is_causal:
+ layers.append(
+ CausalConv2D(
+ in_channels=in_channels,
+ out_channels=conv_channels,
+ kernel_size=self._kernel_size,
+ stride=self._stride,
+ padding=None,
+ )
+ )
+ else:
+ layers.append(
+ torch.nn.Conv2d(
+ in_channels=in_channels,
+ out_channels=conv_channels,
+ kernel_size=self._kernel_size,
+ stride=self._stride,
+ padding=self._left_padding,
+ )
+ )
+ layers.append(activation)
+ in_channels = conv_channels
+
+ elif subsampling == "striding_conv1d":
+ in_channels = feat_in
+
+ self._stride = 2
+ self._kernel_size = 5
+ self._ceil_mode = False
+
+ if self.is_causal:
+ self._left_padding = self._kernel_size - 1
+ self._right_padding = self._stride - 1
+ self._max_cache_len = subsampling_factor + 1
+ else:
+ self._left_padding = (self._kernel_size - 1) // 2
+ self._right_padding = (self._kernel_size - 1) // 2
+ self._max_cache_len = 0
+
+ for i in range(self._sampling_num):
+ if self.is_causal:
+ layers.append(
+ CausalConv1D(
+ in_channels=in_channels,
+ out_channels=feat_out if self._sampling_num == i + 1 else conv_channels,
+ kernel_size=self._kernel_size,
+ stride=self._stride,
+ padding=None,
+ )
+ )
+ else:
+ layers.append(
+ torch.nn.Conv1d(
+ in_channels=in_channels,
+ out_channels=feat_out if self._sampling_num == i + 1 else conv_channels,
+ kernel_size=self._kernel_size,
+ stride=self._stride,
+ padding=self._left_padding,
+ )
+ )
+ layers.append(activation)
+ in_channels = conv_channels
+
+ elif subsampling == "dw_striding_conv1d":
+ in_channels = feat_in
+
+ self._stride = 2
+ self._kernel_size = 5
+ self._ceil_mode = False
+
+ self._left_padding = (self._kernel_size - 1) // 2
+ self._right_padding = (self._kernel_size - 1) // 2
+
+ # Layer 1
+ layers.extend(
+ [
+ torch.nn.Conv1d(
+ in_channels=in_channels,
+ out_channels=in_channels,
+ kernel_size=self._kernel_size,
+ stride=self._stride,
+ padding=self._left_padding,
+ groups=in_channels,
+ ),
+ torch.nn.Conv1d(
+ in_channels=in_channels,
+ out_channels=feat_out if self._sampling_num == 1 else conv_channels,
+ kernel_size=1,
+ stride=1,
+ padding=0,
+ groups=1,
+ ),
+ ]
+ )
+ in_channels = conv_channels
+ layers.append(activation)
+
+ for i in range(self._sampling_num - 1):
+ layers.extend(
+ [
+ torch.nn.Conv1d(
+ in_channels=in_channels,
+ out_channels=in_channels,
+ kernel_size=self._kernel_size,
+ stride=self._stride,
+ padding=self._left_padding,
+ groups=in_channels,
+ ),
+ torch.nn.Conv1d(
+ in_channels=in_channels,
+ out_channels=feat_out if self._sampling_num == i + 2 else conv_channels,
+ kernel_size=1,
+ stride=1,
+ padding=0,
+ groups=1,
+ ),
+ ]
+ )
+ layers.append(activation)
+ in_channels = conv_channels
+
+ else:
+ raise ValueError(f"Not valid sub-sampling: {subsampling}!")
+
+ if subsampling in ["vggnet", "dw_striding", "striding"]:
+ in_length = torch.tensor(feat_in, dtype=torch.float)
+ out_length = calc_length(
+ lengths=in_length,
+ all_paddings=self._left_padding + self._right_padding,
+ kernel_size=self._kernel_size,
+ stride=self._stride,
+ ceil_mode=self._ceil_mode,
+ repeat_num=self._sampling_num,
+ )
+ self.out = torch.nn.Linear(conv_channels * int(out_length), feat_out)
+ self.conv2d_subsampling = True
+ elif subsampling in ["striding_conv1d", "dw_striding_conv1d"]:
+ self.out = None
+ self.conv2d_subsampling = False
+ else:
+ raise ValueError(f"Not valid sub-sampling: {subsampling}!")
+
+ self.conv = torch.nn.Sequential(*layers)
+
+ def get_sampling_frames(self):
+ return [1, self.subsampling_factor]
+
+ def get_streaming_cache_size(self):
+ return [0, self.subsampling_factor + 1]
+
+ def forward(self, x, lengths):
+ lengths = calc_length(
+ lengths,
+ all_paddings=self._left_padding + self._right_padding,
+ kernel_size=self._kernel_size,
+ stride=self._stride,
+ ceil_mode=self._ceil_mode,
+ repeat_num=self._sampling_num,
+ )
+
+ # Unsqueeze Channel Axis
+ if self.conv2d_subsampling:
+ x = x.unsqueeze(1)
+ # Transpose to Channel First mode
+ else:
+ x = x.transpose(1, 2)
+
+ # split inputs if chunking_factor is set
+ if self.subsampling_conv_chunking_factor != -1 and self.conv2d_subsampling:
+ if self.subsampling_conv_chunking_factor == 1:
+ # if subsampling_conv_chunking_factor is 1, we split only if needed
+ # avoiding a bug / feature limiting indexing of tensors to 2**31
+ # see https://github.com/pytorch/pytorch/issues/80020
+ x_ceil = 2**31 / self._conv_channels * self._stride * self._stride
+ if torch.numel(x) > x_ceil:
+ need_to_split = True
+ else:
+ need_to_split = False
+ else:
+ # if subsampling_conv_chunking_factor > 1 we always split
+ need_to_split = True
+
+ if need_to_split:
+ x, success = self.conv_split_by_batch(x)
+ if not success: # if unable to split by batch, try by channel
+ if self._subsampling == "dw_striding":
+ x = self.conv_split_by_channel(x)
+ else:
+ x = self.conv(x) # try anyway
+ else:
+ x = self.conv(x)
+ else:
+ x = self.conv(x)
+
+ # Flatten Channel and Frequency Axes
+ if self.conv2d_subsampling:
+ b, c, t, f = x.size()
+ x = self.out(x.transpose(1, 2).reshape(b, t, -1))
+ # Transpose to Channel Last mode
+ else:
+ x = x.transpose(1, 2)
+
+ return x, lengths
+
+ def reset_parameters(self):
+ # initialize weights
+ if self._subsampling == "dw_striding":
+ with torch.no_grad():
+ # init conv
+ scale = 1.0 / self._kernel_size
+ dw_max = (self._kernel_size**2) ** -0.5
+ pw_max = self._conv_channels**-0.5
+
+ torch.nn.init.uniform_(self.conv[0].weight, -scale, scale)
+ torch.nn.init.uniform_(self.conv[0].bias, -scale, scale)
+
+ for idx in range(2, len(self.conv), 3):
+ torch.nn.init.uniform_(self.conv[idx].weight, -dw_max, dw_max)
+ torch.nn.init.uniform_(self.conv[idx].bias, -dw_max, dw_max)
+ torch.nn.init.uniform_(self.conv[idx + 1].weight, -pw_max, pw_max)
+ torch.nn.init.uniform_(self.conv[idx + 1].bias, -pw_max, pw_max)
+
+ # init fc (80 * 64 = 5120 from https://github.com/kssteven418/Squeezeformer/blob/13c97d6cf92f2844d2cb3142b4c5bfa9ad1a8951/src/models/conformer_encoder.py#L487
+ fc_scale = (self._feat_out * self._feat_in / self._sampling_num) ** -0.5
+ torch.nn.init.uniform_(self.out.weight, -fc_scale, fc_scale)
+ torch.nn.init.uniform_(self.out.bias, -fc_scale, fc_scale)
+
+ def conv_split_by_batch(self, x):
+ """Tries to split input by batch, run conv and concat results"""
+ b, _, _, _ = x.size()
+ if b == 1: # can't split if batch size is 1
+ return x, False
+
+ if self.subsampling_conv_chunking_factor > 1:
+ cf = self.subsampling_conv_chunking_factor
+ logging.debug(f"using manually set chunking factor: {cf}")
+ else:
+ # avoiding a bug / feature limiting indexing of tensors to 2**31
+ # see https://github.com/pytorch/pytorch/issues/80020
+ x_ceil = 2**31 / self._conv_channels * self._stride * self._stride
+ p = math.ceil(math.log(torch.numel(x) / x_ceil, 2))
+ cf = 2**p
+ logging.debug(f"using auto set chunking factor: {cf}")
+
+ new_batch_size = b // cf
+ if new_batch_size == 0: # input is too big
+ return x, False
+
+ logging.debug(f"conv subsampling: using split batch size {new_batch_size}")
+ return torch.cat([self.conv(chunk) for chunk in torch.split(x, new_batch_size, 0)]), True
+
+ def conv_split_by_channel(self, x):
+ """For dw convs, tries to split input by time, run conv and concat results"""
+ x = self.conv[0](x) # full conv2D
+ x = self.conv[1](x) # activation
+
+ for i in range(self._sampling_num - 1):
+ _, c, t, _ = x.size()
+
+ if self.subsampling_conv_chunking_factor > 1:
+ cf = self.subsampling_conv_chunking_factor
+ logging.debug(f"using manually set chunking factor: {cf}")
+ else:
+ # avoiding a bug / feature limiting indexing of tensors to 2**31
+ # see https://github.com/pytorch/pytorch/issues/80020
+ p = math.ceil(math.log(torch.numel(x) / 2**31, 2))
+ cf = 2**p
+ logging.debug(f"using auto set chunking factor: {cf}")
+
+ new_c = int(c // cf)
+ if new_c == 0:
+ logging.warning(f"chunking factor {cf} is too high; splitting down to one channel.")
+ new_c = 1
+
+ new_t = int(t // cf)
+ if new_t == 0:
+ logging.warning(f"chunking factor {cf} is too high; splitting down to one timestep.")
+ new_t = 1
+
+ logging.debug(f"conv dw subsampling: using split C size {new_c} and split T size {new_t}")
+ x = self.channel_chunked_conv(self.conv[i * 3 + 2], new_c, x) # conv2D, depthwise
+
+ # splitting pointwise convs by time
+ x = torch.cat([self.conv[i * 3 + 3](chunk) for chunk in torch.split(x, new_t, 2)], 2) # conv2D, pointwise
+ x = self.conv[i * 3 + 4](x) # activation
+ return x
+
+ def channel_chunked_conv(self, conv, chunk_size, x):
+ """Performs channel chunked convolution"""
+
+ ind = 0
+ out_chunks = []
+ for chunk in torch.split(x, chunk_size, 1):
+ step = chunk.size()[1]
+
+ if self.is_causal:
+ chunk = nn.functional.pad(
+ chunk, pad=(self._kernel_size - 1, self._stride - 1, self._kernel_size - 1, self._stride - 1)
+ )
+ ch_out = nn.functional.conv2d(
+ chunk,
+ conv.weight[ind : ind + step, :, :, :],
+ bias=conv.bias[ind : ind + step],
+ stride=self._stride,
+ padding=0,
+ groups=step,
+ )
+ else:
+ ch_out = nn.functional.conv2d(
+ chunk,
+ conv.weight[ind : ind + step, :, :, :],
+ bias=conv.bias[ind : ind + step],
+ stride=self._stride,
+ padding=self._left_padding,
+ groups=step,
+ )
+ out_chunks.append(ch_out)
+ ind += step
+
+ return torch.cat(out_chunks, 1)
+
+ def change_subsampling_conv_chunking_factor(self, subsampling_conv_chunking_factor: int):
+ if (
+ subsampling_conv_chunking_factor != -1
+ and subsampling_conv_chunking_factor != 1
+ and subsampling_conv_chunking_factor % 2 != 0
+ ):
+ raise ValueError("subsampling_conv_chunking_factor should be -1, 1, or a power of 2")
+ self.subsampling_conv_chunking_factor = subsampling_conv_chunking_factor
+
+
+def calc_length(lengths, all_paddings, kernel_size, stride, ceil_mode, repeat_num=1):
+ """Calculates the output length of a Tensor passed through a convolution or max pooling layer"""
+ add_pad: float = all_paddings - kernel_size
+ one: float = 1.0
+ for i in range(repeat_num):
+ lengths = torch.div(lengths.to(dtype=torch.float) + add_pad, stride) + one
+ if ceil_mode:
+ lengths = torch.ceil(lengths)
+ else:
+ lengths = torch.floor(lengths)
+ return lengths.to(dtype=torch.int)
+
+
+class SubsamplingReductionModule(nn.Module):
+ """Downsamples the audio signal in time dimension."""
+
+ def __init__(self, reduction: str, d_model: int, reduction_factor: int = 2):
+ super().__init__()
+
+ assert reduction in ["pooling", "striding"]
+
+ self.reduction = reduction
+ self.d_model = d_model
+ self._sampling_num = int(math.log(reduction_factor, 2))
+
+ if reduction == "pooling":
+ self.reduction_enc = nn.MaxPool1d(kernel_size=reduction_factor)
+ self.padding = 0
+ self.kernel_size = self.reduction_enc.kernel_size
+ self.stride = self.reduction_enc.stride
+ elif reduction == "striding":
+ self.reduction_enc = ConvSubsampling(
+ subsampling="striding",
+ subsampling_factor=reduction_factor,
+ feat_in=d_model,
+ feat_out=d_model,
+ conv_channels=d_model,
+ activation=nn.ReLU(),
+ is_causal=False,
+ )
+
+ def forward(self, x, lengths):
+ """Shapes:
+ - x: [B, T, C]
+ - lengths: [B]
+ """
+
+ if self.reduction == "striding":
+ x, lengths = self.reduction_enc(x=x, lengths=lengths)
+ else:
+ x = torch.transpose(x, 1, 2) # [B, C, T]
+ lengths = calc_length(
+ lengths=lengths,
+ all_paddings=self.padding,
+ kernel_size=self.kernel_size,
+ stride=self.stride,
+ ceil_mode=False,
+ repeat_num=self._sampling_num,
+ )
+ x = self.reduction_enc(x)
+ x = torch.transpose(x, 1, 2) # [B, T, C]
+
+ return x, lengths
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/compat.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/compat.py
new file mode 100644
index 000000000..f1c648acd
--- /dev/null
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/compat.py
@@ -0,0 +1,644 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat); compatibility shims
+# replacing nemo.core / nemo.utils machinery for dependency-free inference.
+# Each vendored function below is an exact copy from its annotated original source path;
+# only imports were adjusted. Stub classes replace NeMo base classes whose functionality
+# is not needed at inference time.
+
+import importlib
+import logging as _stdlib_logging
+import math
+import warnings
+from contextlib import contextmanager, nullcontext
+from dataclasses import dataclass
+from typing import Dict, Optional
+
+import torch
+import torch.nn as nn
+from omegaconf import OmegaConf
+from safetensors.torch import load_file
+from transformers import AutoConfig, AutoModelForCausalLM
+
+
+# ==============================================================================
+# Logging shim (replaces the nemo.utils logging import)
+# ==============================================================================
+class logging_mode:
+ """Stub of nemo.utils.logging's dedup modes (used as `mode=logging_mode.ONCE`)."""
+
+ EACH = 0
+ ONCE = 1
+
+
+class _NemoLoggerShim:
+ """stdlib-logging based replacement for the nemo.utils logging module.
+
+ NeMo's logger accepts an extra `mode=` kwarg (deduplication); it is accepted
+ and ignored here, with a simple set-based dedup for `mode=logging_mode.ONCE`.
+ """
+
+ def __init__(self, logger):
+ self._logger = logger
+ self._once_seen = set()
+
+ def _log(self, level, msg, args, mode=None):
+ if mode == logging_mode.ONCE:
+ key = (level, str(msg))
+ if key in self._once_seen:
+ return
+ self._once_seen.add(key)
+ self._logger.log(level, msg, *args)
+
+ def debug(self, msg, *args, mode=None, **kwargs):
+ self._log(_stdlib_logging.DEBUG, msg, args, mode=mode)
+
+ def info(self, msg, *args, mode=None, **kwargs):
+ self._log(_stdlib_logging.INFO, msg, args, mode=mode)
+
+ def warning(self, msg, *args, mode=None, **kwargs):
+ self._log(_stdlib_logging.WARNING, msg, args, mode=mode)
+
+ def error(self, msg, *args, mode=None, **kwargs):
+ self._log(_stdlib_logging.ERROR, msg, args, mode=mode)
+
+ def critical(self, msg, *args, mode=None, **kwargs):
+ self._log(_stdlib_logging.CRITICAL, msg, args, mode=mode)
+
+
+logging = _NemoLoggerShim(_stdlib_logging.getLogger("nemo_vendored"))
+
+
+# ==============================================================================
+# typecheck no-op (replaces nemo.core.classes.common.typecheck)
+# ==============================================================================
+class typecheck:
+ """No-op replacement for NeMo's typecheck decorator.
+
+ Supports the two usages found in the vendored files:
+ @typecheck()
+ def forward(...)
+
+ @typecheck.disable_checks()
+ def forward(...)
+
+ with typecheck.disable_checks():
+ ...
+ """
+
+ def __init__(self, *args, **kwargs):
+ pass
+
+ def __call__(self, func):
+ return func
+
+ @staticmethod
+ @contextmanager
+ def disable_checks():
+ yield
+
+
+# ==============================================================================
+# Neural type stubs (replace nemo.core.neural_types)
+# ==============================================================================
+
+
+# ==============================================================================
+# from_config_dict target resolver (replaces hydra-based NeMo Serialization)
+# ==============================================================================
+# Maps the `_target_` strings that appear in the NVIDIA-NemotronLabs-VoiceChat-11B
+# checkpoint config.json to the vendored classes. Extend as needed.
+_TARGET_REGISTRY = {
+ "nemo.collections.asr.modules.ConformerEncoder": (".asr.conformer_encoder", "ConformerEncoder"),
+ "nemo.collections.asr.modules.conformer_encoder.ConformerEncoder": (
+ ".asr.conformer_encoder",
+ "ConformerEncoder",
+ ),
+ "nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor": (
+ ".asr.audio_preprocessing",
+ "AudioToMelSpectrogramPreprocessor",
+ ),
+ "nemo.collections.asr.modules.audio_preprocessing.AudioToMelSpectrogramPreprocessor": (
+ ".asr.audio_preprocessing",
+ "AudioToMelSpectrogramPreprocessor",
+ ),
+ "nemo.collections.speechlm2.modules.perception.IdentityConnector": (".perception", "IdentityConnector"),
+ "nemo.collections.speechlm2.modules.perception.AudioPerceptionModule": (".perception", "AudioPerceptionModule"),
+}
+
+_PACKAGE = __name__.rsplit(".", 1)[0] # ...nemotron_voicechat.nemo_vendored
+
+
+def from_config_dict(config):
+ """Instantiate a vendored class from an OmegaConf/dict config with a `_target_` key.
+
+ Hydra-free replacement for nemo.core.classes.common.Serialization.from_config_dict.
+ Like hydra.utils.instantiate with the default `_convert_="none"`, all keys except
+ `_target_` are passed as keyword arguments (nested containers stay OmegaConf nodes).
+ """
+ cfg = OmegaConf.create(config) if not OmegaConf.is_config(config) else config
+ if "_target_" not in cfg:
+ raise ValueError(f"from_config_dict requires a '_target_' key in the config; got keys: {list(cfg.keys())}")
+ target = cfg["_target_"]
+ if target not in _TARGET_REGISTRY:
+ raise ValueError(
+ f"Unknown _target_ '{target}' — not in the vendored nemo_vendored.compat._TARGET_REGISTRY. "
+ f"Known targets: {sorted(_TARGET_REGISTRY)}. Add a mapping if this class was vendored."
+ )
+ module_name, class_name = _TARGET_REGISTRY[target]
+ cls = getattr(importlib.import_module(module_name, package=_PACKAGE), class_name)
+ kwargs = {str(k): cfg[k] for k in cfg if k != "_target_"}
+ return cls(**kwargs)
+
+
+# ==============================================================================
+# NeuralModule / Exportable / mixin stubs (replace nemo.core.classes.*)
+# ==============================================================================
+class NeuralModule(nn.Module):
+ """Slim replacement of nemo.core.classes.module.NeuralModule."""
+
+ @property
+ def num_weights(self):
+ num: int = 0
+ for p in self.parameters():
+ if p.requires_grad:
+ num += p.numel()
+ return num
+
+ @staticmethod
+ def from_config_dict(config):
+ return from_config_dict(config)
+
+ def unfreeze(self, partial: bool = False) -> None:
+ """Unfreeze all parameters for training. (exact copy of NeuralModule.unfreeze)"""
+ if partial and not hasattr(self, "_frozen_grad_map"):
+ raise ValueError("Cannot unfreeze partially without first freezing the module with `freeze()`")
+ for pname, param in self.named_parameters():
+ if not partial:
+ param.requires_grad = True
+ elif pname in self._frozen_grad_map:
+ param.requires_grad = self._frozen_grad_map[pname]
+ if hasattr(self, "_frozen_grad_map"):
+ delattr(self, "_frozen_grad_map")
+ self.train()
+
+
+class Exportable:
+ """Stub of nemo.core.classes.exportable.Exportable (ONNX export machinery removed)."""
+
+ @property
+ def input_names(self):
+ return []
+
+ @property
+ def output_names(self):
+ return []
+
+
+class AccessMixin:
+ """Stub of nemo.core.classes.mixins.access_mixins.AccessMixin.
+
+ Tensor-capture registry is training/interctc-only; at inference access is
+ always disabled so guarded blocks are skipped.
+ """
+
+ @property
+ def access_cfg(self):
+ return {}
+
+ def is_access_enabled(self, guid=None):
+ return False
+
+ def register_accessible_tensor(self, name, tensor):
+ pass
+
+
+class AttentionAdapterModuleMixin:
+ """Stub of nemo.collections.asr.parts.submodules.adapters.attention_adapter_mixin.
+ Adapters (PEFT) are training-only; report none available."""
+
+ def is_adapter_available(self) -> bool:
+ return False
+
+
+class AdapterModuleMixin(AttentionAdapterModuleMixin):
+ """Stub of nemo.core.classes.mixins.adapter_mixins.AdapterModuleMixin."""
+
+ pass
+
+
+# ==============================================================================
+# StreamingEncoder — exact copy of
+# nemo/collections/asr/parts/mixins/streaming.py (imports adjusted only)
+# ==============================================================================
+from abc import ABC, abstractmethod # noqa: E402
+
+
+class StreamingEncoder(ABC):
+ @abstractmethod
+ def setup_streaming_params(
+ self,
+ max_look_ahead: int = 10000,
+ ):
+ """
+ This function sets the needed values and parameters to perform streaming. The configuration (CacheAwareStreamingConfig) need to be stored in self.streaming_cfg.
+ The streaming configuration is needed to simulate streaming inference. It would set the following
+ """
+ pass
+
+
+# ==============================================================================
+# CacheAwareStreamingConfig — exact copy of the dataclass from
+# nemo/collections/asr/models/configs/asr_models_config.py
+# ==============================================================================
+@dataclass
+class CacheAwareStreamingConfig:
+ chunk_size: int = 0 # the size of each chunk at each step, it can be a list of two integers to specify different chunk sizes for the first step and others
+ shift_size: int = 0 # the size of the shift in each step, it can be a list of two integers to specify different shift sizes for the first step and others
+
+ cache_drop_size: int = 0 # the number of steps to drop from the cache
+ last_channel_cache_size: int = 0 # the size of the needed cache for last channel layers
+
+ valid_out_len: int = (
+ 0 # the number of the steps in the final output which are valid (have the same value as in the offline mode)
+ )
+
+ pre_encode_cache_size: int = 0 # the size of the needed cache for the pre-encoding part of the model to avoid caching inside the pre-encoding layers
+ drop_extra_pre_encoded: int = 0 # the number of steps to get dropped after the pre-encoding layer
+
+ last_channel_num: int = 0 # number of the last channel layers (like MHA layers) which need caching in the model
+ last_time_num: int = 0 # number of the last time layers (like convolutions) which need caching in the model
+
+
+# ==============================================================================
+# avoid_float16_autocast_context / avoid_bfloat16_autocast_context — exact
+# copies of NeMo nemo/utils/cast_utils.py
+# ==============================================================================
+
+
+def avoid_float16_autocast_context():
+ """
+ If the current autocast context is float16, cast it to bfloat16
+ if available (unless we're in jit) or float32
+ """
+
+ if torch.is_autocast_enabled() and torch.get_autocast_gpu_dtype() == torch.float16:
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
+ return torch.amp.autocast("cuda", dtype=torch.float32)
+
+ if torch.cuda.is_bf16_supported():
+ return torch.amp.autocast("cuda", dtype=torch.bfloat16)
+ else:
+ return torch.amp.autocast("cuda", dtype=torch.float32)
+ else:
+ return nullcontext()
+
+
+# ==============================================================================
+# fp32_precision — exact copy of nemo/collections/speechlm2/parts/precision.py
+# ==============================================================================
+@contextmanager
+def fp32_precision():
+ """
+ Workaround for precision related issues when training with bf16-true PyTorch Lightning precision setting.
+ In bf16-true, PTL changes PyTorch's default dtype, which may break implicit assumptions for some models.
+ This context manager restores default float32 precision and runs the computation in float32 autocast context.
+ """
+ default_dtype = torch.get_default_dtype()
+ torch.set_default_dtype(torch.float32)
+ try:
+ with torch.amp.autocast(device_type="cuda" if torch.cuda.is_available() else "cpu", dtype=torch.float32):
+ yield
+ finally:
+ torch.set_default_dtype(default_dtype)
+
+
+# ==============================================================================
+# get_pad_id — exact copy of nemo/collections/speechlm2/data/utils.py
+# ==============================================================================
+def get_pad_id(tokenizer) -> int:
+ pad_id = tokenizer.pad
+ if pad_id is not None:
+ return pad_id
+ pad_id = tokenizer.unk_id
+ if pad_id is not None:
+ return pad_id
+ warnings.warn(
+ "The text tokenizer has no or tokens available, using ID 0 for padding (this may lead to silent bugs)."
+ )
+ return 0
+
+
+# ==============================================================================
+# resample — exact copies of resample/_get_sinc_resample_kernel/_apply_sinc_resample_kernel
+# of nemo/collections/audio/parts/utils/resampling.py (itself vendored from
+# torchaudio 2.6.0, BSD 2-Clause; see that file for the license text).
+# nemo/collections/audio/parts/utils/transforms.py re-exports `resample` from there.
+# ==============================================================================
+def resample(
+ waveform: torch.Tensor,
+ orig_freq: int,
+ new_freq: int,
+ lowpass_filter_width: int = 6,
+ rolloff: float = 0.99,
+ resampling_method: str = "sinc_interp_hann",
+ beta: float | None = None,
+) -> torch.Tensor:
+ r"""Resamples the waveform at the new frequency using bandlimited interpolation. :cite:`RESAMPLE`."""
+
+ if orig_freq <= 0.0 or new_freq <= 0.0:
+ raise ValueError("Original frequency and desired frequency should be positive")
+
+ if orig_freq == new_freq:
+ return waveform
+
+ gcd = math.gcd(int(orig_freq), int(new_freq))
+
+ kernel, width = _get_sinc_resample_kernel(
+ orig_freq,
+ new_freq,
+ gcd,
+ lowpass_filter_width,
+ rolloff,
+ resampling_method,
+ beta,
+ waveform.device,
+ waveform.dtype,
+ )
+ resampled = _apply_sinc_resample_kernel(waveform, orig_freq, new_freq, gcd, kernel, width)
+ return resampled
+
+
+def _get_sinc_resample_kernel(
+ orig_freq: int,
+ new_freq: int,
+ gcd: int,
+ lowpass_filter_width: int = 6,
+ rolloff: float = 0.99,
+ resampling_method: str = "sinc_interp_hann",
+ beta: float | None = None,
+ device: torch.device = "cpu",
+ dtype: torch.dtype | None = None,
+):
+ if not (int(orig_freq) == orig_freq and int(new_freq) == new_freq):
+ raise Exception(
+ "Frequencies must be of integer type to ensure quality resampling computation. "
+ "To work around this, manually convert both frequencies to integer values "
+ "that maintain their resampling rate ratio before passing them into the function. "
+ "Example: To downsample a 44100 hz waveform by a factor of 8, use "
+ "`orig_freq=8` and `new_freq=1` instead of `orig_freq=44100` and `new_freq=5512.5`. "
+ "For more information, please refer to https://github.com/pytorch/audio/issues/1487."
+ )
+
+ if resampling_method not in ["sinc_interp_hann", "sinc_interp_kaiser"]:
+ raise ValueError("Invalid resampling method: {}".format(resampling_method))
+
+ orig_freq = int(orig_freq) // gcd
+ new_freq = int(new_freq) // gcd
+
+ if lowpass_filter_width <= 0:
+ raise ValueError("Low pass filter width should be positive.")
+ base_freq = min(orig_freq, new_freq)
+ # This will perform antialiasing filtering by removing the highest frequencies.
+ # At first I thought I only needed this when downsampling, but when upsampling
+ # you will get edge artifacts without this, as the edge is equivalent to zero padding,
+ # which will add high freq artifacts.
+ base_freq *= rolloff
+
+ # The key idea of the algorithm is that x(t) can be exactly reconstructed from x[i] (tensor)
+ # using the sinc interpolation formula:
+ # x(t) = sum_i x[i] sinc(pi * orig_freq * (i / orig_freq - t))
+ # We can then sample the function x(t) with a different sample rate:
+ # y[j] = x(j / new_freq)
+ # or,
+ # y[j] = sum_i x[i] sinc(pi * orig_freq * (i / orig_freq - j / new_freq))
+
+ # We see here that y[j] is the convolution of x[i] with a specific filter, for which
+ # we take an FIR approximation, stopping when we see at least `lowpass_filter_width` zeros crossing.
+ # But y[j+1] is going to have a different set of weights and so on, until y[j + new_freq].
+ # Indeed:
+ # y[j + new_freq] = sum_i x[i] sinc(pi * orig_freq * ((i / orig_freq - (j + new_freq) / new_freq))
+ # = sum_i x[i] sinc(pi * orig_freq * ((i - orig_freq) / orig_freq - j / new_freq))
+ # = sum_i x[i + orig_freq] sinc(pi * orig_freq * (i / orig_freq - j / new_freq))
+ # so y[j+new_freq] uses the same filter as y[j], but on a shifted version of x by `orig_freq`.
+ # This will explain the F.conv1d after, with a stride of orig_freq.
+ width = math.ceil(lowpass_filter_width * orig_freq / base_freq)
+ # If orig_freq is still big after GCD reduction, most filters will be very unbalanced, i.e.,
+ # they will have a lot of almost zero values to the left or to the right...
+ # There is probably a way to evaluate those filters more efficiently, but this is kept for
+ # future work.
+ idx_dtype = dtype if dtype is not None else torch.float64
+
+ idx = torch.arange(-width, width + orig_freq, dtype=idx_dtype, device=device)[None, None] / orig_freq
+
+ t = torch.arange(0, -new_freq, -1, dtype=dtype, device=device)[:, None, None] / new_freq + idx
+ t *= base_freq
+ t = t.clamp_(-lowpass_filter_width, lowpass_filter_width)
+
+ # we do not use built in torch windows here as we need to evaluate the window
+ # at specific positions, not over a regular grid.
+ if resampling_method == "sinc_interp_hann":
+ window = torch.cos(t * math.pi / lowpass_filter_width / 2) ** 2
+ else:
+ # sinc_interp_kaiser
+ if beta is None:
+ beta = 14.769656459379492
+ beta_tensor = torch.tensor(float(beta))
+ window = torch.i0(beta_tensor * torch.sqrt(1 - (t / lowpass_filter_width) ** 2)) / torch.i0(beta_tensor)
+
+ t *= math.pi
+
+ scale = base_freq / orig_freq
+ kernels = torch.where(t == 0, torch.tensor(1.0).to(t), t.sin() / t)
+ kernels *= window * scale
+
+ if dtype is None:
+ kernels = kernels.to(dtype=torch.float32)
+
+ return kernels, width
+
+
+def _apply_sinc_resample_kernel(
+ waveform: torch.Tensor,
+ orig_freq: int,
+ new_freq: int,
+ gcd: int,
+ kernel: torch.Tensor,
+ width: int,
+):
+ if not waveform.is_floating_point():
+ raise TypeError(f"Expected floating point type for waveform tensor, but received {waveform.dtype}.")
+
+ orig_freq = int(orig_freq) // gcd
+ new_freq = int(new_freq) // gcd
+
+ # pack batch
+ shape = waveform.size()
+ waveform = waveform.view(-1, shape[-1])
+
+ num_wavs, length = waveform.shape
+ waveform = torch.nn.functional.pad(waveform, (width, width + orig_freq))
+ resampled = torch.nn.functional.conv1d(waveform[:, None], kernel, stride=orig_freq)
+ resampled = resampled.transpose(1, 2).reshape(num_wavs, -1)
+ target_length = torch.ceil(torch.as_tensor(new_freq * length / orig_freq)).long()
+ resampled = resampled[..., :target_length]
+
+ # unpack batch
+ resampled = resampled.view(shape[:-1] + resampled.shape[-1:])
+ return resampled
+
+
+# ==============================================================================
+# Swish — exact copy of nemo/collections/asr/parts/utils/activations.py
+# ==============================================================================
+class Swish(nn.SiLU):
+ """
+ Swish activation function introduced in 'https://arxiv.org/abs/1710.05941'
+ Mathematically identical to SiLU. See note in nn.SiLU for references.
+ """
+
+
+# ==============================================================================
+# activation_registry — exact copy of nemo/collections/common/parts/utils.py
+# ==============================================================================
+activation_registry = {
+ "identity": nn.Identity,
+ "hardtanh": nn.Hardtanh,
+ "relu": nn.ReLU,
+ "selu": nn.SELU,
+ "swish": nn.SiLU,
+ "silu": nn.SiLU,
+ "gelu": nn.GELU,
+}
+
+
+# ==============================================================================
+# FusedBatchNorm1d — exact copy of nemo/collections/asr/parts/submodules/batchnorm.py
+# ==============================================================================
+class FusedBatchNorm1d(nn.Module):
+ """
+ Fused BatchNorm to use in Conformer to improve accuracy in finetuning with TTS scenario
+ Drop-in replacement for BatchNorm1d with simple affine projection
+ """
+
+ def __init__(self, num_features: int):
+ """
+ Args:
+ num_features: number of channels, see original BatchNorm1d documentation
+ """
+ super().__init__()
+ self.num_features = num_features
+ self.weight = nn.Parameter(torch.ones(num_features))
+ self.bias = nn.Parameter(torch.zeros(num_features))
+
+ def forward(self, x: torch.Tensor):
+ if x.dim() == 3:
+ return x * self.weight.unsqueeze(-1) + self.bias.unsqueeze(-1)
+ assert x.dim() == 2
+ return x * self.weight + self.bias
+
+
+# ==============================================================================
+# load_pretrained_hf / set_model_dict_for_partial_init / load_checkpoint —
+# exact copies of nemo/collections/speechlm2/parts/pretrained.py
+# (only the functions used by the kept inference paths were vendored)
+# ==============================================================================
+def load_pretrained_hf(
+ model_path_or_name: str, pretrained_weights: bool = True, dtype=torch.float32, trust_remote_code: bool = True
+):
+ """
+ Load pretrained HuggingFace AutoModelForCausalLM.
+
+ Setting ``pretrained_weights=False`` returns a model that has identical architecture with the checkpoint,
+ but is randomly initialized.
+
+ Args:
+ model_path_or_name: Path or name of the model to load
+ pretrained_weights: Whether to load pretrained weights (True) or random init (False)
+ dtype: Data type for the model
+ trust_remote_code: Whether to trust remote code when loading model (needed for some models like Nemotron)
+ """
+ if pretrained_weights:
+ return AutoModelForCausalLM.from_pretrained(
+ model_path_or_name, torch_dtype=dtype, trust_remote_code=trust_remote_code
+ )
+ else:
+ config = AutoConfig.from_pretrained(model_path_or_name, trust_remote_code=trust_remote_code)
+ return AutoModelForCausalLM.from_config(config, torch_dtype=dtype, trust_remote_code=trust_remote_code)
+
+
+def set_model_dict_for_partial_init(
+ pretrained_dict: Dict[str, torch.Tensor], model_dict: Dict[str, torch.Tensor]
+) -> Dict[str, torch.Tensor]:
+ """
+ Partially initialize a model's state dictionary with a pretrained state dictionary.
+ This function safely copies compatible layers from a pretrained model into a new model,
+ ignoring layers with mismatched shapes or missing keys.
+
+ Steps:
+ 1. Remove layers from the pretrained dictionary if their shape does not match the target model.
+ 2. Keep only keys that exist in the target model.
+ 3. Update the model dictionary with the filtered pretrained weights.
+
+ Args:
+ pretrained_dict (Dict[str, torch.Tensor]):
+ The state dictionary of the pretrained model.
+ model_dict (Dict[str, torch.Tensor]):
+ The state dictionary of the target model to be partially initialized.
+
+ Returns:
+ Dict[str, torch.Tensor]:
+ The updated model state dictionary with compatible layers loaded from the pretrained dictionary.
+ """
+ # 1. Remove layers where pretrained shape differs from model shape
+ for k, v in list(pretrained_dict.items()):
+ if k in model_dict and hasattr(model_dict[k], "numel") and v.numel() != model_dict[k].numel():
+ del pretrained_dict[k]
+ logging.info(f" | > Layer with shape mismatch in the model definition: {k}")
+
+ # 2. Keep only keys that exist in the target model
+ pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
+
+ # 3. Update model dictionary with filtered pretrained layers
+ model_dict.update(pretrained_dict)
+ logging.info(f" | > {len(pretrained_dict)} / {len(model_dict)} layers are restored.")
+
+ return model_dict
+
+
+def load_checkpoint(checkpoint_path):
+ """
+ Load a model checkpoint from disk.
+
+ Supports loading checkpoints stored in either PyTorch (`.ckpt`, `.pt`) or
+ SafeTensors (`.safetensors`) formats. All parameters are loaded onto CPU
+ regardless of the original device.
+
+ Args:
+ checkpoint_path (str):
+ Path to the checkpoint file. If the filename contains `.safetensors`,
+ it is loaded using the SafeTensors backend; otherwise, it is assumed
+ to be a PyTorch checkpoint containing a `state_dict` field.
+
+ Returns:
+ dict:
+ A state dictionary mapping parameter names to tensors.
+ """
+ if ".safetensors" in checkpoint_path:
+ checkpoint_state = load_file(checkpoint_path, device="cpu")
+ else:
+ checkpoint_state = torch.load(checkpoint_path, map_location="cpu")["state_dict"]
+ return checkpoint_state
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/duplex_ear_tts.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/duplex_ear_tts.py
new file mode 100644
index 000000000..7c674af5c
--- /dev/null
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/duplex_ear_tts.py
@@ -0,0 +1,1038 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/speechlm2/models/duplex_ear_tts.py; training-only code removed.
+# Reduced for inference: base class changed LightningModule+HFHubMixin -> nn.Module;
+# removed training_step/validation_step/test_step, epoch hooks, log_model_stats,
+# run_evaluation_one_batch, configure_optimizers/backward/configure_model (FSDP/TP),
+# oomptimizer_schema, and the lightning/peft/hydra/metrics imports they required.
+
+import os
+import time
+from collections import Counter
+from contextlib import contextmanager
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from omegaconf import DictConfig
+
+from .compat import (
+ fp32_precision,
+ get_pad_id,
+ load_checkpoint,
+ load_pretrained_hf,
+ logging,
+ set_model_dict_for_partial_init,
+)
+from .ear_tts_model import RVQEARTTSModel
+from .ear_tts_vae_codec import RVQVAEModel
+from .tokenizer import AutoTokenizer
+
+
+class DuplexEARTTS(nn.Module):
+ def __init__(self, cfg: dict) -> None:
+ assert isinstance(cfg, dict), (
+ "You must pass the config to DuplexEARTTS as a Python dict to support hyperparameter serialization "
+ f"in PTL checkpoints (we got: '{type(cfg)=}')."
+ )
+ super().__init__()
+ # convert dict to config
+ cfg = DictConfig(cfg)
+ self.trainer_config = cfg.get("trainer", None)
+ self.data_cfg = cfg.data
+ self.cfg = cfg.model
+ self.target_sample_rate = cfg.data.target_sample_rate
+ self.source_sample_rate = cfg.data.source_sample_rate
+ self.normalize_text = cfg.data.get("normalize_text", False)
+
+ # exp_manager is training-only; keep validation_save_path optional for inference configs.
+ if cfg.get("exp_manager", None) is not None:
+ self.validation_save_path = os.path.join(cfg.exp_manager.explicit_log_dir, "validation_logs")
+ else:
+ self.validation_save_path = None
+
+ # move back text channel by x, in inference it advance the text channel prediction by x frames
+ self.advance_text_channel_by = self.cfg.get("advance_text_channel_by", None)
+
+ # Load ForCausalLM
+ if self.cfg.tts_config.context_hidden_size is not None:
+ self.language_model = self._load_language_model(self.cfg)
+ self.embed_tokens = self._load_embed_tokens(self.cfg)
+ # delete llm because we use it only to get the embedding tokens
+ del self.language_model
+
+ # get codec run precision
+ self.audio_codec_run_dtype = getattr(torch, self.cfg.get("audio_codec_run_dtype", "float32"), torch.float32)
+
+ # Load tokenizer
+ self.tokenizer = AutoTokenizer(
+ self.cfg.pretrained_lm_name,
+ use_fast=True,
+ trust_remote_code=True,
+ bos_token=self.cfg.get("bos_token", None),
+ eos_token=self.cfg.get("eos_token", None),
+ pad_token=self.cfg.get("pad_token", None),
+ ) # Note that we are using fast tokenizer
+
+ # Instantiate TTS model
+ self.tts_model = RVQEARTTSModel(DictConfig(self.cfg.tts_config), tokenizer=self.tokenizer)
+ # Load and initialize audio codec, and bind RVQ embeddings to the TTS model
+ setup_audio_codec(self)
+
+ self._codebook_size = self.tts_model.config.codebook_size
+
+ # compute samples per frame
+ self.source_samples_per_frame = int(self.source_sample_rate * cfg.data.frame_length)
+
+ # get codec silence tokens
+ codec_silence_tokens = self.get_codec_silence_frame()
+ self.register_buffer("codec_silence_tokens", codec_silence_tokens)
+
+ # cached for quicker audio decoding
+ self.register_buffer(
+ "_control_codes",
+ torch.tensor([self.speech_bos_id, self.speech_eos_id, self.speech_pad_id], device=self.device),
+ )
+
+ self._use_fsdp = False
+ self._use_tp = False
+ self.audio_prompt_latents = nn.ParameterDict()
+
+ def get_codec_silence_frame(self):
+ # Generate long zero waveform (silence)
+ audio = torch.zeros(1, 10 * self.target_sample_rate).float().to(self.device)
+ audio_len = torch.tensor([audio.size(-1)]).long()
+ audio, audio_len = self.pad_audio_to_factor(audio, audio_len, self.target_samples_per_frame)
+
+ with ensures_target_precision(self.audio_codec_run_dtype), torch.no_grad():
+ sil_codes, _ = self.audio_codec.encode(audio.unsqueeze(1), audio_len) # [1, T, C]
+ sil_codes = sil_codes[0] # [T, C]
+
+ # Convert each frame (C tokens) into a tuple
+ combos = [tuple(row.tolist()) for row in sil_codes]
+
+ # Count frequencies
+ counter = Counter(combos)
+
+ # Pick the most common combination
+ most_common_combo, freq = counter.most_common(1)[0]
+
+ # Return as tensor [C]
+ return torch.tensor(most_common_combo, device=self.device, dtype=torch.long)
+
+ def _load_embed_tokens(self, cfg) -> nn.Embedding:
+ """Load token embedding layer for RVQ-EAR-TTS."""
+ if self.language_model:
+ assert callable(self.language_model.get_input_embeddings)
+ embed_tokens: nn.Embedding = self.language_model.get_input_embeddings()
+ else:
+ embed_tokens_state_dict = torch.load(
+ cfg.pretrained_lm_embedding_path, map_location="cpu", weights_only=True
+ )
+
+ # Create token embedding layer
+ vocab_size, hidden_size = embed_tokens_state_dict["weight"].size()
+ embed_tokens = nn.Embedding(vocab_size, hidden_size, dtype=torch.bfloat16)
+ embed_tokens.load_state_dict(embed_tokens_state_dict)
+ return embed_tokens
+
+ def _load_language_model(self, cfg):
+ """Load language model for RVQ-EAR-TTS."""
+ if cfg.pretrained_lm_name:
+ language_model = load_pretrained_hf(
+ self.cfg.pretrained_lm_name, pretrained_weights=True, trust_remote_code=True
+ ).eval()
+ else:
+ language_model = None
+ return language_model
+
+ @property
+ def device(self):
+ return next(self.parameters()).device
+
+ @property
+ def speech_vocab_size(self):
+ """Return the size of the audio codec codebook including extra speech BOS and EOS tokens."""
+ if self.use_local_transformer and self.local_transformer_type == "nar": # add extra token for mask
+ return self._codebook_size + 4
+ return self._codebook_size + 3
+
+ @property
+ def speech_bos_id(self) -> int:
+ """Indicates start of utterance generation (not start of inference!)."""
+ if self.cfg.get("custom_speech_bos_id", None):
+ return self.cfg.get("custom_speech_bos_id")
+ return self._codebook_size + 2
+
+ @property
+ def speech_eos_id(self) -> int:
+ """Indicates end of utterance generation."""
+ if self.cfg.get("custom_speech_eos_id", None):
+ return self.cfg.get("custom_speech_eos_id")
+ return self._codebook_size + 1
+
+ @property
+ def speech_pad_id(self) -> int:
+ """Indicates start of inference (the very first frame)."""
+ if self.cfg.get("custom_speech_pad_id", None):
+ return self.cfg.get("custom_speech_pad_id")
+ return self._codebook_size
+
+ @property
+ def text_vocab_size(self):
+ """Return the size of the text tokenizer."""
+ return self.tokenizer.vocab_size
+
+ @property
+ def text_bos_id(self) -> int:
+ return self.tokenizer.bos_id
+
+ @property
+ def text_eos_id(self) -> int:
+ return self.tokenizer.eos_id
+
+ @property
+ def text_pad_id(self) -> int:
+ """
+ Text pad ID is used as a 'blank' for frames when the model is not speaking
+ and for frames where the model is speaking but has already predicted the
+ entire text channel's content.
+
+ Example:
+
+ flow: |---user---||-------assistant--------||-user-|
+ text channel: 0000000000 1xxxxxxx0000000000000002 000000
+
+ Where 0 indicates PAD ID, 1 indicates BOS ID, 2 indacates EOS ID,
+ and x indicates tokens corresponding to actual text
+
+ """
+ return get_pad_id(self.tokenizer)
+
+ def pad_audio_to_factor(self, audio, audio_len, samples_per_frame, downsampling_factor: int = 1):
+ """
+ Zero pad the end of the audio so that we do not have a partial end frame.
+ The output will be zero-padded to have an integer number of frames of
+ length `samples_per_frame * downsampling_factor`.
+
+ Args:
+ audio: input time-domain signal (B, T)
+ audio_len: valid length for each example in the batch (B,)
+ samples_per_frame: number of samples per frame
+ downsampling_factor: how much each frame is downsampled in later processing
+
+ Returns:
+ padded_audio: Padded time-domain signal (B, T')
+ padded_len: Adjusted valid lengths (B,)
+ """
+ with fp32_precision():
+ total_factor = samples_per_frame * downsampling_factor
+ padded_len = total_factor * torch.ceil(audio_len / total_factor).int()
+ max_len = padded_len.max().int().item()
+ num_padding = max_len - audio.shape[1]
+ padded_audio = F.pad(audio, (0, num_padding))
+ return padded_audio, padded_len
+
+ def prepare_inputs(self, batch: dict):
+ """
+ Prepare inputs, extracting audio tokens and padding if needed.
+ """
+ # check if audios has the same batch size
+ assert batch["source_audio"].size(0) == batch["target_audio"].size(0)
+ assert batch["audio_prompt"].size(0) == batch["target_audio"].size(0)
+
+ target_audio = batch["target_audio"]
+ target_audio_lens = batch["target_audio_lens"]
+ target_text_tokens = batch["target_text_tokens"]
+ non_prompt_mask = batch["non_prompt_mask"]
+ aligned_attention_mask = batch["aligned_attention_mask"]
+ aligned_position_ids = batch["aligned_position_ids"]
+
+ if self.training and (self.cfg.get("empty_turn_probability", 0.0) > 0):
+ # Randomly decide whether this batch gets emptied
+ if torch.rand(1).item() < self.cfg.empty_turn_probability:
+ # Zero out audio
+ target_audio = torch.zeros_like(target_audio)
+
+ # Create mask for tokens we want to drop
+ # Keep BOS and EOS, drop the rest.
+ keep_mask = (target_text_tokens == self.text_bos_id) | (target_text_tokens == self.text_eos_id)
+ full_dropout_mask = ~keep_mask # True = positions to replace with PAD
+
+ # Replace all non-BOS/EOS with PAD
+ target_text_tokens = torch.where(
+ full_dropout_mask, torch.full_like(target_text_tokens, self.text_pad_id), target_text_tokens
+ )
+
+ # extract target audio codes
+ target_audio, target_audio_lens = self.pad_audio_to_factor(
+ target_audio, target_audio_lens, self.target_samples_per_frame, 1
+ )
+ with ensures_target_precision(self.audio_codec_run_dtype), torch.no_grad():
+ target_codes, target_codes_lens = self.audio_codec.encode(target_audio.unsqueeze(1), target_audio_lens)
+
+ with fp32_precision():
+ target_len = target_codes.shape[1]
+
+ # Pad or truncate sequence variables
+ def pad_or_truncate(x, pad_value=0):
+ if x.dim() == 2: # [B, T]
+ L = x.shape[1]
+ if L < target_len:
+ return F.pad(x, (0, target_len - L), value=pad_value)
+ else:
+ return x[:, :target_len]
+ return x # leave others for now
+
+ target_text_tokens = pad_or_truncate(target_text_tokens, pad_value=self.text_pad_id)
+ non_prompt_mask = pad_or_truncate(non_prompt_mask, pad_value=0)
+ aligned_position_ids = pad_or_truncate(aligned_position_ids, pad_value=0)
+
+ # Correct attention mask padding/truncation
+ B, H, L1, L2 = aligned_attention_mask.shape
+ new_len = target_len
+ if L1 < new_len or L2 < new_len:
+ pad_rows = new_len - L1
+ pad_cols = new_len - L2
+ aligned_attention_mask = F.pad(aligned_attention_mask, (0, pad_cols, 0, pad_rows))
+ elif L1 > new_len or L2 > new_len:
+ aligned_attention_mask = aligned_attention_mask[:, :, :new_len, :new_len]
+
+ # set the pad token for the first BOS frame
+ target_codes_aligned = target_codes.clone()
+ target_codes_aligned[:, 0] = self.speech_pad_id
+
+ # set special token in the last audio prompt (it will works as a BOS token)
+ pos = non_prompt_mask.float().argmax(dim=1) # shape: [B]
+ row_idx = torch.arange(B, device=self.device)
+ # set the extra self.speech_pad_id at first 1 position in non_prompt_mask
+ target_codes_aligned[row_idx, pos] = self.speech_pad_id
+
+ # EOS dropout to make the model more robust
+ if self.training and self.cfg.get("text_eos_dropout_prob", 0.0) > 0:
+ # Mask EOS positions
+ eos_mask = target_text_tokens == self.text_eos_id
+
+ # Random dropout only on EOS positions
+ dropout_mask = torch.rand(eos_mask.sum(), device=target_text_tokens.device) < self.cfg.text_eos_dropout_prob
+
+ # Scatter dropout decisions into [B, T]
+ full_dropout_mask = torch.zeros_like(target_text_tokens, dtype=torch.bool)
+ full_dropout_mask[eos_mask] = dropout_mask
+
+ # Replace dropped EOS with PAD
+ target_text_tokens = torch.where(
+ full_dropout_mask, torch.full_like(target_text_tokens, self.text_pad_id), target_text_tokens
+ )
+
+ # BOS dropout to make the model more robust
+ if self.training and self.cfg.get("text_bos_dropout_prob", 0.0) > 0:
+ # Mask BOS positions
+ bos_mask = target_text_tokens == self.text_bos_id
+
+ # Random dropout only on BOS positions
+ dropout_mask = torch.rand(bos_mask.sum(), device=target_text_tokens.device) < self.cfg.text_bos_dropout_prob
+
+ # Scatter dropout decisions into [B, T]
+ full_dropout_mask = torch.zeros_like(target_text_tokens, dtype=torch.bool)
+ full_dropout_mask[bos_mask] = dropout_mask
+
+ # Replace dropped BOS with PAD
+ target_text_tokens = torch.where(
+ full_dropout_mask,
+ torch.full_like(target_text_tokens, self.text_pad_id),
+ target_text_tokens,
+ )
+
+ # BOS dropout to make the model more robust
+ if self.training and self.cfg.get("text_bos_dropout_prob", 0.0) > 0:
+ prob = self.cfg.text_bos_dropout_prob # e.g., 0.5
+
+ # Identify all BOS positions [B, T]
+ bos_mask = target_text_tokens == self.text_bos_id
+
+ # Get indices of sequences that actually have a BOS token
+ # We need to know *where* the BOS tokens are to drop them.
+ # tensor of coordinates: [[batch_idx, seq_idx], ...]
+ bos_indices = torch.nonzero(bos_mask)
+ num_bos = bos_indices.shape[0]
+
+ if num_bos > 0:
+ # Create a random dropout decision for each BOS instance
+ drop_decisions = torch.rand(num_bos, device=target_text_tokens.device) < prob
+
+ # Ensure at least one is dropped
+ if drop_decisions.sum() == 0:
+ # Pick one random index from the available BOS locations to drop
+ force_idx = torch.randint(0, num_bos, (1,), device=target_text_tokens.device)
+ drop_decisions[force_idx] = True
+
+ # 5. Apply the dropout
+ # We need to map the decisions back to the full tensor
+ # Create a mask of the same shape as target_text_tokens
+ full_dropout_mask = torch.zeros_like(target_text_tokens, dtype=torch.bool)
+
+ # Set True only at the specific (batch, seq) coordinates we chose to drop
+ # bos_indices[:, 0] are batch indices, bos_indices[:, 1] are seq indices
+ full_dropout_mask[bos_indices[:, 0], bos_indices[:, 1]] = drop_decisions
+
+ # 6. Replace dropped BOS with PAD
+ target_text_tokens = torch.where(
+ full_dropout_mask,
+ torch.full_like(target_text_tokens, self.text_pad_id),
+ target_text_tokens,
+ )
+
+ # shift text tokens
+ subword_ids = F.pad(target_text_tokens[:, 1:], [0, 1])
+ # note that we are using a text mask where we are ignoring the desc + audio prompt but we are keeping 1 until the audio ends to support duplex
+ subword_mask = F.pad(non_prompt_mask[:, 1:], [0, 1])
+
+ # detach embedding as in eartts
+ if self.cfg.tts_config.context_hidden_size is not None:
+ context_hidden_state = self.embed_tokens(target_text_tokens).detach()
+ else:
+ context_hidden_state = None
+
+ if self._use_tp:
+ tp_world_size = self.device_mesh["tensor_parallel"].size()
+ if (remainder := (target_text_tokens.shape[1] - 1) % tp_world_size) != 0:
+ target_text_tokens = target_text_tokens[:, :-remainder]
+ target_codes_aligned = target_codes_aligned[:, :-remainder]
+ target_codes_aligned = target_codes_aligned[:, :-remainder]
+ subword_ids = subword_ids[:, :-remainder]
+ subword_mask = subword_mask[:, :-remainder]
+
+ return {
+ "code": target_codes_aligned,
+ "audio_mask": non_prompt_mask, # set audio_mask as non_prompt_mask to avoid the audio prompt in loss computation
+ "attention_mask": aligned_attention_mask,
+ "position_ids": aligned_position_ids,
+ "subword_ids": subword_ids,
+ "subword_mask": subword_mask,
+ "context_hidden_state": context_hidden_state,
+ "output_lens": target_codes_lens,
+ "non_prompt_mask": non_prompt_mask,
+ "target_text_tokens": target_text_tokens,
+ }
+
+ def _get_generation_config(self, guidance_enabled: bool = False):
+ """Get default generation config for EAR-TTS."""
+ return {
+ "num_iter": 8,
+ "guidance_scale": self.cfg.get("inference_guidance_scale", 0.5) if guidance_enabled else None,
+ "top_p_or_k": self.cfg.get("inference_top_p_or_k", 0.8),
+ "noise_scale": self.cfg.get("inference_noise_scale", 0.8),
+ "eos_threshold": -3.0,
+ }
+
+ def get_audio_prompt_latent(self, name, B):
+ """
+ Retrieve a cached audio prompt latent and adapt it to the requested batch size.
+
+ This fetches a latent previously cached via `set_audio_prompt_latent()` and
+ ensures the returned tensor has batch size `B` by:
+ - returning as-is when the batch already matches,
+ - truncating when the cached batch is larger than `B`,
+ - expanding when the cached batch is smaller (commonly batch=1 warmup).
+
+ Args:
+ name (str):
+ Key of the cached audio prompt latent to retrieve.
+ B (int):
+ Desired batch size.
+
+ Returns:
+ Tensor:
+ Audio prompt latent with batch dimension equal to `B`, moved to `self.device`.
+ Shape: [B, ..., D]
+
+ Raises:
+ KeyError:
+ If `name` does not exist in `self.audio_prompt_latents`.
+ """
+ if not hasattr(self, "audio_prompt_latents") or name not in self.audio_prompt_latents:
+ raise KeyError(f"Unknown audio prompt latent '{name}'. Call set_audio_prompt_latent(...) first.")
+
+ audio_prompt_latent = self.audio_prompt_latents[name] # cached on CPU
+
+ if audio_prompt_latent.shape[0] == B:
+ out = audio_prompt_latent
+ elif audio_prompt_latent.shape[0] >= B:
+ out = audio_prompt_latent[:B]
+ else:
+ out = audio_prompt_latent[:1].expand(B, *audio_prompt_latent.shape[1:])
+
+ return out.to(self.device)
+
+ def set_init_inputs(self, speaker_audio=None, speaker_audio_lens=None, system_prompt=None, speaker_name=None):
+ """
+ Registers and prepares initial input buffers for text/audio prompt and context, to warm up AR inference.
+
+ Args:
+ speaker_audio (torch.Tensor): Batch of prompt audio, (B, T).
+ speaker_audio_lens (torch.Tensor): Lengths for each sample in speaker_audio, (B,).
+ system_prompt (str, optional): System prompt for context.
+ speaker_name (str, optional): Name of a pre-baked speaker latent in audio_prompt_latents.
+ When provided, speaker_audio/speaker_audio_lens are ignored and silent audio is used
+ as the prompt carrier; the actual voice identity comes from the cached latent.
+
+ Returns:
+ dict: Dictionary of input tensors to be passed to inference, with registered buffers.
+ """
+ # compute prompt audio size and slice it
+ with fp32_precision():
+ # compute the exact number of samples for the prompt duration
+ prompt_audio_size = int(
+ ((self.data_cfg.audio_prompt_duration * self.target_sample_rate) // self.target_samples_per_frame)
+ * self.target_samples_per_frame
+ )
+
+ if speaker_name is not None:
+ logging.info(
+ f"set_init_inputs: using pre-baked latent for speaker '{speaker_name}' (silent carrier audio)"
+ )
+ speaker_audio = torch.zeros((1, prompt_audio_size), device=self.device, dtype=torch.float32)
+ speaker_audio_lens = torch.LongTensor([speaker_audio.shape[1]]).to(self.device)
+
+ B, T = speaker_audio.shape
+ device = speaker_audio.device
+ dtype = speaker_audio.dtype
+
+ # allocate result
+ prompt_audio = torch.zeros(B, prompt_audio_size, device=device, dtype=dtype)
+
+ # process each example independently
+ for b in range(B):
+ valid_len = min(speaker_audio_lens[b].item(), T)
+
+ # handle empty
+ if valid_len <= 0:
+ continue
+
+ # valid (non-padded) segment
+ valid_segment = speaker_audio[b, :valid_len]
+
+ if valid_len >= prompt_audio_size:
+ # enough valid audio → crop from start (no silence)
+ prompt_audio[b] = valid_segment[:prompt_audio_size]
+ else:
+ # too short → repeat and crop
+ repeat_factor = (prompt_audio_size + valid_len - 1) // valid_len # ceil division
+ expanded = valid_segment.repeat(repeat_factor)
+ prompt_audio[b] = expanded[:prompt_audio_size]
+
+ # add a silence in the end to smooth the transition between prompt and audio tokens
+ prompt_audio[:, -int(self.target_samples_per_frame * 2) :] = 0
+
+ # get prompt audio size
+ with fp32_precision():
+ prompt_audio_text_pad_size = int(prompt_audio_size // self.target_samples_per_frame)
+
+ # create a eos token id
+ if system_prompt is not None and self.cfg.get("use_system_prompt", None) and system_prompt != "":
+ text_prompt = torch.as_tensor(
+ [self.tokenizer.bos] + self.tokenizer.text_to_ids(system_prompt) + [self.tokenizer.eos],
+ dtype=torch.long,
+ device=self.device,
+ )
+ else:
+ text_prompt = torch.tensor([self.tokenizer.eos], dtype=torch.long, device=self.device)
+
+ # create a padding tensor
+ prompt_audio_text_pad = (
+ torch.ones(prompt_audio_text_pad_size, device=self.device, dtype=text_prompt.dtype) * self.text_pad_id
+ )
+ prompt_audio_text_pad[-1] = self.tokenizer.eos
+
+ # Prepend an initial text EOS token followed by padding tokens that match
+ # the number of audio-prompt frames (in text-token units).
+ target_text_tokens = torch.cat([text_prompt, prompt_audio_text_pad.to(text_prompt.dtype)])
+
+ # create pad audio for the description
+ pad_size = text_prompt.size(-1) * self.target_samples_per_frame
+ pad_audio = (
+ torch.zeros(pad_size, device=prompt_audio.device, dtype=prompt_audio.dtype)
+ .unsqueeze(0)
+ .repeat(prompt_audio.size(0), 1)
+ )
+
+ # repeat to reaches the batch size
+ target_text_tokens = target_text_tokens.unsqueeze(0).repeat(prompt_audio.size(0), 1)
+ target_audio = torch.cat([pad_audio, prompt_audio], dim=1)
+
+ # extract code codes
+ target_audio_len = torch.tensor(
+ [target_audio.size(-1)] * target_audio.size(0), dtype=torch.long, device=self.device
+ )
+ with ensures_target_precision(self.audio_codec_run_dtype), torch.no_grad():
+ code, _ = self.audio_codec.encode(target_audio.unsqueeze(1), target_audio_len)
+
+ # get context hidden
+ if self.cfg.tts_config.context_hidden_size is not None:
+ context_hidden_state = self.embed_tokens(target_text_tokens)
+ else:
+ context_hidden_state = None
+
+ # create masks
+ # non_prompt_mask is all zeros, because all processed is prompt
+ non_prompt_mask = torch.zeros_like(target_text_tokens)
+ non_prompt_mask[:, -2:] = 1 # set last valid prompt frame as 1 to allow the addition of BOS in the right place
+ subword_mask = torch.zeros_like(
+ target_text_tokens
+ ) # subword_mask is almost all zeros because on the warmup there is only the prompt
+ subword_mask[:, -3:] = (
+ 1 # -3 because of the it start right after the first valid prompt token and it is shifted by 1
+ )
+
+ # set the pad token for the first BOS frame
+ code[:, 0] = self.speech_pad_id
+
+ # shift subword_ids
+ subword_ids = F.pad(target_text_tokens[:, 1:], [0, 1], value=0.0)
+
+ # set special token in the last audio prompt (it will works as a BOS token)
+ pos = non_prompt_mask.float().argmax(dim=1) # shape: [B]
+ row_idx = torch.arange(B, device=self.device)
+ # set the extra self.speech_pad_id at first 1 position in non_prompt_mask
+ code[row_idx, pos] = self.speech_pad_id
+
+ init_inputs = {
+ "code": code[:, :-1],
+ "audio_mask": non_prompt_mask.bool()[
+ :, :-1
+ ], # set audio_mask as non_prompt_mask to avoid the audio prompt in loss computation
+ "context_hidden_state": context_hidden_state[:, :-1] if context_hidden_state is not None else None,
+ "subword_ids": subword_ids[:, :-1],
+ "subword_mask": subword_mask.bool()[:, :-1],
+ "non_prompt_mask": non_prompt_mask.bool()[:, :-1],
+ }
+
+ if speaker_name is not None:
+ init_inputs["audio_prompt_latent"] = self.get_audio_prompt_latent(speaker_name, B=1)
+
+ self._init_input_cache = {}
+ for k, v in init_inputs.items():
+ if v is None:
+ self._init_input_cache[k] = None
+ else:
+ # clone → removes inference tensor flag
+ # detach → ensures no graph refs
+ self._init_input_cache[k] = v.detach().clone()
+
+ return init_inputs
+
+ def get_init_inputs(
+ self,
+ B: int,
+ init_inputs_names=[
+ "code",
+ "audio_mask",
+ "context_hidden_state",
+ "subword_ids",
+ "subword_mask",
+ "non_prompt_mask",
+ ],
+ ):
+ """
+ Returns a dictionary of initial inputs for inference, using registered buffers.
+
+ Args:
+ B (int): Required batch size.
+ init_inputs_names (List[str], optional): Names of input buffers to fetch.
+
+ Returns:
+ dict: Each key is name from init_inputs_names, and value is tensor of appropriate shape (B, ...).
+
+ Notes:
+ Expands batch-1 buffers to B if necessary.
+ """
+ if init_inputs_names is None:
+ init_inputs_names = [
+ "code",
+ "audio_mask",
+ "context_hidden_state",
+ "subword_ids",
+ "subword_mask",
+ "non_prompt_mask",
+ ]
+
+ if self._init_input_cache.get("audio_prompt_latent", None) is not None:
+ init_inputs_names = list(init_inputs_names) + ["audio_prompt_latent"]
+
+ init_inputs = {}
+ for name in init_inputs_names:
+ buf = self._init_input_cache.get(name, None)
+
+ if buf is None:
+ init_inputs[name] = None
+ continue
+
+ # Batch already matches
+ if buf.shape[0] == B:
+ init_inputs[name] = buf
+ elif buf.shape[0] >= B:
+ init_inputs[name] = buf[:B]
+ else:
+ # assume batch=1 warmup → expand
+ init_inputs[name] = buf[:1].expand(B, *buf.shape[1:])
+
+ return init_inputs
+
+ @torch.inference_mode()
+ def infer_codes_one_step(
+ self,
+ current_subword_id,
+ prev_subword_id,
+ current_subword_mask,
+ prev_audio_tokens,
+ past_key_values,
+ guidance_enabled=True,
+ generation_config=None,
+ ignore_eos_flag_stop=True,
+ request_id=None,
+ ):
+ """
+ Runs a single autoregressive prediction step to infer audio codec codes.
+
+ Args:
+ current_subword_id (torch.Tensor): Current text token IDs, shape (B, 1).
+ prev_subword_id (torch.Tensor): Previous text token IDs, shape (B, 1).
+ current_subword_mask (torch.Tensor): Current mask, shape (B, 1).
+ prev_audio_tokens (torch.Tensor): Previously generated audio tokens, shape (B, 1, C).
+ past_key_values: Key-value cache for transformer decoder state.
+ guidance_enabled (bool, optional): Enables classifier-free guidance.
+ generation_config (dict, optional): Generation hyperparameters.
+ ignore_eos_flag_stop (bool): If True, ignore EOS flag for stopping.
+
+ Returns:
+ Tuple[torch.Tensor, Any]:
+ - Predicted audio codec token(s), shape (B, 1, C)
+ - Updated past_key_values for the next step.
+ """
+
+ if self.cfg.tts_config.context_hidden_size is not None:
+ # get context_hidden_state it is always one step behind current_subword_id
+ # for the first step uses the last step from warmup
+ context_hidden_state = self.embed_tokens(prev_subword_id)
+ else:
+ context_hidden_state = None
+
+ # force silence as next token
+ if self.cfg.get("inference_force_speech_silence_on_eos", True):
+ silence_codes = self.codec_silence_tokens.view(1, 1, -1).expand(prev_audio_tokens.shape)
+ prev_audio_tokens = torch.where(
+ current_subword_id.unsqueeze(-1) == self.text_eos_id,
+ silence_codes, # silence
+ prev_audio_tokens, # keep original
+ )
+
+ # get subword_ids
+ inputs = {
+ "code": prev_audio_tokens,
+ "context_hidden_state": context_hidden_state,
+ "subword_ids": current_subword_id,
+ "subword_mask": current_subword_mask,
+ "past_key_values": past_key_values,
+ "use_cache": True,
+ "guidance_enabled": guidance_enabled,
+ "generation_config": generation_config,
+ "ignore_eos_flag_stop": ignore_eos_flag_stop,
+ }
+ # request_id is only used by vLLM backend — skip for native PyTorch inference
+ if request_id is not None:
+ inputs["request_id"] = request_id
+
+ outputs = self.tts_model(**inputs)
+
+ return outputs["codes"], outputs["past_key_values"]
+
+ @torch.inference_mode()
+ def decode_one_audio_step(self, gen_audio_codes_history, number_prev_tokens=None):
+ """
+ Decodes one step of generated audio codec tokens to raw waveform.
+
+ Args:
+ gen_audio_codes_history (torch.Tensor): Audio tokens history, shape (B, T, C).
+ number_prev_tokens (int, optional): Number of previous tokens to decode, for incremental decoding.
+
+ Returns:
+ Tuple[torch.Tensor, torch.Tensor]:
+ - audio_pred_cur_step: Latest decoded waveform chunk, shape (B, wav_to_token_ratio).
+ - audio_len: Lengths (number of samples), shape (B,).
+ """
+ with fp32_precision(), torch.no_grad():
+ if number_prev_tokens:
+ gen_audio_codes_history = gen_audio_codes_history[:, -number_prev_tokens:]
+
+ gen_audio_codes_history = replace_control_speech_codes(
+ gen_audio_codes_history, self._control_codes, self.codec_silence_tokens
+ )
+ gen_audio_codes_lens = torch.tensor(
+ [gen_audio_codes_history.size(1)] * gen_audio_codes_history.size(0), device=self.device
+ )
+ audio_pred, audio_len = self.audio_codec.decode(gen_audio_codes_history, gen_audio_codes_lens)
+
+ # return only the current/lastest audio chunk
+ audio_pred_cur_step = audio_pred.squeeze(1)[:, -self.audio_codec.config.wav_to_token_ratio :]
+ audio_len[:] = self.audio_codec.config.wav_to_token_ratio
+ return audio_pred_cur_step, audio_len
+
+ @torch.inference_mode()
+ def offline_inference(
+ self,
+ next_subword_ids: torch.Tensor,
+ init_inputs: dict,
+ formatter: str = "",
+ guidance_enabled: bool = True,
+ generation_config: dict = None,
+ incremental_audio_decoding: bool = False,
+ ) -> dict[str, torch.Tensor]:
+ """
+ Runs offline autoregressive inference for the Duplex EAR-TTS speech decoder.
+
+ This method performs **text-to-speech (TTS)** generation: given subword/text
+ tokens and prompt-initialization states, it autoregressively generates
+ audio codec tokens and decodes them into a waveform.
+
+ Args:
+ next_subword_ids (torch.Tensor):
+ Conditioning subword/text token IDs for the speech decoder.
+ Shape: (B, T_text).
+
+ init_inputs (dict):
+ Dictionary of prompt-dependent initial states produced by
+ ``get_init_inputs()``. May include:
+
+ • "code" — initial audio tokens (e.g., prompt audio)
+ • "audio_mask" — mask for prompt audio positions
+ • "context_hidden_state" — decoder hidden state at t = 0
+ • "subword_ids" — prompt text tokens
+ • "subword_mask" — mask for prompt text
+ • "non_prompt_mask" — mask marking positions to be generated
+
+ ``get_init_inputs()`` automatically expands batch-1 buffers to
+ batch size B.
+
+ formatter (str, optional):
+ Optional formatter identifier used to customize the prompt structure.
+
+ guidance_enabled (bool, optional):
+ Whether classifier-free guidance (CFG) is enabled.
+ If enabled and ``generation_config`` is ``None``, guidance parameters
+ are taken from ``_get_generation_config()``.
+
+ generation_config (dict, optional):
+ Settings controlling autoregressive generation, including sampling
+ strategy, noise scale, refinement iterations, and EOS rules.
+ If ``None``, defaults are taken from
+ ``_get_generation_config(guidance_enabled)``.
+
+ incremental_audio_decoding (bool, optional):
+ If True, codec-to-waveform decoding is performed incrementally during
+ autoregressive generation.
+ If False, waveform decoding occurs only after all audio tokens are produced.
+
+ Returns:
+ dict[str, torch.Tensor]:
+ Contains:
+
+ • **"audio"**:
+ Generated waveform of shape ``(B, T_audio)``, obtained via
+ ``audio_pred.squeeze(1)``.
+
+ • **"audio_len"**:
+ Length of each generated waveform in samples, shape ``(B,)``.
+ """
+ B = next_subword_ids.size(0)
+
+ if generation_config is None:
+ generation_config = self._get_generation_config(guidance_enabled)
+ logging.info(f"Doing inference using the following config: {generation_config} !")
+
+ init_inputs.update({"use_cache": True, "past_key_values": None, "guidance_enabled": guidance_enabled})
+
+ # warmup the model and generate the very first audio token
+ outputs = self.tts_model(**init_inputs)
+
+ if self.cfg.get("inference_skip_first_code_prediction_on_init", True):
+ # use the last token on init, because we are shifthing it in the model forward, so we dont really need to compute it
+ code = init_inputs["code"][:, -1:]
+ else:
+ code, _, _ = self.tts_model.generate_step(outputs.hidden_states[:, -1:], **generation_config)
+
+ past_key_values = outputs["past_key_values"]
+
+ # use the text tokens to stop generation
+ max_steps = next_subword_ids.size(-1)
+ # create variable to store the audios
+ gen_audio_codes = torch.zeros(
+ B, max_steps, self.tts_model.config.num_quantizers, device=self.device, dtype=torch.long
+ )
+
+ # init subwork as all ones
+ subword_mask = torch.ones(B, max_steps, device=self.device, dtype=torch.bool)
+ # get first context subword_id, that is the last subword_ids from the warmup
+ first_context_subword_id = init_inputs["subword_ids"][:, -1].unsqueeze(-1)
+
+ # initialize variables used to save the output audio
+ audio_pred = None
+ audio_pred_len = torch.zeros(B, device=self.device, dtype=torch.long)
+
+ for i in range(max_steps):
+ step_start = time.time()
+ # current subword id is always seem
+ current_subword_id = next_subword_ids[:, i].unsqueeze(-1)
+
+ if i == 0:
+ prev_subword_id = first_context_subword_id
+ else:
+ prev_subword_id = next_subword_ids[:, i - 1].unsqueeze(-1)
+
+ # create subword_mask
+ current_subword_mask = subword_mask[:, i].unsqueeze(-1)
+
+ code, past_key_values = self.infer_codes_one_step(
+ current_subword_id=current_subword_id,
+ prev_subword_id=prev_subword_id,
+ current_subword_mask=current_subword_mask,
+ prev_audio_tokens=code,
+ past_key_values=past_key_values,
+ guidance_enabled=guidance_enabled,
+ generation_config=generation_config,
+ ignore_eos_flag_stop=True,
+ )
+
+ # cache audio tokens
+ gen_audio_codes[:, i] = code.squeeze(1)
+
+ if incremental_audio_decoding:
+ audio_pred_i, audio_pred_i_len = self.decode_one_audio_step(
+ gen_audio_codes[:, : i + 1],
+ number_prev_tokens=self.cfg.get("inference_codec_decoding_prev_tokens_number", None),
+ )
+ if audio_pred is None:
+ audio_pred = audio_pred_i
+ else:
+ audio_pred = torch.cat([audio_pred, audio_pred_i], dim=1)
+ audio_pred_len += audio_pred_i_len
+
+ step_time = time.time() - step_start
+ logging.info(f"Autoregressive inference step: {i} of {max_steps} take around {step_time}s")
+
+ if not incremental_audio_decoding:
+ gen_audio_codes_lens = torch.tensor([gen_audio_codes.shape[1]] * gen_audio_codes.shape[0]).to(self.device)
+ # decode audio. Note that it is not necessary because the prompt is removed, so no special token should be on the output, but lets do it for safety
+ gen_audio_codes = replace_control_speech_codes(
+ gen_audio_codes, self._control_codes, self.codec_silence_tokens
+ )
+ with ensures_target_precision(self.audio_codec_run_dtype), torch.no_grad():
+ audio_pred, audio_pred_len = self.audio_codec.decode(gen_audio_codes, gen_audio_codes_lens)
+
+ return audio_pred.squeeze(1), audio_pred_len
+
+ def maybe_recreate_cached_audio_prompt_latents_structure(self, state_dict):
+ """
+ Recreate audio_prompt_latents ParameterDict structure
+ from keys present in the provided state_dict.
+ """
+ for k, tensor in state_dict.items():
+ if "audio_prompt_latents." in k:
+ name = k.split("audio_prompt_latents.")[1]
+ if name not in self.audio_prompt_latents:
+ self.audio_prompt_latents[name] = nn.Parameter(
+ torch.zeros_like(tensor),
+ requires_grad=False,
+ )
+
+ def load_state_dict(self, state_dict, strict: bool = True):
+ self.maybe_recreate_cached_audio_prompt_latents_structure(state_dict)
+ try:
+ return super().load_state_dict(state_dict, strict=strict)
+ except RuntimeError:
+ logging.info("Error loading model state_dict !! Retrying with partial initialization!")
+ model_dict = set_model_dict_for_partial_init(state_dict, self.state_dict())
+ return super().load_state_dict(model_dict, strict=False)
+
+
+@contextmanager
+def ensures_target_precision(target_dtype):
+ """
+ Workaround for precision related issues when training with bf16-true PyTorch Lightning precision setting.
+ In bf16-true, PTL changes PyTorch's default dtype, which may break implicit assumptions for some models.
+ This context manager restores default float32 precision and runs the computation in float32 autocast context.
+ """
+ default_dtype = torch.get_default_dtype()
+ torch.set_default_dtype(target_dtype)
+ try:
+ with torch.amp.autocast(device_type="cuda" if torch.cuda.is_available() else "cpu", dtype=target_dtype):
+ yield
+ finally:
+ torch.set_default_dtype(default_dtype)
+
+
+def replace_control_speech_codes(
+ speech_codes: torch.Tensor, control_codes: torch.Tensor, silence_tokens: torch.Tensor = None
+) -> torch.Tensor:
+ """
+ Replaces control codes (speech BOS, EOS, etc) in `speech_codes` with the first frame which is
+ assumed to consist of 'valid' codes representing silence.
+ """
+ if silence_tokens is not None:
+ # Expand to [B, 1, 74]
+ silence_tokens_expanded = silence_tokens.unsqueeze(0).unsqueeze(1).expand(speech_codes.shape[0], 1, -1)
+ return torch.where(torch.isin(speech_codes, control_codes), silence_tokens_expanded, speech_codes)
+
+ if torch.isin(speech_codes[:, :1], control_codes).any():
+ return torch.where(torch.isin(speech_codes, control_codes), torch.zeros_like(speech_codes[:, :1]), speech_codes)
+ else:
+ return torch.where(torch.isin(speech_codes, control_codes), speech_codes[:, :1], speech_codes)
+
+
+def setup_audio_codec(model):
+ """
+ Instantiates the RVQ audio codec and injects codec embeddings into the TTS model.
+
+ This function is responsible only for:
+ - Instantiating the codec model (`RVQVAEModel`).
+ - Loading pretrained codec weights (if configured).
+ - Freezing codec parameters.
+ - Registering RVQ embeddings inside the TTS model via `set_rvq_embs`.
+
+ Args:
+ model: Model instance of DuplexEARTTS
+ """
+ with ensures_target_precision(model.audio_codec_run_dtype):
+ model.audio_codec = RVQVAEModel(DictConfig(model.cfg.codec_config))
+ # load pretrained codec checkpoint
+ if model.cfg.get("pretrained_codec_model", None):
+ checkpoint_state = load_checkpoint(model.cfg.pretrained_codec_model)
+ checkpoint_state = set_model_dict_for_partial_init(checkpoint_state, model.audio_codec.state_dict())
+ model.audio_codec.load_state_dict(checkpoint_state, strict=True)
+
+ for p in model.audio_codec.parameters():
+ p.requires_grad = False
+
+ model.audio_codec.eval()
+
+ assert callable(model.tts_model.set_rvq_embs)
+
+ model.tts_model.set_rvq_embs(torch.stack([x.detach() for x in model.audio_codec.prvq.mus_list], 0))
+ model.tts_model.rvq_embs = model.tts_model.rvq_embs.to(next(model.tts_model.parameters()).dtype)
+ # compute target fps
+ model.target_fps = model.target_sample_rate / model.audio_codec.config.wav_to_token_ratio
+ model.target_samples_per_frame = model.audio_codec.config.wav_to_token_ratio
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_commons.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_commons.py
new file mode 100644
index 000000000..1b8b96169
--- /dev/null
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_commons.py
@@ -0,0 +1,260 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/speechlm2/modules/ear_tts_commons.py; training-only code removed.
+
+import argparse
+import glob
+import importlib.machinery
+import json
+import os
+import re
+import shutil
+import subprocess
+import sys
+from collections.abc import Mapping
+from typing import Any
+
+from omegaconf import DictConfig
+from safetensors import safe_open
+from torch import nn
+
+from .compat import logging
+
+# ==============================================================================
+# Constants
+# ==============================================================================
+PYTHON_CONFIG_GETTER_NAME = "get_config"
+CHECKPOINT_FORMAT = "checkpoint_{}/ema.safetensors"
+CONFIG_NAME = "config.json"
+GIT_HASH_NAME = "githash"
+SCRIPT_PLACEHOLDER = "[[[<<>>]]]"
+
+
+# ==============================================================================
+# Configuration Class and Utilities
+# ==============================================================================
+def get_config_from_file(config_path: str) -> DictConfig:
+ """
+ Loads a configuration from a JSON or Python file.
+
+ - For JSON files (`*.json`), it parses the file directly.
+ - For Python files (`*.py`), it imports the file as a module and calls a
+ `get_config()` function within it.
+ - It also supports a special syntax `path/to/config.py:config_name` to select
+ a specific configuration from a Python file that returns a dictionary of configs.
+
+ Args:
+ config_path (str): The path to the configuration file.
+
+ Returns:
+ Config: The loaded configuration object.
+
+ Raises:
+ AssertionError: If the file path is invalid, does not exist, or is not in
+ the expected format.
+ """
+
+ match = re.search(r".+\.((json)|(py)|(py:.+))$", config_path)
+ assert match, f"Only Python (*.py) or JSON (*.json) files are supported, but got {config_path}."
+
+ py_config_name: str | None = None
+ if not (config_path.endswith(".py") or config_path.endswith(".json")):
+ config_path_split = config_path.split(":")
+ config_path = ":".join(config_path_split[:-1])
+ py_config_name = config_path_split[-1]
+
+ assert os.path.isfile(config_path), f"Configuration file not found at: {config_path}"
+
+ if config_path.endswith(".json"):
+ with open(config_path) as f:
+ config = json.load(f)
+ else:
+ config_module = importlib.machinery.SourceFileLoader("_config", config_path).load_module()
+ assert hasattr(config_module, PYTHON_CONFIG_GETTER_NAME), (
+ f"Python config file must define a `{PYTHON_CONFIG_GETTER_NAME}` function."
+ )
+ config = getattr(config_module, PYTHON_CONFIG_GETTER_NAME)(py_config_name)
+ assert isinstance(config, Mapping), f"`{PYTHON_CONFIG_GETTER_NAME}` must return a dictionary-like object."
+ cfg = DictConfig(config)
+ return cfg
+
+
+def get_config() -> DictConfig:
+ """
+ Parses command-line arguments to load the main configuration for a training run.
+
+ This function implements a hierarchical configuration loading strategy:
+ 1. It checks if a `config.json` exists in the specified `--workdir`. If so, it loads it.
+ 2. If a `--config` argument is also provided, it uses that file to update the
+ configuration loaded from the work directory.
+ 3. If no config exists in the work directory, it requires the `--config` argument
+ to be provided as the base configuration.
+
+ This allows for resuming training from a work directory while also being able to
+ override specific parameters for a new run.
+
+ Returns:
+ Config: The final, consolidated configuration object.
+ """
+ parser = argparse.ArgumentParser(description="Load training configuration.")
+ parser.add_argument("-c", "--config", type=str, default=None, help="Path to a Python or JSON configuration file.")
+ parser.add_argument("-w", "--workdir", type=str, required=True, help="Work directory to save logs and checkpoints.")
+
+ args = parser.parse_args()
+ workdir_path = args.workdir
+ config_save_path = os.path.join(workdir_path, CONFIG_NAME)
+
+ if os.path.exists(config_save_path):
+ logging.info(f"Resuming from work directory. Loading configuration from {config_save_path}.")
+ cfg = get_config_from_file(config_save_path)
+ if args.config and args.config != config_save_path:
+ logging.info(f"Updating loaded configuration with parameters from {args.config}.")
+ override_cfg = get_config_from_file(args.config)
+ cfg.update(override_cfg)
+ else:
+ assert args.config is not None, "A configuration file must be specified via `-c` or `--config` for a new run."
+ logging.info(f"Starting a new run. Loading configuration from {args.config}.")
+ cfg = get_config_from_file(args.config)
+ cfg.workdir_path = workdir_path
+
+ return cfg
+
+
+def get_config_from_dir(workdir_path: str) -> DictConfig:
+ """
+ A simple utility to load the configuration directly from a work directory.
+
+ Args:
+ workdir_path (str): The path to the work directory containing a `config.json`.
+
+ Returns:
+ DictConfig: The loaded configuration object.
+ """
+ config_save_path = os.path.join(workdir_path, CONFIG_NAME)
+ cfg = get_config_from_file(config_save_path)
+ cfg.workdir_path = workdir_path
+ return cfg
+
+
+# ==============================================================================
+# Base Model Classes
+# ==============================================================================
+
+
+class PreTrainedModel(nn.Module):
+ config_class = DictConfig
+
+ """
+ A base class for models to handle loading from pretrained checkpoints.
+
+ This class provides a common interface for initializing a model and loading
+ weights from a saved checkpoint, following a pattern similar to libraries
+ like Hugging Face's Transformers.
+
+ Args:
+ config (DictConfig | dict[str, Any]): A configuration object containing model hyperparameters.
+ """
+
+ def __init__(self, config: DictConfig | dict[str, Any], *args, **kwargs):
+ super().__init__()
+ self.config = config if isinstance(config, self.config_class) else self.config_class(config)
+
+ @classmethod
+ def from_pretrained(
+ cls,
+ pretrained_dir: str,
+ cfg: DictConfig | dict[str, Any] | None = None,
+ checkpoint_regex: str = "checkpoint_*/ema.safetensors",
+ strict: bool = False,
+ **model_kwargs,
+ ) -> "PreTrainedModel":
+ """
+ Loads a pretrained model from a directory.
+
+ This method first loads the configuration file from the specified directory,
+ initializes the model with this configuration, and then loads the weights
+ from the latest checkpoint file found in that directory.
+
+ Args:
+ cls (type): The model class to instantiate.
+ pretrained_dir (str): The directory containing the pretrained model
+ config and checkpoint files.
+ cfg (DictConfig | dict[str, Any] | None, optional): An optional config object to override
+ the loaded config. Defaults to None.
+ checkpoint_regex (str, optional): A regex pattern to find the checkpoint
+ file. Defaults to "checkpoint_*/ema.safetensors".
+ strict (bool, optional): Whether to strictly enforce that the keys in
+ the checkpoint match the keys of the model.
+ Defaults to False.
+ **model_kwargs: Additional keyword arguments to pass to the model's
+ constructor.
+
+ Returns:
+ PreTrainedModel: An instance of the model with loaded weights.
+ """
+ pretrained_cfg = get_config_from_dir(pretrained_dir).model
+ if cfg is not None:
+ pretrained_cfg.update(cfg)
+ logging.info(f"The loaded config of the pretrained model is updated to: {pretrained_cfg}")
+ model = cls(
+ pretrained_cfg,
+ **model_kwargs,
+ )
+ model_state_dict = {}
+ with safe_open(latest_checkpoint_path(pretrained_dir, checkpoint_regex), framework="pt", device="cpu") as f:
+ for key in f.keys():
+ model_state_dict[key] = f.get_tensor(key)
+ model.load_state_dict(model_state_dict, strict=strict)
+ return model
+
+
+# ==============================================================================
+# IO and Checkpointing Utilities
+# ==============================================================================
+
+
+def latest_checkpoint_path(dir_path: str, regex: str | None = None) -> str:
+ """
+ Finds the path of the latest checkpoint file or directory in a directory.
+
+ The latest checkpoint is determined by sorting the filenames alphanumerically
+ and picking the last one. This assumes a naming convention like `checkpoint_1000.pt`,
+ `checkpoint_2000.pt`, etc.
+
+ Args:
+ dir_path (str): The directory to search for checkpoints.
+ regex (str | None, optional): A glob pattern to match checkpoint files. If None,
+ a default pattern is used. Defaults to None.
+
+ Returns:
+ str: The full path to the latest checkpoint file.
+
+ Raises:
+ AssertionError: If no files matching the regex are found in the directory.
+ """
+ if regex is None:
+ regex = CHECKPOINT_FORMAT.format("*")
+
+ f_list = glob.glob(os.path.join(dir_path, regex))
+ if not f_list:
+ raise FileNotFoundError(f"No checkpoint files or directories found in {dir_path} matching '{regex}'")
+
+ # Sort files based on the integer values in their names
+ f_list.sort(key=lambda f: int("".join(filter(str.isdigit, f))))
+
+ latest_path = f_list[-1]
+ logging.info(f"Latest checkpoint '{os.path.relpath(latest_path, start=dir_path)}' found in '{dir_path}'.")
+ return latest_path
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_model.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_model.py
new file mode 100644
index 000000000..e83e78cb6
--- /dev/null
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_model.py
@@ -0,0 +1,1620 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/speechlm2/modules/ear_tts_model.py; training-only code removed.
+# NOTE: DuplexEARTTS (duplex_ear_tts.py) imports RVQEARTTSModel from THIS file, not from
+# rvq_ear_tts_model.py — both variants are vendored to match upstream exactly.
+
+import json
+import math
+import os
+from dataclasses import dataclass, fields
+from typing import Any
+
+import torch
+import transformers
+from omegaconf import DictConfig, OmegaConf
+from torch import Tensor, nn
+from torch.nn import functional as F
+from transformers import AutoConfig, AutoModel, AutoModelForTextEncoding, Cache
+from transformers.generation.logits_process import TopKLogitsWarper, TopPLogitsWarper
+
+from .compat import fp32_precision, logging, set_model_dict_for_partial_init
+from .tokenizer import AutoTokenizer
+
+# ==============================================================================
+# MLP module and Norm
+# ==============================================================================
+
+
+class RMSNorm(nn.Module):
+ def __init__(self, dim: int, eps: float = 1e-6):
+ super().__init__()
+ self.eps = eps
+ self.weight = nn.Parameter(torch.zeros(dim))
+
+ def _norm(self, x):
+ return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+ def forward(self, x):
+ output = self._norm(x.float())
+ # Llama does x.to(float16) * w whilst Gemma3 is (x * w).to(float16)
+ output = output * (1.0 + self.weight.float())
+ return output.type_as(x)
+
+ def extra_repr(self):
+ return f"{tuple(self.weight.shape)}, eps={self.eps}"
+
+
+class MLP(nn.Module):
+ def __init__(
+ self,
+ hidden_size: int,
+ intermediate_size: int,
+ ):
+ super().__init__()
+ self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+ self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+ self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+ self.act_fn = nn.GELU(approximate="tanh")
+
+ def forward(self, x: Tensor) -> Tensor:
+ down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+ return down_proj
+
+
+class MLPLayer(nn.Module):
+ def __init__(
+ self,
+ hidden_size: int,
+ intermediate_size: int,
+ eps: float = 1e-6,
+ ):
+ super().__init__()
+ self.pre_norm = RMSNorm(hidden_size, eps=eps)
+ self.mlp = MLP(hidden_size, intermediate_size)
+ self.post_norm = RMSNorm(hidden_size, eps=eps)
+
+ def forward(self, x: Tensor) -> Tensor:
+ y = self.pre_norm(x)
+ y = self.mlp(y)
+ y = self.post_norm(y)
+ x = x + y
+ return x
+
+
+# ==============================================================================
+# Triton-accelerated and Fallback Functions
+# ==============================================================================
+
+TRITON_IMPORTED = False
+try:
+ import triton
+ import triton.language as tl
+
+ TRITON_IMPORTED = True
+except ImportError:
+ TRITON_IMPORTED = False
+
+USE_TRITON = TRITON_IMPORTED and torch.cuda.is_available()
+
+if USE_TRITON:
+ logging.info("Triton available & CUDA detected. Using Triton kernel for batch_matmul.")
+
+ @triton.jit
+ def batch_matmul_kernel(
+ x_ptr,
+ w_ptr,
+ y_ptr,
+ result_ptr,
+ b,
+ d_in,
+ d_out,
+ n,
+ BLOCK_SIZE_DIN: tl.constexpr,
+ BLOCK_SIZE_DOUT: tl.constexpr,
+ ):
+ batch_id = tl.program_id(axis=0)
+ dout_block_id = tl.program_id(axis=1)
+
+ if batch_id >= b:
+ return
+
+ idx = tl.load(y_ptr + batch_id)
+
+ x_offset = x_ptr + batch_id * d_in
+ w_offset = w_ptr + idx * d_out * d_in
+
+ dout_offsets = dout_block_id * BLOCK_SIZE_DOUT + tl.arange(0, BLOCK_SIZE_DOUT)
+ dout_mask = dout_offsets < d_out
+
+ result_block = tl.zeros([BLOCK_SIZE_DOUT], dtype=tl.float32)
+
+ for din_start in range(0, d_in, BLOCK_SIZE_DIN):
+ din_offsets = din_start + tl.arange(0, BLOCK_SIZE_DIN)
+ din_mask = din_offsets < d_in
+
+ x_i = tl.load(x_offset + din_offsets, mask=din_mask, other=0.0)
+
+ w_i_block = tl.load(
+ w_offset + dout_offsets[:, None] * d_in + din_offsets[None, :],
+ mask=(dout_mask[:, None] & din_mask[None, :]),
+ other=0.0,
+ )
+
+ result_block += tl.sum(w_i_block * x_i[None, :], axis=1)
+
+ result_offset = result_ptr + batch_id * d_out + dout_offsets
+ tl.store(result_offset, result_block, mask=dout_mask)
+
+ def batch_matmul_triton(x, w, y, BLOCK_SIZE_DIN: int = 16, BLOCK_SIZE_DOUT: int = 64):
+ assert x.is_contiguous() and w.is_contiguous() and y.is_contiguous()
+
+ b, d_in = x.shape
+ n, d_out, _ = w.shape
+ result = torch.empty(b, d_out, device=x.device, dtype=torch.float32)
+
+ batch_matmul_kernel[lambda meta: (b, triton.cdiv(d_out, meta["BLOCK_SIZE_DOUT"]))](
+ x.float(),
+ w.float(),
+ y,
+ result,
+ b,
+ d_in,
+ d_out,
+ n,
+ BLOCK_SIZE_DIN=BLOCK_SIZE_DIN,
+ BLOCK_SIZE_DOUT=BLOCK_SIZE_DOUT,
+ )
+
+ return result.to(dtype=x.dtype)
+
+ batch_matmul = batch_matmul_triton
+
+else:
+ logging.info("Using PyTorch fallback (Triton unavailable or no CUDA).")
+
+ # Fallback to PyTorch implementation if Triton is not available
+ def batch_matmul_pytorch(x: Tensor, w: Tensor, y: Tensor, *args, **kwargs) -> Tensor:
+ """
+ Performs a batched matrix multiplication using PyTorch's native functions.
+
+ This function serves as a fallback when Triton is not available. It achieves
+ the same result by gathering the appropriate weight matrices and using `torch.bmm`.
+
+ Args:
+ x (Tensor): The input tensor of shape `[batch_size, d_in]`.
+ w (Tensor): The weight tensor of shape `[num_weights, d_out, d_in]`.
+ y (Tensor): The index tensor of shape `[batch_size]`.
+
+ Returns:
+ Tensor: The result of the multiplication, shape `[batch_size, d_out]`.
+ """
+ # w[y] gathers the weight matrices for each item in the batch.
+ # x.unsqueeze(2) reshapes x to [batch_size, d_in, 1] for bmm.
+ # The result is squeezed to remove the trailing dimension of size 1.
+ return torch.bmm(w[y], x.unsqueeze(2)).squeeze(2)
+
+ batch_matmul = batch_matmul_pytorch
+
+
+# ==============================================================================
+# Core Mathematical and Masking Functions
+# ==============================================================================
+
+
+def gumbel_like(tensor: Tensor, eps: float = 1e-8) -> Tensor:
+ """
+ Generates a tensor of Gumbel noise with the same shape as the input tensor.
+
+ This is used for the Gumbel-Max trick, a technique to sample from a categorical
+ distribution in a differentiable way (using a straight-through estimator).
+
+ Args:
+ tensor (torch.Tensor): The input tensor to match the shape of.
+ eps (float): A small epsilon value for numerical stability.
+
+ Returns:
+ torch.Tensor: A tensor containing Gumbel noise.
+ """
+ # Sample from a uniform distribution
+ u = torch.rand_like(tensor)
+ # Apply the inverse CDF of the Gumbel distribution
+ return -torch.log(-torch.log(u + eps) + eps)
+
+
+def sequence_mask(lengths: Tensor, max_length: Tensor | int | None = None) -> Tensor:
+ """
+ Creates a boolean mask from a 1D tensor of sequence lengths.
+
+ This function is useful for masking out padding in sequences. Given a tensor
+ of lengths, it produces a 2D boolean tensor where `mask[i, j]` is `True` if
+ `j < lengths[i]` and `False` otherwise.
+
+ Args:
+ lengths (Long Tensor): A 1D tensor of integer lengths. Shape: `[batch_size]`.
+ max_length (Long Tensor | int | None, optional): The maximum length of the mask. If None,
+ it is inferred from the maximum value
+ in `lengths`. Defaults to None.
+
+ Returns:
+ Tensor: The boolean mask. Shape: `[batch_size, max_length]`.
+ """
+ if max_length is None:
+ max_length = lengths.max()
+
+ # Create a range tensor from 0 to max_length - 1
+ x = torch.arange(max_length, dtype=lengths.dtype, device=lengths.device) # type: ignore[arg-type]
+
+ # Compare each length with the range tensor to create the mask via broadcasting
+ return x.unsqueeze(0) < lengths.unsqueeze(1)
+
+
+def get_masking_rate(rate: Tensor, exponent: float = 3.0) -> Tensor:
+ """
+ Converts a desired token keep rate to a masking rate using a power function.
+
+ This function is part of a scheduling strategy for masking, where the effective
+ masking rate changes non-linearly with the desired keep rate. This function is
+ its own inverse.
+
+ Args:
+ rate (Tensor): The desired rate of tokens to keep (0 to 1).
+ exponent (float, optional): The exponent for the transformation. Defaults to 3.0.
+
+ Returns:
+ Tensor: The corresponding masking rate.
+ """
+ return (1 - rate.pow(exponent)).pow(1 / exponent)
+
+
+# Alias the function for clarity in the inverse context
+get_rate = get_masking_rate
+
+
+def get_mask(
+ code_mask: Tensor,
+ num_masking: Tensor,
+ unmasking: bool = False,
+ validate: bool = False,
+) -> Tensor:
+ """
+ Adjusts a boolean mask by masking or unmasking tokens from the end.
+
+ This function operates on a `code_mask` where `True` values represent valid
+ tokens and are assumed to be contiguous at the start of the sequence. It
+ calculates a new mask by decreasing (masking) or increasing (unmasking)
+ the number of `True` values.
+
+ Args:
+ code_mask (Tensor): The input boolean mask. Shape: `[..., depth]`.
+ num_masking (Tensor): The number of tokens to mask or unmask.
+ Shape matching `code_mask`'s batch dimensions.
+ unmasking (bool, optional): If `True`, increases the number of valid
+ tokens (unmasking). Defaults to `False`.
+ validate (bool, optional): If `True`, asserts that the input `code_mask`
+ is contiguous. This adds a slight overhead and
+ is mainly for debugging. Defaults to `False`.
+
+ Returns:
+ Tensor: A new boolean mask with the adjusted length of valid tokens.
+ Shape is identical to `code_mask`.
+ """
+ depth = code_mask.size(-1)
+ num_valid = code_mask.sum(dim=-1, dtype=torch.long)
+
+ if validate:
+ # Reconstruct the expected contiguous mask and assert equality.
+ expected_mask = sequence_mask(num_valid.view(-1), depth).view_as(code_mask)
+ assert torch.equal(code_mask, expected_mask), (
+ "Input `code_mask` must have contiguous `True` values at the beginning."
+ )
+
+ # Calculate the target number of valid tokens.
+ if not unmasking:
+ # Masking: reduce the number of valid tokens, ensuring it's not negative.
+ num_to_keep = (num_valid - num_masking).clamp_min(0)
+ else:
+ # Unmasking: increase the number of valid tokens, capped by total depth.
+ num_to_keep = (num_valid + num_masking).clamp_max(depth)
+
+ # Generate the new mask using the final number of tokens to keep.
+ return sequence_mask(num_to_keep.view(-1), depth).view_as(code_mask)
+
+
+# ==============================================================================
+# Model and Vocabulary Utilities
+# ==============================================================================
+
+
+@dataclass
+class RVQEARTTSOutput:
+ """
+ Output type for the RVQEARTTSModel, providing a structured way to return model outputs.
+ This class allows accessing outputs by attribute, key, or index.
+ """
+
+ loss: Tensor | None = None
+ lm_loss: Tensor | None = None
+ c_loss: Tensor | None = None
+ k_loss: Tensor | None = None
+
+ hidden_states: Tensor | None = None
+ past_key_values: Tensor | None = None
+ audio_prompt_latent: Tensor | None = None
+
+ codes: Tensor | None = None
+ lm_logits: Tensor | None = None
+ eos_flag: Tensor | None = None
+
+ def __getitem__(self, item: str | int):
+ """Allows for accessing attributes by key or index."""
+ if isinstance(item, str):
+ return getattr(self, item)
+ else:
+ # Access fields in the order they are defined in the dataclass
+ return getattr(self, fields(self)[item].name)
+
+
+def find_and_delete_module(parent_module: nn.Module, target_module: nn.Module, parent_name: str) -> str | None:
+ """
+ Recursively searches for a specific module instance and deletes it from its parent.
+
+ This is useful for dynamically modifying a model's architecture, such as replacing
+ an existing embedding layer with a custom one.
+
+ Args:
+ parent_module (nn.Module): The module to search within.
+ target_module (nn.Module): The exact module instance to find and delete.
+ parent_name (str): The initial name of the parent module for constructing the path.
+
+ Returns:
+ str | None: The full dotted name of the deleted attribute if found, otherwise None.
+ """
+ # Iterate over all direct children of the parent module
+ for name, module in parent_module.named_children():
+ # Use the 'is' operator to check for object identity, not just value equality
+ if module is target_module:
+ # If found, delete the attribute from the parent and return its name
+ delattr(parent_module, name)
+ return f"{parent_name}.{name}"
+
+ # If not found, recurse into the child module
+ found_path = find_and_delete_module(module, target_module, parent_name=f"{parent_name}.{name}")
+ if found_path:
+ return found_path
+ return None
+
+
+def build_vocabs(
+ tokenizer: AutoTokenizer, vocab_dir: str | None = None
+) -> tuple[dict[int, tuple[int, ...]], dict[str, int], int]:
+ """
+ Builds or loads a character-level vocabulary derived from a subword tokenizer.
+
+ This function creates a mapping from each subword in a pretrained tokenizer to a
+ sequence of character IDs. It follows a modern practice of using a directory
+ to save and load vocabulary files, making the process more robust and extensible.
+
+ The primary source of truth is the `char_vocab.json` file. If it exists, it's
+ loaded. Otherwise, it's created from the pretrained tokenizer and saved.
+
+ Args:
+ tokenizer (AutoTokenizer): The pretrained Hugging Face tokenizer class.
+ vocab_dir (str | None, optional): The directory to save or load the character
+ vocabulary from. Defaults to None.
+
+ Returns:
+ tuple[dict[int, tuple[int, ...]], dict[str, int], int]: A tuple containing:
+ - A mapping from subword IDs to tuples of character IDs.
+ - The character-to-ID vocabulary dictionary.
+ - The ID for the subword padding token.
+ """
+
+ def _build_char_vocab() -> dict[str, int]:
+ # Find all single-character tokens in the original tokenizer's vocabulary
+ single_chars = {
+ subword: subword_id for subword, subword_id in tokenizer.tokenizer.vocab.items() if len(subword) == 1
+ }
+ # Create a new, dense character vocabulary sorted by the original token ID
+ sorted_chars = sorted(single_chars.keys(), key=lambda k: single_chars[k])
+ char_vocab = {char: i for i, char in enumerate(sorted_chars)}
+ return char_vocab
+
+ # 1. Load or build the character vocabulary
+ if vocab_dir:
+ from filelock import FileLock
+
+ char_vocab_file = os.path.join(vocab_dir, "char_vocab.json")
+ os.makedirs(vocab_dir, exist_ok=True)
+ with FileLock(char_vocab_file + ".lock", timeout=60):
+ if not os.path.exists(char_vocab_file):
+ char_vocab = _build_char_vocab()
+
+ logging.info(f"Saving character vocabulary to {char_vocab_file}")
+ with open(char_vocab_file, "w", encoding="utf-8") as f:
+ json.dump(char_vocab, f, ensure_ascii=False, indent=2)
+
+ # All processes can now safely load the file.
+ logging.info(f"Loading character vocabulary from {char_vocab_file}")
+ with open(char_vocab_file, encoding="utf-8") as f:
+ char_vocab = json.load(f)
+ else:
+ # No cache directory provided, build in memory.
+ logging.info("Building character vocabulary from tokenizer.")
+ char_vocab = _build_char_vocab()
+
+ # 2. Reconstruct the subword-to-character mapping on the fly
+ subword_id_to_char_ids = {
+ subword_id: tuple(char_vocab[char] for char in subword if char in char_vocab)
+ for subword, subword_id in tokenizer.tokenizer.vocab.items()
+ }
+ # Filter out subwords that contain characters not in our character vocabulary
+ subword_id_to_char_ids = {k: v for k, v in subword_id_to_char_ids.items() if v}
+
+ # 3. Define a padding index for subwords
+ subword_padding_idx = len(tokenizer.vocab)
+ # The padding subword maps to a new character padding ID
+ subword_id_to_char_ids[subword_padding_idx] = (len(char_vocab),)
+ return subword_id_to_char_ids, char_vocab, subword_padding_idx
+
+
+# @torch.compile
+@torch._dynamo.disable
+def depthsum_encoding_step(
+ embs: Tensor,
+ r: Tensor,
+ code: Tensor,
+ depth_str: int = 0,
+ k: int = 72,
+) -> Tensor:
+ for i in range(depth_str, depth_str + k):
+ idx_sel = (
+ embs[i].pow(2).sum(-1) # [g?, v]
+ - 2
+ * (r.unsqueeze(-2) @ embs[i].transpose(-1, -2)).squeeze(-2) # [b, ?, g?, h] , [g?, h, v] -> [b, ?, g?, v]
+ ).argmin(-1)
+
+ emb_i = F.embedding(idx_sel, embs[i])
+ r = r - emb_i
+
+ code[..., i] = idx_sel
+
+ return code
+
+
+class MoGHead(nn.Module):
+ """
+ A Mixture of Gaussians (MoG) prediction head.
+
+ This module takes a hidden state and predicts the parameters for a mixture of
+ Gaussian distributions. It's suitable for modeling continuous, multi-modal data.
+
+ Args:
+ hidden_size (int): The dimensionality of the input hidden state.
+ intermediate_size (int): The dimensionality of the MLP layers.
+ out_size (int): The dimensionality of the output vectors (the mean of each Gaussian).
+ num_layers (int): The number of MLP layers in the stack.
+ num_predictions (int): The number of Gaussian components in the mixture.
+ low_rank (int | None): The dimensionality used for compressing the hidden states.
+ min_log_std (float): The minimum value for the logarithm of the standard deviation.
+ eps (float): A small epsilon value for the RMSNorm layers.
+ """
+
+ def __init__(
+ self,
+ hidden_size: int,
+ intermediate_size: int,
+ out_size: int,
+ num_layers: int,
+ num_predictions: int,
+ low_rank: int | None = 64,
+ min_log_std: float = -4.0,
+ eps: float = 1e-6,
+ ):
+ super().__init__()
+ self.out_size = out_size
+ self.low_rank = low_rank
+ self.num_predictions = num_predictions
+ self.min_log_std = min_log_std
+
+ self.mlp_stack = nn.Sequential(
+ *[MLPLayer(hidden_size, intermediate_size, eps=eps) for _ in range(num_layers)],
+ RMSNorm(hidden_size, eps=eps),
+ )
+
+ if low_rank is None:
+ self.proj_logits = nn.Linear(hidden_size, num_predictions, bias=False) # Predicts mixture weights
+ self.proj_mus = nn.Linear(hidden_size, num_predictions * out_size, bias=False) # Predicts means
+ self.proj_logs = nn.Linear(hidden_size, 1, bias=False) # Predicts log standard deviations
+ else:
+ assert low_rank < out_size
+ self.proj_logits = nn.Linear(hidden_size, num_predictions, bias=False) # Predicts mixture weights
+ self.proj_mus = nn.Linear(hidden_size, num_predictions * low_rank, bias=False) # Predicts means
+ self.proj_logs = nn.Linear(hidden_size, 1, bias=False) # Predicts log standard deviations
+ self.proj_else = nn.Linear(hidden_size, out_size, bias=False)
+ self.low_mat = nn.Parameter(torch.randn(num_predictions, out_size, low_rank) * (low_rank**-0.5))
+
+ def infer(self, x: Tensor, guidance_scale: float = 0.0, top_p_or_k: float | int = 1.0) -> tuple[Tensor, Tensor]:
+ """
+ Performs inference by sampling from the predicted mixture distribution.
+
+ Args:
+ x (Tensor): The input hidden state.
+ guidance_scale (float): The weight for classifier-free guidance.
+ top_p_or_k (float | int): The value for top-p (nucleus) or top-k sampling of the mixture components.
+
+ Returns:
+ tuple[Tensor, Tensor]: A tuple containing the mean of the chosen component,
+ and the log standard deviations.
+ """
+ b, t, _ = x.size()
+ n, d = self.num_predictions, self.low_rank or self.out_size
+ x = self.mlp_stack(x)
+ if guidance_scale > 0:
+ b //= 2
+ x_cond, x_uncond = x.chunk(2, dim=0)
+ x = x_cond + guidance_scale * (x_cond - x_uncond)
+
+ logits = self.proj_logits(x)
+
+ # Apply top-p or top-k filtering to the mixture logits
+ if top_p_or_k is not None:
+ logits = (
+ TopPLogitsWarper(top_p_or_k)(
+ None,
+ logits.view(-1, n),
+ ).view_as(logits)
+ if isinstance(top_p_or_k, float)
+ else TopKLogitsWarper(top_p_or_k)(
+ None,
+ logits.view(-1, n),
+ ).view_as(logits)
+ )
+
+ # Sample a mixture component using the Gumbel-Max trick
+ with fp32_precision():
+ mixture_indices = (F.log_softmax(logits, dim=-1) + gumbel_like(logits)).argmax(-1)
+
+ # Select the mean corresponding to the sampled component
+ mu = batch_matmul(
+ x.view(b * t, -1),
+ self.proj_mus.weight.detach().view(n, d, -1),
+ mixture_indices.view(b * t),
+ ).view(b, t, d)
+ if self.proj_mus.bias is not None:
+ mu += self.proj_mus.bias.detach().view(n, d)[mixture_indices]
+
+ if self.low_rank:
+ assert math.log2(d).is_integer() and math.log2(self.out_size).is_integer()
+ mu = batch_matmul(
+ mu.view(b * t, -1),
+ self.low_mat.detach().view(n, self.out_size, -1),
+ mixture_indices.view(b * t),
+ BLOCK_SIZE_DIN=d,
+ BLOCK_SIZE_DOUT=self.out_size,
+ ).view(b, t, self.out_size)
+
+ mu_res = self.proj_else(x)
+ else:
+ mu_res = torch.zeros((b, t, d), device=x.device)
+
+ logs = self.proj_logs(x).clamp_min(self.min_log_std)
+ return mu * torch.exp(logs) + mu_res, logs
+
+ def forward(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor]:
+ """
+ Performs a forward pass for training.
+
+ Args:
+ x (Tensor): The input hidden state.
+
+ Returns:
+ tuple[Tensor, Tensor, Tensor]: A tuple containing the mixture logits,
+ the means for all components, and the
+ log standard deviations.
+ """
+ b, t, _ = x.size()
+ d = self.low_rank or self.out_size
+ x = self.mlp_stack(x)
+ logits = self.proj_logits(x)
+ mus = self.proj_mus(x).view(b, t, self.num_predictions, d)
+ logs = self.proj_logs(x).clamp_min(self.min_log_std)
+
+ if self.low_rank:
+ mu_res = self.proj_else(x)
+ else:
+ mu_res = torch.zeros((b, t, d), device=x.device)
+ return logits, mus, mu_res, logs
+
+ def dist(self, mus: Tensor, mu: Tensor) -> Tensor:
+ """
+ mus: [b, t, n, d]
+ mu: [b, t, d]
+
+ return: [b, t, n]
+ """
+ if self.low_rank is None:
+ return (mus - mu.unsqueeze(-2)).pow(2).sum(-1)
+ else:
+ low_mat_sq = self.low_mat.transpose(-1, -2) @ self.low_mat
+ x, y = mus, mu
+ wx_sq = (
+ x
+ * torch.einsum(
+ "btni,nij->btnj",
+ x,
+ low_mat_sq.to(x),
+ )
+ ).sum(-1) # [b, t, n]
+ y_sq = y.pow(2).sum(-1, keepdim=True) # [b, t, 1]
+ xwy = (x * torch.einsum("bti,nij->btnj", y, self.low_mat.to(y))).sum(
+ -1
+ ) # [b, t, n, d_l], [n, d_i, d_l], [b, t, d_i] -> [b, t, n]
+
+ dist = wx_sq + y_sq - 2 * xwy
+ return torch.abs(dist)
+
+
+class SubwordFlagEmbedding(nn.Module):
+ """
+ Adds a small continuation embedding for subwords (tokens without word-boundary marker).
+ Automatically adds a custom padding token at index vocab_size.
+ Ignores special tokens (starting with '<') when computing continuation flags.
+ """
+
+ def __init__(self, tokenizer: AutoTokenizer, d_model: int):
+ super().__init__()
+
+ self.tokenizer = tokenizer
+ self.vocab_size = self.tokenizer.vocab_size
+ self.d_model = d_model
+
+ # Custom pad token at vocab_size
+ self.pad_id = self.vocab_size
+ # register pad_id as a tensor buffer to avoid device issues
+ self.register_buffer("pad_tensor", torch.tensor(self.pad_id, dtype=torch.long))
+
+ # Precompute continuation flags
+ tokens = [self.tokenizer.ids_to_tokens(i) for i in range(self.vocab_size)]
+ cont_flags = [
+ 1 if not (tok.startswith("Ġ") or tok.startswith("▁") or tok.startswith("<")) else 0 for tok in tokens
+ ]
+ cont_flags.append(0) # for the custom pad token
+ self.register_buffer("is_continuation", torch.tensor(cont_flags, dtype=torch.long))
+
+ # Continuation embedding
+ init_std = self.d_model**-0.5
+ self.cont_emb = nn.Embedding(2, self.d_model)
+ nn.init.normal_(self.cont_emb.weight, mean=0.0, std=init_std)
+ self.cont_emb.weight.data[0].zero_()
+
+ def forward(self, subword_embeds: torch.Tensor, token_ids: torch.LongTensor):
+ # Replace OOV token IDs with pad_id safely
+ token_ids_clamped = torch.where(token_ids >= self.vocab_size, self.pad_tensor, token_ids)
+ # Continuation flags
+ cont_flags = self.is_continuation[token_ids_clamped]
+ # Add continuation embedding
+ cont_emb = self.cont_emb(cont_flags)
+ return subword_embeds + cont_emb
+
+
+class BOSEOSEmbedding(nn.Module):
+ """
+ Adds independent embeddings for BOS and EOS tokens using a single embedding table.
+ Index 0 = regular token (ignored), 1 = BOS, 2 = EOS.
+ Compatible with Hugging Face tokenizers that may or may not have BOS/EOS.
+ """
+
+ def __init__(self, tokenizer: AutoTokenizer, d_model: int):
+ super().__init__()
+
+ self.tokenizer = tokenizer
+ # vocab size that includes special tokens
+ vocab_dict = self.tokenizer.tokenizer.get_vocab()
+ self.vocab_size = max(vocab_dict.values())
+ self.d_model = d_model
+
+ # Custom pad token for OOVs
+ self.pad_id = self.vocab_size
+ self.register_buffer("pad_tensor", torch.tensor(self.pad_id, dtype=torch.long))
+
+ # Identify BOS and EOS tokens (may be None)
+ tokens = [self.tokenizer.ids_to_tokens(i) for i in range(self.vocab_size)]
+
+ special_flags = []
+ for tok in tokens:
+ if self.tokenizer.bos_token is not None and tok == self.tokenizer.bos_token:
+ special_flags.append(1)
+ elif self.tokenizer.eos_token is not None and tok == self.tokenizer.eos_token:
+ special_flags.append(2)
+ else:
+ special_flags.append(0)
+ special_flags.append(0) # for custom pad token
+ self.register_buffer("special_flags", torch.tensor(special_flags, dtype=torch.long))
+ # Embedding table: 0 = regular, 1 = BOS, 2 = EOS
+ init_std = self.d_model**-0.5
+ self.special_emb = nn.Embedding(3, d_model)
+ nn.init.normal_(self.special_emb.weight, mean=0.0, std=init_std)
+ self.special_emb.weight.data[0].zero_() # regular tokens ignored
+
+ def forward(self, token_embeds: torch.Tensor, token_ids: torch.LongTensor):
+ """
+ token_embeds: (B, T, d_model)
+ token_ids: (B, T)
+ """
+ # Clamp OOVs to custom pad token
+ safe_ids = torch.where(token_ids >= self.vocab_size, self.pad_tensor, token_ids)
+
+ # Lookup flags (0=regular, 1=BOS, 2=EOS)
+ flags = self.special_flags[safe_ids]
+ return token_embeds + self.special_emb(flags)
+
+
+class CharAwareSubwordEncoder(nn.Module):
+ """
+ An encoder that creates subword embeddings from character-level embeddings.
+
+ This module replaces a standard subword embedding layer. It breaks down each
+ subword into its constituent characters, embeds the characters, and then
+ aggregates these character embeddings (e.g., via mean pooling) to form the
+ final subword representation. This allows the model to handle rare or out-of-vocabulary
+ subwords more gracefully.
+
+ Args:
+ out_size (int): The dimensionality of the output embedding vectors.
+ tokenizer (AutoTokenizer): The Hugging Face tokenizer class.
+ vocab_dir (str | None): Directory to save/load the character vocabulary.
+ backbone_type (str | None): The type of backbone model from Hugging Face (e.g., "t5gemma").
+ backbone_model_class (str | None): The class name of the backbone model if not using AutoModel.
+ backbone_config_class (str | None): The class name of the backbone config.
+ backbone_config (DictConfig | None): A configuration for the backbone model.
+ """
+
+ def __init__(
+ self,
+ out_size: int,
+ tokenizer: AutoTokenizer,
+ vocab_dir: str | None = None,
+ backbone_type: str | None = "t5gemma",
+ backbone_model_class: str | None = None,
+ backbone_config_class: str | None = None,
+ backbone_config: DictConfig | None = None,
+ use_subword_flag_emb: bool = True,
+ use_bos_eos_emb: bool = True,
+ use_cumulative_word_emb: bool = False,
+ ):
+ super().__init__()
+
+ # 1. Build or load the character vocabulary
+ self.subword_id_to_char_ids, self.char_vocab, self.subword_padding_idx = build_vocabs(
+ tokenizer,
+ vocab_dir,
+ )
+
+ self.char_padding_idx = len(self.char_vocab)
+ self.use_subword_flag_emb = use_subword_flag_emb
+ self.use_bos_eos_emb = use_bos_eos_emb
+
+ # 2. Initialize the backbone model
+ if backbone_type:
+ config = AutoConfig.for_model(
+ backbone_type, **(OmegaConf.to_container(backbone_config, resolve=True) if backbone_config else {})
+ )
+ self.backbone = AutoModelForTextEncoding.from_config(config)
+ else:
+ assert backbone_model_class and backbone_config_class
+ config_class = getattr(transformers, backbone_config_class)
+ model_class = getattr(transformers, backbone_model_class)
+ config = config_class(**(OmegaConf.to_container(backbone_config, resolve=True) if backbone_config else {}))
+ self.backbone = model_class(config)
+
+ self.hidden_size = self.backbone.get_input_embeddings().weight.size(-1)
+
+ # 3. Delete the original subword embedding layer and replace it with our character embedding layer
+ find_and_delete_module(self.backbone, self.backbone.get_input_embeddings(), "backbone")
+ self.embed_tokens = nn.Embedding(len(self.char_vocab) + 1, self.hidden_size, padding_idx=self.char_padding_idx)
+ self.proj_embedding = nn.Linear(self.hidden_size, out_size, bias=False)
+
+ if self.use_subword_flag_emb:
+ self.subword_flag_emb = SubwordFlagEmbedding(tokenizer, self.hidden_size)
+
+ if self.use_bos_eos_emb:
+ self.bos_eos_emb = BOSEOSEmbedding(tokenizer, self.hidden_size)
+
+ def prepare_inputs(self, subword_ids: Tensor, padding_mask: Tensor) -> tuple[Tensor, Tensor]:
+ """
+ Converts a batch of subword IDs into a padded batch of character IDs.
+
+ Args:
+ subword_ids (Tensor): A tensor of subword IDs. Shape: `[batch, seq_len]`.
+ padding_mask (Tensor): A boolean mask indicating valid (non-padding) subwords.
+
+ Returns:
+ tuple[Tensor, Tensor]: A tuple containing:
+ - Padded character IDs. Shape: `[num_valid_subwords, max_char_len]`.
+ - Lengths of each character sequence. Shape: `[num_valid_subwords]`.
+ """
+ device = subword_ids.device
+ # Select only the valid subword IDs
+ subword_id_list = torch.masked_select(subword_ids, padding_mask).cpu().tolist()
+ # Map each subword ID to its sequence of character IDs
+ char_id_list = [list(self.subword_id_to_char_ids.get(x, ())) for x in subword_id_list]
+
+ char_lengths = torch.tensor([len(x) for x in char_id_list], dtype=torch.long, device=device)
+ batch_size = char_lengths.size(0)
+ max_len = int(char_lengths.max().item()) if batch_size > 0 else 0
+
+ # Create a padded tensor for the character IDs
+ char_ids = torch.full((batch_size, max_len), self.char_padding_idx, dtype=torch.long, device=device)
+ for i, char_seq in enumerate(char_id_list):
+ char_ids[i, : len(char_seq)] = torch.tensor(char_seq, dtype=torch.long, device=device)
+
+ return char_ids, char_lengths
+
+ def forward(self, subword_ids: Tensor, subword_mask: Tensor | None = None) -> Tensor:
+ """
+ Performs the forward pass to get character-aware subword embeddings.
+
+ Args:
+ subword_ids (Tensor): A tensor of subword IDs. Shape: `[batch, seq_len]`.
+ subword_mask (Tensor | None): A boolean mask for padding. Defaults to None.
+
+ Returns:
+ Tensor: The final subword embeddings. Shape: `[batch, seq_len, hidden_size]`.
+ """
+ if subword_mask is None:
+ subword_mask = torch.ones_like(subword_ids, dtype=torch.bool)
+
+ # 1. Convert subword IDs to character IDs
+ char_ids, char_lengths = self.prepare_inputs(subword_ids, subword_mask)
+
+ # char_mask = sequence_mask(char_lengths).float()
+ char_mask = sequence_mask(char_lengths)
+
+ # 2. Get character embeddings and pass them through the backbone
+ char_embeds = self.embed_tokens(char_ids)
+
+ # The backbone model should be able to accept `inputs_embeds`
+ char_hidden_states = self.backbone(inputs_embeds=char_embeds, attention_mask=char_mask).last_hidden_state
+
+ # 3. Aggregate character embeddings to form subword embeddings (mean pooling)
+ # We mask the padding characters before summing to get a correct mean.
+ masked_sum = (char_hidden_states * char_mask.unsqueeze(-1)).sum(dim=1)
+ # Avoid division by zero for empty sequences
+ mean_emb = masked_sum / (char_lengths.unsqueeze(-1).clamp(min=1))
+
+ # 4. Scatter the aggregated embeddings back to the original subword sequence shape
+ out_emb = self.proj_embedding(mean_emb)
+ subword_embeds = torch.zeros(
+ subword_ids.shape + (out_emb.size(-1),), device=subword_ids.device, dtype=out_emb.dtype
+ )
+ subword_embeds[subword_mask] = out_emb
+
+ if self.use_subword_flag_emb:
+ subword_embeds = self.subword_flag_emb(subword_embeds, subword_ids)
+
+ if self.use_bos_eos_emb:
+ subword_embeds = self.bos_eos_emb(subword_embeds, subword_ids)
+
+ return subword_embeds
+
+
+class GatedProjectedSumRMSNorm(nn.Module):
+ def __init__(self, audio_dim, text_dim, hidden_dim, final_norm=True, num_codebooks=31, init_residual_scale=0.5):
+ super().__init__()
+ self.num_codebooks = num_codebooks
+
+ self.audio_proj = nn.Linear(audio_dim, hidden_dim)
+ self.text_proj = nn.Linear(text_dim, hidden_dim)
+
+ nn.init.normal_(self.audio_proj.weight, mean=0.0, std=0.015)
+ nn.init.zeros_(self.audio_proj.bias)
+ nn.init.normal_(self.text_proj.weight, mean=0.0, std=0.015)
+ nn.init.zeros_(self.text_proj.bias)
+
+ # FP32 gate params
+ self.gate = nn.Parameter(torch.zeros(hidden_dim, dtype=torch.float32))
+ self.residual_scale = nn.Parameter(torch.tensor(init_residual_scale, dtype=torch.float32))
+
+ self.final_norm = RMSNorm(hidden_dim) if final_norm else nn.Identity()
+
+ def forward(self, audio_emb, text_emb):
+ audio_emb = audio_emb / self.num_codebooks
+
+ # projections run in model dtype (BF16)
+ audio_h = self.audio_proj(audio_emb)
+ text_h = self.text_proj(text_emb)
+
+ dtype = audio_h.dtype
+
+ with fp32_precision():
+ gate = torch.sigmoid(self.gate) # FP32
+ res = torch.sigmoid(self.residual_scale) # FP32
+
+ h = gate.to(dtype) * audio_h + (1 - gate).to(dtype) * text_h
+ h = res.to(dtype) * h
+ h = self.final_norm(h.float()).to(dtype)
+
+ return h
+
+
+class RVQEARTTSModel(nn.Module):
+ """
+ Main RVQEARTTS model for training and inference.
+
+ This model integrates a character-aware text encoder with a transformer backbone
+ and a Mixture-of-Gaussians (MoG) prediction head.
+
+ The architecture is based on the Streaming TTS model proposed in
+ "Audio Flamingo 3" (https://arxiv.org/abs/2507.08128), with several improvements:
+
+ 1. Gated fusion of text and audio representations
+ (`GatedProjectedSumRMSNorm`).
+
+ 2. Subword-aware embeddings for improved pronunciation of multi-token words
+ (`SubwordFlagEmbedding`).
+
+ 3. Custom BOS and EOS embeddings for duplex interaction support,
+ enabling interruption-aware generation (`BOSEOSEmbedding`).
+
+ Args:
+ config (DictConfig | dict[str, Any]): The configuration object for the model.
+ """
+
+ config_class = DictConfig
+ rvq_embs: Tensor
+
+ def __init__(self, config: DictConfig | dict[str, Any], tokenizer: AutoTokenizer = None):
+ super().__init__()
+ self.config = config
+
+ # Backbone module
+ if self.config.get("pretrained_text_name", None):
+ # Load pretrained backbone from huggingface
+ from .compat import load_pretrained_hf
+
+ llm = load_pretrained_hf(self.config.pretrained_text_name, pretrained_weights=True).train()
+ self.backbone = llm.model # fetch PretrainedBaseModel from model "ForCausalLM"
+ else:
+ if self.config.get("backbone_type", None) is None:
+ assert (
+ self.config.get("backbone_model_class", None) is not None
+ and self.config.get("backbone_config_class", None) is not None
+ )
+ backbone_config = getattr(transformers, self.config.backbone_config_class)(
+ **(
+ OmegaConf.to_container(self.config.backbone_config, resolve=True)
+ if self.config.backbone_config
+ else {}
+ ),
+ )
+ self.backbone = getattr(transformers, self.config.backbone_model_class)(backbone_config)
+ else:
+ backbone_config = AutoConfig.for_model(
+ self.config.backbone_type,
+ **(
+ OmegaConf.to_container(self.config.backbone_config, resolve=True)
+ if self.config.backbone_config
+ else {}
+ ),
+ )
+ self.backbone = AutoModel.from_config(backbone_config)
+
+ self.hidden_size = self.backbone.get_input_embeddings().weight.size(-1)
+ find_and_delete_module(self.backbone, self.backbone.get_input_embeddings(), "backbone")
+
+ # Embedding and projection layers
+ self.bos_emb = nn.Parameter(torch.randn(self.hidden_size))
+ self.null_emb = nn.Parameter(torch.randn(self.hidden_size))
+ if self.config.random_target_masking:
+ self.embed_target_mask = nn.Embedding(self.config.num_quantizers, self.hidden_size)
+
+ self.embed_code = nn.Linear(self.config.latent_size, self.hidden_size, bias=False)
+
+ self.embed_context = (
+ nn.Linear(self.config.context_hidden_size, self.hidden_size, bias=False)
+ if self.config.context_hidden_size
+ else None
+ )
+
+ cas_config = (
+ {k: v for k, v in self.config.cas_config.items() if k != "pretrained_tokenizer_name"}
+ if self.config.cas_config
+ else {}
+ )
+ self.embed_subword = (
+ CharAwareSubwordEncoder(
+ tokenizer=tokenizer,
+ out_size=self.hidden_size,
+ use_subword_flag_emb=self.config.use_subword_flag_emb,
+ use_bos_eos_emb=self.config.use_bos_eos_emb,
+ **cas_config,
+ )
+ if self.config.cas_config
+ else None
+ )
+
+ if self.config.use_gated_fusion_for_text_audio:
+ self.gated_fusion_audio_text = GatedProjectedSumRMSNorm(
+ self.hidden_size, self.hidden_size, self.hidden_size, self.config.num_quantizers
+ )
+
+ if self.config.get("audio_prompt_encoder_config", None):
+ # Dedicated projection for audio prompt (pre-BOS)
+ ape_cfg = OmegaConf.to_container(self.config.audio_prompt_encoder_config, resolve=True)
+ model_type = ape_cfg.pop("type") # remove "type" from kwargs
+ ape_cfg = AutoConfig.for_model(model_type, **ape_cfg)
+ self.audio_prompt_encoder = AutoModel.from_config(ape_cfg)
+
+ if self.config.get("use_audio_prompt_frozen_projection", False):
+ with fp32_precision():
+ U, _ = torch.linalg.qr(torch.randn(self.hidden_size, self.hidden_size))
+ V, _ = torch.linalg.qr(torch.randn(self.hidden_size, self.hidden_size))
+ smin, smax = 0.4, 2.5
+ s = smin + (smax - smin) * torch.rand(self.hidden_size)
+ W = U @ torch.diag(s) @ V.T
+ self.register_buffer("audio_prompt_projection_W", W) # register as buffer to avoid weight update
+
+ # Prediction Heads
+ if not self.config.disable_eos_prediction:
+ self.lm_head = nn.Linear(self.hidden_size, 2, bias=False)
+
+ self.mog_head = MoGHead(
+ hidden_size=self.hidden_size,
+ out_size=self.config.latent_size,
+ **self.config.mog_head_config,
+ )
+
+ def set_rvq_embs(self, rvq_embs: Tensor):
+ self.register_buffer("rvq_embs", rvq_embs.detach().clone())
+
+ def depthsum_embedding(self, code: Tensor) -> Tensor:
+ """
+ code: [b, t, d]
+ rvq_embs: [d, v, h]
+
+ ret: [b, t, h]
+ """
+ b, t, d = code.size()
+ _, v, h = self.rvq_embs.size()
+ device = code.device
+
+ ret = torch.zeros((b, t, h), device=device)
+ embs = F.pad(self.rvq_embs, [0, 0, 0, 1])
+ for i in range(d):
+ emb = embs[i]
+ ret = ret + F.embedding(code[..., i], emb)
+ return ret
+
+ def prepare_training_inputs(self, code: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
+ """Prepares masked and dropped-out versions of the code for training."""
+ b, t, d = code.size()
+ device = code.device
+
+ src_rate = torch.rand((b, t), device=device) * self.config.max_training_rate
+ src_masking_rate = get_masking_rate(src_rate, self.config.exponent)
+ src_num_masking = torch.ceil(src_masking_rate * self.config.num_quantizers).long()
+
+ src_code_mask = torch.ones((b, t, d), dtype=torch.bool, device=device)
+ src_code_mask = get_mask(src_code_mask, src_num_masking)
+ src_masked_code = code * src_code_mask + (torch.zeros_like(code) + self.config.codebook_size) * (~src_code_mask)
+
+ if self.config.random_target_masking:
+ tgt_rate = src_rate + (1.0 - src_rate) * torch.rand((b, t), device=device)
+ tgt_masking_rate = get_masking_rate(tgt_rate, self.config.exponent)
+ tgt_num_masking = torch.floor(tgt_masking_rate * self.config.num_quantizers).long()
+
+ tgt_code_mask = torch.ones((b, t, d), dtype=torch.bool, device=device)
+ tgt_code_mask = get_mask(tgt_code_mask, tgt_num_masking)
+ tgt_masked_code = code * tgt_code_mask + (torch.zeros_like(code) + self.config.codebook_size) * (
+ ~tgt_code_mask
+ )
+ else:
+ tgt_code_mask = torch.ones((b, t, d), dtype=torch.bool, device=device)
+ tgt_masked_code = code
+
+ dropout_mask = torch.where(
+ torch.rand((b, t, 1), device=device) < self.config.quantizer_dropout,
+ (torch.randint(0, self.config.num_quantizers + 1, (b, t, 1), device=device)),
+ self.config.num_quantizers,
+ ) > torch.arange(d, dtype=torch.long, device=device)
+ dropped_code = code * dropout_mask + (torch.zeros_like(code) + self.config.codebook_size) * (~dropout_mask)
+
+ return src_masked_code, src_code_mask, tgt_masked_code, tgt_code_mask, dropped_code
+
+ def _prepare_conditioning(
+ self,
+ context_hidden_state: Tensor | None,
+ subword_ids: Tensor | None,
+ subword_mask: Tensor | None,
+ uncond_dec_flag: Tensor,
+ ) -> Tensor:
+ """Computes the final conditioning tensor by combining all sources."""
+ cond = torch.zeros((1, 1, self.hidden_size), device=uncond_dec_flag.device)
+
+ if self.embed_context is not None and context_hidden_state is not None:
+ cond = cond + self.embed_context(context_hidden_state)
+
+ if self.embed_subword is not None and subword_ids is not None:
+ # Infer subword mask from context if not provided
+ if subword_mask is None and context_hidden_state is not None:
+ subword_mask = torch.any(context_hidden_state != 0, dim=-1)
+ # at least one value should be true, otherwise we can completely skip it to avoid errors
+ if subword_mask is not None and subword_mask.any():
+ cond = cond + self.embed_subword(subword_ids, subword_mask)
+
+ # Replace with null embedding for unconditional generation
+ cond = torch.where(uncond_dec_flag, self.null_emb, cond)
+ return cond
+
+ def _compute_losses(
+ self,
+ code: Tensor,
+ lm_logits: Tensor,
+ mog_logits: Tensor,
+ mog_mus: Tensor,
+ mog_mu_res: Tensor,
+ mog_logs: Tensor,
+ src_code_mask: Tensor,
+ tgt_code_mask: Tensor,
+ audio_mask: Tensor,
+ ) -> tuple[Tensor, Tensor, Tensor]:
+ """Helper to compute all losses for the training step."""
+ with torch.autocast(code.device.type, enabled=False):
+ # 1. LM Loss (predicting discrete tokens)
+ if not self.config.disable_eos_prediction:
+ eos_mask = (~audio_mask) & F.pad(audio_mask[:, :-1], [1, 0])
+ lm_mask = eos_mask | audio_mask
+ lm_target = torch.where(eos_mask, 1, 0)
+ lm_loss = (
+ F.cross_entropy(lm_logits.transpose(1, 2), lm_target, reduction="none") * lm_mask
+ ).sum() / lm_mask.sum().clamp_min(1)
+ else:
+ lm_loss = 0.0
+
+ # 2. Continuous & KL Losses (for the MoG head)
+ target_mask = (~src_code_mask & tgt_code_mask) & audio_mask.unsqueeze(-1)
+ reduced_target_mask = target_mask.any(dim=-1)
+
+ cont_code_target = self.depthsum_embedding(
+ code * target_mask + (torch.zeros_like(code) + self.config.codebook_size) * (~target_mask)
+ )
+ mog_logits = mog_logits.float()
+ mog_mus = mog_mus.float()
+ mog_mu_res = mog_mu_res.float()
+ mog_logs = mog_logs.float()
+ with fp32_precision():
+ # Log probability of the true code under each Gaussian component
+ logp_code = (
+ -0.5 * math.log(2 * math.pi) - mog_logs
+ ) * self.config.latent_size - 0.5 * self.mog_head.dist(
+ mog_mus, (cont_code_target - mog_mu_res) * torch.exp(-mog_logs)
+ )
+
+ # Compute posterior q(k|c)
+ q_kc = (
+ torch.softmax(
+ logp_code,
+ -1,
+ )
+ * (1 - self.config.label_smoothing)
+ + self.config.label_smoothing / self.mog_head.num_predictions
+ ).detach()
+ log_q_kc = torch.log(q_kc + 1e-8).detach()
+
+ # Continuous Loss (negative log-likelihood)
+ c_loss = (-(q_kc * logp_code).sum(-1) * reduced_target_mask).sum() / target_mask.sum().clamp_min(1)
+
+ # KL Divergence Loss
+ k_loss = (
+ (q_kc * (log_q_kc - F.log_softmax(mog_logits, -1))).sum(-1) * reduced_target_mask
+ ).sum() / target_mask.sum().clamp_min(1)
+
+ return lm_loss, c_loss, k_loss
+
+ def forward(
+ self,
+ code: Tensor,
+ attention_mask: Tensor | None = None,
+ position_ids: Tensor | None = None,
+ context_hidden_state: Tensor | None = None,
+ subword_ids: Tensor | None = None,
+ subword_mask: Tensor | None = None,
+ audio_mask: Tensor | None = None,
+ non_prompt_mask: Tensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool = False,
+ training: bool | None = None,
+ guidance_enabled: bool = False,
+ generation_config: dict[str, Any] | None = None,
+ teacher_forcing_inference: bool = False,
+ ignore_eos_flag_stop: bool = False,
+ asr_speech_tokens_emb: Tensor | None = None,
+ audio_prompt_latent: Tensor | None = None,
+ ) -> RVQEARTTSOutput:
+ """
+ Performs a forward pass handling training, generation, or single-step inference.
+
+ Args:
+ code (Tensor): Input audio codes. For training, this is the ground truth.
+ For generation, this is the previously generated code token.
+ attention_mask (Tensor | None): Attention mask for the backbone transformer.
+ position_ids (Tensor | None): Position ids for the backbone transformer.
+ context_hidden_state (Tensor | None): Conditioning from a language model.
+ subword_ids (Tensor | None): Subword token IDs for conditioning.
+ subword_mask (Tensor | None): Mask for subword IDs.
+ audio_mask (Tensor | None): Mask for valid audio positions (for training and inference initialization).
+ past_key_values (Cache | None): Cache for past key-values for fast decoding.
+ use_cache (bool): If True, returns the updated `past_key_values`.
+ training (bool | None): Explicitly set training mode. If `None`, uses `self.training`.
+ guidance_enabled (bool): If True, duplicates inputs internally to run both
+ conditional and unconditional passes
+ generation_config (dict[str, Any] | None): If provided, triggers an iterative code generation.
+
+ Returns:
+ RVQEARTTSOutput: A dataclass containing losses (for training) or generated outputs
+ and the cache (for inference).
+ """
+
+ # Determine operating mode.
+ if training is None:
+ training = self.training
+
+ if audio_mask is not None:
+ if training:
+ (src_masked_code, src_code_mask, tgt_masked_code, tgt_code_mask, dropped_code) = (
+ self.prepare_training_inputs(code)
+ )
+ uncond_dec_flag = torch.rand(code.size(0), 1, 1, device=code.device) < self.config.p_uncond
+ else:
+ dropped_code = code
+ uncond_dec_flag = torch.zeros(code.size(0), 1, 1, device=code.device, dtype=torch.bool)
+
+ # code_embeds = (
+ # self.embed_code(self.depthsum_embedding(F.pad(dropped_code[:, :-1], [0, 0, 1, 0])))
+ # + (audio_mask & (~F.pad(audio_mask[:, :-1], [1, 0]))).unsqueeze(-1) * self.bos_emb
+ # )
+
+ # Shifted code
+ shifted_code = F.pad(dropped_code[:, :-1], [0, 0, 1, 0])
+
+ # Base embeddings
+ code_embed = self.depthsum_embedding(shifted_code)
+
+ # BOS mask
+ bos_mask = audio_mask & (~F.pad(audio_mask[:, :-1], [1, 0])) # [B, T]
+ bos_mask = bos_mask.unsqueeze(-1) # [B, T, 1]
+
+ # Mask for tokens BEFORE BOS
+ pre_bos_mask = bos_mask.cumsum(dim=1) == 0 # [B, T, 1]
+
+ # Apply projection to model size
+ code_embed = self.embed_code(code_embed)
+
+ # Choose projection
+ if self.config.get("audio_prompt_encoder_config", None):
+ # Dedicated projection for audio prompt (pre-BOS)
+ if audio_prompt_latent is None:
+ prompt_attn_mask = pre_bos_mask.squeeze(-1).long()
+ audio_prompt_latent = self.audio_prompt_encoder(
+ inputs_embeds=code_embed,
+ attention_mask=prompt_attn_mask,
+ return_dict=True,
+ ).last_hidden_state
+
+ code_embed = torch.where(
+ pre_bos_mask,
+ audio_prompt_latent,
+ code_embed,
+ )
+
+ if self.config.get("use_audio_prompt_frozen_projection", False):
+ if audio_prompt_latent is None:
+ W = self.audio_prompt_projection_W.to(code_embed.device, code_embed.dtype)
+ audio_prompt_latent = torch.nn.functional.linear(code_embed, W.T)
+
+ code_embed = torch.where(
+ pre_bos_mask,
+ audio_prompt_latent,
+ code_embed,
+ )
+
+ # Add BOS embedding
+ code_embeds = code_embed + bos_mask * self.bos_emb
+
+ else: # Inference
+ code_embeds = self.embed_code(self.depthsum_embedding(code))
+ uncond_dec_flag = torch.zeros(code.size(0), 1, 1, device=code.device, dtype=torch.bool)
+
+ if guidance_enabled:
+ assert not training, "Classifier-free guidance can only be used when `training` is False."
+ code_embeds = torch.cat([code_embeds] * 2, 0)
+ if attention_mask is not None:
+ attention_mask = torch.cat([attention_mask] * 2, 0)
+ if position_ids is not None:
+ position_ids = torch.cat([position_ids] * 2, 0)
+ if context_hidden_state is not None:
+ context_hidden_state = torch.cat([context_hidden_state] * 2, 0)
+ if subword_ids is not None:
+ subword_ids = torch.cat([subword_ids] * 2, 0)
+ if subword_mask is not None:
+ subword_mask = torch.cat([subword_mask] * 2, 0)
+ if asr_speech_tokens_emb is not None:
+ asr_speech_tokens_emb = torch.cat([asr_speech_tokens_emb] * 2, 0)
+
+ uncond_dec_flag = torch.cat([uncond_dec_flag, torch.ones_like(uncond_dec_flag)], 0)
+
+ # Prepare conditioning
+ cond = self._prepare_conditioning(context_hidden_state, subword_ids, subword_mask, uncond_dec_flag)
+
+ if asr_speech_tokens_emb is not None:
+ cond = cond + asr_speech_tokens_emb
+
+ if self.config.use_gated_fusion_for_text_audio:
+ inputs_embeds = self.gated_fusion_audio_text(code_embeds, cond)
+ else:
+ inputs_embeds = code_embeds + cond
+
+ # Main backbone pass
+ backbone_outputs = self.backbone(
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ )
+ hidden_states = backbone_outputs.last_hidden_state
+
+ if audio_mask is not None and training:
+ # --- Training-specific loss computation ---
+ if not self.config.disable_eos_prediction:
+ lm_logits = self.lm_head(hidden_states)
+ else:
+ lm_logits = None
+
+ mog_input_embeds = self.embed_code(self.depthsum_embedding(src_masked_code))
+ if self.config.random_target_masking:
+ mog_input_embeds = mog_input_embeds + self.embed_target_mask((tgt_code_mask.sum(-1) - 1).clamp_min(0))
+ mog_input_embeds = mog_input_embeds + hidden_states
+ mog_logits, mog_mus, mog_mu_res, mog_logs = self.mog_head(mog_input_embeds)
+
+ lm_loss, c_loss, k_loss = self._compute_losses(
+ code, lm_logits, mog_logits, mog_mus, mog_mu_res, mog_logs, src_code_mask, tgt_code_mask, audio_mask
+ )
+ total_loss = lm_loss + c_loss + k_loss
+
+ return RVQEARTTSOutput(
+ loss=total_loss,
+ lm_loss=lm_loss,
+ c_loss=c_loss,
+ k_loss=k_loss,
+ hidden_states=hidden_states,
+ audio_prompt_latent=audio_prompt_latent,
+ )
+ else: # Inference
+ if not generation_config:
+ return RVQEARTTSOutput(
+ hidden_states=hidden_states,
+ past_key_values=backbone_outputs.past_key_values,
+ audio_prompt_latent=audio_prompt_latent,
+ )
+ else:
+ if teacher_forcing_inference:
+ generated_codes, lm_logits, eos_flag = self.generate_teacher_forcing(
+ hidden_states, generation_config
+ )
+ else:
+ generated_codes, lm_logits, eos_flag = self.generate_step(
+ hidden_states, ignore_eos_flag_stop=ignore_eos_flag_stop, **generation_config
+ )
+ return RVQEARTTSOutput(
+ past_key_values=backbone_outputs.past_key_values,
+ codes=generated_codes,
+ lm_logits=lm_logits,
+ eos_flag=eos_flag,
+ hidden_states=hidden_states,
+ )
+
+ @torch.no_grad()
+ def generate_teacher_forcing(self, hidden_states: Tensor, generation_config: dict):
+ """
+ Teacher-forcing wrapper for generate_step, processing all frames in parallel
+ using a per-frame loop internally.
+
+ Args:
+ hidden_states: [B, T, H] hidden states
+ generation_config: kwargs for self.generate_step()
+
+ Returns:
+ generated_codes: [B, T, ...] generated codes per frame
+ lm_logits: [B, T, vocab_size] language model logits
+ eos_flag: [B, T] boolean tensor indicating EOS
+ """
+ B, T, H = hidden_states.shape
+
+ # Preallocate caches
+ generated_codes_cache = []
+ lm_logits_cache = []
+ eos_flag_cache = []
+
+ # Iterate over time steps (frames)
+ for t in range(T):
+ # extract one frame (as the original generate_step expects)
+ frame_hidden = hidden_states[:, t, :] # [B, H]
+
+ # call original generate_step
+ generated_codes, lm_logits, eos_flag = self.generate_step(
+ frame_hidden.unsqueeze(1),
+ **generation_config, # keep batch dim + frame dim
+ )
+ if generated_codes is not None:
+ # store in cache
+ generated_codes_cache.append(generated_codes)
+ lm_logits_cache.append(lm_logits)
+ eos_flag_cache.append(eos_flag)
+
+ # Stack results along time dimension
+ generated_codes = torch.stack(generated_codes_cache, dim=1) # [B, T, ...]
+ if not self.config.disable_eos_prediction:
+ lm_logits = torch.stack(lm_logits_cache, dim=1) # [B, T, vocab_size]
+ eos_flag = torch.stack(eos_flag_cache, dim=1) # [B, T]
+ else:
+ lm_logits = None
+ eos_flag = None
+
+ return generated_codes, lm_logits, eos_flag
+
+ @torch.no_grad()
+ def generate_step(
+ self,
+ hidden_states: Tensor,
+ num_iter: int,
+ guidance_scale: list[float] | float | None = None,
+ top_p_or_k: list[float | int] | float | int | None = None,
+ noise_scale: list[float] | float | None = None,
+ exponent: float | None = None,
+ eos_threshold: float | None = None,
+ ignore_eos_flag_stop: bool = False,
+ ) -> tuple[Tensor | None, Tensor, Tensor]:
+ """
+ Performs the iterative unmasking process for a single generation step.
+
+ This function takes the hidden state from the backbone transformer and generates
+ codes through an iterative unmasking process.
+
+ Args:
+ hidden_states (Tensor): The hidden states from the backbone. If using CFG,
+ this should be the combined [uncond, cond] tensor.
+ num_iter (int): The number of unmasking iterations.
+ guidance_scale (list[float] | float | None): The scale for Classifier-Free Guidance.
+ top_p_or_k (list[float | int] | float | int | None): The value for top-p or top-k sampling.
+ noise_scale (list[float] | float | None): The scale of noise to add during MoG sampling.
+ exponent (float | None): The exponent for the masking schedule.
+ eos_threshold (float | None): The threshold for EOS prediction.
+
+ Returns:
+ tuple[Tensor | None, Tensor, Tensor]: A tuple containing:
+ - the generated codes.
+ - The logits from `lm_head`.
+ - The EOS flag.
+ """
+ # 1. Preparation
+ if guidance_scale is not None:
+ if not isinstance(guidance_scale, list):
+ guidance_scale = [guidance_scale] * (1 + num_iter) # includes one step for `lm_head`
+ assert len(guidance_scale) == 1 + num_iter
+ if top_p_or_k is not None:
+ if not isinstance(top_p_or_k, list):
+ top_p_or_k = [top_p_or_k] * (1 + num_iter) # includes one step for `lm_head`
+ assert len(top_p_or_k) == 1 + num_iter
+ if noise_scale is not None:
+ if not isinstance(noise_scale, list):
+ noise_scale = [noise_scale] * num_iter
+ assert len(noise_scale) == num_iter
+ if exponent is None:
+ exponent = self.config.exponent
+
+ if guidance_scale is not None:
+ # The effective batch size is halved
+ hidden_states, uncond_hidden_states = hidden_states.chunk(2, dim=0)
+ else:
+ uncond_hidden_states = hidden_states[:0, :0, :0]
+
+ b, t, _ = hidden_states.size()
+ d = self.config.num_quantizers
+ device = hidden_states.device
+
+ # 2. Predict the discrete part of the code
+ if not self.config.disable_eos_prediction:
+ if guidance_scale is not None:
+ lm_logits = self.lm_head(hidden_states + guidance_scale[0] * (hidden_states - uncond_hidden_states))
+ else:
+ lm_logits = self.lm_head(hidden_states)
+ if top_p_or_k is not None:
+ lm_logits = (
+ TopPLogitsWarper(top_p_or_k[0])(
+ None,
+ lm_logits.view(-1, lm_logits.size(-1)),
+ ).view_as(lm_logits)
+ if isinstance(top_p_or_k[0], float)
+ else TopKLogitsWarper(top_p_or_k[0])(
+ None,
+ lm_logits.view(-1, lm_logits.size(-1)),
+ ).view_as(lm_logits)
+ )
+
+ with fp32_precision():
+ lm_logits = F.log_softmax(lm_logits, -1)
+
+ if eos_threshold is not None:
+ eos_flag = lm_logits[..., -1] > eos_threshold
+ else:
+ eos_flag = lm_logits.argmax(-1) == 1
+
+ if torch.all(eos_flag) and not ignore_eos_flag_stop:
+ return None, lm_logits, eos_flag
+ else:
+ lm_logits = None
+ eos_flag = None
+
+ # Initialize the full code tensor
+ code = torch.zeros((b, t, d), dtype=torch.long, device=device) + self.config.codebook_size
+
+ # 3. Set up the iterative denoising schedule for the continuous part
+ rates = torch.linspace(0.0, 1.0, num_iter + 1, device=device)[:-1].unsqueeze(-1)
+ masking_rates = get_masking_rate(rates, exponent=exponent)
+ num_maskings = torch.ceil(masking_rates * self.config.num_quantizers).long()
+
+ ks = num_maskings - F.pad(num_maskings[1:], [0, 0, 0, 1])
+
+ # 4. Iteratively unmask the continuous part of the code
+ cnt = 0
+ for i, k in enumerate(ks):
+ if torch.all(k == 0):
+ continue
+
+ # Prepare input for the MoG head
+ guidance_scale_i = guidance_scale[i] if guidance_scale is not None else 0.0
+ top_p_or_k_i = top_p_or_k[i] if top_p_or_k is not None else 1.0
+ noise_scale_i = noise_scale[i] if noise_scale is not None else 1.0
+
+ mog_input_embeds = self.embed_code(self.depthsum_embedding(code))
+ if self.config.random_target_masking:
+ mog_input_embeds += self.embed_target_mask(cnt + k - 1)
+ if guidance_scale_i > 0.0:
+ mog_input_embeds = torch.cat(
+ [mog_input_embeds + hidden_states, mog_input_embeds + uncond_hidden_states], 0
+ )
+ else:
+ mog_input_embeds += hidden_states
+
+ mog_mu, mog_logs = self.mog_head.infer(
+ mog_input_embeds,
+ guidance_scale=guidance_scale_i,
+ top_p_or_k=top_p_or_k_i,
+ )
+ z = mog_mu + torch.exp(mog_logs) * torch.randn_like(mog_mu) * noise_scale_i
+ code = depthsum_encoding_step(self.rvq_embs, z, code, cnt, k[0].item())
+ cnt += k[0].item()
+ return code, lm_logits, eos_flag
+
+ def load_state_dict(self, state_dict, strict: bool = True):
+ try:
+ super().load_state_dict(state_dict, strict=strict)
+ except RuntimeError:
+ logging.info("Error loading model state_dict !! Retrying with partial initialization!")
+ model_dict = set_model_dict_for_partial_init(state_dict, self.state_dict())
+ super().load_state_dict(model_dict, strict=False)
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_vae_codec.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_vae_codec.py
new file mode 100644
index 000000000..9385f0c75
--- /dev/null
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_vae_codec.py
@@ -0,0 +1,874 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/speechlm2/modules/ear_tts_vae_codec.py; training-only code removed.
+
+import functools
+import math
+from collections.abc import Callable
+from contextlib import contextmanager
+from typing import Any, Concatenate
+
+import torch
+from omegaconf import DictConfig
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+
+@contextmanager
+def disable_tf32():
+ prev = torch.backends.cudnn.allow_tf32
+ torch.backends.cudnn.allow_tf32 = False
+ try:
+ yield
+ finally:
+ torch.backends.cudnn.allow_tf32 = prev
+
+
+# ==============================================================================
+# Utility Functions
+# ==============================================================================
+
+
+def zero_module(module: nn.Module) -> nn.Module:
+ """
+ Zeros out the parameters of a PyTorch module in-place.
+
+ This is a utility function that iterates through all parameters of a given
+ `nn.Module` and sets their values to zero. This is often used for specific
+ initialization strategies, for example in diffusion models where some layers
+ are initialized to zero.
+
+ From: https://github.com/openai/guided-diffusion/blob/22e0df8183507e13a7813f8d38d51b072ca1e67c/guided_diffusion/nn.py#L68
+
+ Args:
+ module (nn.Module): The PyTorch module to be zeroed.
+
+ Returns:
+ nn.Module: The same module with its parameters zeroed.
+ """
+ for p in module.parameters():
+ # p.detach().zero_() performs the operation in-place without tracking it in autograd
+ p.detach().zero_()
+ return module
+
+
+def sequence_mask(lengths: Tensor, max_length: int | None = None) -> Tensor:
+ """
+ Creates a boolean mask from a 1D tensor of sequence lengths.
+
+ This function is useful for masking out padding in sequences. Given a tensor
+ of lengths, it produces a 2D boolean tensor where `mask[i, j]` is `True` if
+ `j < lengths[i]` and `False` otherwise.
+
+ Example:
+ >>> lengths = torch.tensor([1, 3, 2])
+ >>> sequence_mask(lengths)
+ tensor([[ True, False, False],
+ [ True, True, True],
+ [ True, True, False]])
+
+ Args:
+ lengths (Tensor): A 1D tensor of integer lengths. Shape: `[batch_size]`.
+ max_length (int | None, optional): The maximum length of the mask. If None,
+ it is inferred from the maximum value
+ in `lengths`. Defaults to None.
+
+ Returns:
+ Tensor: The boolean mask. Shape: `[batch_size, max_length]`.
+ """
+ if max_length is None:
+ # If max_length is not provided, use the longest sequence length in the batch.
+ max_length = int(lengths.max().item())
+
+ # Create a range tensor from 0 to max_length - 1
+ x = torch.arange(max_length, dtype=lengths.dtype, device=lengths.device)
+
+ # Compare each length with the range tensor to create the mask.
+ # `x.unsqueeze(0)` is `[1, max_length]`
+ # `lengths.unsqueeze(1)` is `[batch_size, 1]`
+ # Broadcasting takes care of the comparison.
+ return x.unsqueeze(0) < lengths.unsqueeze(1)
+
+
+# ==============================================================================
+# Signal Processing Functions
+# ==============================================================================
+
+
+def spectrogram(
+ wav: Tensor,
+ n_fft: int,
+ hop_length: int,
+ win_length: int,
+ window_fn: Callable[Concatenate[int, ...], Tensor] = torch.hann_window,
+) -> Tensor:
+ """
+ Computes the Short-Time Fourier Transform (STFT) of a waveform with manual padding.
+
+ This implementation manually applies zero padding before computing the STFT.
+ This is done to center the analysis window at the beginning of the signal
+ without using the `center=True` argument in `torch.stft`, giving more control.
+
+ Args:
+ wav (Tensor): The input audio waveform.
+ Shape: [batch_size?, time_steps], where batch_size? is an
+ optional batch dimension.
+ n_fft (int): The size of the FFT.
+ hop_length (int): The number of samples between adjacent STFT columns.
+ win_length (int): The size of the window function.
+ window_fn (function, optional): The window function to apply.
+ Defaults to torch.hann_window.
+
+ Returns:
+ Tensor: The complex-valued spectrogram.
+ Shape: [batch_size?, n_fft // 2 + 1, num_frames]
+ """
+ # Calculate the padding required on the left and right sides to center the frames.
+ pad_size_l = (n_fft - hop_length) // 2
+ pad_size_r = (n_fft - hop_length) - pad_size_l
+
+ # Use a torch.autocast context to perform STFT in float32 for precision.
+ with torch.autocast(device_type=wav.device.type, enabled=False):
+ # Apply reflection padding to the waveform.
+ wav = F.pad(wav.float(), (pad_size_l, pad_size_r))
+
+ # Create the window tensor on the same device as the waveform.
+ window = window_fn(win_length, dtype=torch.float, device=wav.device)
+
+ # Compute the STFT.
+ # `center=False` because we have already manually padded the signal.
+ spec = torch.stft(
+ wav,
+ n_fft,
+ hop_length=hop_length,
+ win_length=win_length,
+ window=window,
+ center=False,
+ normalized=False,
+ onesided=True,
+ return_complex=True,
+ )
+ return spec
+
+
+def spec_to_wav(
+ spec: Tensor,
+ n_fft: int,
+ hop_length: int,
+ win_length: int,
+ window_fn: Callable[Concatenate[int, ...], Tensor] = torch.hann_window,
+ constrain_value_range: bool = False,
+) -> Tensor:
+ """
+ Converts a spectrogram back into a waveform using the overlap-add method.
+ This function is an approximate inverse of the `spectrogram` function.
+
+ Args:
+ spec (Tensor): The input complex-valued spectrogram.
+ Shape: [batch_size?, dim, time_steps], where batch_size?
+ is an optional batch dimension.
+ n_fft (int): The size of the FFT used to create the spectrogram.
+ hop_length (int): The number of samples between frames in the original signal.
+ win_length (int): The size of the window function used in the original signal.
+ window_fn (function, optional): The window function used. Currently only
+ `torch.hann_window` is supported.
+ constrain_value_range (bool, optional): If True, constrains the IFFT values
+ to be within the range of the window.
+ This ensures that the output values
+ remain within the range of -1.0 to 1.0.
+ Defaults to False.
+
+ Returns:
+ Tensor: The reconstructed waveform.
+ Shape: [batch_size?, time_steps]
+
+ Raises:
+ ValueError: If a window function other than `torch.hann_window` is provided.
+ """
+ with torch.autocast(device_type=spec.device.type, enabled=False):
+ if window_fn != torch.hann_window:
+ raise ValueError(f"`window_fn` should be 'torch.hann_window', but got '{window_fn}'.")
+
+ # Calculate padding and number of frames
+ pad = (win_length - hop_length) // 2
+ T = spec.size(-1)
+ window = window_fn(win_length, device=spec.device)
+
+ # 1. Inverse FFT
+ # Convert from frequency domain back to time domain for each frame.
+ ifft = torch.fft.irfft(spec, n=n_fft, dim=-2, norm="backward")
+ window_unsqz = window.unsqueeze(-1)
+
+ # 2. Optionally constrain values
+ if constrain_value_range:
+ ifft = torch.where(
+ ifft >= 0,
+ torch.minimum(ifft, window_unsqz),
+ torch.maximum(ifft, -window_unsqz),
+ )
+
+ # 3. Apply window to the IFFT result
+ ifft = ifft * window_unsqz
+
+ # 4. Overlap and Add
+ # Use `torch.nn.functional.fold` to perform the overlap-add operation efficiently.
+ # This reconstructs the continuous signal from the windowed frames.
+ output_size = (T - 1) * hop_length + win_length
+ wav = F.fold(
+ ifft,
+ output_size=(1, output_size),
+ kernel_size=(1, win_length),
+ stride=(1, hop_length),
+ )[..., 0, 0, pad:-pad]
+
+ # 5. Calculate the window envelope for normalization
+ # This is necessary to correct for the energy added by overlapping windows.
+ window_sq = window.square().expand(T, -1).transpose(0, 1)
+ window_envelope = F.fold(
+ window_sq,
+ output_size=(1, output_size),
+ kernel_size=(1, win_length),
+ stride=(1, hop_length),
+ ).squeeze()[pad:-pad]
+
+ # 6. Normalize the waveform
+ # Divide by the window envelope to get the final reconstructed signal.
+ assert (window_envelope > 1e-11).all(), "Window envelope has zero values, cannot normalize."
+ wav = wav / window_envelope
+
+ return wav
+
+
+# ==============================================================================
+# Basic Modules
+# ==============================================================================
+
+
+class CausalConv1dCache:
+ """
+ A cache for managing states in causal 1D convolutions.
+
+ This class is used during autoregressive inference to store and update the
+ tail of the input to a causal convolution, which is used as padding for the
+ next time step. This avoids re-computing the entire sequence at each step.
+ """
+
+ def __init__(self) -> None:
+ self.cache: dict[int | str, Tensor] = {}
+
+ def __getitem__(self, layer_id: int | str) -> Tensor:
+ """Retrieves the cached tensor for a given layer."""
+ return self.cache[layer_id]
+
+ def update(
+ self,
+ states: Tensor,
+ layer_id: int | str,
+ padding: int,
+ padding_value: int = 0,
+ flush: bool = False,
+ ) -> Tensor:
+ """
+ Updates the cache for a specific layer and returns the padded input.
+
+ Args:
+ states (Tensor): The new input tensor for the current time step.
+ layer_id (int | str): An identifier for the convolutional layer.
+ padding (int): The amount of left padding required by the convolution.
+ padding_value (int, optional): The value to use for initial padding. Defaults to 0.
+ flush (bool, optional): If True, the cache for this layer is deleted
+ after use. Defaults to False.
+
+ Returns:
+ Tensor: The input states concatenated with the cached padding.
+ """
+ device = states.device
+ dtype = states.dtype
+ b, c, t = states.size()
+
+ if layer_id not in self.cache:
+ # Initialize cache with zero padding if it's the first time step
+ padding_tensor = torch.zeros((b, c, padding), dtype=dtype, device=device) + padding_value
+ else:
+ padding_tensor = self.cache[layer_id]
+ assert padding_tensor.size(2) == padding
+
+ # Concatenate the cached padding with the new states
+ padded_states = torch.cat([padding_tensor, states], dim=2)
+ # Update the cache with the tail of the new padded states
+ self.cache[layer_id] = padded_states[:, :, -padding:]
+
+ if flush:
+ del self.cache[layer_id]
+
+ return padded_states
+
+
+class LayerNormNd(nn.Module):
+ """
+ A LayerNorm module that works for N-dimensional inputs.
+
+ This implementation normalizes over the channel dimension (dim=1), which is
+ a common setup for convolutional networks.
+
+ Args:
+ channels (int): The number of channels of the input tensor.
+ eps (float, optional): A value added to the denominator for numerical
+ stability. Defaults to 1e-6.
+ elementwise_affine (bool, optional): If True, this module has learnable
+ affine parameters (weight and bias).
+ Defaults to True.
+ bias (bool, optional): If True, this module has a learnable bias.
+ Defaults to True.
+ """
+
+ def __init__(self, channels: int, eps=1e-6, elementwise_affine: bool = True, bias: bool = True):
+ super().__init__()
+ self.channels = channels
+ self.eps = eps
+
+ self.weight = nn.Parameter(torch.ones((channels,)), requires_grad=elementwise_affine)
+ self.bias = nn.Parameter(torch.zeros((channels,)), requires_grad=elementwise_affine and bias)
+
+ def forward(self, x: Tensor) -> Tensor:
+ # Calculate mean and reciprocal standard deviation over the channel dimension
+ mean = x.mean(1, keepdim=True)
+ x_shift = x - mean
+ # Using rsqrt for potentially better performance
+ x_rstd = torch.rsqrt(x_shift.pow(2).mean(1, keepdim=True) + self.eps)
+
+ # Reshape weight and bias to be broadcastable with the input tensor
+ shape = [-1 if i == 1 else 1 for i in range(x.ndim)]
+
+ # Apply normalization and affine transformation
+ return (x_shift * x_rstd) * self.weight.view(shape) + self.bias.view(shape)
+
+
+class ConvNeXt1d(nn.Module):
+ """
+ A 1D ConvNeXt block adapted for causal convolutions on audio signals.
+
+ This block is a core component of modern convolutional architectures, featuring
+ a depthwise convolution, layer normalization, and pointwise convolutions to
+ expand and contract the channel dimension, similar to an inverted bottleneck.
+
+ Implementation adapted from: https://github.com/charactr-platform/vocos
+
+ Args:
+ dim (int): Number of input and output channels.
+ intermediate_dim (int): Dimensionality of the intermediate (expanded) layer.
+ kernel_size (int): The kernel size for the causal depthwise convolution.
+ identity_init (bool, optional): If True, the final pointwise convolution
+ is initialized to zero, making the block
+ an identity function at the start of training.
+ Defaults to False.
+ layer_idx (int, optional): An index for this layer, used for caching. Defaults to 0.
+ """
+
+ def __init__(
+ self,
+ dim: int,
+ intermediate_dim: int,
+ kernel_size: int,
+ identity_init: bool = False,
+ layer_idx: int = 0,
+ ):
+ super().__init__()
+ self.layer_idx = layer_idx
+ self.kernel_size = kernel_size
+
+ # Depthwise convolution
+ self.dwconv = nn.Conv1d(dim, dim, kernel_size=kernel_size, groups=dim)
+
+ self.norm = LayerNormNd(dim)
+ self.pwconv1 = nn.Conv1d(dim, intermediate_dim, 1) # Pointwise/1x1 conv for expansion
+ self.act = nn.GELU()
+
+ # Pointwise/1x1 conv for projection
+ if identity_init:
+ self.pwconv2 = zero_module(nn.Conv1d(intermediate_dim, dim, 1))
+ else:
+ self.pwconv2 = nn.Conv1d(intermediate_dim, dim, 1)
+
+ def forward(self, x: Tensor, cache: CausalConv1dCache | None = None, flush: bool = False) -> Tensor:
+ residual = x
+
+ # Apply causal padding, either through a cache or manually
+ if cache is not None:
+ x = cache.update(x, self.layer_idx, self.kernel_size - 1, flush=flush)
+ else:
+ x = F.pad(x, [self.kernel_size - 1, 0]) # Left padding for causality
+
+ # Main ConvNeXt path
+ x = self.dwconv(x)
+ x = self.norm(x)
+ x = self.pwconv1(x)
+ x = self.act(x)
+ x = self.pwconv2(x)
+
+ # Add residual connection
+ x = residual + x
+ return x
+
+
+class PreTrainedEMAVariance(nn.Module):
+ """
+ Exponential Moving Average of Variance
+ """
+
+ def __init__(self, initial_value: float = 1.0):
+ super().__init__()
+ self.variance = nn.Parameter(
+ torch.tensor(initial_value),
+ requires_grad=False,
+ )
+
+ def forward(self) -> Tensor:
+ return self.variance
+
+
+class PreTrainedProbabilisticVQ(nn.Module):
+ def __init__(
+ self,
+ channels: int,
+ num_mixtures: int,
+ depth: int = 1,
+ ):
+ super().__init__()
+ self.channels = channels
+ self.num_mixtures = num_mixtures
+ self.depth = depth
+
+ self.mus_list = nn.ParameterList(
+ [
+ nn.Parameter(
+ F.normalize(torch.randn(num_mixtures, channels), p=2.0, dim=1) * ((depth - i) / depth),
+ requires_grad=False,
+ )
+ for i in range(depth)
+ ]
+ )
+ self._variance_list = nn.ModuleList([PreTrainedEMAVariance() for _ in range(depth)])
+
+ def encode(self, z: Tensor, return_z_q: bool = False) -> list[Tensor] | tuple[list[Tensor], Tensor]:
+ r = z
+ ids_sel = []
+ for i in range(self.depth):
+ mus = self.mus_list[i]
+ idx_sel = self._dist_sq(r, mus).argmin(-1) # [b, ?, h], [v, h] -> [b, ?]
+ r = r - F.embedding(idx_sel, mus)
+ ids_sel.append(idx_sel)
+ if return_z_q:
+ return ids_sel, z - r
+ return ids_sel
+
+ def decode(self, ids_sel: list[Tensor]) -> Tensor:
+ z = torch.zeros((*ids_sel[0].size(), self.channels), device=ids_sel[0].device)
+ for i in range(len(ids_sel)):
+ mus = self.mus_list[i]
+ z = z + F.embedding(ids_sel[i], mus)
+ return z # [b, ?, h]
+
+ def _dist_sq(self, z: Tensor, mus: Tensor) -> Tensor:
+ """
+ z: [b, ?, d?, h]
+ mus: [d?, v, h]
+ """
+ return (
+ z.pow(2).sum(-1, keepdim=True) # [b, ?, d?, 1]
+ + mus.pow(2).sum(-1) # [d?, v]
+ - 2 * (z.unsqueeze(-2) @ mus.transpose(-1, -2)).squeeze(-2) # [b, ?, d?, h] , [d?, h, v] -> [b, ?, d?, v]
+ )
+
+
+class Wav2Latent(nn.Module):
+ """
+ An encoder model that transforms a raw waveform into a latent representation.
+
+ This model first converts the waveform to a spectrogram, then processes it
+ through a series of ConvNeXt blocks and downsampling convolutional layers
+ to produce a compressed latent tensor.
+
+ Args:
+ latent_size (int): The number of channels in the final latent representation.
+ n_fft (int): The FFT size for the initial spectrogram transformation.
+ hop_length (int): The hop length for the STFT.
+ base_hidden_size (int): The base number of channels for the hidden layers.
+ channel_mult (tuple[int, ...]): A tuple of multipliers for the hidden
+ size at each stage of downsampling.
+ rates (tuple[int, ...]): A tuple of downsampling factors (strides) for
+ the convolutional layers.
+ num_blocks (int): The number of ConvNeXt blocks per stage.
+ kernel_size (int): The kernel size for the ConvNeXt blocks.
+ groups (int): The number of groups for the downsampling convolutions.
+ """
+
+ def __init__(
+ self,
+ latent_size: int = 1024,
+ n_fft: int = 32,
+ hop_length: int = 8,
+ base_hidden_size: int = 384,
+ channel_mult: tuple[int, ...] = (1, 2, 4),
+ rates: tuple[int, ...] = (8, 8, 8),
+ num_blocks: int = 3,
+ kernel_size: int = 7,
+ groups: int = 1,
+ ):
+ super().__init__()
+ self.n_fft = n_fft
+ self.hop_length = hop_length
+
+ # Initial projection from spectrogram to hidden size
+ layers: list[nn.Module] = [nn.Conv1d(n_fft + 2, base_hidden_size * channel_mult[0], 1, bias=False)]
+
+ # Downsampling stages
+ for i in range(len(channel_mult)):
+ ch_mult, rate = channel_mult[i], rates[i]
+ hidden_size = base_hidden_size * ch_mult
+ # Add ConvNeXt blocks for this stage
+ for j in range(num_blocks):
+ layers.append(ConvNeXt1d(hidden_size, hidden_size * 4, kernel_size, True, layer_idx=i * num_blocks + j))
+ # Add downsampling convolution
+ next_hidden_size = base_hidden_size * channel_mult[i + 1] if i < len(channel_mult) - 1 else latent_size
+ layers.append(
+ nn.Conv1d(hidden_size, next_hidden_size, kernel_size=rate, stride=rate, bias=False, groups=groups)
+ )
+
+ self.layers = nn.ModuleList(layers)
+
+ def forward(self, x: Tensor, cache=None, flush: bool = False) -> Tensor:
+ if cache is not None:
+ raise NotImplementedError("Caching is not implemented for the encoder.")
+
+ # Convert waveform to spectrogram (magnitude and phase)
+ with torch.autocast(device_type=x.device.type, enabled=False):
+ spec = spectrogram(x.squeeze(1), n_fft=self.n_fft, hop_length=self.hop_length, win_length=self.n_fft)
+ # Split complex spectrogram into real and imaginary, then treat as magnitude and phase
+ mag, ph = torch.view_as_real(spec).chunk(2, dim=-1)
+ x = torch.cat([mag, ph], 1).squeeze(-1)
+
+ # Pass through the network
+ for layer in self.layers:
+ if isinstance(layer, ConvNeXt1d):
+ x = layer(x, cache=cache, flush=flush)
+ else:
+ x = layer(x)
+
+ # Transpose to [batch, time, channels] for compatibility with transformers
+ x = x.transpose(-1, -2)
+ return x
+
+
+class Latent2Wav(nn.Module):
+ """
+ A decoder (vocoder) model that transforms a latent representation back into a raw waveform.
+
+ This model processes a latent tensor through a series of ConvNeXt blocks and
+ upsampling transposed convolutional layers to produce a spectrogram, which is
+ then converted back to a waveform using an inverse STFT.
+
+ Args:
+ latent_size (int): The number of channels in the input latent representation.
+ n_fft (int): The FFT size for the final spectrogram reconstruction.
+ hop_length (int): The hop length for the ISTFT.
+ base_hidden_size (int): The base number of channels for the hidden layers.
+ channel_mult (tuple[int, ...]): A tuple of multipliers for the hidden
+ size at each stage of upsampling.
+ rates (tuple[int, ...]): A tuple of upsampling factors (strides) for
+ the transposed convolutional layers.
+ num_blocks (int): The number of ConvNeXt blocks per stage.
+ kernel_size (int): The kernel size for the ConvNeXt blocks.
+ groups (int): The number of groups for the upsampling convolutions.
+ """
+
+ def __init__(
+ self,
+ latent_size: int = 1024,
+ n_fft: int = 32,
+ hop_length: int = 8,
+ base_hidden_size: int = 384,
+ channel_mult: tuple[int, ...] = (4, 2, 1),
+ rates: tuple[int, ...] = (8, 8, 8),
+ num_blocks: int = 3,
+ kernel_size: int = 7,
+ groups=1,
+ ):
+ super().__init__()
+ self.n_fft = n_fft
+ self.hop_length = hop_length
+ self.spec_cache_idx = (len(channel_mult)) * num_blocks
+
+ layers: list[nn.Module] = []
+
+ # Upsampling stages
+ for i in range(len(channel_mult)):
+ ch_mult, rate = channel_mult[i], rates[i]
+ hidden_size = base_hidden_size * ch_mult
+ # Add upsampling transposed convolution
+ in_size = base_hidden_size * channel_mult[i - 1] if i != 0 else latent_size
+ layers.append(
+ nn.ConvTranspose1d(in_size, hidden_size, kernel_size=rate, stride=rate, bias=False, groups=groups)
+ )
+ # Add ConvNeXt blocks for this stage
+ for j in range(num_blocks):
+ layers.append(ConvNeXt1d(hidden_size, hidden_size * 4, kernel_size, True, layer_idx=i * num_blocks + j))
+
+ # Final projection to spectrogram dimensions (magnitude + phase)
+ layers.append(nn.Conv1d(hidden_size, n_fft + 2, 1, bias=False))
+ self.layers = nn.ModuleList(layers)
+
+ def forward(self, x: Tensor, cache=None, flush: bool = False, constrain_value_range: bool = True) -> Tensor:
+ # Transpose input from [batch, time, channels] to [batch, channels, time]
+ x = x.transpose(-1, -2)
+
+ # Pass through the network
+ for layer in self.layers:
+ if isinstance(layer, ConvNeXt1d):
+ x = layer(x, cache=cache, flush=flush)
+ else:
+ x = layer(x)
+
+ # Convert network output to a complex spectrogram and then to a waveform
+ with torch.autocast(device_type=x.device.type, enabled=False):
+ max_mag = 100.0
+ # Split output into magnitude and phase components
+ mag, ph = x.float().chunk(2, dim=1)
+ # Safeguard to prevent excessively large magnitudes
+ mag = max_mag * torch.exp(-F.softplus(-mag + math.log(max_mag)))
+
+ # Reconstruct the complex spectrogram from magnitude and phase
+ # The DC and Nyquist components are real, so their phase is applied via cosine.
+ mag_dc, mag_mid, mag_nyquist = mag.split([1, mag.size(1) - 2, 1], dim=1)
+ ph_dc, ph_mid, ph_nyquist = torch.cos(ph).split([1, ph.size(1) - 2, 1], dim=1)
+ ph_imag = torch.sin(ph[:, 1:-1, :])
+
+ spec_real = mag_mid * ph_mid
+ spec_imag = mag_mid * ph_imag
+
+ spec = torch.cat([mag_dc * ph_dc, spec_real + 1j * spec_imag, mag_nyquist * ph_nyquist], 1)
+
+ # Handle caching for autoregressive generation of the spectrogram
+ if cache is not None:
+ half_spec_padding = math.ceil(((self.n_fft - self.hop_length) // 2) / self.hop_length)
+ spec = cache.update(spec, self.spec_cache_idx, padding=half_spec_padding * 2, flush=flush)
+ if flush:
+ spec = F.pad(spec, [0, half_spec_padding])
+
+ # Convert spectrogram to waveform
+ x = spec_to_wav(
+ spec, self.n_fft, self.hop_length, self.n_fft, constrain_value_range=constrain_value_range
+ ).unsqueeze(1)
+
+ if cache is not None:
+ # Trim the output to remove the padded region from the start
+ half_wav_padding = half_spec_padding * self.hop_length
+ x = x[:, :, half_wav_padding:-half_wav_padding]
+
+ return x
+
+
+class RVQVAEModel(nn.Module):
+ """
+ Residual Vector-Quantized Variational Autoencoder (RVQ-VAE) model.
+
+ This model learns a discrete representation of audio by encoding a waveform
+ into a latent space and then quantizing the latents into discrete codes.
+ It consists of an encoder, a quantizer, and a decoder.
+
+ Args:
+ config (DictConfig | dict[str, Any]): A configuration object with model hyperparameters.
+ """
+
+ config_class: type[DictConfig] = DictConfig
+
+ def __init__(self, config: DictConfig | dict[str, Any]):
+ super().__init__()
+ self.config = config
+
+ self.encoder = Wav2Latent(
+ latent_size=self.config.latent_size,
+ n_fft=self.config.n_fft,
+ hop_length=self.config.hop_length,
+ base_hidden_size=self.config.base_hidden_size,
+ channel_mult=self.config.channel_mult,
+ rates=self.config.rates,
+ num_blocks=self.config.num_blocks,
+ kernel_size=self.config.kernel_size,
+ groups=self.config.groups,
+ )
+
+ # Layers for quantization
+ self.prvq = PreTrainedProbabilisticVQ(
+ channels=self.config.latent_size,
+ num_mixtures=self.config.codebook_size,
+ depth=self.config.num_quantizers,
+ )
+
+ self.decoder = Latent2Wav(
+ latent_size=self.config.latent_size,
+ n_fft=self.config.n_fft,
+ hop_length=self.config.hop_length,
+ base_hidden_size=self.config.base_hidden_size,
+ channel_mult=tuple(reversed(self.config.channel_mult)),
+ rates=tuple(reversed(self.config.rates)),
+ num_blocks=self.config.num_blocks,
+ kernel_size=self.config.kernel_size,
+ groups=self.config.groups,
+ )
+
+ for p in self.parameters():
+ p.requires_grad = False
+
+ def ae_encode(self, x: Tensor, cache: CausalConv1dCache | None = None, flush: bool = False) -> Tensor:
+ """
+ Runs the encoder part of the autoencoder.
+
+ Args:
+ x (Tensor): Input waveform. Shape: `[batch, 1, time]`.
+ cache (CausalConv1dCache | None): Not implemented for the encoder.
+ flush (bool): Not implemented for the encoder.
+
+ Returns:
+ Tensor: The continuous latent representation. Shape: `[batch, time', channels]`.
+ """
+ assert x.size(1) == 1 and x.dim() == 3, "Input must be a batch of mono audio."
+ assert x.size(2) % self.config.wav_to_token_ratio == 0, (
+ f"Input audio length ({x.size(2)}) must be divisible by the model's "
+ f"wav_to_token_ratio ({self.config.wav_to_token_ratio}). "
+ f"Please pad the input to a compatible length."
+ )
+
+ if cache is not None:
+ raise NotImplementedError("Caching is not supported for the encoder.")
+
+ return self.encoder(x, cache=cache, flush=flush)
+
+ def ae_decode(
+ self,
+ x: Tensor,
+ constrain_value_range: bool = True,
+ cache: CausalConv1dCache | None = None,
+ flush: bool = False,
+ ) -> Tensor:
+ """
+ Runs the decoder part of the autoencoder.
+
+ Args:
+ x (Tensor): The (de-quantized) latent representation. Shape: `[batch, time', channels]`.
+ constrain_value_range (bool): If True, constrains the output of the ISTFT.
+ cache (CausalConv1dCache | None): Cache for autoregressive generation.
+ flush (bool): If True, flushes the cache.
+
+ Returns:
+ Tensor: The reconstructed waveform. Shape: `[batch, 1, time]`.
+ """
+ return self.decoder(x, constrain_value_range=constrain_value_range, cache=cache, flush=flush)
+
+ def encode(self, x: Tensor, x_len: Tensor) -> tuple[Tensor, Tensor]:
+ """
+ Encodes a waveform into discrete codes.
+
+ Args:
+ x (Tensor): Input waveform. Shape: `[batch, 1, time]`.
+ x_len (Tensor): The original lengths of the waveforms in the batch.
+
+ Returns:
+ tuple[Tensor, Tensor]: A tuple containing:
+ - The discrete codes. Shape: `[batch, time', n_quantizers]`.
+ - The lengths of the code sequences.
+ """
+ with disable_tf32():
+ z_e = self.ae_encode(x)
+ code_len = x_len // self.config.wav_to_token_ratio
+ return self.quantize(z_e), code_len
+
+ def decode(
+ self,
+ code: Tensor,
+ code_len: Tensor | None = None,
+ constrain_value_range: bool = True,
+ cache: CausalConv1dCache | None = None,
+ flush: bool = False,
+ ) -> tuple[Tensor, Tensor | None]:
+ """
+ Decodes discrete codes back into a waveform.
+
+ Args:
+ code (Tensor): The discrete codes. Shape: `[batch, time', n_quantizers]`.
+ code_len (Tensor | None): The lengths of the code sequences.
+ constrain_value_range (bool): If True, constrains the output of the ISTFT.
+ cache (CausalConv1dCache | None): Cache for autoregressive generation.
+ flush (bool): If True, flushes the cache.
+
+ Returns:
+ tuple[Tensor, Tensor | None]: A tuple containing:
+ - The reconstructed waveform. Shape: `[batch, 1, time]`.
+ - The lengths of the reconstructed waveforms.
+ """
+ with disable_tf32():
+ z_q = self.dequantize(code)
+ x_hat = self.ae_decode(z_q, constrain_value_range=constrain_value_range, cache=cache, flush=flush)
+ wav_len = code_len * self.config.wav_to_token_ratio if code_len is not None else None
+ return x_hat, wav_len
+
+ def quantize(self, z: Tensor) -> Tensor:
+ """
+ Quantizes a continuous latent tensor into discrete codes.
+
+ Args:
+ z (Tensor): The continuous latent tensor from the encoder.
+ Shape: `[batch, time, channels]`.
+
+ Returns:
+ Tensor: The quantized codes. Shape: `[batch, time, n_quantizers]`.
+ """
+ with disable_tf32():
+ ids_sel = self.prvq.encode(z, return_z_q=False)
+ return torch.stack(ids_sel, -1)
+
+ def dequantize(self, code: Tensor) -> Tensor:
+ """
+ De-quantizes discrete codes back into a continuous latent tensor.
+
+ Args:
+ code (Tensor): The quantized codes. Shape: `[batch, time, n_quantizers]`.
+
+ Returns:
+ Tensor: The de-quantized continuous latent tensor.
+ Shape: `[batch, time, latent_size]`.
+ """
+ ids_sel = [x.squeeze(-1) for x in torch.split(code, 1, -1)]
+ return self.prvq.decode(ids_sel)
+
+ def forward(self, x: Tensor, constrain_value_range: bool = False) -> Tensor:
+ """
+ Performs a full autoencoding pass: encode, quantize, dequantize, and decode.
+
+ Args:
+ x (Tensor): The input waveform. Shape: `[batch, 1, time]`.
+ constrain_value_range (bool): If True, constrains the output of the ISTFT.
+
+ Returns:
+ Tensor: The reconstructed waveform. Shape: `[batch, 1, time]`.
+ """
+
+ with torch.no_grad():
+ z_e = self.ae_encode(x)
+ code = self.quantize(z_e)
+ z_d = self.dequantize(code)
+ x_hat = self.ae_decode(z_d, constrain_value_range=constrain_value_range)
+ return x_hat
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/fusion.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/fusion.py
new file mode 100644
index 000000000..2cf7ba8f4
--- /dev/null
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/fusion.py
@@ -0,0 +1,403 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/speechlm2/parts/fusion.py; training-only code removed.
+
+"""
+Fusion modules for combining multi-modal embeddings in duplex speech models.
+
+Supported fusion methods:
+- "add" or None: Simple additive fusion with per-modality weights
+- "concat": Concatenate (after LayerNorm) and project back to hidden_dim
+- "gated_simple": Per-timestep, per-dimension gates with softmax normalization
+- "gated_gmu": Gated Multimodal Unit (Arevalo et al., 2017) with per-modality sigmoid gates
+
+For gated methods, the channel embeddings, tie_and_roll, and per-modality weights
+are bypassed - the gating mechanism learns the optimal combination.
+"""
+
+import torch
+import torch.nn as nn
+from typing import Optional, Literal, Union
+
+from .compat import logging
+
+FuseMethod = Literal["add", "concat", "gated_simple", "gated_gmu"]
+
+
+class FusionModule(nn.Module):
+ """Base class for fusion modules."""
+
+ def __init__(self, hidden_dim: int):
+ super().__init__()
+ self.hidden_dim = hidden_dim
+
+ def forward(
+ self,
+ agent_text_embeds: torch.Tensor,
+ user_audio_embeds: torch.Tensor,
+ user_text_embeds: torch.Tensor | None = None,
+ function_embeds: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ """
+ Fuse multi-modal embeddings.
+
+ Args:
+ agent_text_embeds: Agent text embeddings (B, T, D)
+ user_audio_embeds: User audio embeddings (B, T, D)
+ user_text_embeds: User text/ASR embeddings (B, T, D) or None
+ function_embeds: Function calling channel embeddings (B, T, D) or None
+
+ Returns:
+ Fused embeddings (B, T, D)
+ """
+ raise NotImplementedError
+
+
+class AddFusion(FusionModule):
+ """
+ Simple additive fusion with optional per-modality weights.
+
+ This is the default behavior matching the original implementation.
+ """
+
+ def __init__(
+ self,
+ hidden_dim: int,
+ agent_text_weight: float = 1.0,
+ user_audio_weight: float = 1.0,
+ user_text_weight: float = 1.0,
+ function_weight: float = 1.0,
+ ):
+ super().__init__(hidden_dim)
+ self.agent_text_weight = agent_text_weight
+ self.user_audio_weight = user_audio_weight
+ self.user_text_weight = user_text_weight
+ self.function_weight = function_weight
+
+ def forward(
+ self,
+ agent_text_embeds: torch.Tensor,
+ user_audio_embeds: torch.Tensor,
+ user_text_embeds: torch.Tensor | None = None,
+ function_embeds: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ output = agent_text_embeds * self.agent_text_weight
+ output = output + user_audio_embeds * self.user_audio_weight
+ if user_text_embeds is not None:
+ output = output + user_text_embeds * self.user_text_weight
+ if function_embeds is not None:
+ output = output + function_embeds * self.function_weight
+ return output
+
+
+class ConcatFusion(FusionModule):
+ """
+ Concatenate embeddings (after LayerNorm) and project back to hidden_dim.
+
+ Supports optional user_text_embeds and function_embeds; zeros are used when
+ not provided. Uses 4 modalities (agent, audio, user_text, function) so that
+ function channel can be optionally present.
+ The linear layer learns to handle this case appropriately.
+ """
+
+ def __init__(
+ self,
+ hidden_dim: int,
+ agent_text_weight: float = 1.0,
+ user_audio_weight: float = 1.0,
+ user_text_weight: float = 1.0,
+ function_weight: float = 1.0,
+ num_modalities: int = 4,
+ ):
+ super().__init__(hidden_dim)
+ self.agent_text_weight = agent_text_weight
+ self.user_audio_weight = user_audio_weight
+ self.user_text_weight = user_text_weight
+ self.function_weight = function_weight
+ self.num_modalities = num_modalities
+
+ # LayerNorm for each modality
+ self.ln_agent = nn.LayerNorm(hidden_dim)
+ self.ln_audio = nn.LayerNorm(hidden_dim)
+ self.ln_text = nn.LayerNorm(hidden_dim)
+ self.ln_function = nn.LayerNorm(hidden_dim)
+
+ # Projection layer (4 modalities; pad with zeros when optional channels missing)
+ self.proj = nn.Linear(hidden_dim * num_modalities, hidden_dim, bias=False)
+
+ def forward(
+ self,
+ agent_text_embeds: torch.Tensor,
+ user_audio_embeds: torch.Tensor,
+ user_text_embeds: torch.Tensor | None = None,
+ function_embeds: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ # Apply layernorm and weights
+ agent_normed = self.ln_agent(agent_text_embeds) * self.agent_text_weight
+ audio_normed = self.ln_audio(user_audio_embeds) * self.user_audio_weight
+
+ if user_text_embeds is not None:
+ text_normed = self.ln_text(user_text_embeds) * self.user_text_weight
+ else:
+ # Zero-pad when ASR embeddings are not available
+ text_normed = torch.zeros_like(agent_text_embeds)
+
+ if function_embeds is not None:
+ function_normed = self.ln_function(function_embeds) * self.function_weight
+ else:
+ function_normed = torch.zeros_like(agent_text_embeds)
+
+ # Concatenate along hidden dim and project
+ concat = torch.cat([agent_normed, audio_normed, text_normed, function_normed], dim=-1)
+ return self.proj(concat)
+
+
+class GatedFusionSimple(FusionModule):
+ """
+ Learns a softmax-normalized gate over 4 streams at each timestep (agent, audio, user_text, function).
+
+ Gate input: LayerNorm(each stream) → concatenate → linear → softmax
+ Output: weighted sum of LayerNorm'd embeddings where weights are input-dependent.
+ Optional channels (user_text, function) use zeros when not provided; their gate is masked to 0.
+
+ Note: This is PER-TIMESTEP gating (one scalar weight per modality per timestep).
+ For PER-DIMENSION gating, use GatedFusionGMU.
+ """
+
+ def __init__(self, hidden_dim: int, num_modalities: int = 4):
+ super().__init__(hidden_dim)
+ self.num_modalities = num_modalities
+
+ # LayerNorm for each stream (normalizes distributions before gate computation)
+ self.ln_agent = nn.LayerNorm(hidden_dim)
+ self.ln_audio = nn.LayerNorm(hidden_dim)
+ self.ln_text = nn.LayerNorm(hidden_dim)
+ self.ln_function = nn.LayerNorm(hidden_dim)
+
+ # Gate network: outputs one gate value per modality per timestep
+ # Input: concatenated normalized modalities (B, T, D*4)
+ # Output: gates (B, T, 4) - one scalar per modality
+ self.gate_net = nn.Linear(hidden_dim * num_modalities, num_modalities, bias=True)
+
+ # Initialize for uniform weighting
+ self._init_uniform()
+
+ def _init_uniform(self):
+ """Initialize gate network for uniform initial weighting."""
+ nn.init.zeros_(self.gate_net.weight)
+ nn.init.zeros_(self.gate_net.bias)
+
+ def forward(
+ self,
+ agent_text_embeds: torch.Tensor,
+ user_audio_embeds: torch.Tensor,
+ user_text_embeds: torch.Tensor | None = None,
+ function_embeds: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ B, T, D = agent_text_embeds.shape
+
+ has_user_text = user_text_embeds is not None
+ has_function = function_embeds is not None
+ text_input = user_text_embeds if has_user_text else torch.zeros_like(agent_text_embeds)
+ function_input = function_embeds if has_function else torch.zeros_like(agent_text_embeds)
+
+ # Apply LayerNorm to each stream before gate computation
+ agent_normed = self.ln_agent(agent_text_embeds)
+ audio_normed = self.ln_audio(user_audio_embeds)
+ text_normed = self.ln_text(text_input)
+ function_normed = self.ln_function(function_input)
+
+ # Concatenate normalized inputs for gate computation
+ concat = torch.cat([agent_normed, audio_normed, text_normed, function_normed], dim=-1)
+
+ # Compute gates: (B, T, D*4) -> (B, T, 4)
+ gates = self.gate_net(concat) # (B, T, num_modalities)
+
+ # Mask out optional channels when not present (before softmax)
+ if not has_user_text:
+ gates = gates.clone()
+ gates[:, :, 2] = float("-inf")
+ if not has_function:
+ gates = gates.clone()
+ gates[:, :, 3] = float("-inf")
+
+ # Softmax over modalities for normalization
+ gates = torch.softmax(gates, dim=-1) # (B, T, 4)
+
+ # Stack normalized modalities for weighted sum: (B, T, 4, D)
+ stacked = torch.stack([agent_normed, audio_normed, text_normed, function_normed], dim=2)
+
+ # Apply gates (broadcast) and sum over modalities
+ fused = (stacked * gates.unsqueeze(-1)).sum(dim=2) # (B, T, D)
+
+ return fused
+
+
+class GatedFusionGMU(FusionModule):
+ """
+ Gated Multimodal Unit inspired by Arevalo et al. (2017).
+
+ PER-DIMENSION gating: each dimension can attend to different modalities.
+ Uses softmax over modalities for numerically stable normalization.
+ Supports 4 modalities (agent, audio, user_text, function); optional channels use zeros and are masked.
+
+ All streams are LayerNorm'd before gate computation and fusion.
+ Gate output shape: (B, T, 4, D) with softmax over the modality dimension.
+ """
+
+ def __init__(self, hidden_dim: int, num_modalities: int = 4):
+ super().__init__(hidden_dim)
+ self.num_modalities = num_modalities
+
+ # LayerNorm for each stream (normalizes distributions before gate computation)
+ self.ln_agent = nn.LayerNorm(hidden_dim)
+ self.ln_audio = nn.LayerNorm(hidden_dim)
+ self.ln_text = nn.LayerNorm(hidden_dim)
+ self.ln_function = nn.LayerNorm(hidden_dim)
+
+ # Per-modality gate projections (output logits, not activations)
+ self.gate_agent = nn.Linear(hidden_dim * num_modalities, hidden_dim)
+ self.gate_audio = nn.Linear(hidden_dim * num_modalities, hidden_dim)
+ self.gate_text = nn.Linear(hidden_dim * num_modalities, hidden_dim)
+ self.gate_function = nn.Linear(hidden_dim * num_modalities, hidden_dim)
+
+ # Initialize for uniform weighting (all logits = 0 → equal softmax weights)
+ self._init_uniform()
+
+ def _init_uniform(self):
+ """Initialize gate networks for uniform initial weighting."""
+ for gate in [self.gate_agent, self.gate_audio, self.gate_text, self.gate_function]:
+ nn.init.zeros_(gate.weight)
+ nn.init.zeros_(gate.bias)
+
+ def forward(
+ self,
+ agent_text_embeds: torch.Tensor,
+ user_audio_embeds: torch.Tensor,
+ user_text_embeds: torch.Tensor | None = None,
+ function_embeds: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ B, T, D = agent_text_embeds.shape
+
+ has_user_text = user_text_embeds is not None
+ has_function = function_embeds is not None
+ text_input = user_text_embeds if has_user_text else torch.zeros_like(agent_text_embeds)
+ function_input = function_embeds if has_function else torch.zeros_like(agent_text_embeds)
+
+ # Apply LayerNorm to each stream before gate computation
+ agent_normed = self.ln_agent(agent_text_embeds)
+ audio_normed = self.ln_audio(user_audio_embeds)
+ text_normed = self.ln_text(text_input)
+ function_normed = self.ln_function(function_input)
+
+ # Concatenate normalized inputs for gate computation
+ concat = torch.cat([agent_normed, audio_normed, text_normed, function_normed], dim=-1)
+
+ # Compute per-modality gate logits: each is (B, T, D)
+ logits_agent = self.gate_agent(concat)
+ logits_audio = self.gate_audio(concat)
+ logits_text = self.gate_text(concat)
+ logits_function = self.gate_function(concat)
+
+ # Stack logits: (B, T, 4, D)
+ stacked_logits = torch.stack([logits_agent, logits_audio, logits_text, logits_function], dim=2)
+
+ # Mask out optional channels when not present (before softmax)
+ if not has_user_text:
+ stacked_logits = stacked_logits.clone()
+ stacked_logits[:, :, 2, :] = float("-inf")
+ if not has_function:
+ stacked_logits = stacked_logits.clone()
+ stacked_logits[:, :, 3, :] = float("-inf")
+
+ # Softmax over modalities (dim=2) for stable normalization
+ gates = torch.softmax(stacked_logits, dim=2) # (B, T, 4, D)
+
+ # Stack normalized embeddings: (B, T, 4, D)
+ stacked_embeds = torch.stack([agent_normed, audio_normed, text_normed, function_normed], dim=2)
+
+ # Fuse with element-wise gating
+ fused = (gates * stacked_embeds).sum(dim=2) # (B, T, D)
+
+ return fused
+
+
+def create_fusion_module(
+ fuse_method: Union[str, FuseMethod] | None,
+ hidden_dim: int,
+ agent_text_weight: float = 1.0,
+ user_audio_weight: float = 1.0,
+ user_text_weight: float = 1.0,
+ function_weight: float = 1.0,
+) -> FusionModule:
+ """
+ Factory function to create the appropriate fusion module.
+
+ Args:
+ fuse_method: One of None, "add", "concat", "gated_simple", "gated_gmu"
+ hidden_dim: Hidden dimension of the embeddings
+ agent_text_weight: Weight for agent text (used by add/concat methods only)
+ user_audio_weight: Weight for user audio (used by add/concat methods only)
+ user_text_weight: Weight for user text (used by add/concat methods only)
+ function_weight: Weight for function channel (used by add/concat only; optional channel)
+
+ Returns:
+ Instantiated fusion module
+
+ Notes:
+ - All fusion modules accept optional function_embeds; when None, that channel
+ contributes nothing (zeros + mask for gated methods).
+ - For gated methods (gated_simple, gated_gmu), the weight parameters are
+ ignored as the gating mechanism learns the optimal combination.
+ - When using gated methods, you should also bypass channel embeddings,
+ tie_and_roll, and any manual embedding weight parameters in the model.
+ """
+ if fuse_method is None or fuse_method == "add":
+ logging.info(
+ f"Using AddFusion with weights: agent={agent_text_weight}, "
+ f"audio={user_audio_weight}, text={user_text_weight}, function={function_weight}"
+ )
+ return AddFusion(
+ hidden_dim=hidden_dim,
+ agent_text_weight=agent_text_weight,
+ user_audio_weight=user_audio_weight,
+ user_text_weight=user_text_weight,
+ function_weight=function_weight,
+ )
+ elif fuse_method == "concat":
+ logging.info(
+ f"Using ConcatFusion with weights: agent={agent_text_weight}, "
+ f"audio={user_audio_weight}, text={user_text_weight}, function={function_weight}"
+ )
+ return ConcatFusion(
+ hidden_dim=hidden_dim,
+ agent_text_weight=agent_text_weight,
+ user_audio_weight=user_audio_weight,
+ user_text_weight=user_text_weight,
+ function_weight=function_weight,
+ )
+ elif fuse_method == "gated_simple":
+ logging.info(
+ f"Using GatedFusionSimple with hidden_dim={hidden_dim} (learned gating, optional function channel)"
+ )
+ return GatedFusionSimple(hidden_dim=hidden_dim)
+ elif fuse_method == "gated_gmu":
+ logging.info(f"Using GatedFusionGMU with hidden_dim={hidden_dim} (learned gating, optional function channel)")
+ return GatedFusionGMU(hidden_dim=hidden_dim)
+ else:
+ raise ValueError(
+ f"Unknown fuse_method: '{fuse_method}'. Expected one of: None, 'add', 'concat', 'gated_simple', 'gated_gmu'"
+ )
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/perception.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/perception.py
new file mode 100644
index 000000000..36df3a0d3
--- /dev/null
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/perception.py
@@ -0,0 +1,135 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/speechlm2/modules/perception.py; training-only code removed.
+
+import torch
+from omegaconf import DictConfig
+from torch import nn
+
+from .compat import (
+ Exportable,
+ NeuralModule,
+ typecheck,
+)
+
+
+class AudioPerceptionModule(NeuralModule, Exportable):
+ """Audio perception module that consists of audio encoder(s) and modality adapter."""
+
+ def __init__(self, cfg: DictConfig):
+ super().__init__()
+ # Initialize components
+ self.cfg = cfg
+ self.preprocessor = self.from_config_dict(cfg.preprocessor)
+ self.encoder = self.from_config_dict(cfg.encoder)
+
+ if "spec_augment" in cfg and cfg.spec_augment is not None:
+ self.spec_augmentation = self.from_config_dict(cfg.spec_augment)
+ else:
+ self.spec_augmentation = None
+ self.modality_adapter = self.from_config_dict(cfg.modality_adapter)
+ if "output_dim" not in cfg.modality_adapter and "d_model" in cfg.modality_adapter: # e.g., conformer encoder
+ self.proj = nn.Linear(cfg.modality_adapter.d_model, cfg.output_dim)
+ else:
+ self.proj = nn.Identity()
+
+ self.modality_adapter_quantizer_levels = cfg.get("modality_adapter_quantizer_levels", None)
+ if self.modality_adapter_quantizer_levels:
+ # FiniteScalarQuantizer (nemo.collections.tts.modules.audio_codec_modules) was not
+ # vendored: `modality_adapter_quantizer_levels` is unset in the
+ # NVIDIA-NemotronLabs-VoiceChat-11B checkpoint config.json.
+ raise RuntimeError(
+ "modality_adapter_quantizer_levels is set in the perception config, but "
+ "FiniteScalarQuantizer was not vendored (it is unused by the "
+ "NVIDIA-NemotronLabs-VoiceChat-11B checkpoint). Vendor "
+ "nemo.collections.tts.modules.audio_codec_modules.FiniteScalarQuantizer to enable it."
+ )
+
+ def maybe_preprocess_audio(
+ self,
+ input_signal=None,
+ input_signal_length=None,
+ processed_signal=None,
+ processed_signal_length=None,
+ ):
+ has_input_signal = input_signal is not None and input_signal_length is not None
+ has_processed_signal = processed_signal is not None and processed_signal_length is not None
+ if (has_input_signal ^ has_processed_signal) is False:
+ raise ValueError(
+ f"{self.__class__} Arguments ``input_signal`` and ``input_signal_length`` are mutually exclusive "
+ " with ``processed_signal`` and ``processed_signal_len`` arguments."
+ )
+
+ if not has_processed_signal:
+ processed_signal, processed_signal_length = self.preprocessor(
+ input_signal=input_signal,
+ length=input_signal_length,
+ )
+ return processed_signal, processed_signal_length
+
+ # disable type checks to avoid type-check errors when using Conformer as modality adapter
+ @typecheck.disable_checks()
+ def forward(
+ self,
+ input_signal=None,
+ input_signal_length=None,
+ processed_signal=None,
+ processed_signal_length=None,
+ return_encoder_emb=False,
+ sample_id=None,
+ ):
+ processed_signal, processed_signal_length = self.maybe_preprocess_audio(
+ input_signal, input_signal_length, processed_signal, processed_signal_length
+ )
+
+ # Spec augment is not applied during evaluation/testing
+ if self.spec_augmentation is not None and self.training:
+ processed_signal = self.spec_augmentation(input_spec=processed_signal, length=processed_signal_length)
+
+ encoder_emb, encoded_len = self.encoder(audio_signal=processed_signal, length=processed_signal_length)
+
+ if self.modality_adapter_quantizer_levels is not None:
+ encoded = self.modality_adapter_quantizer_bottleneck(encoder_emb.transpose(1, 2))
+ encoded, _ = self.modality_adapter_vector_quantizer(inputs=encoded.transpose(1, 2), input_len=None)
+ encoded = self.modality_adapter_quantizer_projection(encoded.transpose(1, 2)).transpose(1, 2)
+ else:
+ encoded = encoder_emb
+
+ encoded, encoded_len = self.modality_adapter(audio_signal=encoded, length=encoded_len)
+
+ # encoded_orig = encoded.clone()
+
+ # b, c, t -> b, t, c
+ encoded = self.proj(encoded.transpose(1, 2))
+
+ if return_encoder_emb:
+ return encoded, encoded_len, encoder_emb.transpose(1, 2)
+ else:
+ return encoded, encoded_len
+
+
+class IdentityConnector(NeuralModule, Exportable):
+ """User to pass encoder's representations as-is to the LLM."""
+
+ def __init__(
+ self,
+ *args,
+ **kwargs,
+ ):
+ super().__init__()
+
+ def forward(self, audio_signal, length=None, *args, **kwargs):
+ return audio_signal, length
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/tokenizer.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/tokenizer.py
new file mode 100644
index 000000000..ed89746b4
--- /dev/null
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/tokenizer.py
@@ -0,0 +1,501 @@
+# SPDX-FileCopyrightText: Copyright (c) 2020 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat)
+# nemo/collections/common/tokenizers/tokenizer_spec.py and
+# nemo/collections/common/tokenizers/huggingface/auto_tokenizer.py; training-only code removed.
+# Only the imports were adjusted (nemo.utils logging -> stdlib shim); the class
+# bodies are exact copies.
+
+from abc import ABC, abstractmethod
+from collections import OrderedDict
+from typing import List, Optional
+
+from transformers import AutoTokenizer as AUTOTOKENIZER
+
+from .compat import logging
+
+__all__ = ["TokenizerSpec", "AutoTokenizer"]
+
+
+class TokenizerSpec(ABC):
+ """
+ Inherit this class to implement a new tokenizer.
+ """
+
+ @abstractmethod
+ def text_to_tokens(self, text):
+ """Converts text into a list of tokens."""
+ pass
+
+ @abstractmethod
+ def tokens_to_ids(self, tokens):
+ """Converts a list of tokens to their corresponding IDs."""
+ pass
+
+ @abstractmethod
+ def ids_to_tokens(self, ids):
+ """Converts a list of token IDs back to tokens."""
+ pass
+
+ @abstractmethod
+ def text_to_ids(self, text):
+ """Converts text directly to token IDs."""
+ pass
+
+ def add_special_tokens(self, special_tokens: List[str]):
+ """Adds special tokens (eos, pad, cls...) to vocab."""
+ raise NotImplementedError("To be implemented")
+
+ def apply_chat_template(self, *args, **kwargs):
+ """Applies chat template and tokenizes results"""
+ raise NotImplementedError("To be implemented")
+
+ @property
+ def name(self):
+ """name of the class"""
+ return type(self).__name__
+
+ @property
+ def cls(self):
+ """Property alias to match MegatronTokenizer; returns cls_id if available."""
+ if hasattr(self, "cls_id"):
+ return self.cls_id
+ raise AttributeError(f"{type(self).__name__} has no attribute 'cls' or 'cls_id'")
+
+ @property
+ def sep(self):
+ """Property alias to match MegatronTokenizer; returns sep_id if available."""
+ if hasattr(self, "sep_id"):
+ return self.sep_id
+ raise AttributeError(f"{type(self).__name__} has no attribute 'sep' or 'sep_id'")
+
+ @property
+ def pad(self):
+ """Property alias to match MegatronTokenizer; returns pad_id if available."""
+ if hasattr(self, "pad_id"):
+ return self.pad_id
+ raise AttributeError(f"{type(self).__name__} has no attribute 'pad' or 'pad_id'")
+
+ @property
+ def eod(self):
+ """Property alias to match MegatronTokenizer; returns eod_id if available."""
+ if hasattr(self, "eod_id"):
+ return self.eod_id
+ if hasattr(self, "eos_id"):
+ # Default to end-of-sentence id if end-of-document is not defined.
+ return self.eos_id
+ raise AttributeError(f"{type(self).__name__} has no attribute 'eod', 'eod_id', 'eos', or 'eos_id'")
+
+ @property
+ def bos(self):
+ """Property alias to match MegatronTokenizer; returns bos_id if available."""
+ if hasattr(self, "bos_id"):
+ return self.bos_id
+ raise AttributeError(f"{type(self).__name__} has no attribute 'bos' or 'bos_id'")
+
+ @property
+ def eos(self):
+ """Property alias to match MegatronTokenizer; returns eos_id if available."""
+ if hasattr(self, "eos_id"):
+ return self.eos_id
+ raise AttributeError(f"{type(self).__name__} has no attribute 'eos' or 'eos_id'")
+
+ @property
+ def mask(self):
+ """Property alias to match MegatronTokenizer; returns mask_id if available."""
+ if hasattr(self, "mask_id"):
+ return self.mask_id
+ raise AttributeError(f"{type(self).__name__} has no attribute 'mask' or 'mask_id'")
+
+
+class AutoTokenizer(TokenizerSpec):
+ """
+ Wrapper of HuggingFace AutoTokenizer
+ https://huggingface.co/docs/transformers/main/en/model_doc/auto#transformers.AutoTokenizer.
+
+ """
+
+ def __init__(
+ self,
+ pretrained_model_name: str,
+ vocab_file: str | None = None,
+ merges_file: str | None = None,
+ mask_token: str | None = None,
+ bos_token: str | None = None,
+ eos_token: str | None = None,
+ pad_token: str | None = None,
+ sep_token: str | None = None,
+ cls_token: str | None = None,
+ unk_token: str | None = None,
+ additional_special_tokens: List | None = [],
+ use_fast: bool | None = True,
+ trust_remote_code: bool | None = False,
+ include_special_tokens: bool = False,
+ ):
+ """
+ Args:
+ pretrained_model_name: corresponds to HuggingFace-AutoTokenizer's 'pretrained_model_name_or_path' input
+ argument. For more details please refer to the documentation of the `from_pretrained` method here:
+ https://huggingface.co/docs/transformers/main/en/model_doc/auto#transformers.AutoTokenizer.
+ The list of all supported models can be found here: https://huggingface.co/models
+ vocab_file: path to file with vocabulary which consists
+ of characters separated by newlines.
+ mask_token: mask token
+ bos_token: the beginning of sequence token
+ eos_token: the end of sequence token. Usually equal to sep_token
+ pad_token: token to use for padding
+ sep_token: token used for separating sequences
+ cls_token: class token. Usually equal to bos_token
+ unk_token: token to use for unknown tokens
+ additional_special_tokens: list of other tokens beside standard special tokens (bos, eos, pad, etc.). For
+ example, sentinel tokens for T5 (, , etc.)
+ use_fast: whether to use fast HuggingFace tokenizer
+ include_special_tokens: when True, converting text to ids will include special tokens / prompt tokens (if
+ any), yielding self.tokenizer(text).input_ids
+ """
+ try:
+ self._initialize_tokenizer(pretrained_model_name, vocab_file, merges_file, use_fast, trust_remote_code)
+ assert self.tokenizer, "tokenizer not initialized"
+ except Exception:
+ try:
+ self._initialize_tokenizer(
+ pretrained_model_name, vocab_file, merges_file, not use_fast, trust_remote_code
+ )
+ assert self.tokenizer, "tokenizer not initialized"
+ except Exception as e:
+ raise ValueError(
+ f"Unable to instantiate HuggingFace AUTOTOKENIZER for {pretrained_model_name}. Exception: {e}"
+ )
+
+ self.include_special_tokens = include_special_tokens
+ self.original_vocab_size = len(self.tokenizer)
+ special_tokens_dict = {}
+
+ # # setting special tokens, by default the default model's special tokens will be preserved
+ # # unless passes new values to the special tokens
+ if unk_token is not None:
+ special_tokens_dict["unk_token"] = unk_token
+ if mask_token is not None:
+ special_tokens_dict["mask_token"] = mask_token
+ if pad_token is not None:
+ special_tokens_dict["pad_token"] = pad_token
+
+ # if the model does not have eos_token but has sep_token,
+ # set eos_token = sep_token, and vice versa
+ if sep_token is not None:
+ special_tokens_dict["sep_token"] = sep_token
+ elif self.tokenizer.sep_token is None and self.tokenizer.eos_token:
+ special_tokens_dict["sep_token"] = self.tokenizer.eos_token
+ if eos_token is not None:
+ special_tokens_dict["eos_token"] = eos_token
+ elif self.tokenizer.eos_token is None and self.tokenizer.sep_token:
+ special_tokens_dict["eos_token"] = self.tokenizer.sep_token
+
+ # if the model does not have bos_token but has cls_token,
+ # set bos_token = cls_token, and vice versa
+ if bos_token is not None:
+ special_tokens_dict["bos_token"] = bos_token
+ elif self.tokenizer.bos_token is None and self.tokenizer.cls_token:
+ special_tokens_dict["bos_token"] = self.tokenizer.cls_token
+ if cls_token is not None:
+ special_tokens_dict["cls_token"] = cls_token
+ elif self.tokenizer.cls_token is None and self.tokenizer.bos_token:
+ special_tokens_dict["cls_token"] = self.tokenizer.bos_token
+
+ # add additional special tokens (not standard special tokens such as bos, eod, sep)
+ if additional_special_tokens is not None:
+ special_tokens_dict["additional_special_tokens"] = additional_special_tokens
+
+ new_tokens_in_vocab = []
+ for token in [mask_token, bos_token, eos_token, pad_token, sep_token, cls_token, unk_token]:
+ if token is not None and token not in self.tokenizer.get_vocab():
+ new_tokens_in_vocab.append(token)
+ for token in additional_special_tokens:
+ if token is not None and token not in self.tokenizer.get_vocab():
+ new_tokens_in_vocab.append(token)
+
+ if len(new_tokens_in_vocab) > 0:
+ """
+ Special tokens that were not previously included in the tokenizer's vocabulary file will be added to
+ the vocabulary and, as a result, the model should be resized, for example:
+
+ # define your model
+ pretrained_model_name = 'roberta-base'
+ model = nemo_nlp.modules.get_lm_model(pretrained_model_name=pretrained_model_name)
+
+ # define pretrained tokenizer
+ tokenizer_default = nemo_nlp.modules.get_tokenizer(tokenizer_name=pretrained_model_name)
+
+ special_tokens = {'bos_token': '',
+ 'cls_token': '',
+ 'additional_special_tokens': ['', '']}
+ tokenizer_default.add_special_tokens(special_tokens_dict=special_tokens)
+
+ # resize your model so that the embeddings for newly added tokens are updated during training/finetuning
+ model.resize_token_embeddings(tokenizer_default.vocab_size)
+
+ See NLP_Tokenizers.ipynb for more details.
+ """
+ logging.warning(
+ f"{new_tokens_in_vocab} \n will be added to the vocabulary.\n"
+ f"Please resize your model accordingly, "
+ f"see NLP_Tokenizers.ipynb for more details."
+ )
+ self.add_special_tokens(special_tokens_dict)
+ self.space_sensitive = self.text_to_tokens("x y") != self.text_to_tokens("x") + self.text_to_tokens("y")
+ self._inv_vocab_dict = {}
+
+ def _initialize_tokenizer(
+ self,
+ pretrained_model_name: str,
+ vocab_file: str | None = None,
+ merges_file: str | None = None,
+ use_fast: bool | None = False,
+ trust_remote_code: bool | None = False,
+ ):
+ # this logic deals with different huggingface tokenizers having different positional args
+ if vocab_file is None:
+ self.tokenizer = AUTOTOKENIZER.from_pretrained(
+ pretrained_model_name_or_path=pretrained_model_name,
+ use_fast=use_fast,
+ trust_remote_code=trust_remote_code,
+ )
+ elif merges_file is None:
+ self.tokenizer = AUTOTOKENIZER.from_pretrained(
+ pretrained_model_name_or_path=pretrained_model_name,
+ vocab_file=vocab_file,
+ use_fast=use_fast,
+ trust_remote_code=trust_remote_code,
+ )
+ else:
+ self.tokenizer = AUTOTOKENIZER.from_pretrained(
+ pretrained_model_name_or_path=pretrained_model_name,
+ vocab_file=vocab_file,
+ merges_file=merges_file,
+ use_fast=use_fast,
+ trust_remote_code=trust_remote_code,
+ )
+
+ @property
+ def vocab_size(self):
+ """
+ Returns the size of the tokenizer's vocabulary.
+
+ Returns:
+ int: The number of tokens in the vocabulary.
+ """
+ return len(self.tokenizer)
+
+ def add_special_tokens(self, special_tokens_dict: dict) -> int:
+ """
+ Adds a dictionary of special tokens (eos, pad, cls...). If special tokens are NOT in the vocabulary, they are
+ added to it (indexed starting from the last index of the current vocabulary).
+
+ Args:
+ special_tokens_dict: dict of string. Keys should be in the list of predefined special attributes:
+ [``bos_token``, ``eos_token``, ``unk_token``, ``sep_token``, ``pad_token``, ``cls_token``,
+ ``mask_token``, ``additional_special_tokens``].
+ Tokens are only added if they are not already in the vocabulary.
+
+ Returns:
+ Number of tokens added to the vocabulary.
+ """
+ num_tokens_added = self.tokenizer.add_special_tokens(special_tokens_dict)
+
+ if num_tokens_added > 0:
+ logging.info(f"{num_tokens_added} special tokens added, resize your model accordingly.")
+ for k in self.tokenizer.SPECIAL_TOKENS_ATTRIBUTES:
+ setattr(self, k, getattr(self.tokenizer, k, None))
+ return num_tokens_added
+
+ def text_to_tokens(self, text):
+ """
+ Converts text into a list of tokens.
+
+ Args:
+ text (str): Input text to be tokenized.
+
+ Returns:
+ List[str]: List of tokens.
+ """
+ tokens = self.tokenizer.tokenize(text)
+ return tokens
+
+ def tokens_to_ids(self, tokens):
+ """
+ Converts a list of tokens to their corresponding IDs.
+
+ Args:
+ tokens (List[str]): List of tokens to convert.
+
+ Returns:
+ List[int]: List of token IDs.
+ """
+ ids = self.tokenizer.convert_tokens_to_ids(tokens)
+ return ids
+
+ def ids_to_tokens(self, ids):
+ """
+ Converts a list of token IDs back to tokens.
+
+ Args:
+ ids (List[int]): List of token IDs to convert.
+
+ Returns:
+ List[str]: List of tokens.
+ """
+ tokens = self.tokenizer.convert_ids_to_tokens(ids)
+ return tokens
+
+ def text_to_ids(self, text):
+ """
+ Converts text directly to token IDs.
+
+ Args:
+ text (str): Input text to be converted to IDs.
+
+ Returns:
+ List[int]: List of token IDs. If include_special_tokens is True, will include special tokens from the
+ tokenizer's configuration.
+ """
+ if self.include_special_tokens:
+ return self.tokenizer(text).input_ids
+ tokens = self.text_to_tokens(text)
+ ids = self.tokens_to_ids(tokens)
+ return ids
+
+ def apply_chat_template(self, *args, **kwargs):
+ """Applies chat template and tokenizes results"""
+ return self.tokenizer.apply_chat_template(*args, **kwargs)
+
+ @property
+ def vocab(self):
+ """
+ Returns the vocabulary as a list where the index corresponds to the token ID.
+
+ Returns:
+ List[str]: List of tokens in the vocabulary.
+ """
+ id2vocab = {v: k for k, v in self.tokenizer.vocab.items()}
+ return [id2vocab[i] for i in range(len(id2vocab))]
+
+ @property
+ def pad_id(self):
+ """
+ Gets the ID of the padding token.
+
+ Returns:
+ int or None: The ID of the padding token if it exists, None otherwise.
+ """
+ if getattr(self, "pad_token") is None:
+ return None
+ return self.tokens_to_ids([getattr(self, "pad_token")])[0]
+
+ @property
+ def bos_id(self):
+ """
+ Gets the ID of the beginning-of-sequence token.
+
+ Returns:
+ int or None: The ID of the BOS token if it exists, None otherwise.
+ """
+ if getattr(self, "bos_token") is None:
+ return None
+ return self.tokens_to_ids([getattr(self, "bos_token")])[0]
+
+ @property
+ def eos_id(self):
+ """
+ Gets the ID of the end-of-sequence token.
+
+ Returns:
+ int or None: The ID of the EOS token if it exists, None otherwise.
+ """
+ if getattr(self, "eos_token") is None:
+ return None
+ return self.tokens_to_ids([getattr(self, "eos_token")])[0]
+
+ @property
+ def eod(self):
+ """
+ Gets the ID of the end-of-document token (same as EOS token). Required for megatron-core compatibility.
+
+ Returns:
+ int: The ID of the EOD/EOS token.
+ """
+ return self.tokens_to_ids([getattr(self, "eos_token")])[0]
+
+ @property
+ def sep_id(self):
+ """
+ Gets the ID of the separator token.
+
+ Returns:
+ int or None: The ID of the separator token if it exists, None otherwise.
+ """
+ if getattr(self, "sep_token") is None:
+ return None
+ return self.tokens_to_ids([getattr(self, "sep_token")])[0]
+
+ @property
+ def cls_id(self):
+ """
+ Gets the ID of the classifier token.
+
+ Returns:
+ int or None: The ID of the classifier token if it exists, None otherwise.
+ """
+ if getattr(self, "cls_token") is None:
+ return None
+ return self.tokens_to_ids([getattr(self, "cls_token")])[0]
+
+ @property
+ def unk_id(self):
+ """
+ Gets the ID of the unknown token.
+
+ Returns:
+ int or None: The ID of the unknown token if it exists, None otherwise.
+ """
+ if getattr(self, "unk_token") is None:
+ return None
+ return self.tokens_to_ids([getattr(self, "unk_token")])[0]
+
+ @property
+ def mask_id(self):
+ """
+ Gets the ID of the mask token.
+
+ Returns:
+ int or None: The ID of the mask token if it exists, None otherwise.
+ """
+ if getattr(self, "mask_token") is None:
+ return None
+ return self.tokens_to_ids([getattr(self, "mask_token")])[0]
+
+ @property
+ def name(self):
+ """
+ Returns the name of the underlying HuggingFace tokenizer class.
+
+ Returns:
+ str: Name of the tokenizer class.
+ """
+ return type(self.tokenizer).__name__
+
+ def save_pretrained(self, save_directory: str):
+ """Saves tokenizer's vocabulary and other artifacts to the specified directory"""
+ return self.tokenizer.save_pretrained(save_directory)
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_code2wav.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_code2wav.py
new file mode 100644
index 000000000..3ceaf93de
--- /dev/null
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_code2wav.py
@@ -0,0 +1,249 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Stage-2 Code2Wav for NemotronVoiceChat: RVQ-VAE codec decode to 22.05 kHz PCM.
+
+``LLM_GENERATION`` stage consuming the talker's per-frame 31-quantizer code
+stacks and decoding them with the vendored NeMo RVQ-VAE codec. The decode runs
+in fp32 (the NeMo reference wraps codec decode in an fp32 context; bf16 collapses
+the reconstruction).
+
+Input contract (mirrors ``PersonaPlexCode2Wav``): the connector payload carries
+``{"codes": {"audio": Tensor[frames, num_quantizers]}}`` per request; the
+scheduler-side ``input_ids`` are placeholders. Codes must lie in
+``[0, codebook_size)`` — out-of-range codes are a hard error (an early bf16-drift
+signal upstream), never clamped.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Mapping
+from typing import Any
+
+import torch
+import torch.nn as nn
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+
+logger = init_logger(__name__)
+
+
+def validate_code_stack(codes: torch.Tensor, num_quantizers: int, codebook_size: int) -> torch.Tensor:
+ """Validate and reshape a code payload to ``[frames, num_quantizers]``.
+
+ Out-of-range codes are a hard error (an early numeric-drift signal
+ upstream), never clamped.
+ """
+ q = num_quantizers
+ if codes.dim() == 1:
+ if codes.numel() % q != 0:
+ raise ValueError(
+ f"NemotronVoiceChat code2wav received a flat code sequence of length "
+ f"{codes.numel()} which is not divisible by num_quantizers={q}."
+ )
+ codes = codes.reshape(-1, q)
+ if codes.dim() != 2 or codes.shape[-1] != q:
+ raise ValueError(f"NemotronVoiceChat code2wav expects codes shaped [frames, {q}], got {tuple(codes.shape)}.")
+ codes = codes.to(dtype=torch.long)
+ bad = (codes < 0) | (codes >= codebook_size)
+ if bool(bad.any()):
+ bad_vals = codes[bad]
+ raise ValueError(
+ "NemotronVoiceChat code2wav received out-of-range codec codes "
+ f"(min={int(bad_vals.min())}, max={int(bad_vals.max())}, count={int(bad.sum())}); "
+ f"valid range is [0, {codebook_size}). Refusing to clamp — this usually "
+ "signals numeric drift or a corrupted stage payload upstream."
+ )
+ return codes
+
+
+class NemotronVoiceChatCode2Wav(nn.Module):
+ """RVQ-VAE codec decoder stage (GenerationModelRunner)."""
+
+ input_modalities = "audio"
+
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+ super().__init__()
+ self.vllm_config = vllm_config
+ self.config = vllm_config.model_config.hf_config
+ codec_cfg = getattr(self.config, "code2wav_config", None)
+ self._num_quantizers = int(getattr(codec_cfg, "num_quantizers", 31))
+ self._codebook_size = int(getattr(codec_cfg, "codebook_size", 1024))
+ self._sample_rate = int(getattr(codec_cfg, "sample_rate", 22050))
+ # Samples per 12.5 Hz frame (streaming chunks slice new audio by frame).
+ self._wav_per_frame = int(getattr(codec_cfg, "wav_to_token_ratio", 1764))
+
+ # Runner-facing capability flags (same contract as PersonaPlexCode2Wav).
+ self.have_multimodal_outputs = True
+ self.has_preprocess = False
+ self.has_postprocess = False
+ self.enable_update_additional_information = True
+ self.requires_raw_input_tokens = True
+
+ # Constructed in load_weights (owns the tts_model.audio_codec.* subtree).
+ self.audio_codec: nn.Module | None = None
+
+ # ------------------------------------------------------------------
+ # Runner-facing placeholder hooks.
+ # ------------------------------------------------------------------
+ def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
+ return torch.zeros((input_ids.shape[0], 1), device=input_ids.device, dtype=torch.float32)
+
+ def compute_logits(self, hidden_states: Any, sampling_metadata: Any = None) -> None:
+ return None
+
+ def get_dummy_runtime_additional_information(self, num_reqs: int) -> list[dict[str, object]]:
+ return [{"meta": {"nemotron_voicechat_dummy_profile": True}} for _ in range(num_reqs)]
+
+ # ------------------------------------------------------------------
+ # Decode.
+ # ------------------------------------------------------------------
+ def _codes_from_runtime_info(self, runtime_info: Mapping[str, Any] | None) -> torch.Tensor | None:
+ """Extract the codes payload; None means "no codes info at all".
+
+ An EMPTY tensor is returned as-is: it is an explicit empty payload
+ (e.g. zero acoustic frames after the prompt trim) and must yield empty
+ audio, NOT a fallback decode of the placeholder input_ids.
+ """
+ if not isinstance(runtime_info, Mapping):
+ return None
+ if "codes" in runtime_info:
+ codes = runtime_info["codes"]
+ if isinstance(codes, Mapping):
+ codes = codes.get("audio")
+ else:
+ codes = runtime_info.get("codes.audio")
+ if isinstance(codes, list | tuple):
+ codes = torch.as_tensor(codes, dtype=torch.long)
+ if isinstance(codes, torch.Tensor):
+ return codes
+ return None
+
+ @torch.no_grad()
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ positions: torch.Tensor | None = None,
+ intermediate_tensors: Any = None,
+ inputs_embeds: torch.Tensor | None = None,
+ runtime_additional_information: list[dict[str, Any]] | None = None,
+ **kwargs: Any,
+ ) -> OmniOutput:
+ sr_tensor = torch.tensor(self._sample_rate, dtype=torch.int32)
+ empty = torch.zeros((0,), dtype=torch.float32)
+
+ infos = runtime_additional_information or []
+ num_req = max(len(infos), 1)
+ audios: list[torch.Tensor] = [empty] * num_req
+ srs = [sr_tensor] * num_req
+
+ if self.audio_codec is None:
+ raise RuntimeError("NemotronVoiceChatCode2Wav.forward called before load_weights().")
+ device = next(self.audio_codec.parameters()).device
+
+ for i in range(num_req):
+ runtime_info = infos[i] if i < len(infos) else None
+ meta = runtime_info.get("meta") if isinstance(runtime_info, Mapping) else None
+ if isinstance(meta, Mapping) and meta.get("nemotron_voicechat_dummy_profile") is True:
+ continue
+ codes = self._codes_from_runtime_info(runtime_info)
+ if codes is None:
+ # No input_ids fallback (unlike PersonaPlexCode2Wav): this
+ # pipeline's scheduler-side prompt is all-zero placeholders, so
+ # decoding it would produce plausible-sounding garbage whenever
+ # n_frames % num_quantizers == 0. A missing payload is a bug.
+ raise ValueError(
+ "NemotronVoiceChat code2wav request carries no 'codes' payload in its "
+ "runtime info; the talker stage must ship the [frames, num_quantizers] "
+ "code stack (placeholder input_ids are never decoded)."
+ )
+ codes = validate_code_stack(codes, self._num_quantizers, self._codebook_size).to(device=device)
+ frames = codes.shape[0]
+ if frames == 0:
+ continue
+ # Async-chunk streaming: the payload carries the CUMULATIVE trimmed
+ # code history plus meta.left_context_size = frames whose audio was
+ # already emitted by earlier chunks. Mirroring NeMo's
+ # decode_one_audio_step (unbounded window), we re-decode the whole
+ # prefix and slice off only the new tail — seam-free by
+ # construction. Sync full-payload chunks carry no left context.
+ left_context_frames = 0
+ if isinstance(meta, Mapping) and meta.get("left_context_size") is not None:
+ left_context_frames = max(int(meta["left_context_size"]), 0)
+ elif isinstance(runtime_info, Mapping) and runtime_info.get("meta.left_context_size") is not None:
+ left_context_frames = max(int(runtime_info["meta.left_context_size"]), 0)
+ if left_context_frames >= frames:
+ continue # terminal empty chunk: all frames already emitted
+ lens = torch.tensor([frames], device=device, dtype=torch.long)
+ # NeMo decodes the codec strictly in fp32.
+ with torch.autocast(device_type=device.type, enabled=False):
+ wav, wav_len = self.audio_codec.decode(codes.unsqueeze(0), lens)
+ wav = wav.reshape(-1)[: int(wav_len.reshape(-1)[0])]
+ if left_context_frames > 0:
+ wav = wav[left_context_frames * self._wav_per_frame :]
+ audios[i] = wav.detach().to(dtype=torch.float32).cpu()
+
+ return OmniOutput(
+ text_hidden_states=None,
+ multimodal_outputs={"model_outputs": audios, "sr": srs},
+ )
+
+ def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput | tuple, **kwargs: Any) -> OmniOutput:
+ if isinstance(model_outputs, OmniOutput):
+ return model_outputs
+ raise TypeError(f"NemotronVoiceChatCode2Wav expected OmniOutput, got {type(model_outputs)}")
+
+ # ------------------------------------------------------------------
+ # Weights.
+ # ------------------------------------------------------------------
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+ """Build the vendored RVQ-VAE codec and load ``tts_model.audio_codec.*``.
+
+ The vLLM safetensors iterator streams the full unified checkpoint; only
+ the codec subtree is materialized here (everything else is skipped
+ without copying). The codec stays fp32.
+ """
+ from vllm_omni.model_executor.models.nemotron_voicechat.nemo_vendored.ear_tts_vae_codec import (
+ RVQVAEModel,
+ )
+
+ codec_cfg = dict(getattr(self.config, "tts_cfg", {}).get("codec_config") or {})
+ if not codec_cfg:
+ raise ValueError(
+ "NemotronVoiceChat checkpoint config lacks "
+ "'model.speech_generation.model.codec_config'; cannot build the RVQ-VAE codec."
+ )
+ from omegaconf import DictConfig
+
+ codec = RVQVAEModel(DictConfig(codec_cfg))
+
+ prefix = "tts_model.audio_codec."
+ state: dict[str, torch.Tensor] = {}
+ for name, tensor in weights:
+ if name.startswith(prefix):
+ state[name.removeprefix(prefix)] = tensor.to(dtype=torch.float32)
+ if not state:
+ raise ValueError(
+ f"No '{prefix}*' tensors found in the checkpoint; the NemotronVoiceChat "
+ "unified model.safetensors is required for the code2wav stage."
+ )
+ missing, unexpected = codec.load_state_dict(state, strict=False)
+ if missing:
+ raise ValueError(f"NemotronVoiceChat code2wav is missing codec weights: {sorted(missing)[:10]} ...")
+ if unexpected:
+ logger.warning(
+ "NemotronVoiceChat code2wav ignoring %d unexpected codec tensors (e.g. %s)",
+ len(unexpected),
+ sorted(unexpected)[:5],
+ )
+ device = self.vllm_config.device_config.device
+ codec = codec.to(device=device, dtype=torch.float32).eval()
+ self.audio_codec = codec
+ logger.info(
+ "NemotronVoiceChat code2wav loaded RVQ-VAE codec: %d tensors, num_quantizers=%d, sr=%d",
+ len(state),
+ self._num_quantizers,
+ self._sample_rate,
+ )
+ return {name for name, _ in self.named_parameters()}
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_talker.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_talker.py
new file mode 100644
index 000000000..f4f98ab00
--- /dev/null
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_talker.py
@@ -0,0 +1,490 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Stage-1 talker for NemotronVoiceChat: frame-locked text -> 31-quantizer codes.
+
+``LLM_AR`` stage driven per frame by the vLLM engine, with the vendored NeMo
+EAR-TTS model (28-layer Gemma3-style backbone + MoG head) executed inside the
+per-request ``preprocess`` hook. Classifier-free guidance (internal batch
+doubling) and MoG sampling happen inside the vendored forward exactly as in
+NeMo — vLLM's sampler never touches the audio codes.
+
+Execution contract (mirrors the NeMo ``NemotronVoiceChat.offline_inference``
+TTS half):
+
+* The request's ``additional_information`` carries the full frame-aligned text
+ timeline (built by ``thinker2talker_token_only``), INCLUDING the PAD prompt
+ region — the talker's KV state at the first acoustic frame depends on having
+ processed it (NeMo trims prompt-region codes only after the loop).
+* vLLM prompt = one placeholder token (timeline position 0; NeMo's loop starts
+ at t=1). Each decode step t runs ``infer_codes_one_step`` with
+ ``current_subword_id=timeline[t]``, ``prev_subword_id=timeline[t-1]`` (the
+ first step uses the warmup's last init subword, as in NeMo), and the previous
+ step's code stack; the resulting ``[1, 31]`` codes accumulate under
+ ``("codes","audio")`` in the request state for the full-payload producer.
+* Warmup: ``set_init_inputs(speaker_name=...)`` (pre-baked "Aria" latent) +
+ one ``tts_model(**init_inputs)`` forward seeds ``past_key_values`` and the
+ initial code, exactly as in the reference.
+* Completion: after consuming timeline position ``T-1`` the stage emits a stop
+ placeholder token (the deploy yaml sets ``stop_token_ids: [1]``); codes for
+ positions ``1..T-1`` are shipped and the producer trims the prompt region.
+
+Heavyweight per-request state (HF ``past_key_values`` etc.) lives in a
+model-side session dict keyed by request id (PersonaPlex pattern), cleaned up
+via ``on_requests_finished``.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+import torch.nn as nn
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+
+from vllm_omni.model_executor.models.nemotron_voicechat.runtime_info import (
+ merge_runtime_info,
+ require_request_id,
+)
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+
+logger = init_logger(__name__)
+
+_CONTINUE_TOKEN = 0
+_STOP_TOKEN = 1
+
+
+def sanitize_tts_model_cfg(tts_model_cfg: dict[str, Any]) -> dict[str, Any]:
+ """Make the vendored DuplexEARTTS construction config-only.
+
+ Every parameter comes from the unified checkpoint streamed in
+ ``load_weights``; the vendored constructors must never download or load
+ external pretrained weights. Returns a deep copy so the shared stage
+ config is never mutated:
+
+ * ``tts_config.pretrained_text_name`` would make ``RVQEARTTSModel``
+ download a full LLM backbone — dropped (the backbone builds from
+ ``backbone_type``/``backbone_config``).
+ * ``pretrained_codec_model`` would make ``setup_audio_codec`` load an
+ external codec checkpoint at construction — dropped (the codec weights
+ arrive with the ``tts_model.audio_codec.*`` subtree, and the RVQ
+ embedding binding is overwritten by the same ``load_state_dict``).
+ * ``tts_config.cas_config.pretrained_tokenizer_name`` is dropped, matching
+ NeMo's own inference override (the CAS vocab is baked into the weights).
+ * ``pretrained_lm_name`` survives ONLY as the tokenizer reference (the
+ context-LM load path is disabled by ``context_hidden_size: null``); the
+ ``NEMOTRON_VOICECHAT_LLM_PATH`` override keeps it air-gap friendly.
+ """
+ import copy
+ import os
+
+ cfg = copy.deepcopy(tts_model_cfg)
+ cfg.pop("pretrained_codec_model", None)
+ tts_config = cfg.get("tts_config")
+ if isinstance(tts_config, dict):
+ tts_config.pop("pretrained_text_name", None)
+ cas_cfg = tts_config.get("cas_config")
+ if isinstance(cas_cfg, dict):
+ cas_cfg.pop("pretrained_tokenizer_name", None)
+ llm_override = os.environ.get("NEMOTRON_VOICECHAT_LLM_PATH")
+ if llm_override:
+ cfg["pretrained_lm_name"] = llm_override
+ return cfg
+
+
+class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module):
+ """vLLM AR stage wrapping the vendored NeMo EAR-TTS per-frame step."""
+
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+ super().__init__()
+ self.vllm_config = vllm_config
+ self.config = vllm_config.model_config.hf_config
+ talker_cfg = getattr(self.config, "talker_config", None)
+ self._hidden = int(getattr(talker_cfg, "hidden_size", 1152))
+ self._vocab = max(int(getattr(talker_cfg, "vocab_size", 1024)), 2)
+ self._dtype = getattr(vllm_config.model_config, "dtype", torch.bfloat16)
+ self._speaker_name = getattr(self.config, "inference_speaker_name", "Aria") or "Aria"
+
+ # Omni AR runner contract flags.
+ self.have_multimodal_outputs = True
+ self.has_preprocess = True
+ self.has_postprocess = False
+ self.requires_full_prefix_cached_hidden_states = False
+
+ # Stateful per-request execution is only validated at batch size 1
+ # (the shipped offline scope); fail fast on silent multi-request use.
+ max_num_seqs = int(getattr(vllm_config.scheduler_config, "max_num_seqs", 1))
+ if max_num_seqs != 1:
+ raise NotImplementedError(
+ f"NemotronVoiceChat talker supports max_num_seqs=1 only (got {max_num_seqs}); "
+ "the vendored EAR-TTS session state is not batch-safe."
+ )
+
+ # Vendored DuplexEARTTS; constructed in load_weights.
+ self.tts: nn.Module | None = None
+ # request_id -> mutable session state (past_key_values, prev code, ...).
+ self._sessions: dict[str, dict[str, Any]] = {}
+
+ # ------------------------------------------------------------------
+ # vLLM model protocol (placeholder token loop).
+ # ------------------------------------------------------------------
+ def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
+ return torch.zeros((input_ids.shape[0], self._hidden), device=input_ids.device, dtype=self._dtype)
+
+ def forward(
+ self,
+ input_ids: torch.Tensor | None,
+ positions: torch.Tensor,
+ intermediate_tensors: Any = None,
+ inputs_embeds: torch.Tensor | None = None,
+ **_: Any,
+ ) -> torch.Tensor:
+ # The per-frame compute happens in preprocess; the "hidden state" only
+ # transports the stop flag (column 0) to compute_logits.
+ if inputs_embeds is not None:
+ return inputs_embeds
+ assert input_ids is not None
+ return self.embed_input_ids(input_ids)
+
+ def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, sampling_metadata: Any = None) -> torch.Tensor:
+ if isinstance(hidden_states, OmniOutput):
+ hidden_states = hidden_states.text_hidden_states
+ hs = hidden_states
+ flag = hs[..., 0].to(torch.float32) # 1.0 => finished
+ logits = torch.full((hs.shape[0], self._vocab), -1.0e4, device=hs.device, dtype=torch.float32)
+ logits[:, _CONTINUE_TOKEN] = (1.0 - flag) * 10.0
+ logits[:, _STOP_TOKEN] = flag * 10.0
+ return logits
+
+ def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput:
+ if isinstance(model_outputs, OmniOutput):
+ return model_outputs
+ hidden = model_outputs
+ info_dicts = kwargs.get("model_intermediate_buffer") or kwargs.get("runtime_additional_information") or []
+ codes_list: list[torch.Tensor] = []
+ prompt_len: int | None = None
+ for info in info_dicts:
+ if not isinstance(info, dict):
+ continue
+ codes = info.get("codes", {})
+ audio = codes.get("audio") if isinstance(codes, dict) else None
+ if isinstance(audio, torch.Tensor):
+ codes_list.append(audio)
+ meta = info.get("meta")
+ reported = meta.get("nvc_logical_prompt_len") if isinstance(meta, dict) else None
+ if reported is None:
+ reported = info.get("meta.nvc_logical_prompt_len")
+ if reported is not None:
+ prompt_len = int(reported)
+ if not codes_list:
+ return OmniOutput(text_hidden_states=hidden, multimodal_outputs={})
+ outputs: dict[str, Any] = {"codes": {"audio": torch.cat(codes_list, dim=0)}}
+ # The prompt length must ride the WIRE payload (multimodal_outputs):
+ # the async-chunk dispatcher hands exactly this dict to the
+ # talker2code2wav producer, while request.additional_information is
+ # overwritten concurrently by arriving thinker chunks.
+ if prompt_len is not None:
+ # 1-D on purpose: the wire-payload builder indexes element.shape[0].
+ outputs["meta"] = {"nvc_logical_prompt_len": torch.tensor([prompt_len], dtype=torch.long)}
+ return OmniOutput(
+ text_hidden_states=hidden,
+ multimodal_outputs=outputs,
+ )
+
+ # ------------------------------------------------------------------
+ # Per-frame vendored TTS step.
+ # ------------------------------------------------------------------
+ def _try_timeline(self, info: dict[str, Any]) -> torch.Tensor | None:
+ """Extract the frame-locked text timeline from either payload shape.
+
+ Sync mode ships ``nvc_text_timeline`` (+ ``nvc_logical_prompt_len``) in
+ additional_information; async-chunk mode replaces the request info with
+ the connector payload each chunk, carrying the CUMULATIVE timeline under
+ ``ids.all`` (and the logical prompt length under
+ ``meta.num_processed_tokens``). Returns None when neither is present.
+ """
+ timeline = info.get("nvc_text_timeline")
+ if timeline is not None and info.get("nvc_logical_prompt_len") is None:
+ raise ValueError(
+ "NemotronVoiceChat talker got 'nvc_text_timeline' without "
+ "'nvc_logical_prompt_len'; the code2wav producer needs the prompt "
+ "length to trim prompt-region codes."
+ )
+ if timeline is None:
+ ids = info.get("ids")
+ timeline = ids.get("all") if isinstance(ids, dict) else None
+ if timeline is None:
+ timeline = info.get("ids.all")
+ if timeline is None:
+ return None
+ if isinstance(timeline, torch.Tensor):
+ return timeline.reshape(-1).to(torch.long)
+ return torch.as_tensor(list(timeline), dtype=torch.long)
+
+ @staticmethod
+ def _logical_prompt_len(info: dict[str, Any]) -> int | None:
+ """Logical prompt length from either payload shape (sync/async)."""
+ value = info.get("nvc_logical_prompt_len")
+ if value is None:
+ meta = info.get("meta")
+ value = meta.get("num_processed_tokens") if isinstance(meta, dict) else None
+ if value is None:
+ value = info.get("meta.num_processed_tokens")
+ return int(value) if value is not None else None
+
+ @staticmethod
+ def _upstream_finished(info: dict[str, Any]) -> bool:
+ """True when the async producer marked its final chunk (meta.finished)."""
+ meta = info.get("meta")
+ flag = meta.get("finished") if isinstance(meta, dict) else None
+ if flag is None:
+ flag = info.get("meta.finished")
+ if isinstance(flag, torch.Tensor):
+ return bool(flag.reshape(-1)[0].item()) if flag.numel() else False
+ return bool(flag)
+
+ def _timeline(self, info: dict[str, Any]) -> torch.Tensor:
+ timeline = self._try_timeline(info)
+ if timeline is None:
+ # Explicit failure — no prompt_token_ids fallback: an implicit
+ # timeline would silently skip the prompt-region trim downstream.
+ raise ValueError(
+ "NemotronVoiceChat talker request is missing its timeline metadata "
+ "('nvc_text_timeline' in additional_information, or 'ids.all' on the "
+ "async-chunk payload); the stage cannot synthesize speech without "
+ "the thinker's frame-locked tokens."
+ )
+ return timeline
+
+ def _init_session(self, request_id: str, info: dict[str, Any], device: torch.device) -> dict[str, Any]:
+ assert self.tts is not None
+ timeline = self._timeline(info).to(device)
+ if timeline.numel() < 2:
+ raise ValueError(
+ f"NemotronVoiceChat talker timeline for request {request_id} has "
+ f"{timeline.numel()} positions; need at least 2 (prompt + one frame)."
+ )
+ # Warmup: pre-baked speaker latent + init prompt forward (NeMo
+ # offline_inference lines around get_init_inputs/first forward).
+ self.tts.set_init_inputs(speaker_name=self._speaker_name)
+ init_inputs = self.tts.get_init_inputs(B=1)
+ generation_config = self.tts._get_generation_config(self._guidance_enabled())
+ init_inputs.update({"use_cache": True, "past_key_values": None, "guidance_enabled": self._guidance_enabled()})
+ with torch.inference_mode():
+ outputs = self.tts.tts_model(**init_inputs)
+ prompt_len = self._logical_prompt_len(info)
+ if prompt_len is None:
+ raise ValueError(
+ "NemotronVoiceChat talker request is missing its logical prompt length "
+ "('nvc_logical_prompt_len' in additional_information, or async "
+ "meta.num_processed_tokens); the code2wav producer cannot trim the "
+ "prompt-region codes without it."
+ )
+ session = {
+ "timeline": timeline,
+ "prompt_len": prompt_len,
+ "step": 1, # NeMo's loop starts at t=1
+ "past_key_values": outputs.past_key_values,
+ "code": init_inputs["code"][:, -1:],
+ "first_context_subword_id": init_inputs["subword_ids"][:, -1].unsqueeze(-1),
+ "tts_initialized": False,
+ "generation_config": generation_config,
+ "codes_rows": [],
+ # Sync mode ships the whole timeline up front via
+ # nvc_text_timeline; async-chunk mode grows it across chunks and
+ # the last chunk sets meta.finished.
+ "sync_mode": info.get("nvc_text_timeline") is not None,
+ "upstream_finished": self._upstream_finished(info) or info.get("nvc_text_timeline") is not None,
+ }
+ self._sessions[request_id] = session
+ return session
+
+ def _refresh_session_timeline(self, session: dict[str, Any], info: dict[str, Any], device: torch.device) -> None:
+ """Adopt a longer timeline from the latest async-chunk payload."""
+ latest = self._try_timeline(info)
+ if latest is not None and latest.numel() > int(session["timeline"].numel()):
+ session["timeline"] = latest.to(device)
+ if self._upstream_finished(info):
+ session["upstream_finished"] = True
+
+ def _guidance_enabled(self) -> bool:
+ return bool(getattr(self.config, "tts_cfg", {}).get("inference_guidance_enabled", True))
+
+ def _step_session(self, session: dict[str, Any]) -> tuple[torch.Tensor, bool]:
+ """One NeMo TTS step; returns (codes_row [1, Q], finished)."""
+ assert self.tts is not None
+ timeline: torch.Tensor = session["timeline"]
+ t = int(session["step"])
+ total = int(timeline.numel())
+ if t >= total:
+ # The drain loop in preprocess never enters here; a direct call
+ # past the received timeline is a programming error. Hard-fail
+ # rather than fabricating a frame — a repeated/guessed subword
+ # would silently desync the audio from the text channel.
+ raise RuntimeError(
+ f"NemotronVoiceChat talker stepped past the received timeline: step {t} "
+ f"needs timeline position {t} but only {total} positions have arrived "
+ f"(upstream_finished={session['upstream_finished']})."
+ )
+ current_subword_id = timeline[t].reshape(1, 1)
+ if not session["tts_initialized"]:
+ prev_subword_id = session["first_context_subword_id"]
+ session["tts_initialized"] = True
+ else:
+ prev_subword_id = timeline[t - 1].reshape(1, 1)
+ current_subword_mask = torch.ones(1, 1, device=timeline.device, dtype=torch.bool)
+ with torch.inference_mode():
+ code, past_key_values = self.tts.infer_codes_one_step(
+ current_subword_id=current_subword_id,
+ prev_subword_id=prev_subword_id,
+ current_subword_mask=current_subword_mask,
+ prev_audio_tokens=session["code"],
+ past_key_values=session["past_key_values"],
+ guidance_enabled=self._guidance_enabled(),
+ generation_config=session["generation_config"],
+ ignore_eos_flag_stop=True,
+ )
+ session["code"] = code
+ session["past_key_values"] = past_key_values
+ session["step"] = t + 1
+ # Async mode: exhausting the received timeline only ends the request
+ # once the upstream marked its final chunk; otherwise more entries are
+ # in flight and the scheduler parks us until they arrive.
+ finished = (t + 1) >= int(session["timeline"].numel()) and bool(session["upstream_finished"])
+ return code.reshape(1, -1).to(torch.long), finished
+
+ def preprocess(
+ self,
+ input_ids: torch.Tensor,
+ input_embeds: torch.Tensor | None,
+ **info_dict: Any,
+ ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
+ info = merge_runtime_info(info_dict)
+ device = input_ids.device
+ span = int(input_ids.shape[0])
+ is_prefill_raw = info.get("_omni_is_prefill")
+ is_prefill = bool(is_prefill_raw) if isinstance(is_prefill_raw, bool) else span > 1
+ request_id = require_request_id(info, "talker")
+
+ embeds = torch.zeros((span, self._hidden), device=device, dtype=self._dtype)
+ if is_prefill or request_id not in self._sessions:
+ # The 1-token vLLM prefill is the session boundary: run the vendored
+ # warmup here so the first decode step is a pure NeMo t=1 step.
+ self._init_session(request_id, info, device)
+ return input_ids, embeds, {}
+
+ session = self._sessions[request_id]
+ self._refresh_session_timeline(session, info, device)
+ # Sync mode: one NeMo TTS step per engine step (the verified parity
+ # path, unchanged). Async-chunk mode: DRAIN every received timeline
+ # position on each wake. The engine tokens are placeholders
+ # (CONTINUE/STOP) and the codes ride the cumulative payload, so
+ # multiple TTS steps per engine step are safe — and necessary:
+ # * the save_async background thread can coalesce thinker steps into
+ # one chunk (fewer wakes than new positions), and
+ # * the prompt-region PAD steps (logical_prompt_len - 1 of them) all
+ # become available with chunk 0 and must not cost one upstream
+ # chunk each.
+ finished = False
+ steps_run = 0
+ while int(session["step"]) < int(session["timeline"].numel()):
+ codes_row, finished = self._step_session(session)
+ session["codes_rows"].append(codes_row.cpu())
+ steps_run += 1
+ if session["sync_mode"]:
+ break
+ if steps_run == 0:
+ # Zero-progress wake (duplicate/terminal-marker chunk): no TTS step
+ # was run, so emit a plain CONTINUE placeholder — never a
+ # fabricated frame. The request finishes only when the upstream is
+ # exhausted AND the timeline is fully consumed.
+ finished = bool(session["upstream_finished"])
+ if finished:
+ embeds[:, 0] = 1.0
+ return input_ids, embeds, {}
+ cumulative = torch.cat(session["codes_rows"], dim=0)
+ if finished:
+ embeds[:, 0] = 1.0 # stop flag -> compute_logits emits the stop token
+ info_update: dict[str, Any] = {
+ "codes": {"audio": cumulative},
+ # nvc_logical_prompt_len rides EVERY step's payload: the request's
+ # additional_information is overwritten both by arriving thinker
+ # chunks and by these model updates, so the async code2wav producer
+ # reads the prompt length from the per-step multimodal_output
+ # instead of racing the request-level dict.
+ "meta": {"nvc_tts_step": int(session["step"]), "nvc_logical_prompt_len": int(session["prompt_len"])},
+ }
+ return input_ids, embeds, info_update
+
+ def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None:
+ for request_id in finished_req_ids:
+ self._sessions.pop(str(request_id), None)
+
+ # ------------------------------------------------------------------
+ # Weights.
+ # ------------------------------------------------------------------
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+ """Build the vendored DuplexEARTTS and load the ``tts_model.*`` subtree.
+
+ The backbone/MoG run in the engine dtype (bf16 by default); the codec
+ submodule inside DuplexEARTTS stays fp32 (it is not used by this stage —
+ decode happens in code2wav — but keeping its dtype canonical avoids
+ surprises if warmup paths touch it).
+ """
+ from vllm_omni.model_executor.models.nemotron_voicechat.nemo_vendored.duplex_ear_tts import (
+ DuplexEARTTS,
+ )
+
+ tts_model_cfg = dict(getattr(self.config, "tts_cfg", {}) or {})
+ if not tts_model_cfg:
+ raise ValueError(
+ "NemotronVoiceChat checkpoint config lacks 'model.speech_generation.model'; cannot build the talker."
+ )
+ tts_model_cfg = sanitize_tts_model_cfg(tts_model_cfg)
+ tts_data = dict(getattr(self.config, "tts_data", {}) or {})
+ tts_data.setdefault("source_sample_rate", 22050)
+ tts_data.setdefault("target_sample_rate", int(getattr(self.config, "target_sample_rate", 22050)))
+ tts_data.setdefault("frame_length", float(getattr(self.config, "frame_length", 0.08)))
+ # DuplexEARTTS consumes the whole speech_generation section layout
+ # ({"data": ..., "model": ...}), matching NeMo's constructor call.
+ tts = DuplexEARTTS({"data": tts_data, "model": tts_model_cfg})
+
+ prefix = "tts_model."
+ state: dict[str, torch.Tensor] = {}
+ for name, tensor in weights:
+ if not name.startswith(prefix):
+ continue
+ target = name.removeprefix(prefix)
+ if target.startswith("audio_codec."):
+ state[target] = tensor.to(dtype=torch.float32)
+ else:
+ state[target] = tensor.to(dtype=self._dtype)
+ if not state:
+ raise ValueError(
+ "No 'tts_model.*' tensors found in the checkpoint; the NemotronVoiceChat "
+ "unified model.safetensors is required for the talker stage."
+ )
+ missing, unexpected = tts.load_state_dict(state, strict=False)
+ real_missing = [m for m in missing if not m.startswith("audio_codec.")]
+ if real_missing:
+ raise ValueError(f"NemotronVoiceChat talker is missing TTS weights: {sorted(real_missing)[:10]} ...")
+ if unexpected:
+ logger.warning(
+ "NemotronVoiceChat talker ignoring %d unexpected TTS tensors (e.g. %s)",
+ len(unexpected),
+ sorted(unexpected)[:5],
+ )
+ device = self.vllm_config.device_config.device
+ tts = tts.to(device=device)
+ tts.tts_model.to(dtype=self._dtype)
+ if hasattr(tts, "audio_codec") and tts.audio_codec is not None:
+ tts.audio_codec.to(dtype=torch.float32)
+ self.tts = tts.eval()
+ logger.info(
+ "NemotronVoiceChat talker loaded EAR-TTS: %d tensors, speaker=%s, guidance=%s",
+ len(state),
+ self._speaker_name,
+ self._guidance_enabled(),
+ )
+ return {name for name, _ in self.named_parameters()}
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_thinker.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_thinker.py
new file mode 100644
index 000000000..9f237a885
--- /dev/null
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_thinker.py
@@ -0,0 +1,596 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Stage-0 thinker for NemotronVoiceChat: 16 kHz speech -> frame-locked text.
+
+``LLM_AR`` stage combining:
+
+* the vendored NeMo FastConformer perception (mel + 24-layer encoder + linear
+ projection to the LLM width, 12.5 Hz frames);
+* the additive channel fusion from the NeMo duplex STT model
+ (``fused[t] = E(text_prev) * w_text + timeline[t] * w_audio +
+ E(func_prev) * w_func`` with the checkpoint's AddFusion weights);
+* the NemotronH (hybrid Mamba) backbone on vLLM paged execution via
+ ``init_vllm_registered_model`` — this is where vLLM earns its keep;
+* the separate ``lm_head`` (text channel, sampled by the vLLM engine) and
+ ``function_head`` (greedy, maintained internally per request).
+
+Timeline contract (verified against NeMo ``_init_inference``/``_step_zero``/
+``_step_inference``; the golden shape test guards it):
+
+* The vLLM prompt covers ``logical_prompt_token_len + 1`` positions: the system
+ prompt tokens (``[bos] + tokens + [eos]``) fused with a PAD text channel
+ (NeMo's ``_get_bos_embedding`` returns the PAD embedding, so position 0 is
+ uniform with the rest of the prompt region), plus acoustic frame 0 as the
+ final prefill position. The first engine-sampled token is therefore NeMo's
+ ``gen_text[prompt_len]`` — the first acoustic-frame prediction.
+* Each decode step t consumes conformer frame t (held in per-request state)
+ fused with the previously sampled text token and the previous greedy
+ function token. Text EOS never terminates generation (``ignore_eos``); the
+ example sets ``max_tokens = acoustic_frame_count``. Never set ``min_tokens``:
+ the tokenizer's eos_token is — the frame-locked PAD/silence
+ token — and min_tokens masks "EOS" to -inf, forbidding silence entirely.
+* The function channel exists for numerical correctness (fusion weight 2.0)
+ even though tool calling is out of scope; the greedy token timeline is
+ exposed for debugging via the request info.
+"""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Iterable
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.model_executor.models.interfaces import HasInnerState, IsHybrid
+from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix
+
+from vllm_omni.model_executor.models.nemotron_voicechat.runtime_info import (
+ merge_runtime_info,
+ require_request_id,
+)
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+
+logger = init_logger(__name__)
+
+_LM_ARCHITECTURE = "NemotronHForCausalLM"
+
+
+def validate_timeline_fits_max_model_len(
+ logical_prompt_token_len: int,
+ acoustic_frame_count: int,
+ max_model_len: int,
+) -> None:
+ """Reject audio whose frame-locked timeline exceeds the stage limit.
+
+ The vLLM request occupies ``logical_prompt_token_len + 1`` prefill
+ positions (the +1 is the acoustic-frame-0 placeholder) plus
+ ``acoustic_frame_count`` sampled tokens, so the binding sequence size is
+ one MORE than NeMo's logical timeline (``prompt + frames``).
+ """
+ logical_total_len = logical_prompt_token_len + acoustic_frame_count
+ vllm_prefill_len = logical_prompt_token_len + 1
+ vllm_total_len = vllm_prefill_len + acoustic_frame_count
+ if vllm_total_len > max_model_len:
+ raise ValueError(
+ "NemotronVoiceChat input exceeds the thinker stage's max_model_len: "
+ f"vllm_prefill_len={vllm_prefill_len} + "
+ f"acoustic_frame_count={acoustic_frame_count} = "
+ f"vllm_total_len={vllm_total_len} > max_model_len={max_model_len} "
+ f"(logical_prompt_token_len={logical_prompt_token_len}, "
+ f"logical_total_len={logical_total_len}). Use shorter audio or raise the "
+ "stage's max_model_len in the deploy yaml."
+ )
+
+
+def compute_acoustic_frame_count(nemo_stt_cfg: dict[str, Any], num_samples: int) -> int:
+ """Exact 12.5 Hz frame count for ``num_samples`` of 16 kHz audio.
+
+ Replicates the vendored mel featurizer length arithmetic followed by the
+ dw-striding conv subsampling length recurrence (NOT ``duration / 0.08``).
+ The thinker validates this against the real perception output at prefill,
+ and the golden shape test pins it against the NeMo reference.
+ """
+ from vllm_omni.model_executor.models.nemotron_voicechat.nemo_vendored.asr.subsampling import (
+ calc_length,
+ )
+
+ perception = nemo_stt_cfg.get("perception") or {}
+ pre = perception.get("preprocessor") or {}
+ encoder = perception.get("encoder") or {}
+ sample_rate = int(pre.get("sample_rate", 16000))
+ window_stride = float(pre.get("window_stride", 0.01))
+ window_size = float(pre.get("window_size", 0.025))
+ n_fft = int(pre.get("n_fft") or 2 ** int(np.ceil(np.log2(window_size * sample_rate))))
+ hop = int(window_stride * sample_rate)
+ pad_amount = n_fft // 2 * 2 # torch.stft center padding on both sides
+ mel_len = (num_samples + pad_amount - n_fft) // hop + 1
+ subsampling_factor = int(encoder.get("subsampling", "dw_striding") and encoder.get("subsampling_factor", 8))
+ conv_kernel = int(encoder.get("subsampling_conv_kernel_size", 3) or 3)
+ stride = 2
+ repeat = int(np.log2(subsampling_factor))
+ # Mirror the vendored ConvSubsampling padding exactly: causal downsampling
+ # (this checkpoint sets encoder.causal_downsampling=true) pads
+ # left=kernel-1 plus right=stride-1, i.e. one more column than the
+ # non-causal symmetric 2*(kernel//2) — which shifts the output length by
+ # one for roughly half of all input durations.
+ if bool(encoder.get("causal_downsampling", False)):
+ all_paddings = (conv_kernel - 1) + (stride - 1)
+ else:
+ all_paddings = 2 * (conv_kernel // 2)
+ out = calc_length(
+ torch.tensor([mel_len]),
+ all_paddings=all_paddings,
+ kernel_size=conv_kernel,
+ stride=stride,
+ ceil_mode=False,
+ repeat_num=repeat,
+ )
+ return int(out[0])
+
+
+class NemotronVoiceChatThinkerForConditionalGeneration(nn.Module, HasInnerState, IsHybrid):
+ """vLLM-native thinker (perception + fusion + NemotronH + dual heads).
+
+ Declares the hybrid-Mamba interfaces (``IsHybrid``/``HasInnerState``) so
+ vLLM configures the mamba state cache for the wrapped NemotronH backbone;
+ the state shape/dtype classmethods mirror ``NemotronHForCausalLM`` but read
+ the NemotronH sub-config (``get_text_config()``) instead of the top-level
+ checkpoint config.
+ """
+
+ @classmethod
+ def get_mamba_state_dtype_from_config(cls, vllm_config: VllmConfig) -> tuple[torch.dtype, torch.dtype]:
+ from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM
+
+ return NemotronHForCausalLM.get_mamba_state_dtype_from_config(vllm_config)
+
+ @classmethod
+ def get_mamba_state_shape_from_config(cls, vllm_config: VllmConfig) -> tuple[tuple[int, int], tuple[int, int, int]]:
+ from vllm.model_executor.layers.mamba.mamba_utils import MambaStateShapeCalculator
+
+ parallel_config = vllm_config.parallel_config
+ hf_config = vllm_config.model_config.hf_config.get_text_config()
+ intermediate_size = hf_config.mamba_num_heads * hf_config.mamba_head_dim
+ return MambaStateShapeCalculator.mamba2_state_shape(
+ intermediate_size=intermediate_size,
+ tp_world_size=parallel_config.tensor_parallel_size,
+ n_groups=hf_config.n_groups,
+ num_heads=hf_config.mamba_num_heads,
+ head_dim=hf_config.mamba_head_dim,
+ state_size=hf_config.ssm_state_size,
+ conv_kernel=hf_config.conv_kernel,
+ num_spec=vllm_config.num_speculative_tokens,
+ )
+
+ @classmethod
+ def get_mamba_state_copy_func(cls):
+ from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM
+
+ return NemotronHForCausalLM.get_mamba_state_copy_func()
+
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+ super().__init__()
+ self.vllm_config = vllm_config
+ config = vllm_config.model_config.hf_config
+ self.config = config
+ stt_cfg: dict[str, Any] = dict(getattr(config, "stt_cfg", {}) or {})
+ if not stt_cfg:
+ raise ValueError("NemotronVoiceChat checkpoint config lacks 'model.stt.model'; cannot build the thinker.")
+ self.stt_cfg = stt_cfg
+ text_cfg = config.get_text_config()
+ self._hidden = int(getattr(text_cfg, "hidden_size", 4480))
+ self._vocab = int(getattr(text_cfg, "vocab_size", 131072))
+ self._dtype = getattr(vllm_config.model_config, "dtype", torch.bfloat16)
+
+ # NemotronH backbone (+ lm_head) on vLLM native layers.
+ self.language_model = init_vllm_registered_model(
+ vllm_config=vllm_config,
+ hf_config=text_cfg,
+ prefix=maybe_prefix(prefix, "language_model"),
+ architectures=[_LM_ARCHITECTURE],
+ )
+
+ # Shared channel embedding + separate function head (both distinct
+ # checkpoints subtrees from the backbone).
+ self.embed_tokens = nn.Embedding(self._vocab, self._hidden)
+ self.function_head = nn.Linear(self._hidden, self._vocab, bias=False)
+
+ # Vendored perception (mel + FastConformer + projection).
+ from vllm_omni.model_executor.models.nemotron_voicechat.nemo_vendored.perception import (
+ AudioPerceptionModule,
+ )
+
+ perception_cfg = stt_cfg.get("perception")
+ if not perception_cfg:
+ raise ValueError("'model.stt.model.perception' missing from the checkpoint config.")
+ from omegaconf import DictConfig
+
+ self.perception = AudioPerceptionModule(DictConfig(perception_cfg))
+
+ # AddFusion weights (the checkpoint uses fuse_method="add" defaults;
+ # the same duplex_*_weight keys double as the fusion weights).
+ use_function_head = bool(stt_cfg.get("use_function_head", True))
+ self._w_text = float(stt_cfg.get("duplex_text_channel_weight", 1.0))
+ self._w_audio = float(stt_cfg.get("duplex_user_channel_weight", 1.0))
+ self._w_func = float(stt_cfg.get("duplex_function_channel_weight", 1.0)) if use_function_head else 0.0
+ self._use_function_head = use_function_head
+ fuse_method = stt_cfg.get("fuse_method", "add") or "add"
+ if fuse_method != "add":
+ raise NotImplementedError(
+ f"NemotronVoiceChat thinker only implements fuse_method='add' "
+ f"(this checkpoint uses it); got {fuse_method!r}."
+ )
+ if bool(stt_cfg.get("predict_user_text", False)):
+ raise NotImplementedError(
+ "NemotronVoiceChat thinker does not implement the user-ASR channel "
+ "(predict_user_text=True); this checkpoint has it disabled."
+ )
+
+ self._special_ids: dict[str, int] | None = None
+
+ # Omni AR runner contract.
+ self.have_multimodal_outputs = True
+ self.has_preprocess = True
+ self.has_postprocess = True # per-step greedy function-channel token
+ self.requires_full_prefix_cached_hidden_states = False
+ # The per-step function token is a scalar; no GPU-resident buffers needed.
+ self.gpu_resident_buffer_keys: set[tuple[str, str]] = set()
+
+ # Stateful per-request execution is only validated at batch size 1
+ # (the shipped offline scope); fail fast on silent multi-request use.
+ max_num_seqs = int(getattr(vllm_config.scheduler_config, "max_num_seqs", 1))
+ if max_num_seqs != 1:
+ raise NotImplementedError(
+ f"NemotronVoiceChat thinker supports max_num_seqs=1 only (got {max_num_seqs}); "
+ "multi-request batching of the per-request fusion state is not implemented."
+ )
+
+ # request_id -> per-request state (audio frames, prefill embeds, counters).
+ self._sessions: dict[str, dict[str, Any]] = {}
+
+ # ------------------------------------------------------------------
+ # Special token ids (from the checkpoint's token strings).
+ # ------------------------------------------------------------------
+ def _resolve_special_ids(self) -> dict[str, int]:
+ if self._special_ids is not None:
+ return self._special_ids
+ from transformers import AutoTokenizer
+
+ ref = os.environ.get("NEMOTRON_VOICECHAT_LLM_PATH") or self.stt_cfg.get(
+ "pretrained_llm", "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
+ )
+ tok = AutoTokenizer.from_pretrained(ref, trust_remote_code=False)
+ ids = {
+ "bos": tok.convert_tokens_to_ids(self.stt_cfg.get("bos_token", "")),
+ "eos": tok.convert_tokens_to_ids(self.stt_cfg.get("eos_token", "")),
+ "pad": tok.convert_tokens_to_ids(self.stt_cfg.get("pad_token", "")),
+ }
+ for name, value in ids.items():
+ if value is None or value < 0:
+ raise ValueError(f"NemotronVoiceChat could not resolve the {name} token id from {ref!r}.")
+ self._special_ids = ids
+ return ids
+
+ # ------------------------------------------------------------------
+ # vLLM model protocol.
+ # ------------------------------------------------------------------
+ def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
+ # Placeholder; preprocess replaces every position with fused embeds.
+ return torch.zeros((input_ids.shape[0], self._hidden), device=input_ids.device, dtype=self._dtype)
+
+ def forward(
+ self,
+ input_ids: torch.Tensor | None,
+ positions: torch.Tensor,
+ intermediate_tensors: Any = None,
+ inputs_embeds: torch.Tensor | None = None,
+ **_: Any,
+ ) -> torch.Tensor:
+ return self.language_model(
+ input_ids=None,
+ positions=positions,
+ intermediate_tensors=intermediate_tensors,
+ inputs_embeds=inputs_embeds,
+ )
+
+ def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, sampling_metadata: Any = None) -> torch.Tensor:
+ if isinstance(hidden_states, OmniOutput):
+ hidden_states = hidden_states.text_hidden_states
+ # vLLM >= 0.25 dropped the sampling_metadata parameter.
+ return self.language_model.compute_logits(hidden_states)
+
+ def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput:
+ if isinstance(model_outputs, OmniOutput):
+ return model_outputs
+ return OmniOutput(text_hidden_states=model_outputs, multimodal_outputs={})
+
+ # ------------------------------------------------------------------
+ # Fusion helpers.
+ # ------------------------------------------------------------------
+ def _fuse(
+ self,
+ text_ids: torch.Tensor,
+ timeline_embeds: torch.Tensor,
+ func_ids: torch.Tensor,
+ ) -> torch.Tensor:
+ """AddFusion over the three channels; shapes [T] ids, [T, H] timeline."""
+ text_e = self.embed_tokens(text_ids)
+ fused = text_e * self._w_text + timeline_embeds.to(text_e.dtype) * self._w_audio
+ if self._use_function_head:
+ fused = fused + self.embed_tokens(func_ids) * self._w_func
+ return fused
+
+ # ------------------------------------------------------------------
+ # Per-request session.
+ # ------------------------------------------------------------------
+ def _load_audio(self, info: dict[str, Any], device: torch.device) -> torch.Tensor:
+ audio = info.get("nvc_audio")
+ if audio is None:
+ raise ValueError(
+ "NemotronVoiceChat thinker request is missing its audio: pass "
+ "additional_information={'nvc_audio': <16 kHz mono float waveform>, "
+ "'nvc_sr': 16000, 'nvc_prompt_token_ids': [...]}."
+ )
+ wav = torch.as_tensor(np.asarray(audio), dtype=torch.float32, device=device).reshape(-1)
+ if wav.numel() == 0:
+ raise ValueError("NemotronVoiceChat thinker received empty audio input.")
+ sr = int(info.get("nvc_sr", 16000))
+ expected_sr = int(getattr(self.config, "source_sample_rate", 16000))
+ if sr != expected_sr:
+ raise ValueError(
+ f"NemotronVoiceChat thinker expects {expected_sr} Hz audio, got {sr} Hz — "
+ "resample client-side (the offline example does this)."
+ )
+ return wav
+
+ def _init_session(self, request_id: str, info: dict[str, Any], device: torch.device) -> dict[str, Any]:
+ pad_id = self._resolve_special_ids()["pad"]
+ wav = self._load_audio(info, device)
+
+ prompt_ids = info.get("nvc_prompt_token_ids")
+ if prompt_ids is None:
+ raise ValueError(
+ "NemotronVoiceChat thinker request is missing 'nvc_prompt_token_ids' "
+ "(the system prompt token ids, [bos] + tokens + [eos])."
+ )
+ prompt_ids = torch.as_tensor(list(prompt_ids), dtype=torch.long, device=device)
+ if prompt_ids.numel() == 0:
+ raise ValueError("NemotronVoiceChat thinker requires a non-empty system prompt.")
+
+ # Perception: mel + FastConformer + projection -> [N, hidden] frames.
+ with torch.inference_mode():
+ encoded, encoded_len = self.perception(
+ input_signal=wav.unsqueeze(0).to(next(self.perception.parameters()).dtype),
+ input_signal_length=torch.tensor([wav.numel()], device=device),
+ )
+ n_frames = int(encoded_len.reshape(-1)[0])
+ audio_embeds = encoded[0, :n_frames].to(self._dtype) # [N, H]
+ expected = info.get("nvc_expected_frames")
+ if expected is not None and int(expected) != n_frames:
+ raise ValueError(
+ f"NemotronVoiceChat frame-count mismatch: the request was sized for "
+ f"{int(expected)} acoustic frames but perception produced {n_frames}. "
+ "The example must size max_tokens with compute_acoustic_frame_count() (never min_tokens)."
+ )
+
+ p = int(prompt_ids.numel())
+ validate_timeline_fits_max_model_len(
+ logical_prompt_token_len=p,
+ acoustic_frame_count=n_frames,
+ max_model_len=int(self.vllm_config.model_config.max_model_len),
+ )
+
+ # Fused prefill embeds for positions 0..P (P prompt slots + frame 0).
+ # All prompt positions use the PAD text channel (NeMo's
+ # _get_bos_embedding returns the PAD embedding) and PAD function channel.
+ timeline = torch.cat([self.embed_tokens(prompt_ids).to(self._dtype), audio_embeds[0:1]], dim=0) # [P+1, H]
+ pad_ids = torch.full((p + 1,), pad_id, dtype=torch.long, device=device)
+ prefill = self._fuse(pad_ids, timeline, pad_ids) # [P+1, H]
+
+ session = {
+ "audio_embeds": audio_embeds,
+ "prefill_embeds": prefill,
+ "prompt_len": p,
+ "n_frames": n_frames,
+ "func_token": int(pad_id),
+ "decode_step": 0,
+ }
+ self._sessions[request_id] = session
+ return session
+
+ # ------------------------------------------------------------------
+ # Omni AR hooks.
+ # ------------------------------------------------------------------
+ def preprocess(
+ self,
+ input_ids: torch.Tensor,
+ input_embeds: torch.Tensor | None,
+ **info_dict: Any,
+ ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
+ info = merge_runtime_info(info_dict)
+ device = input_ids.device
+ span = int(input_ids.shape[0])
+ is_prefill_raw = info.get("_omni_is_prefill")
+ is_prefill = bool(is_prefill_raw) if isinstance(is_prefill_raw, bool) else span > 1
+ request_id = require_request_id(info, "thinker")
+
+ if is_prefill or request_id not in self._sessions:
+ session = self._sessions.get(request_id) or self._init_session(request_id, info, device)
+ offset = max(0, int(info.get("_omni_num_computed_tokens", 0) or 0))
+ prefill: torch.Tensor = session["prefill_embeds"]
+ if offset + span > prefill.shape[0]:
+ raise ValueError(
+ f"NemotronVoiceChat thinker prefill span [{offset}, {offset + span}) exceeds "
+ f"the fused prompt length {prefill.shape[0]}; the request prompt must be "
+ "the system prompt tokens plus exactly one placeholder position."
+ )
+ return input_ids, prefill[offset : offset + span].to(self._dtype), {}
+
+ session = self._sessions[request_id]
+ pad_id = self._resolve_special_ids()["pad"]
+
+ # Function channel: greedy argmax over the PREVIOUS position's hidden,
+ # computed by postprocess right after each forward. At the first decode
+ # step this is the prefill's last position (acoustic frame 0), matching
+ # NeMo's gen_function[prompt_len]. The cumulative debug timeline
+ # (nvc_function_tokens) is owned by postprocess, which also covers the
+ # final acoustic frame (no later preprocess exists for it).
+ if self._use_function_head:
+ prev_func = info.get("nvc_prev_function_token")
+ if prev_func is not None:
+ token = int(prev_func.reshape(-1)[0]) if isinstance(prev_func, torch.Tensor) else int(prev_func)
+ session["func_token"] = token
+
+ # Position-addressed frame index: decode of vLLM position P+k consumes
+ # acoustic frame k (frame 0 was the last prefill position). Deriving it
+ # from _omni_num_computed_tokens keeps preprocess idempotent — if the
+ # runner ever re-enters the same decode position (preemption or
+ # recompute), the audio pointer stays aligned with the text channel
+ # instead of silently advancing. Falls back to a call counter only if
+ # the runner did not inject the key.
+ computed = info.get("_omni_num_computed_tokens")
+ if computed is not None:
+ frame_idx = int(computed) - int(session["prompt_len"])
+ else:
+ frame_idx = int(session["decode_step"]) + 1
+ session["decode_step"] = frame_idx
+ n_frames = int(session["n_frames"])
+ if frame_idx < 1:
+ raise ValueError(
+ f"NemotronVoiceChat thinker derived decode frame index {frame_idx} "
+ f"(computed_tokens={computed}, prompt_len={session['prompt_len']}); "
+ "decode must start at acoustic frame 1."
+ )
+ if frame_idx >= n_frames:
+ raise ValueError(
+ f"NemotronVoiceChat thinker was asked for decode step {frame_idx} but only "
+ f"{n_frames} acoustic frames exist; max_tokens must equal the frame count."
+ )
+ text_prev = input_ids.reshape(-1)[:1].to(torch.long)
+ func_prev = torch.full_like(text_prev, session["func_token"] if self._use_function_head else pad_id)
+ frame = session["audio_embeds"][frame_idx : frame_idx + 1]
+ fused = self._fuse(text_prev, frame, func_prev) # [1, H]
+
+ info_update: dict[str, Any] = {
+ "meta": {"nvc_frame": frame_idx},
+ "nvc_text_pad_id": pad_id,
+ }
+ return input_ids, fused, info_update
+
+ def postprocess(self, hidden_states: torch.Tensor, **kwargs: Any) -> dict[str, Any]:
+ """Greedy function-channel token from this step's last hidden.
+
+ Runs after EVERY forward (prefill and each decode step), so the final
+ acoustic frame's token is covered too — NeMo produces gen_function for
+ every non-prompt position. ``nvc_prev_function_token`` is load-bearing
+ (it feeds the next step's fusion at weight 2.0). The cumulative debug
+ timeline (``nvc_function_tokens``, one entry per acoustic frame) costs
+ a host sync plus a growing concat per step, so it is only accumulated
+ when ``NEMOTRON_VOICECHAT_DEBUG_FUNCTION_TIMELINE=1``.
+ On multi-chunk prefill, only the final chunk (whose forward covers the
+ last prefill position, acoustic frame 0) contributes a token.
+ """
+ if hidden_states is None or hidden_states.numel() == 0 or not self._use_function_head:
+ return {}
+ if kwargs.get("_omni_is_prefill"):
+ computed = int(kwargs.get("_omni_num_computed_tokens", 0) or 0)
+ prompt_len = int(kwargs.get("_omni_prompt_len", 0) or 0)
+ if prompt_len and computed + int(hidden_states.shape[0]) < prompt_len:
+ return {} # intermediate prefill chunk: frame 0 not reached yet
+ with torch.inference_mode():
+ func_logits = self.function_head(hidden_states[-1, :].reshape(1, -1).to(self._dtype))
+ token = func_logits.argmax(dim=-1).reshape(()).detach()
+ update: dict[str, Any] = {"nvc_prev_function_token": token}
+ if os.environ.get("NEMOTRON_VOICECHAT_DEBUG_FUNCTION_TIMELINE", "0") == "1":
+ existing = kwargs.get("nvc_function_tokens")
+ if isinstance(existing, torch.Tensor) and existing.numel() > 0:
+ update["nvc_function_tokens"] = torch.cat([existing.reshape(-1).cpu(), token.reshape(1).cpu()])
+ else:
+ update["nvc_function_tokens"] = token.reshape(1).cpu()
+ return update
+
+ def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None:
+ for request_id in finished_req_ids:
+ self._sessions.pop(str(request_id), None)
+
+ # ------------------------------------------------------------------
+ # Weights.
+ # ------------------------------------------------------------------
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+ """Route the unified checkpoint into the thinker's components.
+
+ * ``stt_model.llm.*`` -> vLLM NemotronH (renamed to ``backbone.*``; the
+ separate ``stt_model.embed_tokens`` doubles as ``backbone.embeddings``
+ since the checkpoint keeps a single shared text embedding).
+ * ``stt_model.lm_head.weight`` -> the backbone's ``lm_head``.
+ * ``stt_model.embed_tokens/function_head`` -> the fusion embedding /
+ function head.
+ * ``stt_model.perception.*`` -> the vendored Conformer.
+ * everything else (``rnnt_*``, ``tts_model.*``) is skipped lazily.
+ """
+ loaded: set[str] = set()
+ lm_weights: list[tuple[str, torch.Tensor]] = []
+ perception_state: dict[str, torch.Tensor] = {}
+ embed_weight: torch.Tensor | None = None
+
+ for name, tensor in weights:
+ if name.startswith("stt_model.llm."):
+ lm_weights.append(("backbone." + name.removeprefix("stt_model.llm."), tensor))
+ elif name == "stt_model.lm_head.weight":
+ lm_weights.append(("lm_head.weight", tensor))
+ elif name == "stt_model.embed_tokens.weight":
+ embed_weight = tensor
+ elif name == "stt_model.function_head.weight":
+ with torch.no_grad():
+ self.function_head.weight.copy_(tensor)
+ loaded.add("function_head.weight")
+ elif name.startswith("stt_model.perception."):
+ perception_state[name.removeprefix("stt_model.perception.")] = tensor
+
+ if embed_weight is None:
+ raise ValueError("Checkpoint is missing 'stt_model.embed_tokens.weight'.")
+ if not lm_weights:
+ raise ValueError("Checkpoint is missing the 'stt_model.llm.*' backbone subtree.")
+ if not perception_state:
+ raise ValueError("Checkpoint is missing the 'stt_model.perception.*' subtree.")
+ if "function_head.weight" not in loaded and self._use_function_head:
+ raise ValueError("Checkpoint is missing 'stt_model.function_head.weight'.")
+
+ with torch.no_grad():
+ self.embed_tokens.weight.copy_(embed_weight)
+ loaded.add("embed_tokens.weight")
+
+ # The vLLM NemotronH model expects an input embedding; the checkpoint's
+ # shared embed_tokens fills that slot (inputs_embeds bypass it at runtime).
+ lm_weights.append(("backbone.embeddings.weight", embed_weight))
+ lm_loaded = self.language_model.load_weights(iter(lm_weights))
+ loaded |= {f"language_model.{n}" for n in lm_loaded}
+
+ missing, unexpected = self.perception.load_state_dict(perception_state, strict=False)
+ if missing:
+ raise ValueError(f"NemotronVoiceChat perception is missing weights: {sorted(missing)[:10]} ...")
+ if unexpected:
+ logger.warning(
+ "NemotronVoiceChat perception ignoring %d unexpected tensors (e.g. %s)",
+ len(unexpected),
+ sorted(unexpected)[:5],
+ )
+ device = self.vllm_config.device_config.device
+ self.perception.to(device=device, dtype=self._dtype).eval()
+ self.embed_tokens.to(device=device, dtype=self._dtype)
+ self.function_head.to(device=device, dtype=self._dtype)
+ loaded |= {f"perception.{n}" for n in perception_state}
+ logger.info(
+ "NemotronVoiceChat thinker loaded: %d backbone tensors, %d perception tensors, "
+ "fusion weights (text=%.1f, audio=%.1f, function=%.1f)",
+ len(lm_weights),
+ len(perception_state),
+ self._w_text,
+ self._w_audio,
+ self._w_func,
+ )
+ return loaded
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/pipeline.py b/vllm_omni/model_executor/models/nemotron_voicechat/pipeline.py
new file mode 100644
index 000000000..4f4834740
--- /dev/null
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/pipeline.py
@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""NemotronVoiceChat pipeline: thinker (speech -> frame-locked text) -> talker
+(text timeline -> 31-quantizer RVQ code stacks) -> code2wav (codes -> 22.05 kHz PCM).
+
+``nvidia/NVIDIA-NemotronLabs-VoiceChat-11B`` is an 11B full-duplex speech-to-speech
+model. Offline it is served as a qwen3-omni-style 3-stage pipeline:
+
+* Stage 0 (``thinker``, LLM_AR): vendored FastConformer perception + additive
+ channel fusion feeding the NemotronH backbone on vLLM; emits the frame-locked
+ agent text channel (also the final text output) and carries the full frame
+ timeline + metadata to the talker as a latent payload.
+* Stage 1 (``talker``, LLM_AR): the vendored Gemma3-based EAR-TTS model stepped
+ per frame (classifier-free guidance + MoG sampling run inside the vendored
+ forward); emits one 31-code stack per frame as latents.
+* Stage 2 (``code2wav``, LLM_GENERATION): the vendored RVQ-VAE codec decoder
+ (fp32) producing the final 22.05 kHz waveform.
+
+The STT->TTS coupling is one-directional and token-ids-only (verified against
+the NeMo reference), which is what makes the offline cascade exact.
+"""
+
+from vllm_omni.config.stage_config import (
+ PipelineConfig,
+ StageExecutionType,
+ StagePipelineConfig,
+)
+
+_PROC = "vllm_omni.model_executor.stage_input_processors.nemotron_voicechat"
+
+NEMOTRON_VOICECHAT_PIPELINE = PipelineConfig(
+ model_type="nemotron_voicechat",
+ model_arch="NemotronVoiceChatThinkerForConditionalGeneration",
+ # Named after the alias key so the deploy-yaml stem substring-matches the
+ # checkpoint directory name (NVIDIA-NemotronLabs-VoiceChat-11B) during
+ # bare-path auto-detection.
+ default_deploy_config_name="nemotron_labs_voicechat.yaml",
+ stages=(
+ StagePipelineConfig(
+ stage_id=0,
+ model_stage="thinker",
+ execution_type=StageExecutionType.LLM_AR,
+ input_sources=(),
+ owns_tokenizer=True,
+ final_output=True,
+ final_output_type="text",
+ engine_output_type="latent",
+ # thinker -> talker sync mode is token-path only: the frame-locked
+ # text timeline rides the talker's prompt_token_ids (built by
+ # thinker2talker_token_only on the talker stage); no connector
+ # payload and no full-payload wait for stage 1. With
+ # ``async_chunk: true`` in the deploy yaml, the timeline instead
+ # streams incrementally over the connector (cumulative ids.all per
+ # thinker step) so the talker starts before the thinker finishes.
+ async_chunk_process_next_stage_input_func=f"{_PROC}.thinker2talker_async_chunk",
+ sampling_constraints={"detokenize": True},
+ ),
+ StagePipelineConfig(
+ stage_id=1,
+ model_stage="talker",
+ execution_type=StageExecutionType.LLM_AR,
+ model_arch="NemotronVoiceChatTalkerForConditionalGeneration",
+ hf_config_name="talker_config",
+ input_sources=(0,),
+ engine_output_type="latent",
+ custom_process_next_stage_input_func=f"{_PROC}.talker2code2wav_full_payload",
+ async_chunk_process_next_stage_input_func=f"{_PROC}.talker2code2wav_async_chunk",
+ sync_process_input_func=f"{_PROC}.thinker2talker_token_only",
+ sampling_constraints={"detokenize": False},
+ ),
+ StagePipelineConfig(
+ stage_id=2,
+ model_stage="code2wav",
+ execution_type=StageExecutionType.LLM_GENERATION,
+ model_arch="NemotronVoiceChatCode2Wav",
+ hf_config_name="code2wav_config",
+ input_sources=(1,),
+ final_output=True,
+ final_output_type="audio",
+ engine_output_type="audio",
+ sync_process_input_func=f"{_PROC}.talker2code2wav_token_only",
+ sampling_constraints={"detokenize": False},
+ ),
+ ),
+)
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/runtime_info.py b/vllm_omni/model_executor/models/nemotron_voicechat/runtime_info.py
new file mode 100644
index 000000000..7f5a897d2
--- /dev/null
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/runtime_info.py
@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Runtime-info helpers shared by the thinker and talker AR stages."""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def merge_runtime_info(info_dict: dict[str, Any]) -> dict[str, Any]:
+ """Flatten ``additional_information`` into the info dict (top-level wins)."""
+ additional = info_dict.get("additional_information")
+ if isinstance(additional, dict):
+ merged = {k: v for k, v in info_dict.items() if k != "additional_information"}
+ for k, v in additional.items():
+ merged.setdefault(k, v)
+ return merged
+ return info_dict
+
+
+def require_request_id(info: dict[str, Any], stage: str) -> str:
+ """Resolve the request id from the runtime info (or its meta), or fail."""
+ request_id = info.get("request_id")
+ if request_id is None:
+ meta = info.get("meta")
+ if isinstance(meta, dict):
+ request_id = meta.get("request_id")
+ if request_id is None:
+ raise RuntimeError(f"NemotronVoiceChat {stage} requires a request_id in the runtime info.")
+ return str(request_id)
diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py
index 3c79ba34c..498f9eb85 100644
--- a/vllm_omni/model_executor/models/registry.py
+++ b/vllm_omni/model_executor/models/registry.py
@@ -176,6 +176,23 @@ _OMNI_MODELS = {
"personaplex_code2wav",
"PersonaPlexCode2Wav",
),
+ ## nemotron_voicechat (NVIDIA-NemotronLabs-VoiceChat-11B, full-duplex S2S;
+ ## 3-stage offline speech->speech pipeline: thinker/talker/code2wav)
+ "NemotronVoiceChatThinkerForConditionalGeneration": (
+ "nemotron_voicechat",
+ "nemotron_voicechat_thinker",
+ "NemotronVoiceChatThinkerForConditionalGeneration",
+ ),
+ "NemotronVoiceChatTalkerForConditionalGeneration": (
+ "nemotron_voicechat",
+ "nemotron_voicechat_talker",
+ "NemotronVoiceChatTalkerForConditionalGeneration",
+ ),
+ "NemotronVoiceChatCode2Wav": (
+ "nemotron_voicechat",
+ "nemotron_voicechat_code2wav",
+ "NemotronVoiceChatCode2Wav",
+ ),
## higgs-audio v2
"HiggsAudioV2ForConditionalGeneration": (
"higgs_audio_v2",
diff --git a/vllm_omni/model_executor/stage_input_processors/nemotron_voicechat.py b/vllm_omni/model_executor/stage_input_processors/nemotron_voicechat.py
new file mode 100644
index 000000000..48de5a380
--- /dev/null
+++ b/vllm_omni/model_executor/stage_input_processors/nemotron_voicechat.py
@@ -0,0 +1,355 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Stage transfer processors for the NemotronVoiceChat 3-stage pipeline.
+
+thinker (0) -> talker (1): token path only. The thinker's frame-locked text
+timeline becomes the talker's ``prompt_token_ids``:
+``[text_pad_id] * logical_prompt_token_len + generated frame tokens``. The
+talker must process the PAD prompt region too (its KV state at the first
+acoustic frame depends on it — NeMo trims prompt-region codes only after the
+loop), so the prompt slots are reconstructed here, and the timeline metadata
+rides ``additional_information``.
+
+talker (1) -> code2wav (2): PersonaPlex-style full payload. The talker
+accumulates one 31-quantizer code stack per frame under ``("codes","audio")``;
+the producer ships the prompt-trimmed ``[frames, 31]`` stack once the request
+finishes, and the sync placeholder builder sizes stage-2's dummy prompt.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from vllm_omni.data_entry_keys import (
+ CodesStruct,
+ IdsStruct,
+ MetaStruct,
+ OmniPayloadStruct,
+)
+
+# Full-sequence tensors are re-shipped whole; the full-payload accumulator must
+# replace (not concatenate) them if the producer fires more than once. The AR
+# runner flattens nested payloads before accumulation, so the effective key is
+# the flattened "codes.audio" (the unflattened root is kept for robustness).
+_FULL_PAYLOAD_REPLACE_KEYS: frozenset[str] = frozenset({"codes", "codes.audio", "meta", "meta.nvc_logical_prompt_len"})
+
+# in the Nemotron-Nano-9B-v2 vocab; the checkpoint config's
+# stt pad_token. The thinker also reports it in its latent metadata, which
+# takes precedence when present.
+_DEFAULT_TEXT_PAD_ID = 12
+
+
+def _info_get(info: Any, key: str) -> Any:
+ if isinstance(info, dict):
+ if key in info:
+ return info[key]
+ meta = info.get("meta")
+ if isinstance(meta, dict):
+ return meta.get(key)
+ return None
+
+
+def thinker2talker_token_only(
+ source_outputs: list,
+ prompt: Any = None,
+ _requires_multimodal_data: bool = False,
+) -> list:
+ """Build the talker's engine inputs from the thinker's finished outputs.
+
+ ``prompt_token_ids`` = the full frame-aligned text timeline including the
+ PAD prompt region. The thinker's vLLM prompt covers
+ ``logical_prompt_token_len + 1`` positions (the +1 is acoustic frame 0), so
+ ``logical_prompt_token_len = len(thinker prompt) - 1`` and the generated
+ tokens are exactly the acoustic-frame text channel.
+ """
+ from vllm_omni.inputs.data import OmniTokensPrompt
+
+ del prompt, _requires_multimodal_data
+ inputs: list = []
+ for thinker_output in source_outputs:
+ if not getattr(thinker_output, "finished", False):
+ continue
+ output = thinker_output.outputs[0]
+ generated = list(getattr(output, "cumulative_token_ids", None) or getattr(output, "token_ids", None) or [])
+ prompt_ids = list(getattr(thinker_output, "prompt_token_ids", None) or [])
+ logical_prompt_len = max(len(prompt_ids) - 1, 0)
+ pad_id = _DEFAULT_TEXT_PAD_ID
+ info = getattr(thinker_output, "additional_information", None)
+ reported_pad = _info_get(info, "nvc_text_pad_id")
+ if reported_pad is not None:
+ pad_id = int(reported_pad)
+ timeline = [pad_id] * logical_prompt_len + [int(t) for t in generated]
+ inputs.append(
+ OmniTokensPrompt(
+ # The vLLM-side prompt is one placeholder (timeline position 0;
+ # the NeMo TTS loop starts at t=1). Each decode step consumes
+ # one timeline position from additional_information, so the
+ # talker needs len(timeline) - 1 decode steps.
+ prompt_token_ids=[0],
+ additional_information={
+ "nvc_text_timeline": timeline,
+ "nvc_logical_prompt_len": logical_prompt_len,
+ "nvc_frame_count": len(generated),
+ "nvc_text_pad_id": pad_id,
+ },
+ multi_modal_data=None,
+ mm_processor_kwargs=None,
+ )
+ )
+ return inputs
+
+
+def thinker2talker_async_chunk(
+ transfer_manager: Any,
+ multimodal_output: Any = None,
+ request: Any = None,
+ is_finished: bool = False,
+ **kwargs: Any,
+) -> OmniPayloadStruct | None:
+ """Streaming producer: ship the CUMULATIVE text timeline each thinker step.
+
+ The AR receive path REPLACES the downstream request's
+ additional_information with each chunk payload, so every chunk must carry
+ the whole timeline so far: ``[pad] * logical_prompt_len + sampled tokens``
+ under ``ids.all`` (the talker adopts any longer timeline it sees). The
+ timeline is a few hundred ints, so cumulative re-ship is cheap and keeps
+ the payload self-contained. ``meta.num_processed_tokens`` carries the
+ logical prompt length for the code2wav producer's prompt-region trim.
+ """
+ del multimodal_output, kwargs
+ request_id = request.external_req_id
+ prompt_ids = list(getattr(request, "prompt_token_ids", None) or [])
+ # vLLM prompt = logical prompt + 1 placeholder (acoustic frame 0).
+ logical_prompt_len = max(len(prompt_ids) - 1, 0)
+ generated = list(getattr(request, "output_token_ids", None) or [])
+ pad_id = _DEFAULT_TEXT_PAD_ID
+ reported_pad = _info_get(getattr(request, "additional_information", None), "nvc_text_pad_id")
+ if reported_pad is not None:
+ pad_id = int(reported_pad)
+ timeline = [pad_id] * logical_prompt_len + [int(t) for t in generated]
+
+ # Skip no-progress wakeups (same length as the last emitted chunk).
+ state = transfer_manager.request_payload.get(request_id)
+ last_len = int(state.get("nvc_timeline_len", -1)) if isinstance(state, dict) else -1
+ if len(timeline) <= last_len and not is_finished:
+ return None
+ transfer_manager.request_payload[request_id] = {"nvc_timeline_len": len(timeline)}
+
+ return OmniPayloadStruct(
+ ids=IdsStruct(all=timeline),
+ meta=MetaStruct(
+ finished=torch.tensor(bool(is_finished), dtype=torch.bool),
+ num_processed_tokens=logical_prompt_len,
+ # The talker's vLLM prompt stays a single placeholder token.
+ next_stage_prompt_len=1,
+ ),
+ )
+
+
+def talker2code2wav_async_chunk(
+ transfer_manager: Any,
+ multimodal_output: Any = None,
+ request: Any = None,
+ is_finished: bool = False,
+ **kwargs: Any,
+) -> OmniPayloadStruct | None:
+ """Streaming producer: ship prompt-trimmed CUMULATIVE code stacks in chunks.
+
+ Follows the NeMo incremental-decode recipe (``decode_one_audio_step``):
+ each chunk carries the FULL trimmed code history and
+ ``meta.left_context_size`` = frames the codec already emitted, so the
+ code2wav stage re-decodes the whole prefix and slices off only the new
+ samples — no seams, no crossfade, and the concatenated stream equals the
+ prefix-decode of the final stack. Chunk cadence comes from the connector
+ config key ``codec_chunk_frames`` (default 13 frames ~= 1 s of audio).
+ """
+ del kwargs
+ request_id = request.external_req_id
+
+ codes = None
+ if isinstance(multimodal_output, dict):
+ nested = multimodal_output.get("codes")
+ codes = nested.get("audio") if isinstance(nested, dict) else None
+ if codes is None:
+ codes = multimodal_output.get("codes.audio")
+ if not isinstance(codes, torch.Tensor) or codes.numel() == 0:
+ if not is_finished:
+ return None
+ return _empty_finished_payload()
+ if codes.ndim == 1:
+ codes = codes.reshape(1, -1)
+
+ # Prompt-region trim (rows are NeMo timeline steps t=1..; keep t >= P).
+ # Read the prompt length from the PER-STEP model payload first: the
+ # request-level additional_information is overwritten both by arriving
+ # thinker chunks and by talker info updates, so it races. Cache the first
+ # sighting in the transfer-manager state as a fallback.
+ state = transfer_manager.request_payload.get(request_id)
+ prompt_len = _info_get(multimodal_output if isinstance(multimodal_output, dict) else None, "nvc_logical_prompt_len")
+ if prompt_len is None and isinstance(multimodal_output, dict):
+ prompt_len = multimodal_output.get("meta.nvc_logical_prompt_len")
+ if prompt_len is None:
+ info = getattr(request, "additional_information", None)
+ prompt_len = _info_get(info, "nvc_logical_prompt_len")
+ if prompt_len is None:
+ prompt_len = _info_get(info, "num_processed_tokens")
+ if prompt_len is None and isinstance(state, dict):
+ prompt_len = state.get("nvc_prompt_len")
+ if prompt_len is None:
+ raise ValueError(
+ "NemotronVoiceChat talker request is missing its logical prompt length "
+ "('nvc_logical_prompt_len' or async meta.num_processed_tokens); cannot "
+ "trim the prompt-region codes for streaming code2wav."
+ )
+ trimmed = codes[max(int(prompt_len) - 1, 0) :]
+
+ frames_sent = int(state.get("nvc_frames_sent", 0)) if isinstance(state, dict) else 0
+ new_frames = int(trimmed.shape[0]) - frames_sent
+ if new_frames <= 0 and not is_finished:
+ return None
+
+ connector = getattr(transfer_manager, "connector", None)
+ raw_cfg = getattr(connector, "config", {}) or {}
+ cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {}
+ chunk_frames = int(cfg.get("codec_chunk_frames", 13))
+ if new_frames < chunk_frames and not is_finished:
+ return None
+ if new_frames <= 0 and is_finished and frames_sent > 0:
+ # Everything already shipped; emit a terminal empty chunk so the
+ # code2wav request observes meta.finished.
+ return _empty_finished_payload()
+
+ transfer_manager.request_payload[request_id] = {
+ "nvc_frames_sent": int(trimmed.shape[0]),
+ "nvc_prompt_len": int(prompt_len),
+ }
+ return OmniPayloadStruct(
+ codes=CodesStruct(audio=trimmed.detach().to(dtype=torch.long, device="cpu")),
+ meta=MetaStruct(
+ left_context_size=frames_sent,
+ finished=torch.tensor(bool(is_finished), dtype=torch.bool),
+ ),
+ )
+
+
+def _empty_finished_payload() -> OmniPayloadStruct:
+ return OmniPayloadStruct(
+ codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)),
+ meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)),
+ )
+
+
+def talker2code2wav_full_payload(
+ transfer_manager: Any = None,
+ pooling_output: Any = None,
+ request: Any = None,
+ is_finished: bool = False,
+ **kwargs: Any,
+) -> OmniPayloadStruct | None:
+ """Producer: ship the talker's prompt-trimmed ``[frames, 31]`` code stack.
+
+ Codes accumulate per frame under ``("codes","audio")`` in the request's
+ additional_information; the prompt-region rows were never emitted by the
+ talker (it only stores codes for generated positions), so no extra trim is
+ needed here. Ships once, at request finish, with replace semantics.
+ """
+ # NOTE: no is_finished gating — in sync full-payload mode the producer only
+ # runs at flush time (request already finished), and the flush-time request
+ # object may still report is_finished()=False. Returning a meta-only
+ # "pending" struct here would be SENT as the real payload and starve the
+ # code2wav stage of its codes.
+ del transfer_manager, is_finished
+
+ def _codes_from(src: Any) -> torch.Tensor | None:
+ if not isinstance(src, dict):
+ return None
+ nested = src.get("codes")
+ audio = nested.get("audio") if isinstance(nested, dict) else None
+ return audio if audio is not None else src.get("codes.audio")
+
+ audio = None
+ for source in (
+ getattr(request, "additional_information", None),
+ getattr(request, "additional_information_cpu", None),
+ pooling_output,
+ kwargs.get("multimodal_output"),
+ ):
+ audio = _codes_from(source)
+ if audio is not None:
+ break
+ if not isinstance(audio, torch.Tensor) or audio.numel() == 0:
+ import logging
+
+ logging.getLogger(__name__).error(
+ "nemotron_voicechat talker full-payload producer found no codes: "
+ "pooling_output keys=%s, request.additional_information keys=%s",
+ sorted(pooling_output.keys()) if isinstance(pooling_output, dict) else type(pooling_output).__name__,
+ sorted(getattr(request, "additional_information", None) or {})
+ if isinstance(getattr(request, "additional_information", None), dict)
+ else type(getattr(request, "additional_information", None)).__name__,
+ )
+ return _empty_finished_payload()
+ if audio.ndim == 1:
+ audio = audio.reshape(1, -1)
+ # Rows correspond to NeMo timeline steps t = 1..T-1. Trim the prompt
+ # region: NeMo keeps gen_codes rows t = P..T-1 (P = logical prompt len),
+ # i.e. row indices P-1 onward here. P >= 1 always (a system prompt is
+ # required), so this never drops acoustic frames.
+ prompt_len = _info_get(getattr(request, "additional_information", None), "nvc_logical_prompt_len")
+ if prompt_len is None:
+ prompt_len = _info_get(getattr(request, "additional_information_cpu", None), "nvc_logical_prompt_len")
+ if prompt_len is None:
+ # Explicit failure, never a silent trim=0 (untrimmed prompt-region codes
+ # would ship to the codec). The connector layer logs producer errors.
+ raise ValueError(
+ "NemotronVoiceChat talker request is missing 'nvc_logical_prompt_len' in its "
+ "additional_information; cannot trim the prompt-region codes for code2wav."
+ )
+ audio = audio[max(int(prompt_len) - 1, 0) :]
+ if audio.shape[0] == 0:
+ return _empty_finished_payload()
+ # 2-D [frames, Q] payload: the chunk/full-payload transfer adapters route
+ # >=2-D tensors through the payload channel (1-D would be coerced into
+ # prompt token ids and dropped from the info merge).
+ return OmniPayloadStruct(
+ codes=CodesStruct(audio=audio.detach().to(dtype=torch.long, device="cpu")),
+ meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)),
+ )
+
+
+def talker2code2wav_token_only(
+ source_outputs: list,
+ prompt: Any = None,
+ _requires_multimodal_data: bool = False,
+) -> list:
+ """Sync placeholder builder for stage 2 (codes arrive via the connector)."""
+ from vllm_omni.inputs.data import OmniTokensPrompt
+
+ del prompt, _requires_multimodal_data
+ inputs: list = []
+ for talker_output in source_outputs:
+ if not getattr(talker_output, "finished", False):
+ continue
+ output = talker_output.outputs[0]
+ token_ids = getattr(output, "cumulative_token_ids", None) or getattr(output, "token_ids", None) or []
+ n_frames = len(token_ids)
+ inputs.append(
+ OmniTokensPrompt(
+ prompt_token_ids=[0] * max(n_frames, 1),
+ additional_information=None,
+ multi_modal_data=None,
+ mm_processor_kwargs=None,
+ )
+ )
+ return inputs
+
+
+__all__ = [
+ "thinker2talker_token_only",
+ "thinker2talker_async_chunk",
+ "talker2code2wav_full_payload",
+ "talker2code2wav_async_chunk",
+ "talker2code2wav_token_only",
+]
diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py
index 0f65489b7..b3c388e3a 100644
--- a/vllm_omni/worker/gpu_ar_model_runner.py
+++ b/vllm_omni/worker/gpu_ar_model_runner.py
@@ -339,6 +339,11 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
"NemotronHForCausalLM",
"DyninOmniForConditionalGeneration",
"IndexTTS2TalkerForConditionalGeneration",
+ # nemotron_voicechat: the talker (stage 1) is the full-payload
+ # producer for code2wav (stage 2); the thinker (stage 0) only
+ # produces over the connector in async-chunk (streaming) mode.
+ "NemotronVoiceChatTalkerForConditionalGeneration",
+ "NemotronVoiceChatThinkerForConditionalGeneration",
}
# The stage-level ``model_arch`` override may be blank so the class
# resolves from the checkpoint's own ``architectures`` (e.g. the Audex
diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py
index 766e84c45..7cf26d005 100644
--- a/vllm_omni/worker/gpu_generation_model_runner.py
+++ b/vllm_omni/worker/gpu_generation_model_runner.py
@@ -69,6 +69,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
"CosyVoice3Model",
"DyninOmniForConditionalGeneration",
"IndexTTS2S2MelDecoder",
+ # nemotron_voicechat: code2wav (stage 2) consumes the talker's
+ # full-payload code stacks.
+ "NemotronVoiceChatCode2Wav",
}
if (
getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS