# AURA Omni deploy profile: Qwen3-ASR -> AURA/Qwen3-VL -> Qwen3-TTS Talker -> Code2Wav.
#
# Pass `--deploy-config vllm_omni/deploy/aura_omni.yaml` and point the parent
# `--model` at any local placeholder repo if needed; each stage overrides its
# own model below.
pipeline: aura_omni
# AURA has a semantic ASR -> AURA -> TTS dependency. Do not prewarm downstream
# stages before AURA text is converted into Qwen3-TTS metadata by aura2tts.
async_chunk: false

connectors:
  connector_of_shared_memory:
    name: SharedMemoryConnector
    extra:
      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
      codec_left_context_frames: 72
      initial_codec_chunk_frames: 1
      decode_cudagraph_capture_sizes: [25, 73, 97, 169, 325]

stages:
  - stage_id: 0
    max_num_seqs: 8
    gpu_memory_utilization: 0.1
    trust_remote_code: true
    enable_prefix_caching: true
    async_scheduling: true
    max_model_len: 8192
    max_num_batched_tokens: 4096
    devices: "0"
    model: Qwen/Qwen3-ASR-1.7B
    default_sampling_params:
      temperature: 0.0
      top_p: 1.0
      top_k: -1
      max_tokens: 256

  - stage_id: 1
    max_num_seqs: 4
    gpu_memory_utilization: 0.4
    trust_remote_code: true
    enable_prefix_caching: true
    async_scheduling: true
    max_model_len: 262144
    max_num_batched_tokens: 15360
    devices: "0"
    model: aurateam/AURA
    default_sampling_params:
      temperature: 0.5
      top_p: 1.0
      top_k: -1
      max_tokens: 256
      repetition_penalty: 1.0

  - stage_id: 2
    max_num_seqs: 8
    gpu_memory_utilization: 0.2
    trust_remote_code: true
    enable_prefix_caching: false
    async_scheduling: true
    max_num_batched_tokens: 512
    max_model_len: 4096
    devices: "0"
    output_connectors:
      to_stage_3: connector_of_shared_memory
    model: Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice
    # model: Qwen/Qwen3-TTS-12Hz-1.7B-Base
    default_sampling_params:
      temperature: 0.9
      top_k: 50
      max_tokens: 4096
      seed: 42
      repetition_penalty: 1.05
    subtalker_sampling_params:
      do_sample: true
      temperature: 0.9
      top_k: 50
      top_p: 1.0

  - stage_id: 3
    max_num_seqs: 8
    gpu_memory_utilization: 0.2
    enforce_eager: false
    trust_remote_code: true
    enable_prefix_caching: false
    async_scheduling: true
    max_num_batched_tokens: 65536
    max_model_len: 65536
    devices: "0"
    input_connectors:
      from_stage_2: connector_of_shared_memory
    model: Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice
    # model: Qwen/Qwen3-TTS-12Hz-1.7B-Base
    default_sampling_params:
      temperature: 0.0
      top_p: 1.0
      top_k: -1
      max_tokens: 65536
      seed: 42
      repetition_penalty: 1.0

platforms:
  npu:
    stages:
      - stage_id: 2
        enforce_eager: true