[New Model] support nemotron Audex (#4976)

Signed-off-by: Yuekai Zhang <zhangyuekai@foxmail.com>
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>

examples/offline_inference/audex/README.md

@@ -0,0 +1,114 @@
+# Audex (Nemotron-Labs-Audex-2B) offline inference
+
+One checkpoint, four tasks — each script builds its own `Omni` engine with
+the right deploy yaml, so they run as plain `python <script>` with no extra
+flags. Pass the HF repo ROOT as `--model` (default
+`nvidia/Nemotron-Labs-Audex-2B`); per-stage subfolders and download subsets
+resolve automatically.
+
+| script | pipeline (`vllm_omni/deploy/<name>.yaml`) | audio in | text out | speech out | general audio out |
+|---|---|---|---|---|---|
+| `text_to_speech.py` | `audex_tts` | ❌ | ❌ | ✅ | ❌ |
+| `text_to_audio.py` | `audex_tta` | ❌ | ❌ | ❌ | ✅ |
+| `audio_qa.py` | `audex_thinker_only` | ✅ | ✅ | ❌ | ❌ |
+| `speech_to_speech.py` | `audex_s2s` | ✅ | ✅ | ✅ | ❌ |
+
+Shared plumbing (corpus/WAV IO, tokenizer loading, the CFG cond/uncond
+pair contract) lives in `common.py`; each task script keeps only its own
+prompt recipe and flow. Note `text_to_audio.py` here drives the Audex
+two-stage AR pipeline — for diffusion text-to-audio models (Stable Audio
+Open, AudioX) use the generic
+`examples/offline_inference/text_to_audio/text_to_audio.py` instead.
+
+## text_to_speech.py — text → speech
+
+Thinker generates `<speechcodec_N>` tokens; the streaming causal decoder
+turns them into 16 kHz WAVs (one file per prompt, slugified filename).
+
+```bash
+# three built-in demo sentences -> results/audex_wavs/
+python examples/offline_inference/audex/text_to_speech.py
+
+# your own texts, with classifier-free guidance (official quality setting)
+python examples/offline_inference/audex/text_to_speech.py \
+    --texts "Hello there." "Nice to meet you." --cfg-scale 1.5 \
+    --output-dir results/my_tts
+```
+
+Key flags: `--texts` / `--texts-file` (TSV `utt_id<TAB>text`),
+`--cfg-scale` (default 1.0 = off; 1.5 recommended — guided requests run
+one at a time, each paired with a length-matched null prompt),
+`--temperature` (overrides the deploy yaml's stage-0 default).
+
+## text_to_audio.py — caption → general audio (sound effects)
+
+Same thinker over the interleaved 4-codebook `<audiocodec_N>` RVQ block,
+decoded by the external XCodec1 checkpoint. CFG is effectively mandatory
+(default scale 3.0). Clips are capped at `--codec-cap 4000` codec tokens
+(10 s).
+
+```bash
+# three built-in demo captions -> results/audex_tta_wavs/
+python examples/offline_inference/audex/text_to_audio.py
+
+python examples/offline_inference/audex/text_to_audio.py \
+    --captions "Thunder rolling across a valley." --output-dir results/my_tta
+```
+
+XCodec1 resolves from `--xcodec1-path`, the `XCODEC1_PATH` env var, or the
+default `hf-audio/xcodec-hubert-general-balanced` repo (downloaded on
+first use).
+
+## audio_qa.py — audio (+ instruction) → text
+
+Single-stage audio understanding on the full checkpoint (NV-Whisper
+encoder + projector + LM). The default question is ASR
+("Transcribe the input speech."); without `--audio-files`, vLLM's public
+`mary_had_lamb` asset is transcribed.
+
+```bash
+# transcribe the built-in demo asset
+python examples/offline_inference/audex/audio_qa.py
+
+# your own clips / free-form QA
+python examples/offline_inference/audex/audio_qa.py \
+    --audio-files a.wav b.wav --question "What language is being spoken?"
+```
+
+## speech_to_speech.py — spoken question → spoken answer
+
+The official three-pass cascade over ONE `audex_s2s` deployment: ASR
+(audio → transcript, text-final) → chat (transcript → answer, text-final)
+→ TTS (answer → speech through the streaming decoder). Only the TTS pass
+touches the codec path. Without `--audio-file`, the `mary_had_lamb` asset
+is used as the spoken input.
+
+```bash
+# full cascade on the built-in demo asset -> results/audex_s2s_answer.wav
+python examples/offline_inference/audex/speech_to_speech.py
+
+python examples/offline_inference/audex/speech_to_speech.py \
+    --audio-file question.wav --output results/answer.wav --cfg-scale 1.5
+```
+
+`--cfg-scale` (default 1.5) applies to the TTS pass only; 1.0 disables it.
+
+Online-serving counterparts (server + HTTP clients) live in
+`examples/online_serving/audex/`.
+
+## 30B-A3B (nvidia/Nemotron-Labs-Audex-30B-A3B)
+
+The same four scripts serve the 30B MoE checkpoint (hybrid Mamba + MoE
+NemotronH thinker; identical token space and decoder). The 30B REQUIRES an
+explicit 30B deploy yaml — the model root's default resolution lands on the
+2B-tuned configs:
+
+    python examples/offline_inference/audex/text_to_speech.py \
+        --model nvidia/Nemotron-Labs-Audex-30B-A3B \
+        --deploy-config vllm_omni/deploy/audex_tts_30b.yaml
+
+Per-mode yamls: `audex_{tts,tta,thinker_only,s2s}_30b.yaml`. Defaults fit a
+single H100 80 GB (thinker + decoder share the card; prefix caching is off —
+hybrid Mamba); if long sequences OOM, set `tensor_parallel_size: 2` on
+stage 0. First run downloads ~60 GB. Verified on a single H100 80 GB: ~61 GiB
+weights + healthy KV headroom; TP2 was not needed.

examples/offline_inference/audex/audio_qa.py

@@ -0,0 +1,93 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline Audex (Nemotron-Labs-Audex-2B) audio understanding example.
+
+Speech (or general audio) + text instruction → text, through the
+single-stage ``audex_thinker_only`` pipeline: NV-Whisper
+encoder + relu2 projector + 2B dense LM on ``checkpoint_folder_full``.
+
+The prompt format mirrors the official audioqa script: a ChatML user turn
+holding the ``<so_embedding>`` placeholder (expanded by the processor to
+750 embedding positions per 30 s clip) plus the instruction, with a closed
+``<think></think>`` priming.
+
+Examples (without --audio-files, vLLM's public ``mary_had_lamb`` asset is
+transcribed):
+
+    # Transcribe WAVs (ASR):
+    python examples/offline_inference/audex/audio_qa.py \\
+        --audio-files a.wav b.wav
+
+    # Free-form audio QA:
+    python examples/offline_inference/audex/audio_qa.py \\
+        --audio-files a.wav --question "What language is being spoken?"
+"""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from common import default_deploy_config, load_audio_file, request_index
+from vllm.assets.audio import AudioAsset
+
+from vllm_omni import Omni
+
+ASR_QUESTION = "Transcribe the input speech."
+
+
+def build_prompt(question: str) -> str:
+    return f"<|im_start|>user\n<so_embedding>\n{question}<|im_end|>\n<|im_start|>assistant\n<think></think>"
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Offline Audex audio understanding")
+    parser.add_argument("--model", type=str, default="nvidia/Nemotron-Labs-Audex-2B")
+    parser.add_argument(
+        "--audio-files",
+        type=str,
+        nargs="+",
+        default=None,
+        help="Input clips. Defaults to vLLM's mary_had_lamb audio asset.",
+    )
+    parser.add_argument("--question", type=str, default=ASR_QUESTION)
+    parser.add_argument(
+        "--deploy-config",
+        type=str,
+        # The model root's default deploy yaml is the TTS pipeline; audio
+        # understanding needs the single-stage thinker-only pipeline.
+        default=default_deploy_config("audex_thinker_only.yaml"),
+        help="Deploy yaml (defaults to the audex_thinker_only pipeline).",
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    engine = Omni(model=args.model, deploy_config=args.deploy_config, trust_remote_code=True)
+
+    prompt_text = build_prompt(args.question)
+    if args.audio_files:
+        clips = [load_audio_file(path) for path in args.audio_files]
+    else:
+        clips = [AudioAsset("mary_had_lamb").audio_and_sample_rate]
+    prompts = []
+    for audio, sr in clips:
+        prompts.append(
+            {
+                "prompt": prompt_text,
+                "multi_modal_data": {"audio": (audio, sr)},
+            }
+        )
+
+    outputs = engine.generate(prompts)
+
+    ordered = sorted(outputs, key=request_index)
+    names = [Path(p).name for p in args.audio_files] if args.audio_files else ["mary_had_lamb"]
+    for name, req_output in zip(names, ordered):
+        text = req_output.outputs[0].text if req_output.outputs else ""
+        print(f"{name}: {text.strip()}")
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/audex/common.py

@@ -0,0 +1,149 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Shared helpers for the offline Audex examples in this directory.
+
+Import style follows the repo's example convention (same-directory module,
+see ``text_to_speech/qwen3_tts/tts_common.py``): scripts run as files, so
+``from common import ...`` resolves via the script's own directory.
+"""
+
+from __future__ import annotations
+
+import copy
+import re
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import soundfile as sf
+import torch
+
+from vllm_omni import Omni
+from vllm_omni.model_executor.models.audex.prompt import build_null_prompt
+
+SAMPLE_RATE = 16_000
+_DEPLOY_DIR = Path(__file__).resolve().parents[3] / "vllm_omni" / "deploy"
+
+# Guided decoding sharpens the distribution, so the unguided temperature
+# (0.1) adds excess sampling noise under CFG. Measured on the en-24 gate
+# corpus at cfg 1.5: temp 0.1 -> CER 7.31%, temp 0.05 -> CER 6.87% (vs the
+# unguided baseline 7.24%).
+GUIDED_TEMPERATURE = 0.05
+
+
+def default_deploy_config(name: str) -> str:
+    """Absolute path of a deploy yaml bundled with vllm_omni."""
+    return str(_DEPLOY_DIR / name)
+
+
+def slugify(text: str, fallback: str = "prompt") -> str:
+    slug = re.sub(r"\s+", "_", text.strip().lower())
+    slug = re.sub(r"[^a-z0-9_]+", "", slug)
+    return slug[:48] or fallback
+
+
+def load_corpus(
+    corpus_file: str | None,
+    texts: list[str] | None,
+    default_texts: tuple[str, ...],
+) -> list[tuple[str, str]]:
+    """(utt_id, text) pairs from a TSV corpus file, CLI texts, or defaults."""
+    if corpus_file:
+        corpus = []
+        for line in Path(corpus_file).read_text().splitlines():
+            if not line.strip():
+                continue
+            utt, text = line.split("\t", 1)
+            corpus.append((utt, text.strip()))
+        return corpus
+    resolved = texts if texts else list(default_texts)
+    return [(slugify(t), t) for t in resolved]
+
+
+def load_audio_file(path: str) -> tuple[np.ndarray, int]:
+    """Mono float32 samples + native sample rate (processors resample)."""
+    audio, sr = sf.read(path, dtype="float32")
+    if audio.ndim == 2:
+        audio = audio.mean(axis=1)
+    return audio, sr
+
+
+def extract_pcm(multimodal_output: dict) -> torch.Tensor:
+    """Flatten one request's audio payload to a 1-D float32 CPU tensor."""
+    audio = multimodal_output.get("model_outputs")
+    if audio is None:
+        audio = multimodal_output.get("audio")
+    if audio is None:
+        raise ValueError(f"no audio key in multimodal_output: {list(multimodal_output.keys())}")
+    if isinstance(audio, list):
+        valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None]
+        if not valid:
+            raise ValueError("audio list is empty")
+        return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0]
+    return torch.as_tensor(audio).float().cpu().reshape(-1)
+
+
+def write_wav(path: Path, pcm: torch.Tensor, sample_rate: int = SAMPLE_RATE) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    arr = (np.clip(pcm.numpy(), -1.0, 1.0) * 32767.0).astype(np.int16)
+    sf.write(str(path), arr, sample_rate, format="WAV", subtype="PCM_16")
+
+
+def load_audex_tokenizer(model: str, profile: str = "tts", folder: str = "checkpoint_folder_audiogen"):
+    """Thinker tokenizer for CFG/TTA prompt building (resolves the snapshot)."""
+    from transformers import AutoTokenizer
+
+    from vllm_omni.model_executor.models.audex.checkpoint import ensure_audex_snapshot
+
+    root = ensure_audex_snapshot(model, profile=profile)
+    return AutoTokenizer.from_pretrained(str(Path(root) / folder), trust_remote_code=True)
+
+
+def request_index(req_output) -> int:
+    """Submission index from the numeric request-id prefix.
+
+    ``Omni.generate`` may return outputs out of submission order, and a
+    lexicographic sort would misplace "10_..." before "2_...".
+    """
+    match = re.search(r"(\d+)", str(req_output.request_id))
+    return int(match.group(1)) if match else 0
+
+
+def guided_stage_params(
+    engine: Omni,
+    cond_prompt: str,
+    tokenizer,
+    cfg_scale: float,
+    pair_id: str,
+    *,
+    temperature: float | None = None,
+    extra_args: dict[str, Any] | None = None,
+    null_prompt_builder: Callable[[str, Any], str] = build_null_prompt,
+) -> list:
+    """Stage sampling params carrying the CFG pair contract for ONE request.
+
+    Clones the engine's resolved defaults (never mutates the shared list),
+    applies ``extra_args`` unconditionally (e.g. the TTA RVQ contract), and
+    attaches the cond/uncond pair contract when ``cfg_scale > 1.0``. Guided
+    requests must go one at a time: each needs its own pair id and a
+    length-matched null prompt built from its own text.
+    """
+    params = copy.deepcopy(engine.resolve_sampling_params_list(None))
+    stage0 = params[0]
+    if temperature is not None:
+        stage0.temperature = temperature
+    if stage0.extra_args is None:
+        stage0.extra_args = {}
+    if extra_args:
+        stage0.extra_args.update(extra_args)
+    if cfg_scale > 1.0:
+        stage0.extra_args.update(
+            {
+                "cfg_scale": float(cfg_scale),
+                "cfg_role": "cond",
+                "cfg_pair_id": pair_id,
+                "cfg_null_prompt": null_prompt_builder(cond_prompt, tokenizer),
+            }
+        )
+    return params

examples/offline_inference/audex/speech_to_speech.py

@@ -0,0 +1,277 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline Audex (Nemotron-Labs-Audex-2B) cascaded speech-to-speech example.
+
+The official three-pass cascade over ONE ``audex_s2s``
+deployment:
+
+  1. ASR pass  (audio + "Transcribe the input speech." → transcript, text)
+  2. chat pass (transcript → answer, text)
+  3. TTS pass  (answer → speech, audio via the streaming causal decoder)
+
+Text passes carry ``modalities: ["text"]`` and finish at stage 0; only the
+TTS pass (``modalities: ["audio"]``) streams codec frames into code2wav.
+
+Example (without --audio-file, vLLM's public ``mary_had_lamb`` asset is
+used as the spoken input):
+
+    python examples/offline_inference/audex/speech_to_speech.py \\
+        --audio-file question.wav --output results/answer.wav
+"""
+
+from __future__ import annotations
+
+import argparse
+import copy
+import re
+from pathlib import Path
+
+import torch
+from common import (
+    SAMPLE_RATE,
+    default_deploy_config,
+    guided_stage_params,
+    load_audex_tokenizer,
+    load_audio_file,
+    write_wav,
+)
+from vllm.assets.audio import AudioAsset
+
+from vllm_omni import Omni
+from vllm_omni.model_executor.models.audex.prompt import build_cond_prompt
+
+ASR_QUESTION = "Transcribe the input speech."
+SYSTEM_PROMPT = "You are a helpful and harmless assistant.\n\nYou are not allowed to use any tools."
+# The official chat-pass recipe below (SYSTEM_PROMPT + TEXT_INSTRUCTION
+# prefix, open <think> priming, temp 1.0/top_p 0.95, seed retry) applies to
+# the 30B-A3B ONLY (see is_30b()/_chat_prompt below and
+# tests/e2e/offline_inference/test_audex_s2s.py, which mirrors this same
+# split). It is load-bearing on the 30B: without the formatting-instruction
+# prefix a bare transcript turn is treated as a TTS request and the model
+# emits literal <speechcodec_N> tokens instead of a text answer (the system
+# prompt alone does not prevent this). The 2B must NOT use this recipe: it
+# is long-verified stable on a bare prompt + engine-default sampling, while
+# the official temp-1.0 seed-0 draw hallucinates off-language text (yet
+# closes </think> cleanly, so the seed retry never fires) and fails the
+# round-trip WER gate — measured, reproduced byte-identically.
+TEXT_INSTRUCTION = (
+    "SYSTEM & FORMATTING INSTRUCTION: You are Audex, created by NVIDIA based on the Nemotron-Cascade-2 architecture. "
+    "You may output your reasoning, followed by your final response. "
+    "You may format your reasoning block however you like. "
+    "However, your final response must strictly follow these rules: "
+    "* Write in plain, unformatted prose like a book or newspaper article. "
+    "* Do not use markdown, bullet points, lists, or headers. "
+    "* Standard numbers, abbreviations, and symbols are acceptable. "
+    "* [CRITICAL] You must press enter after every single sentence, placing each sentence on its own separate line."
+)
+
+
+def is_30b(model: str) -> bool:
+    # Users pass arbitrary repo IDs or local snapshot paths, so this is a
+    # heuristic (not an exact match against the official 30B-A3B repo ID)
+    # — good enough for an example script; keep in sync with the same
+    # branch in examples/online_serving/audex/client.py and
+    # tests/e2e/offline_inference/test_audex_s2s.py.
+    return "30b" in model.lower()
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Offline Audex cascaded S2S")
+    parser.add_argument("--model", type=str, default="nvidia/Nemotron-Labs-Audex-2B")
+    parser.add_argument(
+        "--audio-file",
+        type=str,
+        default=None,
+        help="Spoken question (WAV). Defaults to vLLM's mary_had_lamb audio asset.",
+    )
+    parser.add_argument("--output", type=str, default="results/audex_s2s_answer.wav")
+    parser.add_argument(
+        "--deploy-config",
+        type=str,
+        # The model root's default deploy yaml is the TTS-only pipeline; the
+        # cascade needs the audio-capable 2-stage full pipeline.
+        default=default_deploy_config("audex_s2s.yaml"),
+        help="Deploy yaml (defaults to the audex_s2s pipeline).",
+    )
+    parser.add_argument(
+        "--cfg-scale",
+        type=float,
+        default=1.5,
+        help="CFG strength for the TTS pass (1.0 disables; official setting 1.5).",
+    )
+    return parser.parse_args()
+
+
+def _chat_prompt(user_text: str) -> str:
+    # Official chat-pass recipe: system prompt + formatting instruction
+    # prefixed to the user text, thinking ENABLED via an open <think> block
+    # (the official demo's web UI ships "Enable reasoning" checked, i.e.
+    # build_chat_prompt(..., enable_thinking=True); chat_template.jinja then
+    # renders "<|im_start|>assistant\n<think>\n"). The instruction invites
+    # the model to reason before answering — priming the block open keeps
+    # that reasoning inside <think>...</think> so _extract_final_response
+    # below can drop it; with a closed block the 30B leaks the reasoning as
+    # plain text in front of the answer.
+    return (
+        f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n"
+        f"<|im_start|>user\n{TEXT_INSTRUCTION}\n\n{user_text}<|im_end|>\n"
+        "<|im_start|>assistant\n<think>\n"
+    )
+
+
+def _asr_prompt() -> str:
+    return f"<|im_start|>user\n<so_embedding>\n{ASR_QUESTION}<|im_end|>\n<|im_start|>assistant\n<think></think>"
+
+
+def _extract_final_response(text: str) -> str:
+    # Mirror the official demo's extract_spoken_answer
+    # (cascaded_s2s_web_server.py): keep only what follows the LAST closing
+    # think tag, falling back to the full text when the model skipped
+    # reasoning or never closed the block. Benign no-op for the 2B, whose
+    # answers carry no think block.
+    if "</think>" in text:
+        answer = text.rsplit("</think>", 1)[-1].strip()
+        if answer:
+            return answer
+    return text.strip()
+
+
+def _clean(text: str) -> str:
+    # The full checkpoint prefixes transcripts with a language tag sentence
+    # and quotes the content; extract the quoted payload when present. Span
+    # first-to-last quote (not split on every apostrophe) so contractions
+    # inside the transcript (e.g. "don't") survive intact.
+    text = text.strip()
+    match = re.search(r"'(.*)'", text, flags=re.S)
+    if match and match.group(1).strip():
+        return match.group(1).strip()
+    return text
+
+
+def _text_of(outputs) -> str:
+    (req_output,) = outputs
+    return (req_output.outputs[0].text or "").strip()
+
+
+CHAT_MAX_TOKENS = 2048  # official --text-max-tokens default
+# The official demo samples this pass unseeded; some trajectories deliberate
+# past the token budget without ever closing the think block, so retry a few
+# deterministic seeds until one closes (see the chat pass below).
+CHAT_SEEDS = (0, 1, 2)
+
+
+def _chat_params(engine: Omni, seed: int):
+    # Official demo text-pass sampling (cascaded_s2s_web_server.py defaults:
+    # --text-temperature 1.0, --text-top-p 0.95, --text-max-tokens 2048).
+    # The yaml stage-0 defaults (temperature 0.1, top_k 80) are TTS-thinker
+    # settings; under near-greedy decoding the 30B loops inside its
+    # reasoning block until the token budget runs out and never emits
+    # </think>, so the extracted "answer" would be truncated reasoning.
+    params = copy.deepcopy(engine.resolve_sampling_params_list(None))
+    stage0 = params[0]
+    stage0.temperature = 1.0
+    stage0.top_p = 0.95
+    stage0.top_k = -1
+    stage0.max_tokens = CHAT_MAX_TOKENS
+    stage0.seed = seed
+    return params
+
+
+def _chat_once(engine: Omni, prompt: str, seed: int) -> tuple[str, bool]:
+    (req_output,) = engine.generate([{"prompt": prompt, "modalities": ["text"]}], _chat_params(engine, seed))
+    out = req_output.outputs[0]
+    text = (out.text or "").strip()
+    truncated = getattr(out, "finish_reason", None) == "length" or len(out.token_ids or []) >= CHAT_MAX_TOKENS
+    return text, truncated
+
+
+def main():
+    args = parse_args()
+    engine = Omni(model=args.model, deploy_config=args.deploy_config, trust_remote_code=True)
+
+    tokenizer = None
+    if args.cfg_scale > 1.0:
+        tokenizer = load_audex_tokenizer(args.model, profile="full", folder="checkpoint_folder_full")
+
+    if args.audio_file:
+        audio, sr = load_audio_file(args.audio_file)
+    else:
+        audio, sr = AudioAsset("mary_had_lamb").audio_and_sample_rate
+
+    # Pass 1 — ASR (text-final; never touches the speech decoder).
+    transcript = _clean(
+        _text_of(
+            engine.generate(
+                [
+                    {
+                        "prompt": _asr_prompt(),
+                        "multi_modal_data": {"audio": (audio, sr)},
+                        "modalities": ["text"],
+                    }
+                ]
+            )
+        )
+    )
+    print(f"[1/3] transcript : {transcript}")
+    if not transcript:
+        raise SystemExit("ASR pass produced an empty transcript; aborting cascade")
+
+    # Pass 2 — chat answer (text-final). The recipe is size-conditional (see
+    # the TEXT_INSTRUCTION comment above): the 30B requires the official
+    # recipe, the 2B regresses under it.
+    if is_30b(args.model):
+        # Official recipe. The raw completion is
+        # <reasoning...></think><answer>; speak only the final response.
+        # Retry with the next seed only when the budget ran out before the
+        # reasoning block closed (a natural stop without </think> is a
+        # complete answer); if every seed spins, fall back to the official
+        # demo's behavior of keeping the full text.
+        raw = ""
+        for seed in CHAT_SEEDS:
+            raw, truncated = _chat_once(engine, _chat_prompt(transcript), seed)
+            if "</think>" in raw or not truncated:
+                break
+            print(f"[2/3] chat seed {seed} hit the token budget mid-reasoning; retrying")
+        answer = _extract_final_response(raw)
+    else:
+        # 2B: bare prompt with a closed think block, engine-default sampling
+        # (yaml near-greedy), no retry; the answer is used as-is.
+        chat_prompt = f"<|im_start|>user\n{transcript}<|im_end|>\n<|im_start|>assistant\n<think></think>"
+        answer = _text_of(engine.generate([{"prompt": chat_prompt, "modalities": ["text"]}]))
+    print(f"[2/3] answer     : {answer}")
+    if not answer:
+        raise SystemExit("Chat pass produced an empty answer; aborting cascade")
+
+    # Pass 3 — TTS (audio-final; streams through code2wav). Both final
+    # stages of the pipeline may emit an output record for the request;
+    # collect the audio-bearing one.
+    # The formatting instruction puts every sentence on its own line;
+    # flatten to plain prose for the TTS thinker.
+    tts_prompt = build_cond_prompt(re.sub(r"\s+", " ", answer).strip())
+    outputs = engine.generate(
+        [{"prompt": tts_prompt, "modalities": ["audio"]}],
+        guided_stage_params(engine, tts_prompt, tokenizer, args.cfg_scale, "s2s-tts-pass"),
+    )
+    chunks: list[torch.Tensor] = []
+    for req_output in outputs:
+        mm = getattr(req_output, "multimodal_output", None) or {}
+        if not mm and req_output.outputs:
+            mm = getattr(req_output.outputs[0], "multimodal_output", None) or {}
+        audio_val = mm.get("model_outputs", mm.get("audio"))
+        if audio_val is None:
+            continue
+        if isinstance(audio_val, list):
+            chunks.extend(torch.as_tensor(a).float().cpu().reshape(-1) for a in audio_val if a is not None)
+        else:
+            chunks.append(torch.as_tensor(audio_val).float().cpu().reshape(-1))
+    pcm = torch.cat(chunks) if chunks else torch.empty(0)
+    if pcm.numel() == 0:
+        raise SystemExit("TTS pass produced empty audio")
+
+    out_path = Path(args.output)
+    write_wav(out_path, pcm)
+    print(f"[3/3] speech     : {pcm.numel() / SAMPLE_RATE:.2f}s -> {out_path}")
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/audex/text_to_audio.py

@@ -0,0 +1,157 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline Audex (Nemotron-Labs-Audex-2B) text-to-audio inference example.
+
+Caption → general audio through the vLLM-Omni TTA pipeline: the audiogen
+thinker generates interleaved 4-codebook <audiocodec_N> RVQ tokens under an
+RVQ phase mask, and the external XCodec1 checkpoint decodes them to a 16 kHz
+mono WAV.
+
+(For DIFFUSION text-to-audio models — Stable Audio Open, AudioX — use
+``examples/offline_inference/text_to_audio/text_to_audio.py`` instead; Audex
+TTA is a two-stage autoregressive pipeline with a different interface.)
+
+Classifier-free guidance is effectively mandatory for TTA quality (official
+default scale 3.0), so every request submits a cond/uncond pair; requests
+run one at a time.
+
+Pass the HF repo ROOT as --model. XCodec1 resolves from --xcodec1-path /
+XCODEC1_PATH / the default hf-audio repo.
+
+Example:
+
+    python examples/offline_inference/audex/text_to_audio.py \\
+        --captions "Heavy rain falling on a tin roof." \\
+        --output-dir results/audex_tta_wavs
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import time
+from pathlib import Path
+
+from common import (
+    SAMPLE_RATE,
+    default_deploy_config,
+    extract_pcm,
+    guided_stage_params,
+    load_audex_tokenizer,
+    load_corpus,
+    write_wav,
+)
+
+from vllm_omni import Omni
+from vllm_omni.model_executor.models.audex.prompt import build_tta_cond_prompt, build_tta_null_prompt
+from vllm_omni.model_executor.models.audex.tta import build_tta_phase_token_ids
+
+DEFAULT_CAPTIONS = (
+    "Heavy rain falling on a tin roof.",
+    "A dog barking in the distance while birds chirp.",
+    "Ocean waves crashing on a rocky shore.",
+)
+DEFAULT_CODEC_CAP = 4000
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Offline Audex TTA inference")
+    parser.add_argument("--model", type=str, default="nvidia/Nemotron-Labs-Audex-2B")
+    parser.add_argument("--captions", type=str, nargs="+", default=None, help="Audio captions.")
+    parser.add_argument(
+        "--captions-file",
+        type=str,
+        default=None,
+        help="TSV corpus: one 'utt_id<TAB>caption' per line (overrides --captions).",
+    )
+    parser.add_argument("--output-dir", type=str, default="results/audex_tta_wavs")
+    parser.add_argument(
+        "--deploy-config",
+        type=str,
+        # The model root's default deploy yaml is the TTS pipeline; TTA
+        # prompts and RVQ sampling params require the dedicated TTA pipeline.
+        default=default_deploy_config("audex_tta.yaml"),
+        help="Deploy yaml (defaults to the audex_tta pipeline).",
+    )
+    parser.add_argument(
+        "--xcodec1-path",
+        type=str,
+        default=None,
+        help="Local XCodec1 checkpoint (defaults to $XCODEC1_PATH or the hf-audio repo).",
+    )
+    parser.add_argument(
+        "--cfg-scale",
+        type=float,
+        default=3.0,
+        help="Classifier-free guidance strength (official TTA default 3.0; 1.0 disables).",
+    )
+    parser.add_argument(
+        "--codec-cap",
+        type=int,
+        default=DEFAULT_CODEC_CAP,
+        help="Max generated codec tokens before the phase mask forces <audiogen_end>.",
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    if args.xcodec1_path:
+        os.environ["XCODEC1_PATH"] = args.xcodec1_path
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    corpus = load_corpus(args.captions_file, args.captions, DEFAULT_CAPTIONS)
+
+    tokenizer = load_audex_tokenizer(args.model, profile="tta")
+    phase_token_ids, start_tid, end_tid = build_tta_phase_token_ids(tokenizer)
+    tta_rvq = {
+        "phase_token_ids": phase_token_ids,
+        "start_tid": start_tid,
+        "end_tid": end_tid,
+        "codec_cap": args.codec_cap,
+        # The prompt ends with <audiogen_start>, so generation begins at
+        # RVQ phase 0 immediately.
+        "start_in_prompt": True,
+    }
+
+    engine = Omni(model=args.model, deploy_config=args.deploy_config, trust_remote_code=True)
+
+    print(f"Model       : {args.model}")
+    print(f"Captions    : {len(corpus)}")
+    print(f"CFG scale   : {args.cfg_scale}")
+    print(f"Codec cap   : {args.codec_cap}")
+    print(f"Output dir  : {output_dir}")
+
+    total_elapsed = 0.0
+    total_dur = 0.0
+    for utt, caption in corpus:
+        cond_prompt = build_tta_cond_prompt(caption)
+        sampling = guided_stage_params(
+            engine,
+            cond_prompt,
+            tokenizer,
+            args.cfg_scale,
+            f"tta-{utt}",
+            extra_args={"tta_rvq": tta_rvq},
+            null_prompt_builder=build_tta_null_prompt,
+        )
+        t_start = time.perf_counter()
+        outputs = engine.generate([cond_prompt], sampling)
+        total_elapsed += time.perf_counter() - t_start
+
+        (req_output,) = outputs
+        pcm = extract_pcm(req_output.outputs[0].multimodal_output)
+        dur = pcm.numel() / SAMPLE_RATE
+        if dur <= 0:
+            raise RuntimeError(f"empty audio for {utt}")
+        out_path = output_dir / f"{utt}.wav"
+        write_wav(out_path, pcm)
+        total_dur += dur
+        print(f"  {utt:<48} dur={dur:6.2f}s  -> {out_path}")
+
+    rtf = total_elapsed / total_dur if total_dur > 0 else float("inf")
+    print(f"Total infer : {total_elapsed:.2f}s  total audio: {total_dur:.2f}s  RTF: {rtf:.3f}")
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/audex/text_to_speech.py

@@ -0,0 +1,162 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline Audex (Nemotron-Labs-Audex-2B) TTS inference example.
+
+Runs Stage 0 (thinker) + Stage 1 (streaming causal speech decoder) end-to-end
+through the vLLM-Omni engine and saves a 16 kHz mono WAV per prompt.
+
+Pass the HF repo ROOT as --model; per-stage subfolders resolve automatically.
+
+Examples:
+
+    python examples/offline_inference/audex/text_to_speech.py \\
+        --texts "The weather is so good today." \\
+        --output-dir results/audex_wavs
+
+    # TSV corpus (utt_id<TAB>text per line), all requests in flight at once:
+    python examples/offline_inference/audex/text_to_speech.py \\
+        --texts-file corpus.tsv --output-dir results/audex_wavs
+
+    # One request at a time:
+    python examples/offline_inference/audex/text_to_speech.py \\
+        --texts-file corpus.tsv --sequential --output-dir results/audex_wavs
+"""
+
+from __future__ import annotations
+
+import argparse
+import time
+from pathlib import Path
+
+from common import (
+    GUIDED_TEMPERATURE,
+    SAMPLE_RATE,
+    extract_pcm,
+    guided_stage_params,
+    load_audex_tokenizer,
+    load_corpus,
+    request_index,
+    write_wav,
+)
+
+from vllm_omni import Omni
+from vllm_omni.model_executor.models.audex.prompt import build_cond_prompt
+
+DEFAULT_TEXTS = (
+    "Hello world.",
+    "The quick brown fox jumps over the lazy dog.",
+    "Today is a beautiful day for a walk in the park.",
+)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Offline Audex TTS inference")
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="nvidia/Nemotron-Labs-Audex-2B",
+        help="HF repo root (or local snapshot root).",
+    )
+    parser.add_argument("--texts", type=str, nargs="+", default=None, help="Plain-text prompts.")
+    parser.add_argument(
+        "--texts-file",
+        type=str,
+        default=None,
+        help="TSV corpus: one 'utt_id<TAB>text' per line (overrides --texts).",
+    )
+    parser.add_argument("--output-dir", type=str, default="results/audex_wavs")
+    parser.add_argument("--deploy-config", type=str, default=None, help="Override the deploy config path.")
+    parser.add_argument(
+        "--sequential",
+        action="store_true",
+        help="Submit one request at a time instead of all at once.",
+    )
+    parser.add_argument(
+        "--cfg-scale",
+        type=float,
+        default=1.0,
+        help=(
+            "Classifier-free guidance strength. 1.0 disables guidance (default, "
+            "matches the v1 baseline); 1.5 is the official quality setting. "
+            "Guided runs submit one request at a time (each carries its own "
+            "length-matched null prompt and pair id)."
+        ),
+    )
+    parser.add_argument(
+        "--temperature",
+        type=float,
+        default=None,
+        help="Override the deploy yaml's stage-0 sampling temperature.",
+    )
+    return parser.parse_args()
+
+
+def _write_outputs(outputs, batch: list[tuple[str, str]], output_dir: Path) -> float:
+    ordered = sorted(outputs, key=request_index)
+    if len(ordered) != len(batch):
+        raise RuntimeError(f"expected {len(batch)} outputs, got {len(ordered)}")
+
+    total_dur = 0.0
+    for (utt, _text), req_output in zip(batch, ordered):
+        pcm = extract_pcm(req_output.outputs[0].multimodal_output)
+        dur = pcm.numel() / SAMPLE_RATE
+        if dur <= 0:
+            raise RuntimeError(f"empty audio for {utt}")
+        out_path = output_dir / f"{utt}.wav"
+        write_wav(out_path, pcm)
+        total_dur += dur
+        print(f"  {utt:<40} dur={dur:6.2f}s  -> {out_path}")
+    return total_dur
+
+
+def main():
+    args = parse_args()
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    corpus = load_corpus(args.texts_file, args.texts, DEFAULT_TEXTS)
+
+    engine = Omni(model=args.model, deploy_config=args.deploy_config, trust_remote_code=True)
+
+    cfg_enabled = args.cfg_scale > 1.0
+    tokenizer = load_audex_tokenizer(args.model) if cfg_enabled else None
+
+    print(f"Model       : {args.model}")
+    print(f"Prompts     : {len(corpus)}")
+    print(f"Mode        : {'sequential' if args.sequential or cfg_enabled else 'batched'}")
+    print(f"CFG scale   : {args.cfg_scale}" + (" (guidance off)" if not cfg_enabled else ""))
+    print(f"Output dir  : {output_dir}")
+
+    total_elapsed = 0.0
+    total_dur = 0.0
+    # CFG runs one request at a time: each request needs its own pair id and
+    # text-specific null prompt in the stage-0 sampling params.
+    batches = [[item] for item in corpus] if (args.sequential or cfg_enabled) else [corpus]
+    for batch in batches:
+        prompts = [build_cond_prompt(text) for _utt, text in batch]
+        sampling_params_list = None
+        if cfg_enabled:
+            sampling_params_list = guided_stage_params(
+                engine,
+                prompts[0],
+                tokenizer,
+                args.cfg_scale,
+                f"cfg-{batch[0][0]}",
+                temperature=GUIDED_TEMPERATURE,
+            )
+        if args.temperature is not None:
+            import copy
+
+            if sampling_params_list is None:
+                sampling_params_list = copy.deepcopy(engine.resolve_sampling_params_list(None))
+            sampling_params_list[0].temperature = args.temperature
+        t_start = time.perf_counter()
+        outputs = engine.generate(prompts, sampling_params_list)
+        total_elapsed += time.perf_counter() - t_start
+        total_dur += _write_outputs(outputs, batch, output_dir)
+
+    rtf = total_elapsed / total_dur if total_dur > 0 else float("inf")
+    print(f"Total infer : {total_elapsed:.2f}s  total audio: {total_dur:.2f}s  RTF: {rtf:.3f}")
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/audex/README.md

@@ -0,0 +1,89 @@
+# Audex (Nemotron-Labs-Audex-2B) online serving
+
+One checkpoint, four deployment modes. `run_server.sh` starts the server
+for a given `MODE`; `client.py --mode <mode>` tests it. The capability
+matrix:
+
+| mode (`vllm_omni/deploy/audex_<mode>.yaml`) | audio in | text out | speech out | general audio out | endpoint |
+|---|---|---|---|---|---|
+| `tts` | ❌ | ❌ | ✅ | ❌ | `/v1/audio/speech` |
+| `tta` | ❌ | ❌ | ❌ | ✅ | `/v1/audio/speech` |
+| `thinker_only` | ✅ | ✅ | ❌ | ❌ | `/v1/chat/completions` |
+| `s2s` | ✅ | ✅ | ✅ | ❌ | both |
+
+Audio-input modes of `client.py` fall back to vLLM's public
+`mary_had_lamb` asset when `--audio-file` is omitted; CFG defaults follow
+the official settings per mode (tts/s2s 1.5, tta 3.0; `--cfg-scale 1.0`
+disables).
+
+## tts — text → speech
+
+Text-only thinker (`checkpoint_folder_audiogen`) emits `<speechcodec_N>`
+tokens; the streaming causal decoder produces 16 kHz WAVs. This is also
+the DEFAULT pipeline when the repo root is served without a deploy config.
+
+```bash
+./run_server.sh                       # MODE=tts, port 8097
+
+python client.py --mode tts --text "Hello there." --output hello.wav
+# or raw curl:
+curl -s http://localhost:8097/v1/audio/speech \
+    -H 'Content-Type: application/json' \
+    -d '{"model": "nvidia/Nemotron-Labs-Audex-2B", "input": "Hello there.", "response_format": "wav", "extra_params": {"cfg_scale": 1.5}}' \
+    -o hello.wav
+```
+
+## tta — caption → general audio
+
+Same thinker over the 4-codebook `<audiocodec_N>` RVQ block, decoded by
+the external XCodec1 checkpoint (auto-downloaded; override with
+`XCODEC1_PATH`). Clips are capped at 10 s.
+
+```bash
+MODE=tta ./run_server.sh
+
+python client.py --mode tta --caption "Heavy rain falling on a tin roof." \
+    --output rain.wav
+```
+
+## thinker_only — audio (+ instruction) → text
+
+Single-stage audio understanding on the full checkpoint (NV-Whisper
+encoder + projector + LM). Send the clip as `input_audio` chat content.
+
+```bash
+MODE=thinker_only ./run_server.sh
+
+python client.py --mode thinker_only                      # transcribe the demo asset
+python client.py --mode thinker_only --audio-file a.wav \
+    --question "What language is being spoken?"
+```
+
+## s2s — spoken question → spoken answer
+
+The audio-capable thinker plus the speech decoder in one deployment.
+`client.py --mode s2s` runs the official three passes: ASR (chat with
+`input_audio`) → chat on the transcript → TTS on the answer. Text-final
+passes carry `"modalities": ["text"]`; only the TTS pass streams through
+the decoder.
+
+```bash
+MODE=s2s PORT=8098 ./run_server.sh
+
+python client.py --mode s2s --port 8098 \
+    --audio-file question.wav --output answer.wav
+```
+
+Offline counterparts (no HTTP, one `Omni` engine per script) live in
+`examples/offline_inference/audex/`.
+
+## 30B-A3B (nvidia/Nemotron-Labs-Audex-30B-A3B)
+
+`SIZE=30b` switches run_server.sh to the 30B yamls and model id:
+
+    SIZE=30b ./run_server.sh                 # 30B tts on port 8097
+    SIZE=30b MODE=s2s PORT=8098 ./run_server.sh
+
+Pass `--model nvidia/Nemotron-Labs-Audex-30B-A3B` to the client — the server
+validates the model id and 404s on a mismatch. Single-H100 defaults; see the
+yaml comments for the TP2 fallback. First launch downloads ~60 GB.

examples/online_serving/audex/client.py

@@ -0,0 +1,319 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Audex online client for all four deployment modes.
+
+Pick ``--mode`` to match the server's deploy yaml (see run_server.sh):
+
+  tts           /v1/audio/speech: text -> speech WAV
+  tta           /v1/audio/speech: caption -> general-audio WAV
+  thinker_only  /v1/chat/completions: audio (+ question) -> text
+  s2s           the official three-pass cascade against one audex_s2s
+                server: ASR (chat, input_audio) -> chat -> TTS
+
+Audio-input modes fall back to vLLM's public ``mary_had_lamb`` asset when
+``--audio-file`` is omitted.
+
+Examples:
+
+    python examples/online_serving/audex/client.py --mode tts \\
+        --text "Hello there." --output hello.wav
+    python examples/online_serving/audex/client.py --mode tta \\
+        --caption "Heavy rain falling on a tin roof." --output rain.wav
+    python examples/online_serving/audex/client.py --mode thinker_only
+    python examples/online_serving/audex/client.py --mode s2s \\
+        --audio-file question.wav --output answer.wav
+"""
+
+from __future__ import annotations
+
+import argparse
+import base64
+import re
+from pathlib import Path
+
+import requests
+
+ASR_QUESTION = "Transcribe the input speech."
+SYSTEM_PROMPT = "You are a helpful and harmless assistant.\n\nYou are not allowed to use any tools."
+# The official chat-pass recipe below (SYSTEM_PROMPT + TEXT_INSTRUCTION
+# prefix, reasoning enabled, temp 1.0/top_p 0.95, seed retry) applies to the
+# 30B-A3B ONLY (see is_30b() below and
+# tests/e2e/offline_inference/test_audex_s2s.py, which mirrors this same
+# split). It is load-bearing on the 30B: without the formatting-instruction
+# prefix a bare transcript turn is treated as a TTS request and the model
+# emits literal <speechcodec_N> tokens instead of a text answer (the system
+# prompt alone does not prevent this). The 2B must NOT use this recipe: the
+# official temp-1.0 seed-0 draw measurably hallucinates off-language text on
+# the 2B (yet closes </think> cleanly, so the seed retry never fires) and
+# fails the round-trip WER gate — measured, reproduced byte-identically. The
+# 2B instead keeps its original bare-message / enable_thinking=False /
+# temperature=0.0 recipe.
+TEXT_INSTRUCTION = (
+    "SYSTEM & FORMATTING INSTRUCTION: You are Audex, created by NVIDIA based on the Nemotron-Cascade-2 architecture. "
+    "You may output your reasoning, followed by your final response. "
+    "You may format your reasoning block however you like. "
+    "However, your final response must strictly follow these rules: "
+    "* Write in plain, unformatted prose like a book or newspaper article. "
+    "* Do not use markdown, bullet points, lists, or headers. "
+    "* Standard numbers, abbreviations, and symbols are acceptable. "
+    "* [CRITICAL] You must press enter after every single sentence, placing each sentence on its own separate line."
+)
+DEFAULT_TEXT = "Hello there! This is the Audex text to speech pipeline."
+DEFAULT_CAPTION = "Heavy rain falling on a tin roof."
+# Official quality settings; 1.0 disables guidance.
+DEFAULT_CFG = {"tts": 1.5, "tta": 3.0, "s2s": 1.5}
+
+
+def is_30b(model: str) -> bool:
+    # Users pass arbitrary repo IDs or local snapshot paths, so this is a
+    # heuristic (not an exact match against the official 30B-A3B repo ID)
+    # — good enough for an example script; keep in sync with the same
+    # branch in examples/offline_inference/audex/speech_to_speech.py and
+    # tests/e2e/offline_inference/test_audex_s2s.py.
+    return "30b" in model.lower()
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Audex online client (tts / tta / thinker_only / s2s)")
+    parser.add_argument("--mode", choices=("tts", "tta", "thinker_only", "s2s"), default="s2s")
+    parser.add_argument("--host", type=str, default="127.0.0.1")
+    parser.add_argument("--port", type=int, default=8097)
+    parser.add_argument("--model", type=str, default="nvidia/Nemotron-Labs-Audex-2B")
+    parser.add_argument("--text", type=str, default=DEFAULT_TEXT, help="TTS input text (mode=tts).")
+    parser.add_argument("--caption", type=str, default=DEFAULT_CAPTION, help="TTA caption (mode=tta).")
+    parser.add_argument(
+        "--audio-file",
+        type=str,
+        default=None,
+        help="Input WAV for thinker_only/s2s. Defaults to vLLM's mary_had_lamb asset.",
+    )
+    parser.add_argument("--question", type=str, default=ASR_QUESTION, help="Instruction (mode=thinker_only).")
+    parser.add_argument("--output", type=str, default="results/audex_online.wav")
+    parser.add_argument(
+        "--cfg-scale",
+        type=float,
+        default=None,
+        help="CFG strength; defaults per mode (tts/s2s 1.5, tta 3.0). 1.0 disables.",
+    )
+    parser.add_argument("--max-tokens", type=int, default=512)
+    return parser.parse_args()
+
+
+def _audio_b64(audio_file: str | None) -> str:
+    if audio_file:
+        return base64.b64encode(Path(audio_file).read_bytes()).decode()
+    # Fall back to vLLM's public asset, re-encoded as WAV for the chat API.
+    import io
+
+    import numpy as np
+    import soundfile as sf
+    from vllm.assets.audio import AudioAsset
+
+    audio, sr = AudioAsset("mary_had_lamb").audio_and_sample_rate
+    buf = io.BytesIO()
+    sf.write(buf, np.asarray(audio, dtype="float32"), int(sr), format="WAV")
+    return base64.b64encode(buf.getvalue()).decode()
+
+
+def _chat(
+    base_url: str,
+    model: str,
+    messages: list[dict],
+    max_tokens: int,
+    enable_thinking: bool = False,
+    temperature: float = 0.0,
+    top_p: float = 1.0,
+    seed: int = 0,
+) -> tuple[str, str]:
+    resp = requests.post(
+        f"{base_url}/v1/chat/completions",
+        json={
+            "model": model,
+            "messages": messages,
+            "max_tokens": max_tokens,
+            "temperature": temperature,
+            "top_p": top_p,
+            "seed": seed,
+            # Text-final passes must say so explicitly: without "modalities"
+            # the server falls back to the deployment's configured output
+            # modalities, which for the s2s pipeline include the audio final
+            # stage — routing ASR/chat through code2wav instead of stopping
+            # at stage 0.
+            "modalities": ["text"],
+            # The checkpoint's chat template thinking toggle. The official
+            # cascaded demo transcribes with enable_thinking=False but runs
+            # the chat pass with reasoning ENABLED (its web UI ships the
+            # "Enable reasoning" checkbox checked), priming the generation
+            # prompt with an open "<think>\n" so the model's reasoning stays
+            # inside the think block; _extract_final_response then drops it.
+            # With thinking disabled the 30B leaks its reasoning as plain
+            # text in front of the final response.
+            "chat_template_kwargs": {"enable_thinking": enable_thinking},
+        },
+        timeout=300,
+    )
+    resp.raise_for_status()
+    choice = resp.json()["choices"][0]
+    return choice["message"]["content"].strip(), choice.get("finish_reason") or ""
+
+
+def _speech(base_url: str, model: str, text: str, cfg_scale: float, out_path: Path) -> None:
+    payload: dict = {"model": model, "input": text, "response_format": "wav"}
+    if cfg_scale > 1.0:
+        payload["extra_params"] = {"cfg_scale": cfg_scale}
+    resp = requests.post(f"{base_url}/v1/audio/speech", json=payload, timeout=600)
+    resp.raise_for_status()
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    out_path.write_bytes(resp.content)
+
+
+def _audio_question(audio_b64: str, question: str) -> list[dict]:
+    return [
+        {
+            "role": "user",
+            "content": [
+                {"type": "input_audio", "input_audio": {"data": audio_b64, "format": "wav"}},
+                {"type": "text", "text": question},
+            ],
+        }
+    ]
+
+
+def _tts_text(answer: str) -> str:
+    # The formatting instruction's CRITICAL rule puts each final-response
+    # sentence on its own line, and the official demo TTS-es those sentence
+    # segments (split_sentence_segments). Mirror that per-line segmentation
+    # but keep a single /v1/audio/speech request: speak the leading whole
+    # sentences of the CLEANED final response up to ~200 chars, never
+    # cutting mid-sentence. Markdown markup is stripped as before (it makes
+    # the TTS thinker stop almost immediately).
+    sentences = []
+    for line in answer.splitlines():
+        line = re.sub(r"[*#_`]+", "", line)
+        line = re.sub(r"\s+", " ", line).strip()
+        if line:
+            sentences.append(line)
+    picked: list[str] = []
+    total = 0
+    for sentence in sentences:
+        if picked and total + len(sentence) > 200:
+            break
+        picked.append(sentence)
+        total += len(sentence) + 1
+    return " ".join(picked)
+
+
+def _extract_final_response(text: str) -> str:
+    # Mirror the official demo's extract_spoken_answer
+    # (cascaded_s2s_web_server.py): the chat pass reasons inside the primed
+    # <think> block, so keep only what follows the LAST closing tag; fall
+    # back to the full text when the model skipped reasoning or never
+    # closed the block. Benign no-op for the 2B's clean answers.
+    if "</think>" in text:
+        answer = text.rsplit("</think>", 1)[-1].strip()
+        if answer:
+            return answer
+    return text.strip()
+
+
+def _clean_transcript(text: str) -> str:
+    # The full checkpoint prefixes transcripts with a language-tag sentence
+    # and quotes the content; extract the quoted payload when present. Span
+    # first-to-last quote (not split on every apostrophe) so contractions
+    # inside the transcript (e.g. "don't") survive intact.
+    text = text.strip()
+    match = re.search(r"'(.*)'", text, flags=re.S)
+    if match and match.group(1).strip():
+        return match.group(1).strip()
+    return text
+
+
+def main():
+    args = parse_args()
+    base_url = f"http://{args.host}:{args.port}"
+    cfg_scale = args.cfg_scale if args.cfg_scale is not None else DEFAULT_CFG.get(args.mode, 1.0)
+    out_path = Path(args.output)
+
+    if args.mode == "tts":
+        _speech(base_url, args.model, args.text, cfg_scale, out_path)
+        print(f"speech: {out_path.stat().st_size} bytes -> {out_path}")
+        return
+
+    if args.mode == "tta":
+        _speech(base_url, args.model, args.caption, cfg_scale, out_path)
+        print(f"audio : {out_path.stat().st_size} bytes -> {out_path}")
+        return
+
+    audio_b64 = _audio_b64(args.audio_file)
+
+    if args.mode == "thinker_only":
+        content, _ = _chat(base_url, args.model, _audio_question(audio_b64, args.question), args.max_tokens)
+        print(f"answer: {_extract_final_response(content)}")
+        return
+
+    # s2s — the official three-pass cascade.
+    asr_content, _ = _chat(base_url, args.model, _audio_question(audio_b64, ASR_QUESTION), args.max_tokens)
+    transcript = _clean_transcript(asr_content)
+    print(f"[1/3] transcript : {transcript}")
+    if not transcript:
+        raise SystemExit("ASR pass produced an empty transcript; aborting cascade")
+
+    # The chat-pass recipe is size-conditional (see the TEXT_INSTRUCTION
+    # comment above): the 30B requires the official recipe, the 2B keeps its
+    # original bare-message recipe.
+    if is_30b(args.model):
+        # Official recipe: system prompt + formatting instruction prefixed
+        # to the transcript, reasoning enabled, and the official demo's
+        # text-pass sampling (its defaults: --text-temperature 1.0,
+        # --text-top-p 0.95, --text-max-tokens 2048). Near-greedy decoding
+        # makes the 30B loop inside its reasoning block until the token
+        # budget runs out without ever closing </think>. Even at the
+        # official temperature some trajectories spin past the budget, so
+        # retry a few deterministic seeds; a natural stop without </think>
+        # is a complete answer, and if every seed spins we keep the full
+        # text like the official demo does.
+        chat_messages = [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": f"{TEXT_INSTRUCTION}\n\n{transcript}"},
+        ]
+        raw = ""
+        for seed in (0, 1, 2):
+            raw, finish_reason = _chat(
+                base_url,
+                args.model,
+                chat_messages,
+                max(args.max_tokens, 2048),
+                enable_thinking=True,
+                temperature=1.0,
+                top_p=0.95,
+                seed=seed,
+            )
+            if "</think>" in raw or finish_reason != "length":
+                break
+            print(f"[2/3] chat seed {seed} hit the token budget mid-reasoning; retrying")
+        answer = _extract_final_response(raw)
+    else:
+        # 2B: bare user message (no system/instruction prefix), reasoning
+        # disabled, temperature 0.0, no retry — its original recipe. The
+        # official temp-1.0 seed-0 recipe measurably makes the 2B
+        # hallucinate off-language answers (see the module-level comment).
+        chat_messages = [{"role": "user", "content": transcript}]
+        raw, _ = _chat(
+            base_url,
+            args.model,
+            chat_messages,
+            args.max_tokens,
+            enable_thinking=False,
+            temperature=0.0,
+        )
+        answer = _extract_final_response(raw)
+    print(f"[2/3] answer     : {answer}")
+    if not answer:
+        raise SystemExit("Chat pass produced an empty answer; aborting cascade")
+
+    _speech(base_url, args.model, _tts_text(answer), cfg_scale, out_path)
+    print(f"[3/3] speech     : {out_path.stat().st_size} bytes -> {out_path}")
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/audex/run_server.sh

@@ -0,0 +1,57 @@
+#!/bin/bash
+# Launch the vLLM-Omni server for Audex (Nemotron-Labs-Audex-2B).
+#
+# MODE picks the deployment (see README.md for the capability matrix):
+#   tts           text -> speech            /v1/audio/speech   (default)
+#   tta           caption -> general audio  /v1/audio/speech
+#   thinker_only  audio -> text             /v1/chat/completions
+#   s2s           cascaded speech-to-speech (both endpoints)
+#
+# Pass the HF repo ROOT as MODEL — per-stage subfolders resolve
+# automatically.
+#
+# Usage:
+#   ./run_server.sh                       # 2B tts on port 8097, GPU 0
+#   SIZE=30b ./run_server.sh              # 30B-A3B tts (explicit 30B yaml)
+#   SIZE=30b MODE=s2s PORT=8098 ./run_server.sh
+#   MODE=tta ./run_server.sh              # needs XCodec1 (auto-downloaded)
+#   MODEL=/path/to/local/snapshot ./run_server.sh
+
+set -e
+
+MODE="${MODE:-${1:-tts}}"
+SIZE="${SIZE:-2b}"
+if [ "$SIZE" = "30b" ]; then
+    MODEL="${MODEL:-nvidia/Nemotron-Labs-Audex-30B-A3B}"
+    YAML_SUFFIX="_30b"
+elif [ "$SIZE" = "2b" ]; then
+    MODEL="${MODEL:-nvidia/Nemotron-Labs-Audex-2B}"
+    YAML_SUFFIX=""
+else
+    echo "Unknown SIZE '$SIZE'; expected 2b|30b" >&2
+    exit 1
+fi
+PORT="${PORT:-8097}"
+GPUS="${GPUS:-0}"
+
+REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)"
+DEPLOY_YAML="$REPO_ROOT/vllm_omni/deploy/audex_${MODE}${YAML_SUFFIX}.yaml"
+if [ ! -f "$DEPLOY_YAML" ]; then
+    echo "Unknown MODE '$MODE' / SIZE '$SIZE' (no $DEPLOY_YAML); expected MODE=tts|tta|thinker_only|s2s, SIZE=2b|30b" >&2
+    exit 1
+fi
+
+echo "Starting Audex server"
+echo "  MODE=$MODE ($DEPLOY_YAML)"
+echo "  SIZE=$SIZE"
+echo "  MODEL=$MODEL"
+echo "  PORT=$PORT"
+echo "  CUDA_VISIBLE_DEVICES=$GPUS"
+
+CUDA_VISIBLE_DEVICES="$GPUS" \
+vllm-omni serve "$MODEL" \
+    --host 0.0.0.0 \
+    --port "$PORT" \
+    --trust-remote-code \
+    --stage-configs-path "$DEPLOY_YAML" \
+    --omni