# stage config for running dynin_omni with a 3-stage architecture.
# this config is intended for e2e smoke tests.

stage_args:
  - stage_id: 0
    stage_type: llm
    runtime:
      process: true
      devices: "0"
      max_batch_size: 1
    engine_args:
      model_stage: token2text
      model_arch: DyninOmniForConditionalGeneration
      worker_type: generation
      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
      engine_output_type: latent
      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2text_to_token2image_full_payload
      trust_remote_code: true
      gpu_memory_utilization: 0.5
      enforce_eager: true
      enable_prefix_caching: false
      async_scheduling: false
      max_num_batched_tokens: 4096
    is_comprehension: true
    final_output: true
    final_output_type: text

  - stage_id: 1
    stage_type: llm
    runtime:
      process: true
      devices: "0"
      max_batch_size: 1
    engine_args:
      model_stage: token2image
      model_arch: DyninOmniForConditionalGeneration
      worker_type: generation
      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
      engine_output_type: latent
      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2image_to_token2audio_full_payload
      trust_remote_code: true
      gpu_memory_utilization: 0.2
      enforce_eager: true
      enable_prefix_caching: false
      async_scheduling: false
      max_num_batched_tokens: 4096
    engine_input_source: [0]
    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2text_to_token2image
    final_output: true
    final_output_type: image

  - stage_id: 2
    stage_type: llm
    runtime:
      process: true
      devices: "0"
      max_batch_size: 1
    engine_args:
      model_stage: token2audio
      model_arch: DyninOmniForConditionalGeneration
      worker_type: generation
      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
      engine_output_type: latent
      trust_remote_code: true
      gpu_memory_utilization: 0.2
      enforce_eager: true
      enable_prefix_caching: false
      async_scheduling: false
      max_num_batched_tokens: 4096
    engine_input_source: [1]
    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2image_to_token2audio
    final_output: true
    final_output_type: audio

runtime:
  enabled: true
  edges:
    - from: 0
      to: 1
    - from: 1
      to: 2