# Audex (Nemotron-Labs-Audex-2B) full deploy: cascaded speech-to-speech.
#
# One audio-capable thinker (checkpoint_folder_full) serves the official
# three-pass cascade — (1) ASR: audio -> transcript, (2) chat: transcript ->
# answer, (3) TTS: answer -> speech — plus the streaming causal speech
# decoder as stage 1. Per-request output modalities route the passes: text
# passes finish at stage 0; only audio-modality requests stream codec frames
# into code2wav. CFG applies to the TTS pass via extra_params.cfg_scale
# (recommended 1.5; 1.0 = off).
pipeline: audex_s2s
async_chunk: true

connectors:
  connector_of_shared_memory:
    name: SharedMemoryConnector
    extra:
      shm_threshold_bytes: 65536
      codec_streaming: true
      connector_get_sleep_s: 0.01
      connector_get_max_wait_first_chunk: 3000
      connector_get_max_wait: 300
      codec_chunk_frames: 25
      initial_codec_chunk_frames: 1
      codec_token_offset: 131077
      codec_vocab_size: 65536
      decoder_emit_chunk_frames: 5

stages:
  - stage_id: 0
    max_num_seqs: 8
    gpu_memory_utilization: 0.5
    trust_remote_code: true
    enable_prefix_caching: true
    async_scheduling: true
    max_num_batched_tokens: 8192
    # Sound embeddings (up to 30 clips x 750) + speech-token generation.
    max_model_len: 32768
    dtype: bfloat16
    devices: "0"
    logits_processors:
      - "vllm_omni.model_executor.models.audex.cfg:AudexCFGLogitsProcessor"
    output_connectors:
      to_stage_1: connector_of_shared_memory
    default_sampling_params:
      temperature: 0.1
      top_k: 80
      max_tokens: 2048
      seed: 0

  - stage_id: 1
    max_num_seqs: 8
    gpu_memory_utilization: 0.25
    enforce_eager: true
    trust_remote_code: true
    enable_prefix_caching: false
    async_scheduling: false
    max_num_batched_tokens: 8192
    max_model_len: 8192
    devices: "0"
    input_connectors:
      from_stage_0: connector_of_shared_memory
    default_sampling_params:
      temperature: 0.0
      max_tokens: 8192
      seed: 0