diff --git a/docs/cli/serve.md b/docs/cli/serve.md index 9d1747e0b..ff9e561e0 100644 --- a/docs/cli/serve.md +++ b/docs/cli/serve.md @@ -8,8 +8,7 @@ The stage-based CLI is designed for deployments that require launching each pipe - For **migrated models** that utilize the bundled deployment YAML configurations located in `vllm_omni/deploy/`, the `--deploy-config` flag is only required to override the default configuration. By default, executing `vllm serve MODEL --omni ...` automatically loads the bundled deployment configuration. -- For **legacy models** utilizing configuration files located in - `vllm_omni/model_executor/stage_configs/`, the `--stage-configs-path` parameter remains mandatory. +- For custom legacy `stage_args` YAMLs, pass the file with `--stage-configs-path`. Example: Initializing Stage 0 (Orchestrator and API Server): The commands below show a common device mapping where Stage 0 uses GPU 0 and diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index 1f14d21fd..1aa4ca1b3 100644 --- a/docs/configuration/stage_configs.md +++ b/docs/configuration/stage_configs.md @@ -3,7 +3,7 @@ In vLLM-Omni, a model's `PipelineConfig` defines its fixed stage topology, while a deploy configuration controls how those stages run. Models that have not migrated to this split still use legacy stage configurations with a `stage_args` schema. !!! note - Default deploy config YAMLs (for example, `vllm_omni/deploy/qwen2_5_omni.yaml`, `vllm_omni/deploy/qwen3_omni_moe.yaml`, and `vllm_omni/deploy/qwen3_tts.yaml`) are bundled and loaded automatically when neither `--stage-configs-path` nor `--deploy-config` is provided — the model registry resolves the right pipeline + deploy YAML by `model_type`. The bundled defaults have been verified on 1xH100 for Qwen2.5-Omni and 2xH100 for Qwen3-Omni. Models that have not yet migrated to the new schema continue to use the legacy `vllm_omni/model_executor/stage_configs/.yaml` files via `--stage-configs-path`. + Default deploy config YAMLs (for example, `vllm_omni/deploy/qwen2_5_omni.yaml`, `vllm_omni/deploy/qwen3_omni_moe.yaml`, and `vllm_omni/deploy/qwen3_tts.yaml`) are bundled and loaded automatically when neither `--stage-configs-path` nor `--deploy-config` is provided — the model registry resolves the right pipeline + deploy YAML by `model_type`. Custom legacy `stage_args` YAMLs are still accepted via `--stage-configs-path`. ## Deploy configuration schema diff --git a/docs/contributing/model/adding_omni_model.md b/docs/contributing/model/adding_omni_model.md index c20e94eda..5032c9805 100644 --- a/docs/contributing/model/adding_omni_model.md +++ b/docs/contributing/model/adding_omni_model.md @@ -34,6 +34,7 @@ vllm_omni/model_executor/models/ └── your_model_name/ # Model directory (e.g., qwen3_omni) ├── __init__.py # Exports main model class ├── your_model.py # Main unified model class + ├── pipeline.py # Pipeline topology registry entry ├── your_model_stage1_implementation.py # Stage 1 implementation (e.g., thinker) ├── your_model_stage2_implementation.py # Stage 2 implementation (e.g., talker) └── your_model_stage3_implementation.py # Stage 3 implementation (e.g., code2wav) @@ -42,10 +43,14 @@ vllm_omni/model_executor/models/ vllm_omni/model_executor/stage_input_processors/ └── your_model_name.py # Stage transition processors -vllm_omni/model_executor/stage_configs/ -└── your_model_name.yaml # Stage configuration file +vllm_omni/deploy/ +└── your_model_name.yaml # Deployment configuration file ``` +New in-tree models should register their `PipelineConfig` in +`vllm_omni/config/pipeline_registry.py`; use `--stage-configs-path` only for +custom legacy `stage_args` YAMLs. + ## Step-by-Step Implementation ### Step 1: Create the Model Directory diff --git a/docs/contributing/model/adding_tts_model.md b/docs/contributing/model/adding_tts_model.md index 3e5ae30df..b9bb9484f 100644 --- a/docs/contributing/model/adding_tts_model.md +++ b/docs/contributing/model/adding_tts_model.md @@ -175,15 +175,16 @@ vllm_omni/model_executor/models/ your_model_name/ __init__.py your_model.py # Unified class (stage dispatch) + pipeline.py # Pipeline topology registry entry your_model_ar_stage.py # Stage 0: AR stage your_model_decoder.py # Stage 1: audio decoder vllm_omni/model_executor/stage_input_processors/ your_model_name.py # Stage 0 -> Stage 1 transition -vllm_omni/model_executor/stage_configs/ - your_model_name.yaml # Batch mode config - your_model_name_async_chunk.yaml # Streaming mode config +vllm_omni/deploy/ + your_model_name.yaml # Batch/deploy config + your_model_name_async_chunk.yaml # Streaming/deploy config ``` ### Example placement @@ -445,6 +446,13 @@ Key configuration fields: | `final_output` | Whether this stage produces the final user-facing output | | `final_output_type` | Type of final output (`audio`, `text`, etc.) | +!!! note + New in-tree models should define frozen topology in `models//pipeline.py`, + register it in `vllm_omni/config/pipeline_registry.py`, and put deployment + knobs in `vllm_omni/deploy/.yaml`. Legacy `stage_args` YAMLs are still + accepted for custom configs via `--stage-configs-path`, but should not be used + for new bundled defaults. + ### Batch mode ```yaml @@ -771,8 +779,9 @@ directly from the model's own generator. vllm_omni/model_executor/models/your_model_name/ __init__.py modeling_your_model_name.py # unified class: load_weights + forward + streaming + pipeline.py # single-stage pipeline topology -vllm_omni/model_executor/stage_configs/your_model_name.yaml +vllm_omni/deploy/your_model_name.yaml ``` No stage input processor is needed. diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py index 2d1e70697..d64367bc2 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/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 diff --git a/examples/offline_inference/mimo_audio/end2end.py b/examples/offline_inference/mimo_audio/end2end.py index 2fe46c972..139351aab 100644 --- a/examples/offline_inference/mimo_audio/end2end.py +++ b/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?"}, diff --git a/examples/offline_inference/step_audio2/README.md b/examples/offline_inference/step_audio2/README.md index 84d0ecdde..abd21f371 100644 --- a/examples/offline_inference/step_audio2/README.md +++ b/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 diff --git a/examples/offline_inference/step_audio2/end2end.py b/examples/offline_inference/step_audio2/end2end.py index c67f44926..8cc7e2a26 100644 --- a/examples/offline_inference/step_audio2/end2end.py +++ b/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", diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index 0f9bb8c8e..addf7cb7a 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/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: diff --git a/examples/online_serving/step_audio2/README.md b/examples/online_serving/step_audio2/README.md index e93f58ccc..25e2a6785 100644 --- a/examples/online_serving/step_audio2/README.md +++ b/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", diff --git a/examples/online_serving/step_audio2/openai_chat_completion_client.py b/examples/online_serving/step_audio2/openai_chat_completion_client.py index 2b91dacb2..94acb5533 100644 --- a/examples/online_serving/step_audio2/openai_chat_completion_client.py +++ b/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( diff --git a/examples/online_serving/step_audio2/openai_speech_client.py b/examples/online_serving/step_audio2/openai_speech_client.py index 740b36805..762ca72a6 100644 --- a/examples/online_serving/step_audio2/openai_speech_client.py +++ b/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( diff --git a/examples/online_serving/step_audio2/run_curl.sh b/examples/online_serving/step_audio2/run_curl.sh index 3811308e7..3a3b1c9f4 100755 --- a/examples/online_serving/step_audio2/run_curl.sh +++ b/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 @- < token2wav. +# +# Sequential mode waits for the full Thinker output before running Token2Wav. +async_chunk: false +trust_remote_code: true +enable_prefix_caching: false + +stages: + - stage_id: 0 + devices: "0,1" + tensor_parallel_size: 2 + max_num_seqs: 1 + gpu_memory_utilization: 0.7 + enforce_eager: true + async_scheduling: false + max_num_batched_tokens: 8192 + default_sampling_params: + temperature: 0.7 + top_p: 0.9 + top_k: -1 + max_tokens: 1024 + seed: 42 + repetition_penalty: 1.05 + + - stage_id: 1 + devices: "1" + max_num_seqs: 1 + gpu_memory_utilization: 0.2 + enforce_eager: true + async_scheduling: false + max_num_batched_tokens: 4096 + hf_overrides: + architectures: ["StepAudio2Token2WavForConditionalGeneration"] + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 1 + seed: 42 diff --git a/vllm_omni/deploy/step_audio_2_asr.yaml b/vllm_omni/deploy/step_audio_2_asr.yaml new file mode 100644 index 000000000..0eb0a8523 --- /dev/null +++ b/vllm_omni/deploy/step_audio_2_asr.yaml @@ -0,0 +1,22 @@ +# Step-Audio2 ASR deploy: thinker only. +pipeline: step_audio_2_asr +async_chunk: false +trust_remote_code: true +enable_prefix_caching: false + +stages: + - stage_id: 0 + devices: "0,1" + tensor_parallel_size: 2 + max_num_seqs: 1 + gpu_memory_utilization: 0.8 + enforce_eager: true + async_scheduling: false + max_num_batched_tokens: 2048 + default_sampling_params: + temperature: 0.7 + top_p: 0.9 + top_k: -1 + max_tokens: 1024 + seed: 42 + repetition_penalty: 1.05 diff --git a/vllm_omni/deploy/step_audio_2_async_chunk.yaml b/vllm_omni/deploy/step_audio_2_async_chunk.yaml new file mode 100644 index 000000000..75b5cf901 --- /dev/null +++ b/vllm_omni/deploy/step_audio_2_async_chunk.yaml @@ -0,0 +1,37 @@ +# Step-Audio2 deploy: thinker -> token2wav with async chunk streaming. +async_chunk: true +trust_remote_code: true +enable_prefix_caching: false + +stages: + - stage_id: 0 + devices: "0,1" + tensor_parallel_size: 2 + max_num_seqs: 1 + gpu_memory_utilization: 0.7 + enforce_eager: true + async_scheduling: false + max_num_batched_tokens: 1024 + default_sampling_params: + temperature: 0.7 + top_p: 0.9 + top_k: -1 + max_tokens: 1024 + seed: 42 + repetition_penalty: 1.05 + + - stage_id: 1 + devices: "1" + max_num_seqs: 1 + gpu_memory_utilization: 0.2 + enforce_eager: true + async_scheduling: false + max_num_batched_tokens: 2048 + hf_overrides: + architectures: ["StepAudio2Token2WavForConditionalGeneration"] + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 1 + seed: 42 diff --git a/vllm_omni/deploy/wan2_2_ti2v.yaml b/vllm_omni/deploy/wan2_2_ti2v.yaml new file mode 100644 index 000000000..bdb5cfecb --- /dev/null +++ b/vllm_omni/deploy/wan2_2_ti2v.yaml @@ -0,0 +1,17 @@ +# Wan2.2 TI2V-5B deploy. +# +# Single GPU by default. Increase tensor_parallel_size and devices for A14B MoE variants. +async_chunk: false +trust_remote_code: true +distributed_executor_backend: mp + +stages: + - stage_id: 0 + devices: "0" + max_num_seqs: 1 + tensor_parallel_size: 1 + enforce_eager: true + # Keep VAE tiling on by default for higher resolutions / longer clips. + vae_use_tiling: true + default_sampling_params: + seed: 42 diff --git a/vllm_omni/model_executor/models/hunyuan_video/__init__.py b/vllm_omni/model_executor/models/hunyuan_video/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/vllm_omni/model_executor/models/hunyuan_video/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm_omni/model_executor/models/hunyuan_video/pipeline.py b/vllm_omni/model_executor/models/hunyuan_video/pipeline.py new file mode 100644 index 000000000..341d14e5d --- /dev/null +++ b/vllm_omni/model_executor/models/hunyuan_video/pipeline.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""HunyuanVideo-1.5 pipeline topology.""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +HUNYUAN_VIDEO_15_PIPELINE = PipelineConfig( + model_type="hunyuan_video_15", + model_arch="HunyuanVideo15Pipeline", + diffusers_class_name="HunyuanVideo15Pipeline", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="dit", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(), + final_output=True, + final_output_type="video", + model_arch="HunyuanVideo15Pipeline", + ), + ), +) diff --git a/vllm_omni/model_executor/models/step_audio2/__init__.py b/vllm_omni/model_executor/models/step_audio2/__init__.py index 0b6f07e28..e69de29bb 100644 --- a/vllm_omni/model_executor/models/step_audio2/__init__.py +++ b/vllm_omni/model_executor/models/step_audio2/__init__.py @@ -1,3 +0,0 @@ -from .step_audio2 import StepAudio2ForConditionalGeneration - -__all__ = ["StepAudio2ForConditionalGeneration"] diff --git a/vllm_omni/model_executor/models/step_audio2/pipeline.py b/vllm_omni/model_executor/models/step_audio2/pipeline.py new file mode 100644 index 000000000..7ba6d9ee1 --- /dev/null +++ b/vllm_omni/model_executor/models/step_audio2/pipeline.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Step-Audio2 pipeline topologies (frozen).""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +_PROC = "vllm_omni.model_executor.stage_input_processors.step_audio2" +_HF_ARCHITECTURES = ( + "StepAudio2ForCausalLM", + "StepAudio2ForConditionalGeneration", +) + + +STEP_AUDIO2_PIPELINE = PipelineConfig( + model_type="step_audio_2", + model_arch="StepAudio2ForConditionalGeneration", + hf_architectures=_HF_ARCHITECTURES, + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="step_audio2_thinker", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + final_output=True, + final_output_type="text", + owns_tokenizer=True, + requires_multimodal_data=True, + engine_output_type="latent", + model_arch="StepAudio2ThinkerForConditionalGeneration", + async_chunk_process_next_stage_input_func=f"{_PROC}.thinker2token2wav_async_chunk", + sampling_constraints={"detokenize": True}, + ), + StagePipelineConfig( + stage_id=1, + model_stage="token2wav", + execution_type=StageExecutionType.LLM_GENERATION, + input_sources=(0,), + final_output=True, + final_output_type="audio", + engine_output_type="audio", + model_arch="StepAudio2Token2WavModel", + sync_process_input_func=f"{_PROC}.thinker2token2wav", + sampling_constraints={"detokenize": False}, + ), + ), +) + + +STEP_AUDIO2_ASR_PIPELINE = PipelineConfig( + model_type="step_audio_2_asr", + model_arch="StepAudio2ForConditionalGeneration", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="step_audio2_thinker", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + final_output=True, + final_output_type="text", + owns_tokenizer=True, + requires_multimodal_data=True, + engine_output_type="text", + model_arch="StepAudio2ThinkerForConditionalGeneration", + sampling_constraints={"detokenize": True}, + ), + ), +) diff --git a/vllm_omni/model_executor/models/wan2_2/__init__.py b/vllm_omni/model_executor/models/wan2_2/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/vllm_omni/model_executor/models/wan2_2/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm_omni/model_executor/models/wan2_2/pipeline.py b/vllm_omni/model_executor/models/wan2_2/pipeline.py new file mode 100644 index 000000000..7059845fa --- /dev/null +++ b/vllm_omni/model_executor/models/wan2_2/pipeline.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Wan2.2 TI2V pipeline topology.""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +WAN2_2_TI2V_PIPELINE = PipelineConfig( + model_type="wan2_2_ti2v", + model_arch="WanPipeline", + diffusers_class_name="WanPipeline", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="dit", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(), + final_output=True, + final_output_type="video", + model_arch="WanPipeline", + ), + ), +) diff --git a/vllm_omni/model_executor/stage_configs/__init__.py b/vllm_omni/model_executor/stage_configs/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/vllm_omni/model_executor/stage_configs/hunyuan_video_15_dit_fp8.yaml b/vllm_omni/model_executor/stage_configs/hunyuan_video_15_dit_fp8.yaml deleted file mode 100644 index 097afc020..000000000 --- a/vllm_omni/model_executor/stage_configs/hunyuan_video_15_dit_fp8.yaml +++ /dev/null @@ -1,38 +0,0 @@ -# Stage config for running HunyuanVideo-1.5 DiT with ModelOpt FP8 auto-detect. -# Single GPU. Bump `tensor_parallel_size` and `devices` for multi-GPU TP. -# -# Use with a ModelOpt FP8 checkpoint (e.g. produced by -# examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py). - -stage_args: - - stage_id: 0 - stage_type: diffusion - runtime: - devices: "0" - max_batch_size: 1 - engine_args: - model_stage: dit - model_class_name: HunyuanVideo15Pipeline - max_num_seqs: 1 - enforce_eager: true - trust_remote_code: true - distributed_executor_backend: "mp" - # HunyuanVideo-1.5 peak VRAM is dominated by the single-shot VAE decode of - # the full video (the DiT contributes little), so tiling the decode is the - # main memory lever — far more than FP8 (e.g. 480x832x33: ~117 -> ~47 GiB). - # Default it on here; override with --vae-use-tiling/false if you have headroom. - vae_use_tiling: true - parallel_config: - tensor_parallel_size: 1 - - final_output: true - final_output_type: video - is_comprehension: false - default_sampling_params: - seed: 42 - -runtime: - enabled: true - defaults: - window_size: -1 - max_inflight: 1 diff --git a/vllm_omni/model_executor/stage_configs/mimo_audio.yaml b/vllm_omni/model_executor/stage_configs/mimo_audio.yaml deleted file mode 100644 index b81f27fe4..000000000 --- a/vllm_omni/model_executor/stage_configs/mimo_audio.yaml +++ /dev/null @@ -1,61 +0,0 @@ - -# MiMo Audio deploy: 2-stage thinker+talker → code2wav. -# -# Default mode is async-chunk streaming on a single GPU (both stages on -# device 0) through SharedMemoryConnector. For the legacy 2-GPU sync -# pipeline (stage 1 on a second card), pass ``--no-async-chunk -# --stage-overrides '{"1": {"devices": "1", "max_model_len": 18192, "max_num_batched_tokens": 18192}}'`` -# to the serve CLI. -async_chunk: true -dtype: bfloat16 - -connectors: - connector_of_shared_memory: - name: SharedMemoryConnector - extra: - shm_threshold_bytes: 365536 - codec_streaming: true - connector_get_sleep_s: 0.001 - connector_get_max_wait_first_chunk: 3000 - connector_get_max_wait: 300 - codec_chunk_frames: 30 - codec_left_context_frames: 40 - -stages: - - stage_id: 0 - max_num_seqs: 1 - gpu_memory_utilization: 0.5 - enforce_eager: true - trust_remote_code: true - enable_prefix_caching: false - max_num_batched_tokens: 8192 - max_model_len: 8192 - devices: "0" - output_connectors: - to_stage_1: connector_of_shared_memory - default_sampling_params: - temperature: 0.6 - top_p: 0.95 - top_k: 50 - max_tokens: 18192 - seed: 42 - repetition_penalty: 1.1 - - - stage_id: 1 - max_num_seqs: 1 - gpu_memory_utilization: 0.4 - 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 - top_p: 1.0 - top_k: -1 - max_tokens: 18192 - seed: 42 diff --git a/vllm_omni/model_executor/stage_configs/step_audio_2.yaml b/vllm_omni/model_executor/stage_configs/step_audio_2.yaml deleted file mode 100644 index b560fea9b..000000000 --- a/vllm_omni/model_executor/stage_configs/step_audio_2.yaml +++ /dev/null @@ -1,122 +0,0 @@ -# Stage config for Step-Audio2 with vLLM-Omni -# -# Two-stage pipeline: -# Stage 0 (Thinker): Audio understanding → Text + Audio tokens -# Stage 1 (Token2Wav): Audio tokens → Waveform (24kHz) -# -# Hardware requirement: 1-2 GPUs (can run on single GPU with devices: "0" for both) - -# Model configuration reference (for documentation) -# Actual values are set in hf_overrides below -# -# Token configuration: -# text_max: 151688 # Maximum text token ID (inclusive) -# audio_start: 151696 # First audio token ID in absolute vocabulary -# audio_vocab_size: 6562 # Total number of audio tokens -# audio_eos: 6561 # Audio EOS token ID (relative to audio_start) -# audio_patch_token_id: 151690 # placeholder token ID -# -# Encoder configuration: -# n_mels: 128 # Number of mel frequency bins -# n_audio_ctx: 1500 # Audio context length -# n_audio_state: 512 # Audio encoder hidden state dimension -# n_audio_head: 8 # Number of attention heads -# n_audio_layer: 6 # Number of encoder layers -# kernel_size: 3 # Convolution kernel size for adapter -# adapter_stride: 2 # Stride for adapter convolution - -stage_args: - # Stage 0: Thinker (Audio Understanding + Token Generation) - - stage_id: 0 - runtime: - process: true # Run in separate process - devices: "0,1,2,3" # Use 3 GPUs for tensor parallelism - max_batch_size: 1 # Single sample at a time - engine_args: - model_stage: step_audio2_thinker - model_arch: StepAudio2ThinkerForConditionalGeneration - tensor_parallel_size: 4 # Split model across 3 GPUs - # NOTE: All default values are defined in step_audio2_constants.py (DEFAULT_MODEL_CONFIG) - # Only override here if your model differs from the defaults - # - # Default values (for reference, from constants.py): - # hidden_size: 4096 # Qwen2-7B (change to 1536 for 1.5B, 896 for 0.5B) - # text_max: 151688 - # audio_start: 151696 - # audio_vocab_size: 6562 - # audio_eos: 6561 - # audio_patch_token_id: 151690 - # n_mels: 128 - # n_audio_ctx: 1500 - # n_audio_state: 512 - # n_audio_head: 8 - # n_audio_layer: 6 - # kernel_size: 3 - # adapter_stride: 2 - # - # Uncomment and modify only if needed: - # hidden_size: 4096 - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.8 - enforce_eager: true # Required for now - trust_remote_code: true - enable_prefix_caching: false - max_num_batched_tokens: 8192 - engine_output_type: text # Outputs text + audio tokens - is_comprehension: true # This stage does comprehension - final_output: true # Can be final output (for ASR mode) - final_output_type: text - default_sampling_params: - temperature: 0.7 - top_p: 0.9 - top_k: -1 - max_tokens: 1024 - seed: 42 - detokenize: true - repetition_penalty: 1.05 - - # Stage 1: Token2Wav (Audio Synthesis) - - stage_id: 1 - runtime: - process: true - devices: "3" # Use GPU 3 (separate from Stage 0's 0-2) - max_batch_size: 1 - engine_args: - model_stage: token2wav - model_arch: StepAudio2Token2WavModel - hf_overrides: - architectures: ["StepAudio2Token2WavForConditionalGeneration"] - worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker - scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler - - gpu_memory_utilization: 0.3 # Token2Wav needs less memory - enforce_eager: true - trust_remote_code: true - async_scheduling: false - enable_prefix_caching: false - max_num_batched_tokens: 4096 - engine_output_type: audio # Outputs waveform - engine_input_source: [0] # Takes input from Stage 0 - custom_process_input_func: vllm_omni.model_executor.stage_input_processors.step_audio2.thinker2token2wav - final_output: true # Final output stage - final_output_type: audio - default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 - max_tokens: 1 - seed: 42 - detokenize: false - -# Runtime configuration -runtime: - enabled: true - defaults: - window_size: -1 # Wait for full input before triggering - max_inflight: 1 # Process one at a time - - edges: - - from: 0 # Thinker → Token2Wav - to: 1 - window_size: -1 # Wait for complete output from Thinker diff --git a/vllm_omni/model_executor/stage_configs/step_audio_2_asr.yaml b/vllm_omni/model_executor/stage_configs/step_audio_2_asr.yaml deleted file mode 100644 index 1cecfddc5..000000000 --- a/vllm_omni/model_executor/stage_configs/step_audio_2_asr.yaml +++ /dev/null @@ -1,35 +0,0 @@ -# Stage config for Step-Audio2 ASR (speech-to-text only) -# -# Single-stage pipeline: only Thinker, no Token2Wav. -# Use this for audio_to_text tasks where no audio output is needed. - -stage_args: - - stage_id: 0 - stage_type: llm - runtime: - process: true - devices: "0,1" - max_batch_size: 1 - engine_args: - model_stage: step_audio2_thinker - model_arch: StepAudio2ThinkerForConditionalGeneration - worker_type: ar - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.8 - enforce_eager: true - trust_remote_code: true - enable_prefix_caching: false - max_num_batched_tokens: 2048 - tensor_parallel_size: 2 - engine_output_type: text - is_comprehension: true - final_output: true - final_output_type: text - default_sampling_params: - temperature: 0.7 - top_p: 0.9 - top_k: -1 - max_tokens: 1024 - seed: 42 - detokenize: true - repetition_penalty: 1.05 diff --git a/vllm_omni/model_executor/stage_configs/step_audio_2_async_chunk.yaml b/vllm_omni/model_executor/stage_configs/step_audio_2_async_chunk.yaml deleted file mode 100644 index 98c484072..000000000 --- a/vllm_omni/model_executor/stage_configs/step_audio_2_async_chunk.yaml +++ /dev/null @@ -1,78 +0,0 @@ -# Stage config for Step-Audio2 with async_chunk support -# -# Two-stage pipeline with streaming between stages: -# Stage 0 (Thinker): Audio understanding → Text + Audio tokens (streamed in chunks) -# Stage 1 (Token2Wav): Audio tokens → Waveform (24kHz) -# -# Key difference from step_audio_2.yaml: -# - async_chunk: true enables streaming from Thinker to Token2Wav -# - Thinker uses custom_process_next_stage_input_func to send audio token chunks -# - Token2Wav receives chunks incrementally instead of waiting for full output -# -# Hardware requirement: 1-2 GPUs (can run on single GPU with devices: "0" for both) - -async_chunk: true -stage_args: - # Stage 0: Thinker (Audio Understanding + Token Generation) - - stage_id: 0 - stage_type: llm - runtime: - process: true - devices: "0,1" - max_batch_size: 1 - engine_args: - model_stage: step_audio2_thinker - model_arch: StepAudio2ThinkerForConditionalGeneration - worker_type: ar - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.8 - enforce_eager: true - trust_remote_code: true - enable_prefix_caching: false - max_num_batched_tokens: 1024 - tensor_parallel_size: 2 - engine_output_type: latent - custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.step_audio2.thinker2token2wav_async_chunk - is_comprehension: true - final_output: true - final_output_type: text - default_sampling_params: - temperature: 0.7 - top_p: 0.9 - top_k: -1 - max_tokens: 1024 - seed: 42 - detokenize: true - repetition_penalty: 1.05 - - # Stage 1: Token2Wav (Audio Synthesis) - - stage_id: 1 - stage_type: llm - runtime: - process: true - devices: "3" - max_batch_size: 1 - engine_args: - model_stage: token2wav - model_arch: StepAudio2Token2WavModel - worker_type: generation - hf_overrides: - architectures: ["StepAudio2Token2WavForConditionalGeneration"] - scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler - gpu_memory_utilization: 0.3 - enforce_eager: true - trust_remote_code: true - async_scheduling: false - enable_prefix_caching: false - max_num_batched_tokens: 2048 - engine_output_type: audio - engine_input_source: [0] - final_output: true - final_output_type: audio - default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 - max_tokens: 1 - seed: 42 - detokenize: false diff --git a/vllm_omni/model_executor/stage_configs/wan2_2_ti2v_dit_fp8.yaml b/vllm_omni/model_executor/stage_configs/wan2_2_ti2v_dit_fp8.yaml deleted file mode 100644 index 1b1c40537..000000000 --- a/vllm_omni/model_executor/stage_configs/wan2_2_ti2v_dit_fp8.yaml +++ /dev/null @@ -1,41 +0,0 @@ -# Stage config for running Wan2.2 TI2V-5B DiT with ModelOpt FP8 auto-detect. -# Single GPU (TI2V-5B fits 80GB BF16; FP8 drops by ~half). -# For the A14B MoE variants, bump `tensor_parallel_size` and `devices`. -# -# Use with a ModelOpt FP8 checkpoint (e.g. produced by -# examples/quantization/quantize_wan2_2_modelopt_fp8.py). - -stage_args: - - stage_id: 0 - stage_type: diffusion - runtime: - devices: "0" - max_batch_size: 1 - engine_args: - model_stage: dit - # Wan2.2 TI2V-5B is served by the base Wan22Pipeline, which auto-detects - # expand_timesteps (TI2V-5B mode) from model_index.json. (The standalone - # Wan22TI2VPipeline was consolidated into Wan22Pipeline upstream.) - model_class_name: Wan22Pipeline - max_num_seqs: 1 - enforce_eager: true - trust_remote_code: true - distributed_executor_backend: "mp" - # Tile the VAE decode by default. Wan2.2's 480p peak is DiT-dominated (so - # FP8 is the main lever here), but tiling adds headroom for higher - # resolutions / longer clips at negligible cost. Override with --vae-use-tiling/false. - vae_use_tiling: true - parallel_config: - tensor_parallel_size: 1 - - final_output: true - final_output_type: video - is_comprehension: false - default_sampling_params: - seed: 42 - -runtime: - enabled: true - defaults: - window_size: -1 - max_inflight: 1 diff --git a/vllm_omni/model_executor/stage_input_processors/mimo_audio.py b/vllm_omni/model_executor/stage_input_processors/mimo_audio.py index 3a1a4441f..cd9ca7bf2 100644 --- a/vllm_omni/model_executor/stage_input_processors/mimo_audio.py +++ b/vllm_omni/model_executor/stage_input_processors/mimo_audio.py @@ -17,8 +17,8 @@ logger = init_logger(__name__) # sequence fed to stage-1 must not exceed this, otherwise gpu_input_batch # add_request will fail with a broadcast error when copying prompt_token_ids # into token_ids_cpu. Keep in sync with the stage-1 ``max_model_len`` in -# ``vllm_omni/model_executor/stage_configs/mimo_audio.yaml`` and the offline -# example ``examples/offline_inference/mimo_audio/end2end.py``. +# ``vllm_omni/deploy/mimo_audio.yaml`` and the offline example +# ``examples/offline_inference/mimo_audio/end2end.py``. MAX_CODE2WAV_TOKENS = 18192 # Minimum safe values for codec streaming parameters. diff --git a/vllm_omni/model_executor/stage_input_processors/step_audio2.py b/vllm_omni/model_executor/stage_input_processors/step_audio2.py index 5e015aa7c..c0c0932b2 100644 --- a/vllm_omni/model_executor/stage_input_processors/step_audio2.py +++ b/vllm_omni/model_executor/stage_input_processors/step_audio2.py @@ -7,6 +7,7 @@ import torch from vllm.inputs import TextPrompt from vllm.logger import init_logger +from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.inputs.data import OmniTokensPrompt from vllm_omni.model_executor.models.step_audio2.step_audio2_constants import ( @@ -35,14 +36,14 @@ def _ensure_list(x): def thinker2token2wav_async_chunk( transfer_manager: Any, - pooling_output: dict[str, Any], + multimodal_output: dict[str, Any] | None, request: OmniEngineCoreRequest, is_finished: bool = False, -) -> dict[str, Any] | None: +) -> OmniPayloadStruct | None: """ Async chunk processor: stream audio tokens from Thinker to Token2Wav. - Unlike Qwen3-Omni which passes hidden states via pooling_output, + Unlike Qwen3-Omni which passes hidden states via multimodal_output, Step Audio2 generates audio tokens as part of the LLM autoregressive output (token IDs >= audio_start). This processor extracts those audio tokens from the running token stream and sends them in chunks. @@ -62,14 +63,13 @@ def thinker2token2wav_async_chunk( Args: transfer_manager: OmniChunkTransferAdapter instance (has ``code_prompt_token_ids`` defaultdict for state tracking). - pooling_output: Pooler output dict (unused — audio tokens are in + multimodal_output: Multimodal output dict (unused — audio tokens are in the token ID stream, not in hidden states). request: Current engine request with access to all_token_ids. is_finished: Whether the upstream request has finished generating. Returns: - dict with framework-compatible fields, or None if the chunk is - not yet ready. + Structured connector payload, or None if the chunk is not yet ready. """ audio_start = DEFAULT_TOKEN_CONFIG.audio_start audio_eos = DEFAULT_TOKEN_CONFIG.audio_eos @@ -100,18 +100,22 @@ def thinker2token2wav_async_chunk( # Last chunk: send all remaining tokens if available <= 0: # No audio tokens at all (text-only response) — send EOF marker - return { - "code_predictor_codes": [], - "left_context_size": 1, # 1 = last chunk - "finished": torch.tensor(True, dtype=torch.bool), - } + return OmniPayloadStruct( + codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)), + meta=MetaStruct( + left_context_size=1, # 1 = last chunk + finished=torch.tensor(True, dtype=torch.bool), + ), + ) remaining_tokens = audio_tokens[consumed:] transfer_manager.code_prompt_token_ids[request_id].extend(remaining_tokens) - return { - "code_predictor_codes": remaining_tokens, - "left_context_size": 1, # 1 = last chunk - "finished": torch.tensor(True, dtype=torch.bool), - } + return OmniPayloadStruct( + codes=CodesStruct(audio=torch.tensor(remaining_tokens, dtype=torch.long)), + meta=MetaStruct( + left_context_size=1, # 1 = last chunk + finished=torch.tensor(True, dtype=torch.bool), + ), + ) else: # Non-last chunk: need chunk_size + pre_lookahead_len tokens required = chunk_size + pre_lookahead_len @@ -120,11 +124,13 @@ def thinker2token2wav_async_chunk( chunk_tokens = audio_tokens[consumed : consumed + required] # Only consume chunk_size tokens; the lookahead portion is re-used transfer_manager.code_prompt_token_ids[request_id].extend(audio_tokens[consumed : consumed + chunk_size]) - return { - "code_predictor_codes": chunk_tokens, - "left_context_size": 0, # 0 = not last chunk - "finished": torch.tensor(False, dtype=torch.bool), - } + return OmniPayloadStruct( + codes=CodesStruct(audio=torch.tensor(chunk_tokens, dtype=torch.long)), + meta=MetaStruct( + left_context_size=0, # 0 = not last chunk + finished=torch.tensor(False, dtype=torch.bool), + ), + ) def thinker2token2wav( diff --git a/vllm_omni/platforms/cuda/platform.py b/vllm_omni/platforms/cuda/platform.py index a05bb289a..7be135352 100644 --- a/vllm_omni/platforms/cuda/platform.py +++ b/vllm_omni/platforms/cuda/platform.py @@ -36,7 +36,7 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase): @classmethod def get_default_stage_config_path(cls) -> str: - return "vllm_omni/model_executor/stage_configs" + return "vllm_omni/deploy" @classmethod def has_flash_attn_package(cls) -> bool: diff --git a/vllm_omni/platforms/musa/platform.py b/vllm_omni/platforms/musa/platform.py index 92a194d59..818338da4 100644 --- a/vllm_omni/platforms/musa/platform.py +++ b/vllm_omni/platforms/musa/platform.py @@ -32,7 +32,7 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase): @classmethod def get_default_stage_config_path(cls) -> str: - return "vllm_omni/model_executor/stage_configs" + return "vllm_omni/deploy" @classmethod def has_flash_attn_package(cls) -> bool: diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index fe49455d5..5133fa98e 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -75,7 +75,7 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform): @classmethod def get_default_stage_config_path(cls) -> str: - return "vllm_omni/platforms/npu/stage_configs" + return "vllm_omni/deploy" @classmethod def get_diffusion_model_impl_qualname(cls, op_name: str) -> str: diff --git a/vllm_omni/platforms/rocm/platform.py b/vllm_omni/platforms/rocm/platform.py index f905fd73c..38a14c325 100644 --- a/vllm_omni/platforms/rocm/platform.py +++ b/vllm_omni/platforms/rocm/platform.py @@ -122,7 +122,7 @@ class RocmOmniPlatform(OmniPlatform, RocmPlatform): @classmethod def get_default_stage_config_path(cls) -> str: - return "vllm_omni/platforms/rocm/stage_configs" + return "vllm_omni/deploy" @classmethod def get_torch_device(cls, local_rank: int | None = None) -> torch.device: