Audex (Nemotron-Labs-Audex-2B) offline inference¶

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

# 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

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

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