# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Audex (Nemotron-Labs-Audex-2B) TTS pipeline topology. Stage 0: Thinker — ChatML text prompt → tokens (LLM AR). Stage 1: Code2Wav — streaming causal speech decoder → 16 kHz waveform. Users pass the HF repo ROOT (its manifest config.json carries ``model_type: nemotron_labs_audex``); per-stage checkpoints resolve through ``model_subdir``/``tokenizer_subdir``. The decoder subfolder has no tokenizer files, so stage 1's tokenizer also points at the thinker checkpoint. """ from vllm_omni.config.stage_config import ( PipelineConfig, StageExecutionType, StagePipelineConfig, ) _PROC = "vllm_omni.model_executor.stage_input_processors.audex" # token id in checkpoint_folder_audiogen's tokenizer; pinned # by a unit test against the real tokenizer fixture. AUDEX_SPEECHGEN_END_TOKEN_ID = 131076 # stops TTA generation (pinned in tests via models/audex/tta.py). AUDEX_AUDIOGEN_END_TOKEN_ID = 131074 # No pipeline-level model_arch: stage-0 classes resolve from each # checkpoint's own ``architectures`` (dense on the 2B, NemotronH on the # 30B-A3B); forcing one name via hf_overrides would misbind the other # size. Decoder stages keep their explicit per-stage model_arch. AUDEX_TTS_PIPELINE = PipelineConfig( model_type="audex_tts", default_deploy_config_name="audex_tts.yaml", stages=( StagePipelineConfig( stage_id=0, model_stage="audex_thinker", execution_type=StageExecutionType.LLM_AR, input_sources=(), owns_tokenizer=True, engine_output_type="latent", model_subdir="checkpoint_folder_audiogen", tokenizer_subdir="checkpoint_folder_audiogen", # Emits the unconditional CFG companion for cfg_scale > 1.0 # requests; a no-op for unguided requests. prompt_expand_func=f"{_PROC}.expand_cfg_prompts", async_chunk_process_next_stage_input_func=(f"{_PROC}.thinker2code2wav_async_chunk"), custom_process_next_stage_input_func=f"{_PROC}.thinker2code2wav_full_payload", sampling_constraints={ "detokenize": False, "stop_token_ids": [AUDEX_SPEECHGEN_END_TOKEN_ID], }, ), StagePipelineConfig( stage_id=1, model_stage="audex_code2wav", execution_type=StageExecutionType.LLM_GENERATION, input_sources=(0,), final_output=True, final_output_type="audio", engine_output_type="audio", model_arch="AudexCode2Wav", model_subdir="audex_causal_speech_decoder", tokenizer_subdir="checkpoint_folder_audiogen", sync_process_input_func=f"{_PROC}.thinker2code2wav_token_only", sampling_constraints={"detokenize": True}, ), ), ) # Text-to-audio (caption → general audio): same audiogen thinker checkpoint, # but the output rides the interleaved 4-codebook RVQ block # and decodes through the EXTERNAL XCodec1 checkpoint. XCodec1 is a CNN codec # decoded over the full sequence, so this pipeline is sync full-payload only # (deploy yaml sets async_chunk: false). CFG is effectively mandatory for TTA # quality (official default scale 3.0). # Cascaded speech-to-speech: the audio-capable full checkpoint serves all # three official passes (ASR → chat → TTS) in one deployment; only the TTS # pass (output modality "audio", final_stage_id 1) traverses the streaming # speech decoder — text-modality passes finish at stage 0 and never reach # the codec path. The cascade itself is orchestrated by the client/example. AUDEX_S2S_PIPELINE = PipelineConfig( model_type="audex_s2s", default_deploy_config_name="audex_s2s.yaml", stages=( StagePipelineConfig( stage_id=0, model_stage="audex_omni", execution_type=StageExecutionType.LLM_AR, input_sources=(), final_output=True, final_output_type="text", owns_tokenizer=True, requires_multimodal_data=True, engine_output_type="latent", model_subdir="checkpoint_folder_full", tokenizer_subdir="checkpoint_folder_full", prompt_expand_func=f"{_PROC}.expand_cfg_prompts", async_chunk_process_next_stage_input_func=(f"{_PROC}.thinker2code2wav_async_chunk"), custom_process_next_stage_input_func=f"{_PROC}.thinker2code2wav_full_payload", # additionally stops the TTS pass; text passes # stop at their natural eos first. detokenize stays on for the # text-final passes. sampling_constraints={ "detokenize": True, "stop_token_ids": [AUDEX_SPEECHGEN_END_TOKEN_ID], }, ), StagePipelineConfig( stage_id=1, model_stage="audex_code2wav", execution_type=StageExecutionType.LLM_GENERATION, input_sources=(0,), final_output=True, final_output_type="audio", engine_output_type="audio", model_arch="AudexCode2Wav", model_subdir="audex_causal_speech_decoder", # The decoder folder has no tokenizer files. Borrow the audiogen # tokenizer (not checkpoint_folder_full): this stage resolves its # snapshot with the "tts" download profile, which guarantees # checkpoint_folder_audiogen but not the full folder's tokenizer, # so per-stage / clean-cache launches would otherwise fail. tokenizer_subdir="checkpoint_folder_audiogen", sync_process_input_func=f"{_PROC}.thinker2code2wav_token_only", sampling_constraints={"detokenize": True}, ), ), ) # Audio understanding (ASR / audio QA): single-stage deployment of the # audio-capable full checkpoint — speech (+ text instruction) in, text out. # Modeled on the ming_flash_omni / qwen2_5_omni thinker-only pipelines. AUDEX_THINKER_ONLY_PIPELINE = PipelineConfig( model_type="audex_thinker_only", default_deploy_config_name="audex_thinker_only.yaml", stages=( StagePipelineConfig( stage_id=0, model_stage="audex_omni", execution_type=StageExecutionType.LLM_AR, input_sources=(), final_output=True, final_output_type="text", owns_tokenizer=True, requires_multimodal_data=True, engine_output_type="text", model_subdir="checkpoint_folder_full", tokenizer_subdir="checkpoint_folder_full", sampling_constraints={"detokenize": True}, ), ), ) AUDEX_TTA_PIPELINE = PipelineConfig( model_type="audex_tta", default_deploy_config_name="audex_tta.yaml", stages=( StagePipelineConfig( stage_id=0, model_stage="audex_tta_thinker", execution_type=StageExecutionType.LLM_AR, input_sources=(), owns_tokenizer=True, engine_output_type="latent", model_subdir="checkpoint_folder_audiogen", tokenizer_subdir="checkpoint_folder_audiogen", prompt_expand_func=f"{_PROC}.expand_cfg_prompts", custom_process_next_stage_input_func=f"{_PROC}.thinker2xcodec_full_payload", sampling_constraints={ "detokenize": False, "stop_token_ids": [AUDEX_AUDIOGEN_END_TOKEN_ID], }, ), StagePipelineConfig( stage_id=1, model_stage="audex_xcodec", execution_type=StageExecutionType.LLM_GENERATION, input_sources=(0,), final_output=True, final_output_type="audio", engine_output_type="audio", model_arch="AudexXCodec1", # Length placeholder for the stage-1 request; the bulk codec # payload ships via thinker2xcodec_full_payload. sync_process_input_func=f"{_PROC}.thinker2code2wav_token_only", sampling_constraints={"detokenize": True}, ), ), )