[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:
41 files changed
+520
-516
No files matched your search
@@ -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
|
||||
|
||||
@@ -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",
|
||||
|
||||
Reference in new issue
Block a user