# Audex (Nemotron-Labs-Audex-2B) text-to-audio deploy: TTA thinker → XCodec1.
#
# Notable choices:
#   * Sync full-payload (async_chunk: false): XCodec1 is a CNN codec the
#     official flow decodes over the full sequence; there is no streaming
#     reference and chunked decode risks boundary artifacts.
#   * Stage 0 sampling defaults follow the official TTA script (temperature
#     1.0, top_k 80; stop on <audiogen_end>; codec cap 4000 tokens = 1000
#     frames generated, decode truncates at 500 frames ≈ 10 s).
#   * Classifier-free guidance is effectively mandatory for TTA quality; the
#     official default is cfg_scale 3.0. The RVQ phase mask is required for
#     validity — both logits processors are registered here and driven per
#     request via extra_args (cfg_* and tta_rvq).
#   * Stage 1's model is the EXTERNAL hf-audio XCodec1 checkpoint (override
#     with a local path here or via XCODEC1_PATH); its tokenizer is borrowed
#     from the thinker checkpoint automatically.
pipeline: audex_tta
async_chunk: false

connectors:
  connector_of_shared_memory:
    name: SharedMemoryConnector
    extra:
      shm_threshold_bytes: 65536
      connector_get_sleep_s: 0.01
      connector_get_max_wait_first_chunk: 3000
      connector_get_max_wait: 300
      audiocodec_token_offset: 196613
      audiocodec_vocab_size: 8192

stages:
  - stage_id: 0
    max_num_seqs: 8
    gpu_memory_utilization: 0.4
    trust_remote_code: true
    enable_prefix_caching: true
    async_scheduling: true
    max_num_batched_tokens: 512
    max_model_len: 8192
    dtype: bfloat16
    devices: "0"
    logits_processors:
      - "vllm_omni.model_executor.models.audex.cfg:AudexCFGLogitsProcessor"
      - "vllm_omni.model_executor.models.audex.tta:TTARVQPhaseMaskLogitsProcessor"
    output_connectors:
      to_stage_1: connector_of_shared_memory
    default_sampling_params:
      temperature: 1.0
      top_k: 80
      max_tokens: 4200
      seed: 0

  - stage_id: 1
    model: hf-audio/xcodec-hubert-general-balanced
    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