From 5014fea4a0652c61feaf72672b587fcf10d259ab Mon Sep 17 00:00:00 2001 From: Yuekai Zhang Date: Tue, 28 Jul 2026 11:58:22 +0800 Subject: [PATCH] [New Model] support nemotron Audex (#4976) Signed-off-by: Yuekai Zhang Signed-off-by: Yueqian Lin Co-authored-by: Claude Fable 5 Co-authored-by: Yueqian Lin Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> --- examples/offline_inference/audex/README.md | 114 ++++ examples/offline_inference/audex/audio_qa.py | 93 +++ examples/offline_inference/audex/common.py | 149 +++++ .../audex/speech_to_speech.py | 277 +++++++++ .../offline_inference/audex/text_to_audio.py | 157 +++++ .../offline_inference/audex/text_to_speech.py | 162 +++++ examples/online_serving/audex/README.md | 89 +++ examples/online_serving/audex/client.py | 319 ++++++++++ examples/online_serving/audex/run_server.sh | 57 ++ recipes/NVIDIA/Nemotron-Labs-Audex.md | 142 +++++ recipes/README.md | 1 + tests/config/test_omni_config.py | 5 + .../sched/test_omni_scheduling_coordinator.py | 4 + .../test_invalid_audio_speech.py | 46 +- tests/e2e/offline_inference/test_audex_s2s.py | 262 ++++++++ tests/e2e/offline_inference/test_audex_s2t.py | 106 ++++ tests/e2e/offline_inference/test_audex_tta.py | 145 +++++ tests/e2e/offline_inference/test_audex_tts.py | 161 +++++ tests/e2e/online_serving/test_audex_tta.py | 84 +++ tests/e2e/online_serving/test_audex_tts.py | 230 +++++++ .../openai_api/test_audex_serving_guards.py | 376 ++++++++++++ tests/helpers/media.py | 21 + vllm_omni/config/config_factory.py | 45 +- vllm_omni/config/model.py | 5 +- vllm_omni/config/omni_config.py | 9 + vllm_omni/config/pipeline_registry.py | 14 + vllm_omni/config/stage_config.py | 2 +- .../core/sched/omni_scheduling_coordinator.py | 6 + vllm_omni/deploy/audex_s2s.yaml | 65 ++ vllm_omni/deploy/audex_s2s_30b.yaml | 60 ++ vllm_omni/deploy/audex_thinker_only.yaml | 25 + vllm_omni/deploy/audex_thinker_only_30b.yaml | 22 + vllm_omni/deploy/audex_tta.yaml | 69 +++ vllm_omni/deploy/audex_tta_30b.yaml | 57 ++ vllm_omni/deploy/audex_tts.yaml | 83 +++ vllm_omni/deploy/audex_tts_30b.yaml | 76 +++ vllm_omni/engine/async_omni_engine.py | 17 +- vllm_omni/engine/stage_engine_core_proc.py | 9 +- vllm_omni/engine/stage_init_utils.py | 65 ++ vllm_omni/entrypoints/cli/serve.py | 4 +- .../entrypoints/openai/serving_speech.py | 184 ++++++ .../openai/tts_adapters/__init__.py | 2 + .../entrypoints/openai/tts_adapters/audex.py | 67 +++ .../openai/tts_adapters/audex_tta.py | 69 +++ vllm_omni/entrypoints/utils.py | 26 +- .../model_executor/models/audex/__init__.py | 2 + .../models/audex/audex_code2wav.py | 227 +++++++ .../model_executor/models/audex/audex_omni.py | 322 ++++++++++ .../models/audex/audex_omni_processing.py | 260 ++++++++ .../models/audex/audex_thinker.py | 414 +++++++++++++ .../models/audex/audex_xcodec.py | 178 ++++++ vllm_omni/model_executor/models/audex/cfg.py | 567 ++++++++++++++++++ .../model_executor/models/audex/checkpoint.py | 186 ++++++ .../model_executor/models/audex/pipeline.py | 195 ++++++ .../model_executor/models/audex/prompt.py | 83 +++ .../models/audex/speech_decoder/__init__.py | 22 + ...nfiguration_audex_causal_speech_decoder.py | 50 ++ .../modeling_audex_causal_speech_decoder.py | 478 +++++++++++++++ .../audex/speech_decoder/streaming_utils.py | 140 +++++ vllm_omni/model_executor/models/audex/tta.py | 293 +++++++++ vllm_omni/model_executor/models/registry.py | 25 + .../stage_input_processors/audex.py | 356 +++++++++++ vllm_omni/worker/gpu_ar_model_runner.py | 18 +- .../worker/gpu_generation_model_runner.py | 2 + 64 files changed, 7760 insertions(+), 39 deletions(-) create mode 100644 examples/offline_inference/audex/README.md create mode 100644 examples/offline_inference/audex/audio_qa.py create mode 100644 examples/offline_inference/audex/common.py create mode 100644 examples/offline_inference/audex/speech_to_speech.py create mode 100644 examples/offline_inference/audex/text_to_audio.py create mode 100644 examples/offline_inference/audex/text_to_speech.py create mode 100644 examples/online_serving/audex/README.md create mode 100644 examples/online_serving/audex/client.py create mode 100755 examples/online_serving/audex/run_server.sh create mode 100644 recipes/NVIDIA/Nemotron-Labs-Audex.md create mode 100644 tests/e2e/offline_inference/test_audex_s2s.py create mode 100644 tests/e2e/offline_inference/test_audex_s2t.py create mode 100644 tests/e2e/offline_inference/test_audex_tta.py create mode 100644 tests/e2e/offline_inference/test_audex_tts.py create mode 100644 tests/e2e/online_serving/test_audex_tta.py create mode 100644 tests/e2e/online_serving/test_audex_tts.py create mode 100644 tests/entrypoints/openai_api/test_audex_serving_guards.py create mode 100644 vllm_omni/deploy/audex_s2s.yaml create mode 100644 vllm_omni/deploy/audex_s2s_30b.yaml create mode 100644 vllm_omni/deploy/audex_thinker_only.yaml create mode 100644 vllm_omni/deploy/audex_thinker_only_30b.yaml create mode 100644 vllm_omni/deploy/audex_tta.yaml create mode 100644 vllm_omni/deploy/audex_tta_30b.yaml create mode 100644 vllm_omni/deploy/audex_tts.yaml create mode 100644 vllm_omni/deploy/audex_tts_30b.yaml create mode 100644 vllm_omni/entrypoints/openai/tts_adapters/audex.py create mode 100644 vllm_omni/entrypoints/openai/tts_adapters/audex_tta.py create mode 100644 vllm_omni/model_executor/models/audex/__init__.py create mode 100644 vllm_omni/model_executor/models/audex/audex_code2wav.py create mode 100644 vllm_omni/model_executor/models/audex/audex_omni.py create mode 100644 vllm_omni/model_executor/models/audex/audex_omni_processing.py create mode 100644 vllm_omni/model_executor/models/audex/audex_thinker.py create mode 100644 vllm_omni/model_executor/models/audex/audex_xcodec.py create mode 100644 vllm_omni/model_executor/models/audex/cfg.py create mode 100644 vllm_omni/model_executor/models/audex/checkpoint.py create mode 100644 vllm_omni/model_executor/models/audex/pipeline.py create mode 100644 vllm_omni/model_executor/models/audex/prompt.py create mode 100644 vllm_omni/model_executor/models/audex/speech_decoder/__init__.py create mode 100644 vllm_omni/model_executor/models/audex/speech_decoder/configuration_audex_causal_speech_decoder.py create mode 100644 vllm_omni/model_executor/models/audex/speech_decoder/modeling_audex_causal_speech_decoder.py create mode 100644 vllm_omni/model_executor/models/audex/speech_decoder/streaming_utils.py create mode 100644 vllm_omni/model_executor/models/audex/tta.py create mode 100644 vllm_omni/model_executor/stage_input_processors/audex.py diff --git a/examples/offline_inference/audex/README.md b/examples/offline_inference/audex/README.md new file mode 100644 index 000000000..793b69eb9 --- /dev/null +++ b/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