Skip to content

Audex (Nemotron-Labs-Audex-2B) online serving

Source https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/audex.

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.

./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.

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.

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.

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.

Example materials

client.py

Large file omitted from the rendered docs. View it on GitHub: https://github.com/vllm-project/vllm-omni/blob/main/examples/online_serving/audex/client.py.

run_server.sh
#!/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