[Refactor] Migrate Legacy Stage Configs to the Pipeline Registry (#5031)

Signed-off-by: zwhzzz0821 <2831474076@qq.com>
Co-authored-by: zwhzzz0821 <zwhzzz0821@users.noreply.github.com>
Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

examples/offline_inference/image_to_video/image_to_video.py

@@ -102,6 +102,11 @@ def parse_args() -> argparse.Namespace:
         default=None,
         help="Override model class name (e.g., LTX2ImageToVideoPipeline or LTX23ImageToVideoPipeline).",
     )
+    parser.add_argument(
+        "--deploy-config",
+        default=None,
+        help="Optional deploy config YAML to use for pipeline-backed runs.",
+    )
     parser.add_argument("--image", help="Path to the first-frame or source image.")
     parser.add_argument("--last-image", help="Path to a last-frame condition (used by models such as VACE).")
     parser.add_argument("--mask-image", help="Path to an inpainting mask (used by models such as VACE).")
@@ -453,6 +458,8 @@ def main():
         enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
         profiler_config=args.profiler_config,
     )
+    if args.deploy_config:
+        omni_kwargs["deploy_config"] = args.deploy_config
     if args.quantization is not None:
         omni_kwargs["quantization"] = args.quantization
     # Cosmos3 loads its (gated) guardrail models at build time, so the guardrails

examples/offline_inference/mimo_audio/end2end.py

@@ -219,31 +219,31 @@ def main(args):
     # Build query result based on query type
     # Notice: The audio files used in this example are available at: https://github.com/XiaomiMiMo/MiMo-Audio/tree/main/examples
     if args.query_type == "tts_sft":
-        # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH}  --query-type tts_sft
+        # python3 -u end2end.py --deploy-config ${config_file} --model-name ${MODEL_PATH} --query-type tts_sft
         query_result = query_func(text=text, read_text_only=True)
     elif args.query_type == "tts_sft_with_instruct":
-        # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH}  --query-type tts_sft_with_instruct --instruct "Speak happily in a child's voice"
+        # python3 -u end2end.py --deploy-config ${config_file} --model-name ${MODEL_PATH} --query-type tts_sft_with_instruct --instruct "Speak happily in a child's voice"
         query_result = query_func(text=text, instruct=instruct, read_text_only=True)
     elif args.query_type == "tts_sft_with_audio":
-        # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH}  --query-type tts_sft_with_audio --audio_path "./spoken_dialogue_assistant_turn_1.wav"
+        # python3 -u end2end.py --deploy-config ${config_file} --model-name ${MODEL_PATH} --query-type tts_sft_with_audio --audio-path "./spoken_dialogue_assistant_turn_1.wav"
         audio_list = [get_audio_data(audio_path)]
         query_result = query_func(text=text, read_text_only=True, prompt_speech=audio_path, audio_list=audio_list)
     elif args.query_type == "tts_sft_with_natural_instruction":
-        # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH}  --query-type tts_sft_with_natural_instruction --text "In a panting young male voice, he said: I can't run anymore, wait for me!"
+        # python3 -u end2end.py --deploy-config ${config_file} --model-name ${MODEL_PATH} --query-type tts_sft_with_natural_instruction --text "In a panting young male voice, he said: I can't run anymore, wait for me!"
         query_result = query_func(text=text, read_text_only=False)
     elif args.query_type == "audio_trancribing_sft":
-        # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH}  --query-type audio_trancribing_sft --audio_path "./spoken_dialogue_assistant_turn_1.wav"
+        # python3 -u end2end.py --deploy-config ${config_file} --model-name ${MODEL_PATH} --query-type audio_trancribing_sft --audio-path "./spoken_dialogue_assistant_turn_1.wav"
         audio_path = "spoken_dialogue_assistant_turn_1.wav"
         text = "Please transcribe this audio and repeat it once."
         query_result = query_func(text=text, audio_path=audio_path, use_sostm=True)
     elif args.query_type == "audio_understanding_sft":
-        # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH}  --query-type audio_understanding_sft --text "Summarize the audio." --audio_path "./spoken_dialogue_assistant_turn_1.wav"
+        # python3 -u end2end.py --deploy-config ${config_file} --model-name ${MODEL_PATH} --query-type audio_understanding_sft --text "Summarize the audio." --audio-path "./spoken_dialogue_assistant_turn_1.wav"
         query_result = query_func(text=text, audio_path=audio_path)
     elif args.query_type == "audio_understanding_sft_with_thinking":
-        # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH}  --query-type audio_understanding_sft_with_thinking --text "Summarize the audio." --audio_path "./spoken_dialogue_assistant_turn_1.wav"
+        # python3 -u end2end.py --deploy-config ${config_file} --model-name ${MODEL_PATH} --query-type audio_understanding_sft_with_thinking --text "Summarize the audio." --audio-path "./spoken_dialogue_assistant_turn_1.wav"
         query_result = query_func(text=text, audio_path=audio_path, thinking=True)
     elif args.query_type == "spoken_dialogue_sft_multiturn":
-        # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH}  --query-type spoken_dialogue_sft_multiturn  --audio_path "./prompt_speech_zh_m.wav"
+        # python3 -u end2end.py --deploy-config ${config_file} --model-name ${MODEL_PATH} --query-type spoken_dialogue_sft_multiturn --audio-path "./prompt_speech_zh_m.wav"
         first_turn_text_response = "我没办法获取实时的天气信息。不过呢，你可以试试几个方法来查看今天的天气。首先，你可以用手机自带的天气功能，比如苹果手机的天气应用，或者直接在系统设置里查看。其次，你也可以用一些专业的天气服务，像是国外的AccuWeather、Weather.com，或者国内的中国天气网、墨迹天气等等。再有就是，你还可以在谷歌或者百度里直接搜索你所在的城市加上天气这两个字。如果你能告诉我你所在的城市，我也可以帮你分析一下历史天气趋势，不过最新的数据还是需要你通过官方渠道去获取哦。"
         audio_list = []
         s1_audio_path = "weather_of_today.mp3"
@@ -261,7 +261,7 @@ def main(args):
         ]
         query_result = query_func(message_list, system_prompt=None, ref_audio_path=audio_path, audio_list=audio_list)
     elif args.query_type == "speech2text_dialogue_sft_multiturn":
-        # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH}  --query-type speech2text_dialogue_sft_multiturn
+        # python3 -u end2end.py --deploy-config ${config_file} --model-name ${MODEL_PATH} --query-type speech2text_dialogue_sft_multiturn
         s1_audio_path = "weather_of_today.mp3"
         s2_audio_path = "beijing.mp3"
         audio_list = []
@@ -277,7 +277,7 @@ def main(args):
         ]
         query_result = query_func(message_list, thinking=True, audio_list=audio_list)
     elif args.query_type == "text_dialogue_sft_multiturn":
-        # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH}  --query-type text_dialogue_sft_multiturn
+        # python3 -u end2end.py --deploy-config ${config_file} --model-name ${MODEL_PATH} --query-type text_dialogue_sft_multiturn
         message_list = [
             {"role": "user", "content": "Could you recommend some tourist attractions in China?"},
             {"role": "assistant", "content": "Hello, which city would you like to travel to?"},

examples/offline_inference/step_audio2/README.md

@@ -51,7 +51,7 @@ Single request latency comparison between vLLM-Omni and official Step-Audio2 imp
 
 **Benchmark environment:**
 - GPU: NVIDIA H100 80GB (single card)
-- Model: Step-Audio2-mini
+- Model: Step-Audio-2-mini
 - Warmup: 1 run, Measured: 3 runs (averaged)
 
 ### Async Chunk Streaming Performance
@@ -73,7 +73,7 @@ Comparison between sequential (non-async) and async chunk modes via `/v1/audio/s
 
 **Benchmark environment:**
 - GPU: 4x NVIDIA RTX 3090 24GB (TP=2 for Thinker, 1 GPU for Token2Wav)
-- Model: Step-Audio2-mini
+- Model: Step-Audio-2-mini
 - Endpoint: `/v1/audio/speech` (10 prompts, concurrency=1)
 - Measured via `bench_tts_serve.py`
 
@@ -91,23 +91,22 @@ pip install step-audio2
 
 ## Model Setup
 
+Step-Audio2 ships a custom Transformers configuration, so this model-specific
+example enables remote code loading when it initializes the model.
+
 ### Option 1: Auto-download from HuggingFace (Recommended)
 
-The script will **automatically download** the model on first run:
+The script will **automatically download** the model on first run when a
+HuggingFace model id is passed:
 
 ```bash
-# Just run without specifying --model, it will auto-download stepfun-ai/Step-Audio2-mini
-python end2end.py --query-type audio_to_text
-
-# Or explicitly specify the HuggingFace model
-python end2end.py --query-type audio_to_text --model stepfun-ai/Step-Audio2-mini
+python end2end.py --query-type audio_to_text --model stepfun-ai/Step-Audio-2-mini
 ```
 
 Models will be cached in `~/.cache/huggingface/hub/` for future use.
 
-**Available models**:
-- `stepfun-ai/Step-Audio2-mini` (smaller, faster)
-- `stepfun-ai/Step-Audio2-7B` (larger, better quality)
+**Supported model**:
+- `stepfun-ai/Step-Audio-2-mini`
 
 ### Option 2: Manual Download (for offline use)
 
@@ -115,15 +114,15 @@ Download and use locally:
 
 ```bash
 # Download from HuggingFace
-huggingface-cli download stepfun-ai/Step-Audio2-mini --local-dir ./models/Step-Audio2-mini
+hf download stepfun-ai/Step-Audio-2-mini --local-dir ./models/Step-Audio-2-mini
 
 # Then use the local path
-python end2end.py --query-type audio_to_text --model ./models/Step-Audio2-mini
+python end2end.py --query-type audio_to_text --model ./models/Step-Audio-2-mini
 ```
 
 Ensure the model directory contains:
 ```
-Step-Audio2-mini/
+Step-Audio-2-mini/
 ├── config.json
 ├── model.safetensors (or pytorch_model.bin)
 ├── tokenizer.json
@@ -144,20 +143,18 @@ Transcribe audio to text:
 
 ```bash
 # Quick start - Using default model and test audio
-python end2end.py --query-type audio_to_text
+python end2end.py --query-type audio_to_text \
+    --model stepfun-ai/Step-Audio-2-mini
 
 # Using your own audio file (model will auto-download)
-python end2end.py --query-type audio_to_text \
-    --audio-path /path/to/input.wav
-
-# With specific model
 python end2end.py --query-type audio_to_text \
     --audio-path /path/to/input.wav \
-    --model stepfun-ai/Step-Audio2-7B
+    --model stepfun-ai/Step-Audio-2-mini
 
 # With custom question
 python end2end.py --query-type audio_to_text \
     --audio-path input.wav \
+    --model stepfun-ai/Step-Audio-2-mini \
     --question "What is the speaker saying?"
 ```
 
@@ -170,12 +167,8 @@ Convert text to speech:
 ```bash
 # Basic TTS (model auto-downloads)
 python end2end.py --query-type text_to_audio \
-    --text "Hello, this is a test of Step Audio 2 synthesis."
-
-# With specific model
-python end2end.py --query-type text_to_audio \
-    --text "Hello, this is a test." \
-    --model stepfun-ai/Step-Audio2-7B
+    --text "Hello, this is a test of Step Audio 2 synthesis." \
+    --model stepfun-ai/Step-Audio-2-mini
 ```
 
 **Note**: Speaker voice is controlled by the `STEP_AUDIO2_DEFAULT_PROMPT_WAV` environment variable or the default prompt wav bundled with the model.
@@ -191,12 +184,8 @@ Process input audio and generate output audio:
 ```bash
 # Basic voice conversion (model auto-downloads)
 python end2end.py --query-type audio_to_audio \
-    --audio-path /path/to/source_audio.wav
-
-# With specific model
-python end2end.py --query-type audio_to_audio \
-    --audio-path source.wav \
-    --model stepfun-ai/Step-Audio2-7B
+    --audio-path /path/to/source_audio.wav \
+    --model stepfun-ai/Step-Audio-2-mini
 ```
 
 This mode:
@@ -210,8 +199,14 @@ This mode:
 ```bash
 # Use custom stage configuration
 python end2end.py --query-type audio_to_text \
+    --model stepfun-ai/Step-Audio-2-mini \
     --stage-configs-path /path/to/custom_config.yaml
 
+# Use custom deploy configuration
+python end2end.py --query-type audio_to_text \
+    --model stepfun-ai/Step-Audio-2-mini \
+    --deploy-config /path/to/step_audio_2_asr.yaml
+
 # Multiple prompts (for batch testing)
 python end2end.py --query-type audio_to_text \
     --audio-path input.wav \
@@ -246,14 +241,14 @@ python end2end.py --query-type text_to_audio \
 
 ## Configuration
 
-### Stage Configuration
+### Deploy Configuration
 
-The default configuration (`step_audio_2.yaml`) uses:
+The default configuration (`vllm_omni/deploy/step_audio_2.yaml`) uses:
 
-- **Stage 0 (Thinker)**: GPU 0, 80% memory
-- **Stage 1 (Token2Wav)**: GPU 1, 30% memory
+- **Stage 0 (Thinker)**: GPUs 0-1 with tensor parallel size 2, 70% memory
+- **Stage 1 (Token2Wav)**: GPU 1, 20% memory
 
-For **single GPU** setup, edit the config to use `devices: "0"` for both stages.
+For **single GPU** setup, edit a deploy config copy to use `devices: "0"` for both stages.
 
 ### Sampling Parameters
 
@@ -351,7 +346,7 @@ Complete example from audio to final output:
 # 1. ASR: Transcribe audio
 python end2end.py --query-type audio_to_text \
     --audio-path interview.wav \
-    --model ./models/Step-Audio2-7B \
+    --model ./models/Step-Audio-2-mini \
     --output-dir ./outputs
 
 # 2. Check the transcription
@@ -361,7 +356,7 @@ cat ./outputs/00000_text.txt
 STEP_AUDIO2_DEFAULT_PROMPT_WAV=./speaker_samples/female_voice.wav \
 python end2end.py --query-type text_to_audio \
     --text "The quick brown fox jumps over the lazy dog" \
-    --model ./models/Step-Audio2-7B \
+    --model ./models/Step-Audio-2-mini \
     --output-dir ./outputs
 
 # 4. Listen to the result

examples/offline_inference/step_audio2/end2end.py

@@ -232,32 +232,43 @@ def main(args):
         raise ValueError(f"Unknown query type: {args.query_type}")
 
     # Initialize vLLM-Omni with Step-Audio2
-    # Resolve stage config path
+    # Resolve deploy config path. Keep explicit stage_configs_path for legacy
+    # custom configs, but use bundled deploy configs by default.
+    deploy_config_path = None
+    stage_config_path = None
+    if args.stage_configs_path and args.deploy_config:
+        raise ValueError("--stage-configs-path and --deploy-config are mutually exclusive")
     if args.stage_configs_path:
         stage_config_path = args.stage_configs_path
+    elif args.deploy_config:
+        deploy_config_path = args.deploy_config
     else:
-        configs_dir = Path(__file__).parent.parent.parent.parent / "vllm_omni/model_executor/stage_configs"
+        configs_dir = Path(__file__).parent.parent.parent.parent / "vllm_omni/deploy"
         if args.query_type == "audio_to_text":
             # ASR only needs Thinker (Stage 0), no Token2Wav
-            stage_config_path = str(configs_dir / "step_audio_2_asr.yaml")
-            if not os.path.exists(stage_config_path):
-                stage_config_path = str(configs_dir / "step_audio_2.yaml")
+            deploy_config_path = str(configs_dir / "step_audio_2_asr.yaml")
         else:
             # TTS/S2ST need both stages
-            stage_config_path = str(configs_dir / "step_audio_2.yaml")
+            deploy_config_path = str(configs_dir / "step_audio_2.yaml")
 
-    omni_llm = Omni(
-        model=model_name,
-        log_stats=args.enable_stats,
-        log_file=("step_audio2_pipeline.log" if args.enable_stats else None),
-        init_sleep_seconds=args.init_sleep_seconds,
-        batch_timeout=args.batch_timeout,
-        init_timeout=args.init_timeout,
-        shm_threshold_bytes=args.shm_threshold_bytes,
-        stage_configs_path=stage_config_path,
-        worker_backend=args.worker_backend,
-        ray_address=args.ray_address,
-    )
+    omni_kwargs = {
+        "model": model_name,
+        "log_stats": args.enable_stats,
+        "log_file": ("step_audio2_pipeline.log" if args.enable_stats else None),
+        "init_sleep_seconds": args.init_sleep_seconds,
+        "batch_timeout": args.batch_timeout,
+        "init_timeout": args.init_timeout,
+        "shm_threshold_bytes": args.shm_threshold_bytes,
+        "worker_backend": args.worker_backend,
+        "ray_address": args.ray_address,
+        "trust_remote_code": True,
+    }
+    if stage_config_path is not None:
+        omni_kwargs["stage_configs_path"] = stage_config_path
+    else:
+        omni_kwargs["deploy_config"] = deploy_config_path
+
+    omni_llm = Omni(**omni_kwargs)
 
     # Configure sampling parameters for each stage
     # Stage 0 (Thinker): Generate text and audio tokens
@@ -405,16 +416,21 @@ def parse_args():
         help=(
             "Model path (required). Can be:\n"
             "  - Local path: /path/to/step-audio-2\n"
-            "  - HuggingFace ID: stepfun-ai/Step-Audio2-mini"
+            "  - HuggingFace ID: stepfun-ai/Step-Audio-2-mini"
         ),
     )
     parser.add_argument(
         "--stage-configs-path",
         type=str,
         default=None,
-        help="Path to stage config YAML file (default: use built-in step_audio_2.yaml)",
+        help="Path to a legacy stage config YAML file (default: use bundled deploy config)",
+    )
+    parser.add_argument(
+        "--deploy-config",
+        type=str,
+        default=None,
+        help="Path to a deploy config YAML file (default: use bundled Step-Audio2 deploy config)",
     )
-
     # Query configuration
     parser.add_argument(
         "--query-type",

examples/offline_inference/text_to_video/text_to_video.py

@@ -124,6 +124,11 @@ def parse_args() -> argparse.Namespace:
         default=None,
         help="Override model class name (e.g., LTX2TwoStagesVideoPipeline).",
     )
+    parser.add_argument(
+        "--deploy-config",
+        default=None,
+        help="Optional deploy config YAML to use for pipeline-backed runs.",
+    )
     parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.")
     parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
     parser.add_argument(
@@ -375,6 +380,8 @@ def main():
         enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
         profiler_config=args.profiler_config,
     )
+    if args.deploy_config:
+        omni_kwargs["deploy_config"] = args.deploy_config
     if args.boundary_ratio is not None:
         omni_kwargs["boundary_ratio"] = args.boundary_ratio
     if args.flow_shift is not None:

examples/online_serving/step_audio2/README.md

@@ -10,15 +10,15 @@ Please refer to [README.md](../../../README.md)
 
 ```bash
 # Async chunk mode (recommended — lower first-packet latency for TTS)
-vllm serve stepfun-ai/Step-Audio2-mini --omni --port 8092 \
-    --stage-configs-path vllm_omni/model_executor/stage_configs/step_audio_2_async_chunk.yaml \
+vllm serve stepfun-ai/Step-Audio-2-mini --omni --port 8092 \
+    --deploy-config vllm_omni/deploy/step_audio_2_async_chunk.yaml \
     --trust-remote-code --enforce-eager
 ```
 
 Sequential mode:
 ```bash
-vllm serve stepfun-ai/Step-Audio2-mini --omni --port 8092 \
-    --stage-configs-path vllm_omni/model_executor/stage_configs/step_audio_2.yaml \
+vllm serve stepfun-ai/Step-Audio-2-mini --omni --port 8092 \
+    --deploy-config vllm_omni/deploy/step_audio_2.yaml \
     --trust-remote-code --enforce-eager
 ```
 
@@ -51,7 +51,7 @@ Or via curl:
 ```bash
 curl -X POST http://localhost:8092/v1/audio/speech \
     -H "Content-Type: application/json" \
-    -d '{"model":"stepfun-ai/Step-Audio2-mini","input":"你好世界","voice":"default"}' \
+    -d '{"model":"stepfun-ai/Step-Audio-2-mini","input":"你好世界","voice":"default"}' \
     --output output.wav
 ```
 
@@ -139,7 +139,7 @@ Step-Audio2 uses the OpenAI-compatible chat completions API:
 
 ```json
 {
-  "model": "stepfun-ai/Step-Audio2-mini",
+  "model": "stepfun-ai/Step-Audio-2-mini",
   "messages": [
     {
       "role": "system",

examples/online_serving/step_audio2/openai_chat_completion_client.py

@@ -9,7 +9,7 @@ Step-Audio2 is a two-stage model:
 
 Usage examples:
   # Start server first:
-  vllm serve stepfun-ai/Step-Audio2-mini --omni --port 8092
+  vllm serve stepfun-ai/Step-Audio-2-mini --omni --port 8092
 
   # Audio to text (ASR)
   python openai_chat_completion_client.py --query-type audio_to_text
@@ -311,7 +311,7 @@ def parse_args():
         "--model",
         "-m",
         type=str,
-        default="stepfun-ai/Step-Audio2-mini",
+        default="stepfun-ai/Step-Audio-2-mini",
         help="Model name",
     )
     parser.add_argument(

examples/online_serving/step_audio2/openai_speech_client.py

@@ -87,7 +87,7 @@ def parse_args():
         "--model",
         "-m",
         type=str,
-        default="stepfun-ai/Step-Audio2-mini",
+        default="stepfun-ai/Step-Audio-2-mini",
         help="Model name/path",
     )
     parser.add_argument(

examples/online_serving/step_audio2/run_curl.sh

@@ -110,7 +110,7 @@ output=$(curl -sS -X POST "${API_BASE}/v1/chat/completions" \
     -H "Content-Type: application/json" \
     -d @- <<EOF
 {
-  "model": "stepfun-ai/Step-Audio2-mini",
+  "model": "stepfun-ai/Step-Audio-2-mini",
   "sampling_params_list": $sampling_params_list,
   "messages": [
     {