[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>
This commit is contained in:
authored and GitHub committed 2026-07-18 09:34:33 +00:00
1 parent 793d3e9ab7
commit 4fbbdfd250
41 files changed
+520 -516

No files matched your search

@@ -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",