diff --git a/.buildkite/test-merge.yml b/.buildkite/test-merge.yml index 7b920cd61..84c93a2f3 100644 --- a/.buildkite/test-merge.yml +++ b/.buildkite/test-merge.yml @@ -164,6 +164,28 @@ steps: - group: ":card_index_dividers: E2E Test" depends_on: upload-merge-pipeline steps: + - label: "TTS · Step-Audio2 E2E Test" + timeout_in_minutes: 20 + commands: + - | + timeout 20m bash -c " + export VLLM_LOGGING_LEVEL=DEBUG + export VLLM_WORKER_MULTIPROC_METHOD=spawn + pytest -s -v tests/e2e/offline_inference/test_step_audio2.py -m 'advanced_model' --run-level 'advanced_model' + " + agents: + queue: "gpu_4_queue" + plugins: + - docker#v5.2.0: + image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + always-pull: true + propagate-environment: true + shm-size: "8gb" + environment: + - "HF_HOME=/fsx/hf_cache" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" + - label: "TTS · Qwen3-TTS CustomVoice Test" source_file_dependencies: - tests/e2e/online_serving/test_qwen3_tts_customvoice.py diff --git a/examples/offline_inference/step_audio2/README.md b/examples/offline_inference/step_audio2/README.md new file mode 100644 index 000000000..2ba69ceb5 --- /dev/null +++ b/examples/offline_inference/step_audio2/README.md @@ -0,0 +1,375 @@ +# Step-Audio2 Offline Inference Examples + +This directory contains examples for running offline inference with Step-Audio2 using vLLM-Omni. + +## Model Overview + +Step-Audio2 is a two-stage audio model: + +- **Stage 0 (Thinker)**: Audio understanding → Text + Audio tokens + - Input: Audio (16kHz) + - Output: Text transcription + Audio tokens for synthesis + +- **Stage 1 (Token2Wav)**: Audio synthesis + - Input: Audio tokens + Speaker prompt wav + - Output: Synthesized audio waveform (24kHz) + +## Hardware Requirements + +| Mode | GPU Configuration | VRAM Required | +|------|-------------------|---------------| +| ASR (S2T) | 1x GPU | ~20-25GB | +| TTS/S2ST (single GPU) | 1x GPU | ~40-50GB | +| TTS/S2ST (multi GPU) | 2x GPU | GPU0: ~28GB, GPU1: ~22GB | + +**Tested on:** +- 1x NVIDIA H100 80GB (single-card S2ST) +- 2x NVIDIA A10 40GB (multi-card S2ST) + +**Notes:** +- Single GPU mode requires high VRAM due to both stages sharing memory +- Multi GPU mode separates Stage 0 (Thinker) and Stage 1 (Token2Wav) across GPUs +- VRAM usage can be adjusted via `gpu_memory_utilization` in stage config + +## Performance Benchmark + +### vLLM-Omni vs Official Step-Audio2 + +Single request latency comparison between vLLM-Omni and official Step-Audio2 implementation. + +| Task | Tokens | vllm-omni | Step-Audio2 | Speedup | +|------|--------|-----------|-------------|---------| +| S2ST | ~85 | 5.45s | 7.36s | **1.35x** | +| S2ST | ~160 | 6.67s | 13.92s | **2.09x** | +| S2ST | ~315 | 9.42s | 31.50s | **3.34x** | +| TTS | ~1024 | 16.65s | ~87s | **~5.2x** | + +**Key observations:** +- Speedup increases with sequence length due to vLLM's efficient KV cache management +- TTS (pure generation) shows the largest speedup (~5x) +- S2ST benefits from optimized multi-stage pipeline + +**Benchmark environment:** +- GPU: NVIDIA H100 80GB (single card) +- Model: Step-Audio2-mini +- Warmup: 1 run, Measured: 3 runs (averaged) + +### Async Chunk Streaming Performance + +Comparison between sequential (non-async) and async chunk modes via `/v1/audio/speech` TTS endpoint. + +| Mode | Mean TTFP | Mean E2E | Mean RTF | Audio Throughput | +|------|-----------|----------|----------|-----------------| +| Sequential | 4316ms | 4316ms | 0.938 | 1.07x realtime | +| **Async Chunk** | **1437ms** | 4362ms | 0.949 | 1.06x realtime | +| **Improvement** | **-67% (3x faster)** | ~same | ~same | ~same | + +**Key observations:** +- Async chunk reduces **time-to-first-audio (TTFP) by 67%** (4.3s → 1.4s) +- E2E latency remains comparable — async chunk overlaps Thinker decode with Token2Wav synthesis +- RTF < 1 in both modes (real-time capable) +- Sequential mode: TTFP ≈ E2E (must wait for all audio tokens before synthesis starts) +- Async chunk mode: Token2Wav starts after first 28 tokens (chunk_size=25 + lookahead=3) + +**Benchmark environment:** +- GPU: 4x NVIDIA RTX 3090 24GB (TP=2 for Thinker, 1 GPU for Token2Wav) +- Model: Step-Audio2-mini +- Endpoint: `/v1/audio/speech` (10 prompts, concurrency=1) +- Measured via `bench_tts_serve.py` + +## Installation + +Make sure you have installed vLLM-Omni and all required dependencies: + +```bash +# Install vLLM-Omni +pip install vllm-omni + +# Install Step-Audio2 (REQUIRED for Token2Wav stage) +pip install step-audio2 +``` + +## Model Setup + +### Option 1: Auto-download from HuggingFace (Recommended) + +The script will **automatically download** the model on first run: + +```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 +``` + +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) + +### Option 2: Manual Download (for offline use) + +Download and use locally: + +```bash +# Download from HuggingFace +huggingface-cli download stepfun-ai/Step-Audio2-mini --local-dir ./models/Step-Audio2-mini + +# Then use the local path +python end2end.py --query-type audio_to_text --model ./models/Step-Audio2-mini +``` + +Ensure the model directory contains: +``` +Step-Audio2-mini/ +├── config.json +├── model.safetensors (or pytorch_model.bin) +├── tokenizer.json +├── tokenizer_config.json +└── token2wav/ # Token2Wav models (REQUIRED) + ├── speech_tokenizer_v2_25hz.onnx # Audio tokenizer + ├── campplus.onnx # Speaker encoder + ├── flow.yaml # Flow model config + ├── flow.pt # Flow model weights + └── hift.pt # HiFT vocoder weights +``` + +## Usage Examples + +### 1. Audio to Text (ASR - Speech Recognition) + +Transcribe audio to text: + +```bash +# Quick start - Using default model and test audio +python end2end.py --query-type audio_to_text + +# 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 + +# With custom question +python end2end.py --query-type audio_to_text \ + --audio-path input.wav \ + --question "What is the speaker saying?" +``` + +**Output**: Text transcription saved to `output_step_audio2/00000_text.txt` + +### 2. Text to Audio (TTS - Speech Synthesis) + +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 +``` + +**Note**: Speaker voice is controlled by the `STEP_AUDIO2_DEFAULT_PROMPT_WAV` environment variable or the default prompt wav bundled with the model. + +**Output**: +- Text: `output_step_audio2/00000_text.txt` +- Audio: `output_step_audio2/00000_output.wav` (24kHz) + +### 3. Audio to Audio (Voice Conversion) + +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 +``` + +This mode: +1. Understands the content in `--audio-path` (source) +2. Generates audio output with the default voice + +**Note**: To use a custom speaker voice, set the `STEP_AUDIO2_DEFAULT_PROMPT_WAV` environment variable. + +### Advanced Options + +```bash +# Use custom stage configuration +python end2end.py --query-type audio_to_text \ + --stage-configs-path /path/to/custom_config.yaml + +# Multiple prompts (for batch testing) +python end2end.py --query-type audio_to_text \ + --audio-path input.wav \ + --num-prompts 5 + +# Custom output directory +python end2end.py --query-type text_to_audio \ + --text "Test synthesis" \ + --output-dir ./my_outputs + +# Enable detailed logging +python end2end.py --query-type audio_to_text \ + --audio-path input.wav \ + --enable-stats + +# Adjust generation parameters +python end2end.py --query-type audio_to_text \ + --audio-path input.wav \ + --max-tokens 2048 + +# Use custom speaker voice via environment variable +STEP_AUDIO2_DEFAULT_PROMPT_WAV=/path/to/speaker.wav python end2end.py \ + --query-type text_to_audio \ + --text "Hello world" + +# Use Ray backend for distributed processing +python end2end.py --query-type text_to_audio \ + --text "Hello world" \ + --worker-backend ray \ + --ray-address "auto" +``` + +## Configuration + +### Stage Configuration + +The default configuration (`step_audio_2.yaml`) uses: + +- **Stage 0 (Thinker)**: GPU 0, 80% memory +- **Stage 1 (Token2Wav)**: GPU 1, 30% memory + +For **single GPU** setup, edit the config to use `devices: "0"` for both stages. + +### Sampling Parameters + +- **Thinker (Stage 0)**: + - Temperature: 0.7 (balanced creativity) + - Top-p: 0.9 + - Max tokens: 1024 (configurable) + +- **Token2Wav (Stage 1)**: + - Temperature: 0.0 (deterministic) + - Operates in generation mode (not sampling) + +## Common Issues + +### 1. ImportError: No module named 's3tokenizer' + +**Solution**: Install Step-Audio2 package: +```bash +pip install step-audio2 +``` + +### 2. FileNotFoundError: prompt_wav file not found + +**Solution**: Set the `STEP_AUDIO2_DEFAULT_PROMPT_WAV` environment variable to a valid audio file: +```bash +export STEP_AUDIO2_DEFAULT_PROMPT_WAV=/path/to/speaker.wav +python end2end.py --query-type text_to_audio --text "Hello" +``` +Or ensure the default prompt wav (`default_female.wav`) exists in your model directory. + +### 3. FileNotFoundError: token2wav models not found + +**Solution**: Ensure your model directory has the complete `token2wav/` subdirectory with all ONNX and PyTorch models. + +### 4. CUDA Out of Memory + +**Solutions**: +- Use single GPU mode (set both stages to `devices: "0"`) +- Reduce `gpu_memory_utilization` in config +- Reduce `max_num_batched_tokens` +- Process fewer prompts at once + +### 5. Model not found in registry + +**Solution**: Ensure you're using vLLM-Omni's entry point with `--omni` flag or install vllm-omni properly: +```bash +pip install vllm-omni +``` + +## Output Files + +The script generates files in the output directory (default: `output_step_audio2/`): + +``` +output_step_audio2/ +├── 00000_text.txt # Text output from Thinker stage +├── 00000_output.wav # Audio output from Token2Wav stage (24kHz) +├── 00001_text.txt # (if multiple prompts) +└── 00001_output.wav +``` + +## Performance Tips + +1. **First run is slow**: Stage initialization takes 20-60 seconds +2. **Single GPU**: Set both stages to `devices: "0"` in config +3. **Multiple prompts**: Use `--num-prompts N` for batch testing +4. **Ray backend**: For multi-node or advanced scheduling +5. **Logging**: Use `--enable-stats` to debug performance issues + +## Speaker Voice Configuration + +The Token2Wav stage requires a speaker prompt wav for voice conditioning. It is automatically resolved in this order: + +1. `STEP_AUDIO2_DEFAULT_PROMPT_WAV` environment variable (if set) +2. `{model_dir}/assets/default_female.wav` +3. `{model_dir}/default_female.wav` + +If none are found, set the environment variable explicitly: +```bash +export STEP_AUDIO2_DEFAULT_PROMPT_WAV=/path/to/speaker.wav +``` + +**Guidelines for custom speaker prompt:** + +- **Duration**: 3-10 seconds recommended +- **Quality**: Clean audio, minimal background noise +- **Format**: WAV, MP3, FLAC (will be resampled internally) +- **Content**: Clear speech, representative of target voice + +## Example Workflow + +Complete example from audio to final output: + +```bash +# 1. ASR: Transcribe audio +python end2end.py --query-type audio_to_text \ + --audio-path interview.wav \ + --model ./models/Step-Audio2-7B \ + --output-dir ./outputs + +# 2. Check the transcription +cat ./outputs/00000_text.txt + +# 3. TTS: Synthesize new speech (with custom voice) +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 \ + --output-dir ./outputs + +# 4. Listen to the result +# Audio saved to: ./outputs/00000_output.wav +``` + +## References + +- [Step-Audio2 Paper](https://arxiv.org/abs/2507.16632) +- [vLLM-Omni Documentation](https://vllm-omni.readthedocs.io) +- [Model on HuggingFace](https://huggingface.co/stepfun-ai/Step-Audio-2-mini) diff --git a/examples/offline_inference/step_audio2/end2end.py b/examples/offline_inference/step_audio2/end2end.py new file mode 100644 index 000000000..c67f44926 --- /dev/null +++ b/examples/offline_inference/step_audio2/end2end.py @@ -0,0 +1,527 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +This example shows how to use vLLM-Omni for running offline inference +with Step-Audio2 model for audio understanding and generation. + +Step-Audio2 is a two-stage model: + - Stage 0 (Thinker): Audio understanding → Text + Audio tokens + - Stage 1 (Token2Wav): Audio tokens → Waveform (24kHz) + +Usage examples: + # Audio to text (ASR) + python end2end.py --query-type audio_to_text --audio-path input.wav --model-path "$MODEL_PATH" + + # Text to audio (TTS with default voice) + python end2end.py --query-type text_to_audio --model-path "$MODEL_PATH" + + # Audio to audio (Voice conversion with default voice) + python end2end.py --query-type audio_to_audio --audio-path input.wav --model-path "$MODEL_PATH" +""" + +import os +import sys +from pathlib import Path +from typing import NamedTuple + +import librosa # noqa: TID251 +import numpy as np +import soundfile as sf +import torch +from vllm.assets.audio import AudioAsset +from vllm.sampling_params import SamplingParams +from vllm.utils.argparse_utils import FlexibleArgumentParser + +from vllm_omni.entrypoints.omni import Omni + +SEED = 42 + + +class QueryResult(NamedTuple): + inputs: dict + limit_mm_per_prompt: dict[str, int] + + +def get_audio_to_text_query( + audio_path: str | None = None, + question: str | None = None, + sampling_rate: int = 16000, +) -> QueryResult: + """ + Audio Speech Recognition (ASR) - Convert audio to text. + + Args: + audio_path: Path to audio file + question: Question about the audio (optional) + sampling_rate: Target sampling rate (16kHz for Step-Audio2) + + Returns: + QueryResult with prompt and audio data + """ + if question is None: + question = "请记录下你所听到的语音内容。" + + prompt = ( + f"<|im_start|>system\n{question}<|im_end|>\n<|im_start|>user\n<|im_end|>\n<|im_start|>assistant\n" + ) + + # Load audio + if audio_path: + if not os.path.exists(audio_path): + raise FileNotFoundError(f"Audio file not found: {audio_path}") + audio_signal, sr = librosa.load(audio_path, sr=sampling_rate) + audio_data = (audio_signal.astype(np.float32), sr) + else: + # Use default test audio + audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate + # Resample to 16kHz if needed + if audio_data[1] != sampling_rate: + audio_signal = librosa.resample(audio_data[0], orig_sr=audio_data[1], target_sr=sampling_rate) + audio_data = (audio_signal.astype(np.float32), sampling_rate) + + return QueryResult( + inputs={ + "prompt": prompt, + "multi_modal_data": { + "audio": audio_data, + }, + }, + limit_mm_per_prompt={"audio": 1}, + ) + + +def get_text_to_audio_query( + text: str | None = None, + sampling_rate: int = 16000, +) -> QueryResult: + """ + Text-to-Speech (TTS) - Convert text to audio. + + Args: + text: Text to synthesize + sampling_rate: Target sampling rate (16kHz for Step-Audio2) + + Returns: + QueryResult with text prompt and speaker wav path + + Note: + For audio generation, the prompt must end with "" + to signal the model to generate audio tokens instead of just text. + """ + if text is None: + text = "Hello, this is a test of Step Audio 2 text to speech synthesis." + + prompt = ( + "<|im_start|>system\n" + "请逐字朗读用户提供的文本，不要改写或补充。<|im_end|>\n" + "<|im_start|>user\n" + f"{text}<|im_end|>\n" + "<|im_start|>assistant\n" + "" + ) + + inputs = { + "prompt": prompt, + } + + return QueryResult( + inputs=inputs, + limit_mm_per_prompt={}, + ) + + +def get_audio_to_audio_query( + audio_path: str | None = None, + question: str | None = None, + sampling_rate: int = 16000, +) -> QueryResult: + """ + Audio-to-Audio - Voice conversion/cloning. + + Args: + audio_path: Path to source audio file + question: Question/instruction about the audio + sampling_rate: Target sampling rate (16kHz for Step-Audio2) + + Returns: + QueryResult with audio input + + Note: + This mode processes input audio and generates output audio. + Speaker voice is controlled by STEP_AUDIO2_DEFAULT_PROMPT_WAV env var. + """ + if question is None: + question = "请仔细聆听这段语音，然后复述其内容。" + + prompt = ( + "<|im_start|>system\n" + f"{question}<|im_end|>\n" + "<|im_start|>user\n" + "<|im_end|>\n" + "<|im_start|>assistant\n" + "" + ) + + # Load source audio + if audio_path: + if not os.path.exists(audio_path): + raise FileNotFoundError(f"Audio file not found: {audio_path}") + audio_signal, sr = librosa.load(audio_path, sr=sampling_rate) + audio_data = (audio_signal.astype(np.float32), sr) + + else: + audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate + if audio_data[1] != sampling_rate: + audio_signal = librosa.resample(audio_data[0], orig_sr=audio_data[1], target_sr=sampling_rate) + audio_data = (audio_signal.astype(np.float32), sampling_rate) + + return QueryResult( + inputs={ + "prompt": prompt, + "multi_modal_data": { + "audio": audio_data, + }, + }, + limit_mm_per_prompt={"audio": 1}, + ) + + +query_map = { + "audio_to_text": get_audio_to_text_query, + "text_to_audio": get_text_to_audio_query, + "audio_to_audio": get_audio_to_audio_query, +} + + +def main(args): + # Use local model path or HuggingFace model name + if not args.model: + sys.exit(1) + + model_name = args.model + + # Check if it's a local path + if os.path.isdir(model_name): + # Verify essential files exist + config_path = os.path.join(model_name, "config.json") + if not os.path.exists(config_path): + sys.exit(1) + + # Get query configuration + query_func = query_map[args.query_type] + + # Build query based on type + if args.query_type == "audio_to_text": + query_result = query_func( + audio_path=args.audio_path, + question=args.question, + sampling_rate=args.sampling_rate, + ) + elif args.query_type == "text_to_audio": + query_result = query_func( + text=args.text, + sampling_rate=args.sampling_rate, + ) + elif args.query_type == "audio_to_audio": + query_result = query_func( + audio_path=args.audio_path, + question=args.question, + sampling_rate=args.sampling_rate, + ) + else: + raise ValueError(f"Unknown query type: {args.query_type}") + + # Initialize vLLM-Omni with Step-Audio2 + # Resolve stage config path + if args.stage_configs_path: + stage_config_path = args.stage_configs_path + else: + configs_dir = Path(__file__).parent.parent.parent.parent / "vllm_omni/model_executor/stage_configs" + 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") + else: + # TTS/S2ST need both stages + stage_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, + ) + + # Configure sampling parameters for each stage + # Stage 0 (Thinker): Generate text and audio tokens + # Note: For audio generation tasks, we need to allow the model to generate audio tokens + # Audio tokens are in range [151696, 158257], so we should not stop too early + # Qwen EOS token is typically 151645 (<|im_end|>), but we want to allow audio generation + # So we don't set stop_token_ids for audio tasks, but we do set it for text-only tasks + stop_token_ids = None + if args.query_type == "audio_to_text": + # For ASR, stop at EOS token + stop_token_ids = [151645] # Qwen EOS token <|im_end|> + + # Adjust sampling parameters based on query type + # For audio generation tasks, we need higher repetition_penalty to avoid text loops + # and allow the model to transition to audio token generation + repetition_penalty = 1.1 if args.query_type in ["text_to_audio", "audio_to_audio"] else 1.05 + + thinker_sampling_params = SamplingParams( + temperature=0.7, + top_p=0.9, + top_k=-1, + max_tokens=args.max_tokens, + seed=SEED, + detokenize=True, + repetition_penalty=repetition_penalty, + stop_token_ids=stop_token_ids, + ) + + # Stage 1 (Token2Wav): Convert audio tokens to waveform + # Note: This stage doesn't actually sample, it's deterministic generation + token2wav_sampling_params = SamplingParams( + temperature=0.0, + top_p=1.0, + top_k=-1, + max_tokens=1, + seed=SEED, + detokenize=False, + ) + + if args.query_type == "audio_to_text": + # ASR: single stage, only Thinker + sampling_params_list = [thinker_sampling_params] + else: + # TTS/S2ST: two stages + sampling_params_list = [thinker_sampling_params, token2wav_sampling_params] + + # Prepare prompts + if args.num_prompts > 1: + prompts = [query_result.inputs for _ in range(args.num_prompts)] + else: + prompts = [query_result.inputs] + + # Run inference + omni_outputs = omni_llm.generate(prompts, sampling_params_list) + + # Create output directory + output_dir = args.output_dir + os.makedirs(output_dir, exist_ok=True) + + def _get_request_index(req_id: str, fallback: int) -> int: + """Parse request index from vLLM request_id. + + vLLM may return IDs like "0_". Use the leading integer if present, + otherwise fall back to the loop index. + """ + try: + return int(req_id) + except (TypeError, ValueError): + if isinstance(req_id, str): + head = req_id.split("_", 1)[0] + if head.isdigit(): + return int(head) + return fallback + + # Process outputs from each stage + for stage_idx, stage_outputs in enumerate(omni_outputs): + # request_output may be a single object or a list + raw_output = stage_outputs.request_output + outputs_list = raw_output if isinstance(raw_output, list) else [raw_output] + + if stage_outputs.final_output_type == "text": + # Stage 0 (Thinker) text output + for i, output in enumerate(outputs_list): + request_id = _get_request_index(output.request_id, i) + text_output = output.outputs[0].text + + # Save to file + out_txt = os.path.join(output_dir, f"{request_id:05d}_text.txt") + with open(out_txt, "w", encoding="utf-8") as f: + f.write(f"Prompt:\n{prompts[request_id]['prompt']}\n\n") + f.write(f"Output:\n{text_output}\n") + print(f"Text output saved to: {out_txt}") + print(f"Text: {text_output[:200]}") + + elif stage_outputs.final_output_type == "audio": + # Stage 1 (Token2Wav) audio output + for i, output in enumerate(outputs_list): + request_id = _get_request_index(output.request_id, i) + + # Get audio tensor (24kHz) + mm_out = {} + if hasattr(output, "multimodal_output") and output.multimodal_output: + mm_out = output.multimodal_output + elif output.outputs and hasattr(output.outputs[0], "multimodal_output"): + mm_out = output.outputs[0].multimodal_output or {} + audio_tensor = None + for key in ("audio", "wav", "waveform", "audio_pcm", "pcm", "model_outputs"): + if key in mm_out and mm_out[key] is not None: + audio_tensor = mm_out[key] + break + + if audio_tensor is not None: + # Save audio file + output_wav = os.path.join(output_dir, f"{request_id:05d}_output.wav") + if isinstance(audio_tensor, list): + chunks = [] + for chunk in audio_tensor: + if chunk is None: + continue + if isinstance(chunk, torch.Tensor): + chunk_np = chunk.detach().cpu().numpy() + else: + chunk_np = np.asarray(chunk) + if chunk_np.size > 0: + chunks.append(chunk_np) + if not chunks: + continue + audio_numpy = np.concatenate(chunks, axis=-1) + elif isinstance(audio_tensor, torch.Tensor): + audio_numpy = audio_tensor.detach().cpu().numpy() + else: + audio_numpy = np.asarray(audio_tensor) + sf.write(output_wav, audio_numpy, samplerate=24000) + + +def parse_args(): + parser = FlexibleArgumentParser(description="Demo on using vLLM-Omni for offline inference with Step-Audio2") + + # Model and config + parser.add_argument( + "--model", + "-m", + type=str, + required=True, + help=( + "Model path (required). Can be:\n" + " - Local path: /path/to/step-audio-2\n" + " - HuggingFace ID: stepfun-ai/Step-Audio2-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)", + ) + + # Query configuration + parser.add_argument( + "--query-type", + "-q", + type=str, + default="audio_to_text", + choices=query_map.keys(), + help="Query type: audio_to_text (ASR), text_to_audio (TTS), audio_to_audio (voice conversion)", + ) + parser.add_argument( + "--audio-path", + "-a", + type=str, + default=None, + help="Path to input audio file (for audio_to_text and audio_to_audio)", + ) + parser.add_argument( + "--text", + "-t", + type=str, + default=None, + help="Text to synthesize (for text_to_audio mode)", + ) + parser.add_argument( + "--question", + type=str, + default=None, + help="Question/instruction about the audio", + ) + parser.add_argument( + "--sampling-rate", + type=int, + default=16000, + help="Audio sampling rate for input (default: 16000 Hz)", + ) + + # Generation parameters + parser.add_argument( + "--max-tokens", + type=int, + default=1024, + help="Maximum tokens to generate in Thinker stage (default: 1024)", + ) + parser.add_argument( + "--num-prompts", + type=int, + default=1, + help="Number of times to repeat the prompt (for testing)", + ) + + # System configuration + parser.add_argument( + "--enable-stats", + action="store_true", + default=False, + help="Enable detailed statistics logging", + ) + parser.add_argument( + "--init-sleep-seconds", + type=int, + default=20, + help="Sleep seconds after starting each stage (default: 20)", + ) + parser.add_argument( + "--batch-timeout", + type=int, + default=5, + help="Batch timeout in seconds (default: 5)", + ) + parser.add_argument( + "--init-timeout", + type=int, + default=300, + help="Initialization timeout in seconds (default: 300)", + ) + parser.add_argument( + "--shm-threshold-bytes", + type=int, + default=65536, + help="Shared memory threshold in bytes (default: 65536)", + ) + parser.add_argument( + "--worker-backend", + type=str, + default="multi_process", + choices=["multi_process", "ray"], + help="Worker backend (default: multi_process)", + ) + parser.add_argument( + "--ray-address", + type=str, + default=None, + help="Ray cluster address (for ray backend)", + ) + + # Output + parser.add_argument( + "--output-dir", + "-o", + type=str, + default="output_step_audio2", + help="Output directory for generated files (default: output_step_audio2)", + ) + + return parser.parse_args() + + +if __name__ == "__main__": + args = parse_args() + main(args) diff --git a/examples/online_serving/step_audio2/README.md b/examples/online_serving/step_audio2/README.md new file mode 100644 index 000000000..3a04f4565 --- /dev/null +++ b/examples/online_serving/step_audio2/README.md @@ -0,0 +1,193 @@ +# Step-Audio2 Online Serving + +This directory contains examples for running Step-Audio2 with vLLM-Omni's online serving API. + +## Installation + +Please refer to [README.md](../../../README.md) + +## Launch the Server + +```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 \ + --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 \ + --trust-remote-code --enforce-eager +``` + +With local model: +```bash +vllm serve /path/to/Step-Audio-2-mini --omni --port 8092 \ + --trust-remote-code --enforce-eager +``` + +## Send Requests + +### TTS via `/v1/audio/speech` (Recommended) + +```bash +cd examples/online_serving/step_audio2 + +# Python client +python openai_speech_client.py --text "你好世界" + +# With custom system prompt +python openai_speech_client.py --text "Hello, how are you?" \ + --instructions "You are a friendly assistant." + +# Save to specific file +python openai_speech_client.py --text "你好世界" -o output.wav +``` + +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"}' \ + --output output.wav +``` + +This endpoint bypasses the chat template and directly triggers TTS mode. It supports async chunk streaming for low first-packet latency. + +**Note**: Speaker voice is controlled by `STEP_AUDIO2_DEFAULT_PROMPT_WAV` env var on the server side. + +### Chat Completions (ASR / S2ST) + +```bash +# Audio to Text (ASR) +python openai_chat_completion_client.py --query-type audio_to_text + +# Audio to Audio (S2ST) +python openai_chat_completion_client.py --query-type audio_to_audio --audio-path /path/to/input.wav +``` + +| Argument | Description | +|----------|-------------| +| `--query-type`, `-q` | Query type: `audio_to_text`, `text_to_audio`, `audio_to_audio` | +| `--audio-path`, `-a` | Path to input audio file (local or URL) | +| `--text`, `-t` | Text to synthesize (for TTS mode) | +| `--prompt`, `-p` | Custom prompt/question | +| `--output-dir`, `-o` | Output directory for audio files (default: `output_online`) | +| `--api-base` | API base URL (default: `http://localhost:8092/v1`) | + +### Curl (Chat Completions) + +```bash +# Audio to Text +bash run_curl.sh audio_to_text + +# Text to Audio +bash run_curl.sh text_to_audio + +# Audio to Audio +bash run_curl.sh audio_to_audio +``` + +## Query Types + +### 1. Audio to Text (ASR) + +Transcribe audio to text. + +```bash +python openai_chat_completion_client.py \ + --query-type audio_to_text \ + --audio-path /path/to/speech.wav \ + --prompt "Transcribe this audio." +``` + +### 2. Text to Audio (TTS) + +Convert text to speech: + +```bash +# Via speech endpoint (recommended, returns WAV directly) +python openai_speech_client.py --text "Hello, welcome to Step-Audio2." + +# Via chat completions +python openai_chat_completion_client.py \ + --query-type text_to_audio \ + --text "Hello, welcome to Step-Audio2." +``` + +### 3. Audio to Audio (S2ST) + +Process input audio and generate text transcription + audio output. + +```bash +python openai_chat_completion_client.py \ + --query-type audio_to_audio \ + --audio-path /path/to/source.wav +``` + +## Output + +- **Text output**: Printed to console +- **Audio output**: Saved to `output_online/audio_0.wav` (24kHz WAV) + +## API Format + +Step-Audio2 uses the OpenAI-compatible chat completions API: + +```json +{ + "model": "stepfun-ai/Step-Audio2-mini", + "messages": [ + { + "role": "system", + "content": [{"type": "text", "text": "Transcribe the audio."}] + }, + { + "role": "user", + "content": [ + {"type": "audio_url", "audio_url": {"url": "..."}}, + {"type": "text", "text": "Please transcribe."} + ] + } + ], + "sampling_params_list": [ + {"temperature": 0.7, "max_tokens": 1024}, + {"temperature": 0.0, "max_tokens": 1} + ] +} +``` + +## Performance + +### Async Chunk vs Sequential + +Benchmark via `/v1/audio/speech` (4x RTX 3090, 10 prompts, concurrency=1): + +| Mode | Mean TTFP | Mean E2E | Mean RTF | +|------|-----------|----------|----------| +| Sequential | 4316ms | 4316ms | 0.938 | +| **Async Chunk** | **1437ms** | 4362ms | 0.949 | + +Async chunk reduces TTFP by **67%** by streaming audio token chunks from Thinker to Token2Wav as they are generated. RTF < 1 in both modes (real-time capable). + +## Troubleshooting + +### Server not responding +- Check if the server is running: `curl http://localhost:8092/v1/models` +- Verify the port number matches + +### FileNotFoundError: prompt_wav file not found +- Ensure `default_female.wav` exists at `{model_dir}/assets/default_female.wav` +- Or set `STEP_AUDIO2_DEFAULT_PROMPT_WAV` environment variable when launching the server + +### Audio not generated +- For TTS, use the `/v1/audio/speech` endpoint (recommended) or `openai_speech_client.py` +- For chat completions TTS, ensure the prompt ends with `` +- Check server logs for errors + +### Out of memory +- Reduce `gpu_memory_utilization` in stage configs +- Use a smaller batch size diff --git a/examples/online_serving/step_audio2/openai_chat_completion_client.py b/examples/online_serving/step_audio2/openai_chat_completion_client.py new file mode 100644 index 000000000..2b91dacb2 --- /dev/null +++ b/examples/online_serving/step_audio2/openai_chat_completion_client.py @@ -0,0 +1,342 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +OpenAI-compatible client for Step-Audio2 online serving. + +Step-Audio2 is a two-stage model: + - Stage 0 (Thinker): Audio understanding -> Text + Audio tokens + - Stage 1 (Token2Wav): Audio tokens -> Waveform (24kHz) + +Usage examples: + # Start server first: + vllm serve stepfun-ai/Step-Audio2-mini --omni --port 8092 + + # Audio to text (ASR) + python openai_chat_completion_client.py --query-type audio_to_text + + # Text to audio (TTS) + python openai_chat_completion_client.py --query-type text_to_audio --text "Hello world" + + # Audio to audio (Voice conversion) + python openai_chat_completion_client.py --query-type audio_to_audio --audio-path input.wav +""" + +import base64 +import logging +import os + +import requests +from openai import OpenAI +from vllm.assets.audio import AudioAsset +from vllm.utils.argparse_utils import FlexibleArgumentParser + +logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") +logger = logging.getLogger(__name__) + +SEED = 42 + + +def encode_base64_content_from_file(file_path: str) -> str: + """Encode a local file to base64 format.""" + with open(file_path, "rb") as f: + return base64.b64encode(f.read()).decode("utf-8") + + +def encode_base64_content_from_url(url: str) -> str: + """Encode content from URL to base64 format.""" + with requests.get(url) as response: + response.raise_for_status() + return base64.b64encode(response.content).decode("utf-8") + + +def get_audio_url(audio_path: str | None = None) -> str: + """Convert audio path to URL format for the API. + + Args: + audio_path: Local file path or URL. If None, uses default test audio. + + Returns: + Audio URL (either original URL or base64 data URL) + """ + if not audio_path: + return AudioAsset("mary_had_lamb").url + + if audio_path.startswith(("http://", "https://")): + return audio_path + + if not os.path.exists(audio_path): + raise FileNotFoundError(f"Audio file not found: {audio_path}") + + # Detect MIME type + audio_path_lower = audio_path.lower() + if audio_path_lower.endswith(".wav"): + mime_type = "audio/wav" + elif audio_path_lower.endswith((".mp3", ".mpeg")): + mime_type = "audio/mpeg" + elif audio_path_lower.endswith(".ogg"): + mime_type = "audio/ogg" + elif audio_path_lower.endswith(".flac"): + mime_type = "audio/flac" + else: + mime_type = "audio/wav" + + audio_base64 = encode_base64_content_from_file(audio_path) + return f"data:{mime_type};base64,{audio_base64}" + + +def get_audio_to_text_query( + audio_path: str | None = None, + custom_prompt: str | None = None, +) -> dict: + """Build query for Audio-to-Text (ASR) mode.""" + question = custom_prompt or "Please transcribe the audio content." + audio_url = get_audio_url(audio_path) + + return { + "role": "user", + "content": [ + {"type": "audio_url", "audio_url": {"url": audio_url}}, + {"type": "text", "text": question}, + ], + } + + +def get_text_to_audio_query( + text: str | None = None, + custom_prompt: str | None = None, +) -> dict: + """Build query for Text-to-Audio (TTS) mode. + + Returns a list of messages: user message + assistant message with . + The assistant message triggers TTS mode and must be used with + continue_final_message=True so the chat template does not append <|im_end|>. + """ + text_to_speak = text or "Hello, this is a test of Step Audio 2 text to speech." + + return [ + { + "role": "user", + "content": [{"type": "text", "text": text_to_speak}], + }, + { + "role": "assistant", + "content": "", + }, + ] + + +def get_audio_to_audio_query( + audio_path: str | None = None, + custom_prompt: str | None = None, +) -> dict: + """Build query for Audio-to-Audio (S2ST) mode. + + Returns a list of messages: user message (with audio) + assistant message + with . Must be used with continue_final_message=True. + """ + question = custom_prompt or "Please listen to this audio and repeat its content." + audio_url = get_audio_url(audio_path) + + return [ + { + "role": "user", + "content": [ + {"type": "audio_url", "audio_url": {"url": audio_url}}, + {"type": "text", "text": question}, + ], + }, + { + "role": "assistant", + "content": "", + }, + ] + + +def get_system_prompt(query_type: str) -> dict: + """Get system prompt based on query type.""" + if query_type == "audio_to_text": + system_text = "You are a speech recognition assistant. Transcribe the audio accurately." + elif query_type == "text_to_audio": + system_text = "You are a text-to-speech assistant. Read the text aloud exactly as provided." + else: # audio_to_audio + system_text = "You are an audio processing assistant. Listen and repeat the audio content." + + return { + "role": "system", + "content": [{"type": "text", "text": system_text}], + } + + +QUERY_MAP = { + "audio_to_text": get_audio_to_text_query, + "text_to_audio": get_text_to_audio_query, + "audio_to_audio": get_audio_to_audio_query, +} + + +def run_step_audio2(args) -> None: + """Run Step-Audio2 inference via OpenAI-compatible API.""" + + # Initialize OpenAI client + client = OpenAI( + api_key="EMPTY", + base_url=args.api_base, + ) + + model_name = args.model + + # Sampling parameters for each stage + thinker_sampling_params = { + "temperature": 0.7, + "top_p": 0.9, + "top_k": -1, + "max_tokens": args.max_tokens, + "seed": SEED, + "detokenize": True, + "repetition_penalty": 1.1 if args.query_type != "audio_to_text" else 1.05, + } + + # For ASR, stop at EOS token + if args.query_type == "audio_to_text": + thinker_sampling_params["stop_token_ids"] = [151645] + + token2wav_sampling_params = { + "temperature": 0.0, + "top_p": 1.0, + "top_k": -1, + "max_tokens": 1, + "seed": SEED, + "detokenize": False, + } + + sampling_params_list = [ + thinker_sampling_params, + token2wav_sampling_params, + ] + + # Build query based on type + query_func = QUERY_MAP[args.query_type] + if args.query_type == "audio_to_text": + user_message = query_func(audio_path=args.audio_path, custom_prompt=args.prompt) + elif args.query_type == "text_to_audio": + user_message = query_func(text=args.text, custom_prompt=args.prompt) + else: # audio_to_audio + user_message = query_func(audio_path=args.audio_path, custom_prompt=args.prompt) + + system_message = get_system_prompt(args.query_type) + + logger.info(f"Query type: {args.query_type}") + logger.info(f"Model: {model_name}") + logger.info("Sending request to server...") + + # Build messages list. + # TTS and S2ST return a list of messages (user + assistant with ); + # ASR returns a single user message dict. + if isinstance(user_message, list): + messages = [system_message] + user_message + else: + messages = [system_message, user_message] + + # For TTS/S2ST, use continue_final_message so in the + # assistant turn is not followed by <|im_end|>. + extra_body: dict = {"sampling_params_list": sampling_params_list} + if args.query_type in ("text_to_audio", "audio_to_audio"): + extra_body["continue_final_message"] = True + extra_body["add_generation_prompt"] = False + + # Send request + chat_completion = client.chat.completions.create( + messages=messages, + model=model_name, + extra_body=extra_body, + ) + + # Process response + os.makedirs(args.output_dir, exist_ok=True) + + audio_count = 0 + for choice in chat_completion.choices: + if choice.message.audio: + # Audio response + audio_data = base64.b64decode(choice.message.audio.data) + audio_file = os.path.join(args.output_dir, f"audio_{audio_count}.wav") + with open(audio_file, "wb") as f: + f.write(audio_data) + logger.info(f"Audio saved to: {audio_file}") + audio_count += 1 + + if choice.message.content: + # Text response + logger.info(f"Text output: {choice.message.content}") + + if audio_count == 0 and args.query_type != "audio_to_text": + logger.warning("No audio output received. Check server logs for details.") + + logger.info("Done!") + + +def parse_args(): + parser = FlexibleArgumentParser(description="OpenAI-compatible client for Step-Audio2 online serving") + + parser.add_argument( + "--query-type", + "-q", + type=str, + default="audio_to_text", + choices=QUERY_MAP.keys(), + help="Query type: audio_to_text (ASR), text_to_audio (TTS), audio_to_audio", + ) + parser.add_argument( + "--audio-path", + "-a", + type=str, + default=None, + help="Path to input audio file (local path or URL)", + ) + parser.add_argument( + "--text", + "-t", + type=str, + default=None, + help="Text to synthesize (for text_to_audio mode)", + ) + parser.add_argument( + "--prompt", + "-p", + type=str, + default=None, + help="Custom prompt/question", + ) + parser.add_argument( + "--model", + "-m", + type=str, + default="stepfun-ai/Step-Audio2-mini", + help="Model name", + ) + parser.add_argument( + "--api-base", + type=str, + default="http://localhost:8092/v1", + help="API base URL", + ) + parser.add_argument( + "--max-tokens", + type=int, + default=1024, + help="Maximum tokens for Thinker stage", + ) + parser.add_argument( + "--output-dir", + "-o", + type=str, + default="output_online", + help="Output directory for audio files", + ) + + return parser.parse_args() + + +if __name__ == "__main__": + args = parse_args() + run_step_audio2(args) diff --git a/examples/online_serving/step_audio2/openai_speech_client.py b/examples/online_serving/step_audio2/openai_speech_client.py new file mode 100644 index 000000000..740b36805 --- /dev/null +++ b/examples/online_serving/step_audio2/openai_speech_client.py @@ -0,0 +1,130 @@ +"""OpenAI-compatible client for Step-Audio2 TTS via /v1/audio/speech endpoint. + +Examples: + # Basic TTS + python openai_speech_client.py --text "你好世界" + + # With custom system prompt + python openai_speech_client.py --text "Hello, how are you?" \ + --instructions "You are a friendly assistant." + + # Save to specific file + python openai_speech_client.py --text "你好世界" -o output.wav +""" + +import argparse + +import httpx + +DEFAULT_API_BASE = "http://localhost:8092" +DEFAULT_API_KEY = "EMPTY" + + +def run_tts_generation(args) -> None: + """Run TTS generation via /v1/audio/speech API.""" + payload = { + "model": args.model, + "input": args.text, + "voice": args.voice, + "response_format": args.response_format, + } + + if args.instructions: + payload["instructions"] = args.instructions + + print(f"Model: {args.model}") + print(f"Text: {args.text}") + print(f"Voice: {args.voice}") + print("Generating audio...") + + api_url = f"{args.api_base}/v1/audio/speech" + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {args.api_key}", + } + + with httpx.Client(timeout=300.0) as client: + response = client.post(api_url, json=payload, headers=headers) + + if response.status_code != 200: + print(f"Error: {response.status_code}") + print(response.text) + return + + try: + text = response.content.decode("utf-8") + if text.startswith('{"error"'): + print(f"Error: {text}") + return + except UnicodeDecodeError: + pass + + output_path = args.output or "tts_output.wav" + with open(output_path, "wb") as f: + f.write(response.content) + print(f"Audio saved to: {output_path}") + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Step-Audio2 TTS client via /v1/audio/speech", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=__doc__, + ) + parser.add_argument( + "--api-base", + type=str, + default=DEFAULT_API_BASE, + help=f"API base URL (default: {DEFAULT_API_BASE})", + ) + parser.add_argument( + "--api-key", + type=str, + default=DEFAULT_API_KEY, + help="API key (default: EMPTY)", + ) + parser.add_argument( + "--model", + "-m", + type=str, + default="stepfun-ai/Step-Audio2-mini", + help="Model name/path", + ) + parser.add_argument( + "--text", + type=str, + required=True, + help="Text to synthesize", + ) + parser.add_argument( + "--voice", + type=str, + default="default", + help="Voice name (default: default)", + ) + parser.add_argument( + "--instructions", + type=str, + default=None, + help="System prompt for the Thinker stage", + ) + parser.add_argument( + "--response-format", + type=str, + default="wav", + choices=["wav", "pcm"], + help="Audio output format (default: wav)", + ) + parser.add_argument( + "--output", + "-o", + type=str, + default=None, + help="Output audio file path (default: tts_output.wav)", + ) + return parser.parse_args() + + +if __name__ == "__main__": + args = parse_args() + run_tts_generation(args) diff --git a/examples/online_serving/step_audio2/run_curl.sh b/examples/online_serving/step_audio2/run_curl.sh new file mode 100755 index 000000000..3811308e7 --- /dev/null +++ b/examples/online_serving/step_audio2/run_curl.sh @@ -0,0 +1,149 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Step-Audio2 curl client for online serving +# Usage: bash run_curl.sh [audio_to_text|text_to_audio|audio_to_audio] + +QUERY_TYPE="${1:-audio_to_text}" +API_BASE="${API_BASE:-http://localhost:8092}" + +# Validate query type +if [[ ! "$QUERY_TYPE" =~ ^(audio_to_text|text_to_audio|audio_to_audio)$ ]]; then + echo "Error: Invalid query type '$QUERY_TYPE'" + echo "Usage: $0 [audio_to_text|text_to_audio|audio_to_audio]" + echo " audio_to_text: Speech recognition (ASR)" + echo " text_to_audio: Text-to-speech (TTS)" + echo " audio_to_audio: Voice conversion" + exit 1 +fi + +SEED=42 + +# Default test audio URL +MARY_HAD_LAMB_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/mary_had_lamb.ogg" + +# Sampling parameters for Thinker stage +thinker_sampling_params='{ + "temperature": 0.7, + "top_p": 0.9, + "top_k": -1, + "max_tokens": 1024, + "seed": 42, + "detokenize": true, + "repetition_penalty": 1.05 +}' + +# Sampling parameters for Token2Wav stage +token2wav_sampling_params='{ + "temperature": 0.0, + "top_p": 1.0, + "top_k": -1, + "max_tokens": 1, + "seed": 42, + "detokenize": false +}' + +# Build request based on query type +case "$QUERY_TYPE" in + audio_to_text) + system_content='[{"type": "text", "text": "You are a speech recognition assistant. Transcribe the audio accurately."}]' + user_content='[ + {"type": "audio_url", "audio_url": {"url": "'"$MARY_HAD_LAMB_URL"'"}}, + {"type": "text", "text": "Please transcribe this audio."} + ]' + # Add stop token for ASR + thinker_sampling_params='{ + "temperature": 0.7, + "top_p": 0.9, + "top_k": -1, + "max_tokens": 1024, + "seed": 42, + "detokenize": true, + "repetition_penalty": 1.05, + "stop_token_ids": [151645] + }' + ;; + text_to_audio) + system_content='[{"type": "text", "text": "You are a text-to-speech assistant. Read the text aloud exactly as provided."}]' + user_content='[ + {"type": "text", "text": "Hello, this is a test of Step Audio 2 text to speech synthesis."} + ]' + thinker_sampling_params='{ + "temperature": 0.7, + "top_p": 0.9, + "top_k": -1, + "max_tokens": 1024, + "seed": 42, + "detokenize": true, + "repetition_penalty": 1.1 + }' + ;; + audio_to_audio) + system_content='[{"type": "text", "text": "You are an audio processing assistant. Listen and repeat the audio content."}]' + user_content='[ + {"type": "audio_url", "audio_url": {"url": "'"$MARY_HAD_LAMB_URL"'"}}, + {"type": "text", "text": "Please listen to this audio and repeat its content."} + ]' + thinker_sampling_params='{ + "temperature": 0.7, + "top_p": 0.9, + "top_k": -1, + "max_tokens": 1024, + "seed": 42, + "detokenize": true, + "repetition_penalty": 1.1 + }' + ;; +esac + +sampling_params_list='[ + '"$thinker_sampling_params"', + '"$token2wav_sampling_params"' +]' + +echo "Query type: $QUERY_TYPE" +echo "API base: $API_BASE" +echo "Sending request..." +echo "" + +output=$(curl -sS -X POST "${API_BASE}/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -d @- < tuple[np.ndarray, int]: + """Create a dummy audio signal (sine wave) for testing.""" + t = np.linspace(0, duration_sec, int(sample_rate * duration_sec), dtype=np.float32) + audio = 0.5 * np.sin(2 * np.pi * 440 * t) + return audio, sample_rate + + +class StepAudio2OmniRunner(OmniRunner): + """Test runner specialized for Step-Audio2 model.""" + + def get_omni_inputs( + self, + prompts: list[str] | str, + system_prompt: str | None = None, + audios=None, + **kwargs, + ) -> list[dict]: + """Construct Omni input format for Step-Audio2.""" + if system_prompt is None: + system_prompt = "你是一个语音对话助手，能够理解音频输入并生成语音回复。" + + audio_placeholder = "" + + if isinstance(prompts, str): + prompts = [prompts] + + num_prompts = len(prompts) + + if audios is None: + audios_list = [None] * num_prompts + elif isinstance(audios, list): + audios_list = audios if len(audios) == num_prompts else [audios[0]] * num_prompts + else: + audios_list = [audios] * num_prompts + + omni_inputs = [] + for i, prompt_text in enumerate(prompts): + user_content = "" + multi_modal_data = {} + + audio = audios_list[i] + if audio is not None: + user_content += audio_placeholder + multi_modal_data["audio"] = audio + + user_content += prompt_text + + full_prompt = ( + f"<|im_start|>system\n{system_prompt}<|im_end|>\n" + f"<|im_start|>user\n{user_content}<|im_end|>\n" + f"<|im_start|>assistant\n" + ) + + input_dict = {"prompt": full_prompt} + if multi_modal_data: + input_dict["multi_modal_data"] = multi_modal_data + + omni_inputs.append(input_dict) + + return omni_inputs + + +@pytest.fixture(scope="session") +def step_audio2_runner(): + return StepAudio2OmniRunner + + +@pytest.mark.core_model +def test_processor_loads_correctly(): + """Test that Step-Audio2 processor can be imported and initialized.""" + from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import ( + StepAudio2MultiModalProcessor, + StepAudio2ProcessingInfo, + StepAudio2Processor, + ) + + assert StepAudio2MultiModalProcessor is not None + assert StepAudio2ProcessingInfo is not None + assert StepAudio2Processor is not None + + +@pytest.mark.core_model +def test_audio_preprocessing(): + """Test that audio preprocessing produces correct tensor shapes.""" + from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import ( + log_mel_spectrogram, + padding_mels, + ) + + audio, sample_rate = create_dummy_audio(sample_rate=16000, duration_sec=1.0) + + mel = log_mel_spectrogram(audio) + + assert mel.ndim == 2, f"Expected 2D tensor, got {mel.ndim}D" + assert mel.shape[0] == 128, f"Expected 128 mel bins, got {mel.shape[0]}" + + mels = [mel, mel] + padded_mels, mel_lens = padding_mels(mels) + + assert padded_mels.ndim == 3, f"Expected 3D tensor, got {padded_mels.ndim}D" + assert padded_mels.shape[0] == 2, f"Expected batch size 2, got {padded_mels.shape[0]}" + assert padded_mels.shape[1] == 128, f"Expected 128 mel bins, got {padded_mels.shape[1]}" + + +@pytest.mark.core_model +def test_feature_length_calculation(): + """Test that audio feature length calculation is correct.""" + from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import ( + calculate_audio_feature_length, + ) + + assert calculate_audio_feature_length(1000) == 125 + + assert calculate_audio_feature_length(100) > 0 + + assert calculate_audio_feature_length(1) >= 1 + + +@pytest.mark.core_model +def test_token_config_constants(): + """Test that token configuration constants are set correctly.""" + from vllm_omni.model_executor.models.step_audio2.step_audio2_constants import ( + DEFAULT_TOKEN_CONFIG, + STEP_AUDIO2_AUDIO_END, + STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID, + STEP_AUDIO2_AUDIO_START, + STEP_AUDIO2_AUDIO_VOCAB_SIZE, + STEP_AUDIO2_TEXT_MAX, + ) + + assert STEP_AUDIO2_TEXT_MAX < STEP_AUDIO2_AUDIO_START, "Text tokens should come before audio tokens" + assert STEP_AUDIO2_AUDIO_END == STEP_AUDIO2_AUDIO_START + STEP_AUDIO2_AUDIO_VOCAB_SIZE - 1 + assert STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID > STEP_AUDIO2_TEXT_MAX + assert STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID < STEP_AUDIO2_AUDIO_START + + assert DEFAULT_TOKEN_CONFIG.text_max == STEP_AUDIO2_TEXT_MAX + assert DEFAULT_TOKEN_CONFIG.audio_start == STEP_AUDIO2_AUDIO_START + + +@pytest.mark.core_model +def test_mm_field_config_structure(): + """Test that multimodal field config is properly structured for vLLM batching.""" + from vllm.multimodal.inputs import MultiModalFieldConfig + + audio_lens = torch.tensor([100, 200, 150], dtype=torch.int32) + + field_config = MultiModalFieldConfig.flat_from_sizes("audio", audio_lens) + + assert field_config is not None + + +@pytest.mark.core_model +def test_audio_encoder_output_shape(): + """Test that audio encoder produces correct output shapes.""" + from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import ( + Adaptor, + AudioEncoder, + ) + + encoder = AudioEncoder(n_mels=128, n_ctx=1500, n_state=512, n_head=8, n_layer=6) + + batch_size = 2 + n_mels = 128 + time_steps = 400 + x = torch.randn(batch_size, n_mels, time_steps) + x_len = torch.tensor([time_steps, time_steps // 2], dtype=torch.int32) + + encoded, encoded_lens = encoder(x, x_len) + + assert encoded.ndim == 3, f"Expected 3D tensor, got {encoded.ndim}D" + assert encoded.shape[0] == batch_size + assert encoded.shape[2] == 512 + + adapter = Adaptor(n_state=512, n_hidden=4096, kernel_size=3, stride=2) + adapted = adapter(encoded) + + assert adapted.ndim == 3 + assert adapted.shape[0] == batch_size + assert adapted.shape[2] == 4096 + + +@pytest.mark.core_model +def test_token_separation(): + """Test that token separation works correctly.""" + from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import ( + StepAudio2ThinkerForConditionalGeneration, + ) + + token_ids = [100, 200, 151700, 300, 151800, 400] + + text_tokens, audio_tokens = StepAudio2ThinkerForConditionalGeneration.separate_tokens(token_ids) + + assert text_tokens == [100, 200, 300, 400] + assert audio_tokens == [4, 104] + + +@pytest.mark.core_model +def test_has_audio_output(): + """Test detection of audio tokens in output.""" + from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import ( + StepAudio2ThinkerForConditionalGeneration, + ) + + text_only = [100, 200, 300] + assert not StepAudio2ThinkerForConditionalGeneration.has_audio_output(text_only) + + with_audio = [100, 200, 151700, 300] + assert StepAudio2ThinkerForConditionalGeneration.has_audio_output(with_audio) + + +@pytest.mark.advanced_model +@pytest.mark.parametrize("test_config", test_params) +def test_audio_to_text_and_audio(step_audio2_runner: type[StepAudio2OmniRunner], test_config: tuple[str, str]) -> None: + """Test processing audio input and generating text + audio output.""" + model, stage_config_path = test_config + + with step_audio2_runner(model, seed=42, stage_configs_path=stage_config_path) as runner: + audio = create_dummy_audio(sample_rate=16000, duration_sec=2.0) + + outputs = runner.generate_multimodal( + prompts="请复述这段音频的内容。", + audios=audio, + ) + + assert len(outputs) > 0 + + text_output = None + for stage_output in outputs: + if stage_output.final_output_type == "text": + text_output = stage_output + break + + assert text_output is not None + assert text_output.request_output is not None + text_content = text_output.request_output.outputs[0].text + assert text_content is not None + assert len(text_content.strip()) > 0 + + audio_output = None + for stage_output in outputs: + if stage_output.final_output_type == "audio": + audio_output = stage_output + break + + if audio_output is not None: + assert audio_output.request_output is not None + audio_tensor = audio_output.request_output.outputs[0].multimodal_output.get("audio") + if audio_tensor is not None: + assert audio_tensor.numel() > 0 + + +@pytest.mark.advanced_model +@pytest.mark.parametrize("test_config", test_params) +def test_text_only_input(step_audio2_runner: type[StepAudio2OmniRunner], test_config: tuple[str, str]) -> None: + """Test processing text-only input (no audio).""" + model, stage_config_path = test_config + + with step_audio2_runner(model, seed=42, stage_configs_path=stage_config_path) as runner: + outputs = runner.generate_multimodal( + prompts="你好，请用中文回答。", + audios=None, + ) + + assert len(outputs) > 0 + + text_output = None + for stage_output in outputs: + if stage_output.final_output_type == "text": + text_output = stage_output + break + + assert text_output is not None + assert text_output.request_output is not None + text_content = text_output.request_output.outputs[0].text + assert text_content is not None + assert len(text_content.strip()) > 0 diff --git a/tests/e2e/online_serving/stage_configs/step_audio2_ci.yaml b/tests/e2e/online_serving/stage_configs/step_audio2_ci.yaml new file mode 100644 index 000000000..703a7cdcc --- /dev/null +++ b/tests/e2e/online_serving/stage_configs/step_audio2_ci.yaml @@ -0,0 +1,81 @@ +# Stage config for Step-Audio2 CI tests +# Stage 0: Thinker (audio understanding → text + audio tokens) +# Stage 1: Token2Wav (audio tokens → waveform) + +# The following config has been verified on 2x 24GB GPU (L4/RTX3090/RTX4090). +stage_args: + - stage_id: 0 + runtime: + process: true + devices: "0" + max_batch_size: 1 + engine_args: + model_stage: thinker + model_arch: StepAudio2ThinkerForConditionalGeneration + tensor_parallel_size: 1 + hf_overrides: + architectures: ["StepAudio2ThinkerForConditionalGeneration"] + + worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + max_model_len: 1024 + max_num_batched_tokens: 1024 + max_num_seqs: 1 + gpu_memory_utilization: 0.7 + skip_mm_profiling: true + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + engine_output_type: text # Output text + audio tokens + 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: 256 + seed: 42 + detokenize: true + repetition_penalty: 1.05 + + - stage_id: 1 + runtime: + process: true + devices: "0" + max_batch_size: 1 + engine_args: + + model_stage: token2wav + model_arch: StepAudio2Token2WavModel + worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker + scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + hf_overrides: + architectures: ["StepAudio2Token2WavModel"] + gpu_memory_utilization: 0.2 + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + max_num_batched_tokens: 1024 + engine_output_type: audio # Output waveform + engine_input_source: [0] + custom_process_input_func: vllm_omni.model_executor.stage_input_processors.step_audio2.thinker2token2wav + 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 + +runtime: + enabled: true + defaults: + window_size: -1 + max_inflight: 1 + edges: + - from: 0 + to: 1 + window_size: -1 diff --git a/tests/e2e/online_serving/test_step_audio2.py b/tests/e2e/online_serving/test_step_audio2.py new file mode 100644 index 000000000..a23aaa0ff --- /dev/null +++ b/tests/e2e/online_serving/test_step_audio2.py @@ -0,0 +1,368 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +E2E tests for Step-Audio2 online serving with concurrent requests. + +Tests: +1. Single request via API +2. Concurrent throughput (1/2/4/8 concurrent requests) +3. Latency percentiles (p50/p95/p99) +""" + +import asyncio +import os +import socket +import subprocess +import sys +import time +from pathlib import Path + +import aiohttp +import numpy as np +import pytest +from vllm.utils.network_utils import get_open_port + +MODEL = "stepfun-ai/Step-Audio-2-mini" +STAGE_CONFIG = str(Path(__file__).parent / "stage_configs" / "step_audio2_ci.yaml") + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + + +class OmniServer: + """Omniserver for vLLM-Omni tests.""" + + def __init__( + self, + model: str, + serve_args: list[str], + *, + env_dict: dict[str, str] | None = None, + ) -> None: + self.model = model + self.serve_args = serve_args + self.env_dict = env_dict + self.proc: subprocess.Popen | None = None + self.host = "127.0.0.1" + self.port = get_open_port() + + def _start_server(self) -> None: + """Start the vLLM-Omni server subprocess.""" + env = os.environ.copy() + env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + if self.env_dict is not None: + env.update(self.env_dict) + + cmd = [ + sys.executable, + "-m", + "vllm_omni.entrypoints.cli.main", + "serve", + self.model, + "--omni", + "--host", + self.host, + "--port", + str(self.port), + ] + self.serve_args + + print(f"Launching OmniServer with: {' '.join(cmd)}") + self.proc = subprocess.Popen( + cmd, + env=env, + cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + ) + + max_wait = 600 # 10 minutes + start_time = time.time() + while time.time() - start_time < max_wait: + try: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.settimeout(1) + result = sock.connect_ex((self.host, self.port)) + if result == 0: + print(f"Server ready on {self.host}:{self.port}") + return + except Exception: + pass + time.sleep(2) + + raise RuntimeError(f"Server failed to start within {max_wait} seconds") + + def __enter__(self): + self._start_server() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + if self.proc: + self.proc.terminate() + try: + self.proc.wait(timeout=30) + except subprocess.TimeoutExpired: + self.proc.kill() + self.proc.wait() + + +def create_dummy_audio_base64(duration_sec: float = 5.0, sample_rate: int = 16000) -> str: + """Create dummy audio as base64.""" + import base64 + import io + import wave + + num_samples = int(duration_sec * sample_rate) + t = np.linspace(0, duration_sec, num_samples, dtype=np.float32) + audio = (0.5 * np.sin(2 * np.pi * 440 * t) * 32767).astype(np.int16) + + buffer = io.BytesIO() + with wave.open(buffer, "wb") as wav_file: + wav_file.setnchannels(1) + wav_file.setsampwidth(2) + wav_file.setframerate(sample_rate) + wav_file.writeframes(audio.tobytes()) + + return base64.b64encode(buffer.getvalue()).decode("utf-8") + + +def create_s2st_request(audio_base64: str) -> dict: + """Create S2ST request payload.""" + # Use audio_url format with base64 data URL + audio_url = f"data:audio/wav;base64,{audio_base64}" + return { + "model": MODEL, + "messages": [ + {"role": "system", "content": "请仔细聆听这段语音，然后复述其内容。"}, + { + "role": "user", + "content": [ + { + "type": "audio_url", + "audio_url": {"url": audio_url}, + }, + { + "type": "text", + "text": "", + }, + ], + }, + ], + "max_tokens": 512, + "temperature": 0.7, + "top_p": 0.9, + } + + +async def send_request(session: aiohttp.ClientSession, url: str, payload: dict) -> tuple[float, int, str]: + """Send single request and return (latency, num_tokens, response_text).""" + start = time.perf_counter() + async with session.post(url, json=payload) as response: + result = await response.json() + latency = time.perf_counter() - start + + if "choices" in result and len(result["choices"]) > 0: + text = result["choices"][0].get("message", {}).get("content", "") + tokens = result.get("usage", {}).get("completion_tokens", len(text.split())) + else: + text = str(result) + tokens = 0 + + return latency, tokens, text + + +async def benchmark_concurrent( + url: str, + payload: dict, + num_concurrent: int, + num_requests: int, + warmup: int = 2, +) -> dict: + """Run concurrent benchmark with true sustained concurrency. + + Uses semaphore to maintain exactly num_concurrent requests in flight at all times, + rather than batch-and-wait approach. Measures real wall clock time. + """ + semaphore = asyncio.Semaphore(num_concurrent) + + async def bounded_request(session: aiohttp.ClientSession): + async with semaphore: + return await send_request(session, url, payload) + + async with aiohttp.ClientSession() as session: + if warmup > 0: + print(f" Warming up ({warmup} requests)...", end=" ", flush=True) + for _ in range(warmup): + try: + await send_request(session, url, payload) + except Exception as e: + print(f"Warmup failed: {e}") + print("done") + + print(f" Running {num_requests} requests with concurrency={num_concurrent}...", end=" ", flush=True) + wall_start = time.perf_counter() + + tasks = [bounded_request(session) for _ in range(num_requests)] + results = await asyncio.gather(*tasks, return_exceptions=True) + + wall_time = time.perf_counter() - wall_start + print(f"done ({wall_time:.2f}s)") + + all_latencies = [] + all_tokens = [] + errors = 0 + + for result in results: + if isinstance(result, Exception): + print(f" Request failed: {result}") + errors += 1 + continue + latency, tokens, _ = result + all_latencies.append(latency) + all_tokens.append(tokens) + + if not all_latencies: + return {"error": "All requests failed"} + + latencies = np.array(all_latencies) + total_tokens = sum(all_tokens) + successful_requests = len(all_latencies) + + return { + "num_concurrent": num_concurrent, + "num_requests": num_requests, + "successful_requests": successful_requests, + "errors": errors, + "wall_time": wall_time, + "p50_latency": float(np.percentile(latencies, 50)), + "p95_latency": float(np.percentile(latencies, 95)), + "p99_latency": float(np.percentile(latencies, 99)), + "mean_latency": float(np.mean(latencies)), + "min_latency": float(np.min(latencies)), + "max_latency": float(np.max(latencies)), + "throughput_req_per_sec": successful_requests / wall_time, + "throughput_tokens_per_sec": total_tokens / wall_time, + } + + +class TestStepAudio2OnlineServing: + """Online serving tests for Step-Audio2.""" + + @pytest.fixture(scope="class") + def server(self): + """Start vLLM-Omni server.""" + with OmniServer(MODEL, ["--stage-configs-path", STAGE_CONFIG]) as omni_server: + yield f"http://{omni_server.host}:{omni_server.port}/v1/chat/completions" + + @pytest.mark.asyncio + async def test_single_request(self, server): + """Test single request works.""" + audio_base64 = create_dummy_audio_base64(duration_sec=2.0) + payload = create_s2st_request(audio_base64) + + async with aiohttp.ClientSession() as session: + latency, tokens, text = await send_request(session, server, payload) + + assert latency > 0 + assert tokens >= 0 + print(f"Single request: latency={latency:.2f}s, tokens={tokens}") + + @pytest.mark.asyncio + @pytest.mark.parametrize("num_concurrent", [1, 2, 4]) + async def test_concurrent_throughput(self, server, num_concurrent): + """Test concurrent throughput at different concurrency levels.""" + audio_base64 = create_dummy_audio_base64(duration_sec=5.0) + payload = create_s2st_request(audio_base64) + + num_requests = num_concurrent * 3 + + stats = await benchmark_concurrent(server, payload, num_concurrent, num_requests) + + assert "error" not in stats + print(f"\nConcurrency={num_concurrent}:") + print(f" p50={stats['p50_latency']:.2f}s, p95={stats['p95_latency']:.2f}s") + print(f" throughput={stats['throughput_req_per_sec']:.2f} req/s") + + +async def run_throughput_benchmark( + server_url: str, + audio_duration: float = 5.0, + concurrency_levels: list[int] = [1, 2, 4, 8], + requests_per_level: int = 10, + warmup: int = 2, +): + """Run full throughput benchmark.""" + audio_base64 = create_dummy_audio_base64(duration_sec=audio_duration) + payload = create_s2st_request(audio_base64) + + print(f"{'=' * 70}") + print("Step-Audio2 Online Serving Throughput Benchmark") + print(f"{'=' * 70}") + print(f"Server: {server_url}") + print(f"Audio duration: {audio_duration}s") + print(f"Requests per level: {requests_per_level}") + print(f"Warmup requests: {warmup}") + print(f"Concurrency levels: {concurrency_levels}") + print(f"{'=' * 70}") + + results = [] + for num_concurrent in concurrency_levels: + print(f"\n[Concurrency={num_concurrent}]") + stats = await benchmark_concurrent(server_url, payload, num_concurrent, requests_per_level, warmup) + results.append(stats) + + if "error" not in stats: + print( + f" Result: p50={stats['p50_latency']:.2f}s, " + f"p95={stats['p95_latency']:.2f}s, " + f"throughput={stats['throughput_req_per_sec']:.2f} req/s, " + f"{stats['throughput_tokens_per_sec']:.1f} tok/s" + ) + + print(f"\n{'=' * 80}") + print("SUMMARY") + print(f"{'=' * 80}") + print( + f"{'Conc':<6} {'Reqs':<6} {'Err':<5} {'Wall(s)':<9} {'p50':<8} {'p95':<8} {'p99':<8} {'Req/s':<8} {'Tok/s':<8}" + ) + print("-" * 80) + + for stats in results: + if "error" in stats: + print(f"{stats.get('num_concurrent', '?'):<6} ERROR: {stats['error']}") + continue + print( + f"{stats['num_concurrent']:<6} " + f"{stats['successful_requests']:<6} " + f"{stats['errors']:<5} " + f"{stats['wall_time']:<9.2f} " + f"{stats['p50_latency']:<8.2f} " + f"{stats['p95_latency']:<8.2f} " + f"{stats['p99_latency']:<8.2f} " + f"{stats['throughput_req_per_sec']:<8.2f} " + f"{stats['throughput_tokens_per_sec']:<8.1f}" + ) + + print("=" * 80) + return results + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description="Step-Audio2 Online Serving Benchmark") + parser.add_argument( + "--server-url", type=str, default="http://localhost:8000/v1/chat/completions", help="vLLM server URL" + ) + parser.add_argument("--audio-duration", type=float, default=5.0, help="Input audio duration in seconds") + parser.add_argument("--concurrency", type=str, default="1,2,4,8", help="Comma-separated concurrency levels to test") + parser.add_argument("--requests-per-level", type=int, default=10, help="Number of requests per concurrency level") + parser.add_argument("--warmup", type=int, default=2, help="Number of warmup requests before each benchmark") + + args = parser.parse_args() + concurrency_levels = [int(x) for x in args.concurrency.split(",")] + + asyncio.run( + run_throughput_benchmark( + args.server_url, + args.audio_duration, + concurrency_levels, + args.requests_per_level, + args.warmup, + ) + ) diff --git a/tests/model_executor/models/step_audio2/test_step_audio2_token2wav_async_chunk.py b/tests/model_executor/models/step_audio2/test_step_audio2_token2wav_async_chunk.py new file mode 100644 index 000000000..050a5b724 --- /dev/null +++ b/tests/model_executor/models/step_audio2/test_step_audio2_token2wav_async_chunk.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +# Token2Wav imports the optional `step-audio2` extras (flashcosyvoice / s3tokenizer +# / hyperpyyaml) at module load. Skip instead of erroring at collection time on CI +# images that do not install them (e.g. the CPU core_model runner); a genuine +# import error (e.g. vLLM API drift) still surfaces. See +# examples/offline_inference/step_audio2/README.md. +try: + from vllm_omni.model_executor.models.step_audio2.step_audio2_token2wav import ( + StepAudio2Token2WavForConditionalGeneration, + ) +except ModuleNotFoundError as exc: + pytest.skip( + f"Step-Audio2 Token2Wav extras not installed ({exc.name})", + allow_module_level=True, + ) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _DummyModelConfig: + def __init__(self) -> None: + self.hf_config = SimpleNamespace( + token2wav_path="/tmp/fake-token2wav", + token2wav_float16=False, + ) + self.model = "/tmp/fake-model" + + def get_hidden_size(self) -> int: + return 8 + + +def _build_model() -> StepAudio2Token2WavForConditionalGeneration: + vllm_config = SimpleNamespace( + model_config=_DummyModelConfig(), + device_config=SimpleNamespace(device="cpu"), + ) + return StepAudio2Token2WavForConditionalGeneration(vllm_config=vllm_config) + + +def test_step_audio2_token2wav_sync_path_not_misdetected_by_empty_runtime_info(monkeypatch): + model = _build_model() + + monkeypatch.setattr("os.path.exists", lambda _: True) + monkeypatch.setattr( + model.token2wav, + "forward", + lambda *args, **kwargs: torch.zeros(1, 16, dtype=torch.float32), + ) + + out = model.forward( + input_ids=torch.tensor([1, 2, 3]), + positions=torch.tensor([0, 1, 2]), + runtime_additional_information=[{}], + ) + + audio = out.multimodal_outputs["model_outputs"] + # sync path returns the full waveform wrapped in a single-element list, + # same contract as the async-chunk path (see tests below) + assert isinstance(audio, list) + assert len(audio) == 1 + assert isinstance(audio[0], torch.Tensor) + assert model._stream_states == [] + + +def test_step_audio2_token2wav_async_chunk_batch_guard(): + model = _build_model() + + with pytest.raises(RuntimeError, match="only supports batch=1"): + model._forward_async_chunk( + input_ids=torch.tensor([1, 2, 3]), + runtime_additional_information=[ + {"left_context_size": 0}, + {"left_context_size": 1}, + ], + ) + + +def test_step_audio2_token2wav_async_chunk_last_chunk_resets_state(monkeypatch): + model = _build_model() + + monkeypatch.setattr("os.path.exists", lambda _: True) + calls = {"setup": 0, "stream": 0, "reset": 0} + + def _setup_stream_for(prompt_wav, state): + calls["setup"] += 1 + state.setup_done = True + state.stream_cache = {} + state.hift_cache_dict = { + "mel": torch.zeros(1, 1, 0), + "source": torch.zeros(1, 1, 0), + "speech": torch.zeros(1, 0), + } + + def _stream_chunk_for(audio_tokens, prompt_wav, last_chunk, state): + calls["stream"] += 1 + return torch.tensor([0.1, 0.2, 0.3], dtype=torch.float32) + + def _reset_stream_for(state): + calls["reset"] += 1 + state.setup_done = False + state.finished = True + state.stream_cache = None + state.hift_cache_dict = {} + + monkeypatch.setattr(model.token2wav, "setup_stream_for", _setup_stream_for) + monkeypatch.setattr(model.token2wav, "stream_chunk_for", _stream_chunk_for) + monkeypatch.setattr(model.token2wav, "reset_stream_for", _reset_stream_for) + + out = model.forward( + input_ids=torch.tensor([10, 11, 12]), + positions=torch.tensor([0, 1, 2]), + runtime_additional_information=[{"left_context_size": 1}], + ) + + audio_list = out.multimodal_outputs["model_outputs"] + assert isinstance(audio_list, list) + assert len(audio_list) == 1 + assert torch.allclose(audio_list[0], torch.tensor([0.1, 0.2, 0.3], dtype=torch.float32)) + assert calls == {"setup": 1, "stream": 1, "reset": 1} + + +def test_step_audio2_token2wav_async_chunk_empty_eof_returns_zero_chunk(): + model = _build_model() + + out = model._forward_async_chunk( + input_ids=torch.tensor([], dtype=torch.int64), + runtime_additional_information=[{"left_context_size": 1}], + ) + + audio_list = out.multimodal_outputs["model_outputs"] + assert isinstance(audio_list, list) + assert len(audio_list) == 1 + assert torch.equal(audio_list[0], torch.zeros(1, dtype=torch.float32)) + assert model._stream_states + assert model._stream_states[0].finished is True diff --git a/tests/model_executor/stage_input_processors/test_step_audio2_async_chunk.py b/tests/model_executor/stage_input_processors/test_step_audio2_async_chunk.py new file mode 100644 index 000000000..186f00890 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_step_audio2_async_chunk.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections import defaultdict +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.models.step_audio2.step_audio2_constants import ( + DEFAULT_STREAM_CONFIG, + DEFAULT_TOKEN_CONFIG, +) +from vllm_omni.model_executor.stage_input_processors.step_audio2 import ( + thinker2token2wav_async_chunk, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _req(external_req_id: str, *, prompt_token_ids: list[int], all_token_ids: list[int], finished: bool): + return SimpleNamespace( + external_req_id=external_req_id, + prompt_token_ids=prompt_token_ids, + all_token_ids=all_token_ids, + is_finished=lambda: finished, + ) + + +def test_step_audio2_async_chunk_uses_decode_only_tokens_not_prompt_history(): + audio_start = DEFAULT_TOKEN_CONFIG.audio_start + audio_eos = DEFAULT_TOKEN_CONFIG.audio_eos + transfer_manager = SimpleNamespace(code_prompt_token_ids=defaultdict(list)) + + prompt = [11, 12, audio_start + 5] # historical prompt audio token + generated = [99, audio_start + 1, audio_start + 2, audio_start + audio_eos] + request = _req( + "rid-decode-only", + prompt_token_ids=prompt, + all_token_ids=prompt + generated, + finished=True, + ) + + payload = thinker2token2wav_async_chunk( + transfer_manager=transfer_manager, + pooling_output=None, + request=request, + ) + + assert payload is not None + assert payload["code_predictor_codes"] == [1, 2] + assert payload["left_context_size"] == 1 + assert payload["finished"].item() is True + + +def test_step_audio2_async_chunk_returns_none_when_not_enough_tokens(): + audio_start = DEFAULT_TOKEN_CONFIG.audio_start + required = DEFAULT_STREAM_CONFIG.chunk_size + DEFAULT_STREAM_CONFIG.pre_lookahead_len + transfer_manager = SimpleNamespace(code_prompt_token_ids=defaultdict(list)) + + generated = [audio_start + i for i in range(required - 1)] + request = _req( + "rid-not-ready", + prompt_token_ids=[1, 2, 3], + all_token_ids=[1, 2, 3] + generated, + finished=False, + ) + + payload = thinker2token2wav_async_chunk( + transfer_manager=transfer_manager, + pooling_output=None, + request=request, + ) + + assert payload is None + + +def test_step_audio2_async_chunk_emits_non_last_chunk_and_advances_consumed_by_chunk_size(): + audio_start = DEFAULT_TOKEN_CONFIG.audio_start + chunk_size = DEFAULT_STREAM_CONFIG.chunk_size + required = chunk_size + DEFAULT_STREAM_CONFIG.pre_lookahead_len + transfer_manager = SimpleNamespace(code_prompt_token_ids=defaultdict(list)) + + generated = [audio_start + i for i in range(required + 10)] + request = _req( + "rid-ready", + prompt_token_ids=[7, 8], + all_token_ids=[7, 8] + generated, + finished=False, + ) + + payload = thinker2token2wav_async_chunk( + transfer_manager=transfer_manager, + pooling_output=None, + request=request, + ) + + assert payload is not None + assert payload["left_context_size"] == 0 + assert payload["finished"].item() is False + assert payload["code_predictor_codes"] == list(range(required)) + assert transfer_manager.code_prompt_token_ids["rid-ready"] == list(range(chunk_size)) + + +def test_step_audio2_async_chunk_emits_eof_when_finished_with_no_remaining_audio(): + transfer_manager = SimpleNamespace(code_prompt_token_ids=defaultdict(list)) + transfer_manager.code_prompt_token_ids["rid-eof"] = [1, 2, 3] + + request = _req( + "rid-eof", + prompt_token_ids=[10, 11], + all_token_ids=[ + 10, + 11, + DEFAULT_TOKEN_CONFIG.audio_start + 1, + DEFAULT_TOKEN_CONFIG.audio_start + 2, + DEFAULT_TOKEN_CONFIG.audio_start + 3, + ], + finished=True, + ) + + payload = thinker2token2wav_async_chunk( + transfer_manager=transfer_manager, + pooling_output=None, + request=request, + ) + + assert payload == { + "code_predictor_codes": [], + "left_context_size": 1, + "finished": torch.tensor(True, dtype=torch.bool), + } diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 99df2332b..92ebb0f9f 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -80,6 +80,7 @@ _MOSS_TTS_FULL_MODEL_STAGES = {"moss_tts", "moss_tts_codec"} _HIGGS_AUDIO_V2_TTS_MODEL_STAGES = {"higgs_audio_v2"} _HIGGS_V3_TTS_MODEL_STAGES = {"higgs_audio_v3"} _GLM_TTS_MODEL_STAGES = {"glm_tts"} +_STEP_AUDIO2_TTS_MODEL_STAGES = {"step_audio2_thinker"} _TTS_MODEL_STAGES: set[str] = ( _VOXTRAL_TTS_MODEL_STAGES | _QWEN3_TTS_MODEL_STAGES @@ -94,6 +95,7 @@ _TTS_MODEL_STAGES: set[str] = ( | _MOSS_TTS_MODEL_STAGES | _MOSS_TTS_FULL_MODEL_STAGES | _GLM_TTS_MODEL_STAGES + | _STEP_AUDIO2_TTS_MODEL_STAGES ) _SAMPLING_MAX_TOKENS_TTS_MODEL_TYPES = { "fish_tts", @@ -655,6 +657,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return "higgs_audio_v3" if model_stage in _GLM_TTS_MODEL_STAGES: return "glm_tts" + if model_stage in _STEP_AUDIO2_TTS_MODEL_STAGES: + return "step_audio2" return None def _get_custom_voice_dir(self) -> str | None: @@ -1356,6 +1360,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): def _validate_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None: """Validate TTS request parameters. Returns error message or None.""" + if self._tts_model_type == "step_audio2": + # StepAudio2 only requires non-empty input text + if not request.input or not request.input.strip(): + return "Input text cannot be empty" + return None if self._tts_model_type == "voxtral_tts": return self._validate_voxtral_tts_request(request) if self._tts_model_type == "fish_tts": @@ -2858,6 +2867,33 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): "multi_modal_data": {"audio": [(audio.audio_array, audio.sampling_rate)]}, } + # ---- Step-Audio2 helpers ---- + + def _build_step_audio2_prompt( + self, + request: OpenAICreateSpeechRequest, + ) -> dict[str, Any]: + """Build prompt for Step-Audio2 TTS. + + Constructs the chat prompt with ```` as the last token + of the assistant turn (without ``<|im_end|>``), so the thinker + continues generating audio tokens. + + Prompt format:: + <|im_start|>system\\n{system_prompt}<|im_end|>\\n + <|im_start|>user\\n{input_text}<|im_end|>\\n + <|im_start|>assistant\\n + """ + system_prompt = getattr(request, "instructions", None) or "You are a voice assistant. Read the text aloud." + text = request.input + + raw_prompt = ( + f"<|im_start|>system\n{system_prompt}<|im_end|>\n" + f"<|im_start|>user\n{text}<|im_end|>\n" + f"<|im_start|>assistant\n" + ) + return {"prompt": raw_prompt} + # ---- Fish Speech helpers ---- def _build_fish_speech_prompt( @@ -3202,7 +3238,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if err: raise ValueError(err) - if self._is_fish_speech: + if self._tts_model_type == "step_audio2": + if not request.input or not request.input.strip(): + raise ValueError("Input text cannot be empty") + prompt = self._build_step_audio2_prompt(request) + tts_params = {} + elif self._is_fish_speech: validation_error = self._validate_fish_tts_request(request) if validation_error: raise ValueError(validation_error) @@ -3413,7 +3454,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): tts_params = {} prompt = {"prompt": request.input} - if self._is_fish_speech: + if self._tts_model_type == "step_audio2": + model_type = "step_audio2" + elif self._is_fish_speech: model_type = "fish_speech" elif self._tts_model_type == "covo_audio": model_type = "covo_audio" diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 5368c801e..9906fa4e0 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -53,6 +53,32 @@ _OMNI_MODELS = { "qwen3_omni_code2wav", "Qwen3OmniMoeCode2Wav", ), + # Step-Audio2 models + "StepAudio2ForCausalLM": ( + "step_audio2", + "step_audio2", + "StepAudio2ForConditionalGeneration", + ), + "StepAudio2ForConditionalGeneration": ( + "step_audio2", + "step_audio2", + "StepAudio2ForConditionalGeneration", + ), + "StepAudio2ThinkerForConditionalGeneration": ( + "step_audio2", + "step_audio2_thinker", + "StepAudio2ThinkerForConditionalGeneration", + ), + "StepAudio2Token2WavModel": ( + "step_audio2", + "step_audio2_token2wav", + "StepAudio2Token2WavForConditionalGeneration", + ), + "StepAudio2Token2WavForConditionalGeneration": ( + "step_audio2", + "step_audio2_token2wav", + "StepAudio2Token2WavForConditionalGeneration", + ), "CosyVoice3Model": ( "cosyvoice3", "cosyvoice3", diff --git a/vllm_omni/model_executor/models/step_audio2/__init__.py b/vllm_omni/model_executor/models/step_audio2/__init__.py new file mode 100644 index 000000000..0b6f07e28 --- /dev/null +++ b/vllm_omni/model_executor/models/step_audio2/__init__.py @@ -0,0 +1,3 @@ +from .step_audio2 import StepAudio2ForConditionalGeneration + +__all__ = ["StepAudio2ForConditionalGeneration"] diff --git a/vllm_omni/model_executor/models/step_audio2/assets/default_female.wav b/vllm_omni/model_executor/models/step_audio2/assets/default_female.wav new file mode 100644 index 000000000..04283a734 Binary files /dev/null and b/vllm_omni/model_executor/models/step_audio2/assets/default_female.wav differ diff --git a/vllm_omni/model_executor/models/step_audio2/assets/default_male.wav b/vllm_omni/model_executor/models/step_audio2/assets/default_male.wav new file mode 100644 index 000000000..c65f363b3 Binary files /dev/null and b/vllm_omni/model_executor/models/step_audio2/assets/default_male.wav differ diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2.py b/vllm_omni/model_executor/models/step_audio2/step_audio2.py new file mode 100644 index 000000000..82e5609ff --- /dev/null +++ b/vllm_omni/model_executor/models/step_audio2/step_audio2.py @@ -0,0 +1,195 @@ +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP +from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.sequence import IntermediateTensors + +from .step_audio2_thinker import ( + StepAudio2DummyInputsBuilder, + StepAudio2MultiModalProcessor, + StepAudio2ProcessingInfo, +) + +logger = init_logger(__name__) + + +@MULTIMODAL_REGISTRY.register_processor( + StepAudio2MultiModalProcessor, + info=StepAudio2ProcessingInfo, + dummy_inputs=StepAudio2DummyInputsBuilder, +) +class StepAudio2ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP): + """ + Step-Audio2 Main Controller + + Manages two-stage inference pipeline: + - Stage 1 (Thinker): Audio understanding and token generation + - Stage 2 (Token2Wav): Audio token to waveform synthesis + + Usage: + # Stage 1: Thinker + model = StepAudio2ForConditionalGeneration( + vllm_config=config, + model_stage="thinker" + ) + + # Stage 2: Token2Wav + model = StepAudio2ForConditionalGeneration( + vllm_config=config, + model_stage="token2wav" + ) + """ + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.have_multimodal_outputs = True + + config = vllm_config.model_config.hf_config + multimodal_config = vllm_config.model_config.multimodal_config + + self.config = config + self.multimodal_config = multimodal_config + self.vllm_config = vllm_config + + raw_model_stage = vllm_config.model_config.model_stage + # Normalise: "step_audio2_thinker" → "thinker" + self.model_stage = "thinker" if raw_model_stage in ("thinker", "step_audio2_thinker") else raw_model_stage + + if self.model_stage == "thinker": + self.thinker = init_vllm_registered_model( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "thinker"), + hf_config=config, + architectures=["StepAudio2ThinkerForConditionalGeneration"], + ) + self.model = self.thinker + self.token2wav = None + + logger.info("Initialized Step-Audio2 Thinker (Stage 0)") + + elif self.model_stage == "token2wav": + self.thinker = None + self.token2wav = init_vllm_registered_model( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "token2wav"), + hf_config=config, + architectures=["StepAudio2Token2WavForConditionalGeneration"], + ) + self.model = self.token2wav + + logger.info("Initialized Step-Audio2 Token2Wav (Stage 1)") + + else: + raise ValueError( + f"Invalid model_stage: {self.model_stage}. Must be 'thinker'/'step_audio2_thinker' or 'token2wav'" + ) + + self.make_empty_intermediate_tensors = ( + self.thinker.make_empty_intermediate_tensors if self.model_stage == "thinker" else lambda: None + ) + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + """Get placeholder string for a modality + + Returns: + For audio: "" (matches processor's audio_token) + """ + if modality == "audio": + return "" + return None + + def get_language_model(self) -> nn.Module: + """Get the underlying language model.""" + if self.model_stage == "thinker": + return self.thinker.get_language_model() + # Token2Wav is not a real language model. Return the wrapper itself so + # vLLM interface helpers resolve embed_input_ids() on the adapter layer. + return self.token2wav + + def get_multimodal_embeddings(self, **kwargs): + """Get multimodal embeddings - only used in Thinker stage.""" + if self.model_stage == "thinker": + return self.thinker.get_multimodal_embeddings(**kwargs) + return None + + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: + """Compatibility helper used by older call sites.""" + if self.model_stage == "thinker": + return self.model.get_input_embeddings(input_ids) + return self.model.embed_input_ids(input_ids) + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings=None, + is_multimodal=None, + ) -> torch.Tensor: + """Explicit vLLM embedding hook for both stages.""" + if self.model_stage == "token2wav": + return self.model.embed_input_ids(input_ids) + return self.model.get_input_embeddings(input_ids, multimodal_embeddings) + + def embed_multimodal(self, **kwargs): + """Delegate multimodal embedding to thinker stage only.""" + if self.model_stage == "thinker": + return self.model.embed_multimodal(**kwargs) + return None + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ): + """ + Forward pass through the model + + For Thinker: + Returns hidden states/logits + For Token2Wav: + Returns waveform + """ + return self.model.forward( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **kwargs, + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + """Compute logits from hidden states""" + return self.model.compute_logits(hidden_states) + + def load_weights(self, weights): + """Load weights""" + return self.model.load_weights(weights) + + def move_submodules_to_devices( + self, + *, + thinker_device: str | torch.device | None = None, + token2wav_device: str | torch.device | None = None, + ) -> None: + """ + Optionally move thinker/token2wav to different devices + + Example: + model.move_submodules_to_devices( + thinker_device='cuda:0', + token2wav_device='cuda:1', + ) + """ + if thinker_device is not None and self.thinker is not None: + self.thinker.to(torch.device(thinker_device)) + logger.info(f"Moved Thinker to {thinker_device}") + + if token2wav_device is not None and self.token2wav is not None: + self.token2wav.to(torch.device(token2wav_device)) + logger.info(f"Moved Token2Wav to {token2wav_device}") diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2_constants.py b/vllm_omni/model_executor/models/step_audio2/step_audio2_constants.py new file mode 100644 index 000000000..802cd63be --- /dev/null +++ b/vllm_omni/model_executor/models/step_audio2/step_audio2_constants.py @@ -0,0 +1,181 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Step-Audio2 configuration constants - Single Source of Truth.""" + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class StepAudio2TokenConfig: + """ + Step-Audio2 token ID ranges and vocabulary configuration. + + Token ID Layout (based on vocab_size=158720): + - Text tokens: 0 - 151688 (standard Qwen tokenizer) + - Special tokens: 151689 - 151695 (reserved) + - Audio tokens: 151696 - 158257 (vocab size 6562) + + These values are architecture constants for Step-Audio2 and do not + change between model variants (mini/7B). + """ + + # Text token range + text_max: int = 151688 + """Maximum text token ID (inclusive)""" + + # Audio token range (absolute IDs in full vocabulary) + audio_start: int = 151696 + """First audio token ID in absolute vocabulary""" + + audio_vocab_size: int = 6562 + """Total number of audio tokens""" + + audio_eos: int = 6561 + """Audio EOS token ID (relative to audio_start, used for padding)""" + + # Special tokens + audio_patch_token_id: int = 151690 + """ placeholder token ID""" + + @property + def audio_end(self) -> int: + """Last audio token ID in absolute vocabulary""" + return self.audio_start + self.audio_vocab_size - 1 + + +@dataclass(frozen=True) +class StepAudio2EncoderConfig: + """ + Step-Audio2 audio encoder architecture configuration. + + These values define the encoder structure and should match your model weights. + """ + + # Mel spectrogram parameters + n_mels: int = 128 + """Number of mel frequency bins""" + + # Encoder architecture + n_audio_ctx: int = 1500 + """Audio context length (max sequence length)""" + + n_audio_state: int = 512 + """Audio encoder hidden state dimension""" + + n_audio_head: int = 8 + """Number of attention heads in encoder""" + + n_audio_layer: int = 6 + """Number of encoder layers""" + + # Adapter parameters + kernel_size: int = 3 + """Convolution kernel size for adapter""" + + adapter_stride: int = 2 + """Stride for adapter convolution (downsampling factor)""" + + +@dataclass(frozen=True) +class StepAudio2Token2WavConfig: + """Token2Wav synthesis configuration.""" + + n_timesteps: int = 10 + """Number of diffusion timesteps for flow model inference""" + + sample_rate: int = 24000 + """Output audio sample rate""" + + +@dataclass(frozen=True) +class StepAudio2StreamConfig: + """Streaming (chunked) inference configuration for Token2Wav.""" + + chunk_size: int = 25 + """Number of audio tokens consumed per chunk.""" + + pre_lookahead_len: int = 3 + """Conformer encoder lookahead tokens beyond the chunk boundary.""" + + up_rate: int = 2 + """Token-to-mel upsampling rate inside the flow model.""" + + mel_cache_len: int = 8 + """HiFT overlap mel frames (160 ms at 50 Hz mel rate).""" + + n_timesteps: int = 10 + """Flow model ODE solver steps.""" + + estimator_cache_keep: int = 100 + """Attention-cache pruning window for the flow estimator.""" + + +@dataclass(frozen=True) +class StepAudio2ModelConfig: + """ + Step-Audio2 complete model configuration - Single Source of Truth. + + All architecture constants are defined here. The stage YAML and model config.json + can override these values if needed, but this serves as the default. + """ + + # LLM configuration (Qwen2) + hidden_size: int = 4096 + """LLM hidden dimension (Qwen2-7B: 4096, Qwen2-1.5B: 1536, Qwen2-0.5B: 896)""" + + # Token configuration + text_max: int = 151688 + audio_start: int = 151696 + audio_vocab_size: int = 6562 + audio_eos: int = 6561 + audio_patch_token_id: int = 151690 + + # Encoder configuration + n_mels: int = 128 + n_audio_ctx: int = 1500 + n_audio_state: int = 512 + n_audio_head: int = 8 + n_audio_layer: int = 6 + kernel_size: int = 3 + adapter_stride: int = 2 + + +# Default configuration instances +DEFAULT_TOKEN_CONFIG = StepAudio2TokenConfig() +DEFAULT_ENCODER_CONFIG = StepAudio2EncoderConfig() +DEFAULT_TOKEN2WAV_CONFIG = StepAudio2Token2WavConfig() +DEFAULT_STREAM_CONFIG = StepAudio2StreamConfig() +DEFAULT_MODEL_CONFIG = StepAudio2ModelConfig() + +STREAM_SOURCE_CACHE_LEN = DEFAULT_STREAM_CONFIG.mel_cache_len * 480 # 3840 samples + +# Export constants for backward compatibility +STEP_AUDIO2_TEXT_MAX = DEFAULT_TOKEN_CONFIG.text_max +STEP_AUDIO2_AUDIO_START = DEFAULT_TOKEN_CONFIG.audio_start +STEP_AUDIO2_AUDIO_VOCAB_SIZE = DEFAULT_TOKEN_CONFIG.audio_vocab_size +STEP_AUDIO2_AUDIO_EOS = DEFAULT_TOKEN_CONFIG.audio_eos +STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID = DEFAULT_TOKEN_CONFIG.audio_patch_token_id +STEP_AUDIO2_AUDIO_END = DEFAULT_TOKEN_CONFIG.audio_end + +# Default prompt wav path for Token2Wav (can be overridden by env var) +STEP_AUDIO2_DEFAULT_PROMPT_WAV = "default_female.wav" + + +__all__ = [ + "StepAudio2TokenConfig", + "StepAudio2EncoderConfig", + "StepAudio2Token2WavConfig", + "StepAudio2StreamConfig", + "StepAudio2ModelConfig", + "DEFAULT_TOKEN_CONFIG", + "DEFAULT_ENCODER_CONFIG", + "DEFAULT_TOKEN2WAV_CONFIG", + "DEFAULT_STREAM_CONFIG", + "DEFAULT_MODEL_CONFIG", + "STREAM_SOURCE_CACHE_LEN", + "STEP_AUDIO2_TEXT_MAX", + "STEP_AUDIO2_AUDIO_START", + "STEP_AUDIO2_AUDIO_VOCAB_SIZE", + "STEP_AUDIO2_AUDIO_EOS", + "STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID", + "STEP_AUDIO2_AUDIO_END", +] diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py b/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py new file mode 100644 index 000000000..7fd0bf4fb --- /dev/null +++ b/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py @@ -0,0 +1,898 @@ +""" +Step-Audio2 Thinker - Stage 1 LLM for Audio Understanding + +This file contains: +- Audio encoder components (AudioEncoder, Adaptor) +- Audio preprocessing utilities (mel-spectrogram extraction) +- Multi-modal processing (StepAudio2MultiModalProcessor) +- Main thinker model (StepAudio2ThinkerForConditionalGeneration) +""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping, Sequence +from typing import TYPE_CHECKING, Any, TypedDict + +import librosa # noqa: TID251 +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn.utils.rnn import pad_sequence +from transformers import AutoTokenizer +from transformers.processing_utils import ProcessorMixin +from vllm.config import VllmConfig +from vllm.inputs import MultiModalDataDict +from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP +from vllm.model_executor.models.utils import ( + _merge_multimodal_embeddings, + flatten_bn, + init_vllm_registered_model, + maybe_prefix, +) +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import ( + MultiModalFieldConfig, + MultiModalKwargsItems, + NestedTensors, +) +from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser +from vllm.multimodal.processing import ( + BaseDummyInputsBuilder, + BaseMultiModalProcessor, + BaseProcessingInfo, + PromptReplacement, + PromptUpdate, +) +from vllm.sequence import IntermediateTensors + +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.model_executor.models.step_audio2.step_audio2_constants import ( + DEFAULT_MODEL_CONFIG, + DEFAULT_TOKEN_CONFIG, + STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID, +) + +if TYPE_CHECKING: + pass + + +# ============================================================================ +# Audio Preprocessing Utilities +# ============================================================================ + + +def _mel_filters(n_mels: int) -> torch.Tensor: + """Generate mel filter banks""" + assert n_mels in {80, 128}, f"Unsupported n_mels: {n_mels}" + return torch.from_numpy(librosa.filters.mel(sr=16000, n_fft=400, n_mels=n_mels)) + + +def _normalize_audio(audio: Any) -> torch.Tensor: + """Normalize audio to torch.Tensor""" + if isinstance(audio, tuple): + audio = audio[0] + if isinstance(audio, np.ndarray): + audio = torch.from_numpy(audio) + elif not torch.is_tensor(audio): + audio = torch.tensor(audio, dtype=torch.float32) + return audio.float() + + +def log_mel_spectrogram(audio: Any, n_mels: int = 128, padding: int = 479) -> torch.Tensor: + """ + Convert audio to log mel-spectrogram + + Args: + audio: Audio waveform (any format) + n_mels: Number of mel frequency bins (default: 128) + padding: Right padding for STFT (default: 479) + + Returns: + Log mel-spectrogram tensor of shape (n_mels, T) + """ + audio = _normalize_audio(audio) + if padding > 0: + audio = F.pad(audio, (0, padding)) + window = torch.hann_window(400, device=audio.device) + stft = torch.stft(audio, 400, 160, window=window, return_complex=True) + magnitudes = stft[..., :-1].abs() ** 2 + filters = _mel_filters(n_mels).to(audio.device) + mel_spec = filters @ magnitudes + + log_spec = torch.clamp(mel_spec, min=1e-10).log10() + log_spec = torch.maximum(log_spec, log_spec.max() - 8.0) + log_spec = (log_spec + 4.0) / 4.0 + return log_spec + + +def padding_mels(data: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]: + """ + Pad mel spectrograms to same length + + Args: + data: List of mel spectrograms, each of shape (n_mels, T) + + Returns: + padded_feats: Padded tensor of shape (batch, n_mels, max_T) + feats_lengths: Lengths of each sequence + """ + feats_lengths = torch.tensor([s.size(1) - 2 for s in data], dtype=torch.int32) + feats = [s.t() for s in data] + padded_feats = pad_sequence(feats, batch_first=True, padding_value=0) + return padded_feats.transpose(1, 2), feats_lengths + + +# ============================================================================ +# Encoder Utilities +# ============================================================================ + + +def make_non_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor: + """Create non-padding mask from lengths""" + batch_size = lengths.size(0) + max_len = max_len if max_len > 0 else lengths.max().item() + seq_range = torch.arange(0, max_len, dtype=torch.int64, device=lengths.device) + seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len) + seq_length_expand = lengths.unsqueeze(-1) + mask = seq_range_expand >= seq_length_expand + return ~mask + + +def mask_to_bias(mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + """Convert mask to attention bias""" + assert mask.dtype == torch.bool + assert dtype in [torch.float32, torch.bfloat16, torch.float16] + mask = mask.to(dtype) + mask = (1.0 - mask) * -1.0e10 + return mask + + +# ============================================================================ +# Audio Encoder Components +# ============================================================================ + + +class MultiHeadAttention(nn.Module): + """Multi-head attention for audio encoder""" + + def __init__(self, n_state: int, n_head: int): + super().__init__() + self.n_head = n_head + self.query = nn.Linear(n_state, n_state) + self.key = nn.Linear(n_state, n_state, bias=False) + self.value = nn.Linear(n_state, n_state) + self.out = nn.Linear(n_state, n_state) + + def forward( + self, + x: torch.Tensor, + mask: torch.Tensor | None = None, + ): + q = self.query(x) + k = self.key(x) + v = self.value(x) + + wv, qk = self.qkv_attention(q, k, v, mask) + return self.out(wv), qk + + def qkv_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor | None = None): + B, T, D = q.shape + head_dim = D // self.n_head + + # Reshape: (B, T, D) -> (B, n_head, T, head_dim) + q = q.view(B, T, self.n_head, head_dim).transpose(1, 2) + k = k.view(B, T, self.n_head, head_dim).transpose(1, 2) + v = v.view(B, T, self.n_head, head_dim).transpose(1, 2) + + # Use F.scaled_dot_product_attention + # mask is additive bias (0 or -1e10), pass directly as attn_mask + out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask) + + # Reshape back: (B, n_head, T, head_dim) -> (B, T, D) + out = out.transpose(1, 2).reshape(B, T, D) + return out, None + + +class ResidualAttentionBlock(nn.Module): + """Residual attention block for audio encoder""" + + def __init__(self, n_state: int, n_head: int): + super().__init__() + + self.attn = MultiHeadAttention(n_state, n_head) + self.attn_ln = nn.LayerNorm(n_state) + + n_mlp = n_state * 4 + self.mlp = nn.Sequential(nn.Linear(n_state, n_mlp), nn.GELU(), nn.Linear(n_mlp, n_state)) + self.mlp_ln = nn.LayerNorm(n_state) + + def forward( + self, + x: torch.Tensor, + mask: torch.Tensor | None = None, + ): + x = x + self.attn(self.attn_ln(x), mask=mask)[0] + x = x + self.mlp(self.mlp_ln(x)) + return x + + +class AudioEncoder(nn.Module): + """ + Step-Audio2 Audio Encoder + + Lightweight audio encoder (6 layers, 512 hidden) + Optimized for 25s audio chunks at 25Hz + """ + + def __init__(self, n_mels: int, n_ctx: int, n_state: int, n_head: int, n_layer: int): + super().__init__() + self.conv1 = nn.Conv1d(n_mels, n_state, kernel_size=3, padding=1) + self.conv2 = nn.Conv1d(n_state, n_state, kernel_size=3, stride=2, padding=1) + self.positional_embedding = nn.Embedding(n_ctx, n_state) + + self.blocks: Iterable[ResidualAttentionBlock] = nn.ModuleList( + [ResidualAttentionBlock(n_state, n_head) for _ in range(n_layer)] + ) + self.avg_pooler = nn.AvgPool1d(2, stride=2) + self.after_norm = nn.LayerNorm(n_state) + + def forward(self, x: torch.Tensor, x_len: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """ + Args: + x: mel spectrogram, shape (batch_size, n_mels, T) + x_len: length of each audio, shape (batch_size,) + + Returns: + x: encoded features, shape (batch_size, T//4, n_state) + x_len: updated lengths, shape (batch_size,) + """ + T = x.size(-1) + x = F.gelu(self.conv1(x)) + x = F.gelu(self.conv2(x)) + x = x.permute(0, 2, 1) # (B, T // 2, n_state) + + # Create attention mask + mask = make_non_pad_mask(x_len, T).unsqueeze(1) # (B, 1, T) + mask = mask_to_bias(mask[:, :, (T + 1) % 2 :: 2], x.dtype) # (B, 1, T // 2) + + # Add positional embedding + x = (x + self.positional_embedding.weight[: x.shape[1], :]).to(x.dtype).contiguous() + + # Apply attention blocks + for block in self.blocks: + x = block(x, mask.unsqueeze(1)) + + # Pool and normalize + x = x.permute(0, 2, 1) + x = self.avg_pooler(x) + x = x.permute(0, 2, 1) + x_len = (x_len + 1) // 2 // 2 + x = self.after_norm(x) + + return x, x_len + + +class Adaptor(nn.Module): + """ + Adaptor to project audio features to LLM dimension + + Maps from n_state (512) to n_hidden (LLM dimension, e.g., 4096) + with optional downsampling via convolution + """ + + def __init__( + self, n_state: int = 512, n_hidden: int = 4096, kernel_size: int = 3, stride: int = 2, adapter_state: int = 2048 + ): + super().__init__() + self.stride = stride + + if self.stride != -1: + self.conv = nn.Conv1d(n_state, n_state, kernel_size, stride, padding=1) + + self.linear1 = nn.Linear(n_state, adapter_state) + self.relu = nn.ReLU() + self.linear2 = nn.Linear(adapter_state, n_hidden) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Args: + x: audio features, shape (batch_size, T, n_state) + + Returns: + x: projected features, shape (batch_size, T//stride, n_hidden) + """ + if self.stride != -1: + x = x.permute(0, 2, 1) # (B, n_state, T) + x = F.gelu(self.conv(x)) + x = x.permute(0, 2, 1) # (B, T//stride, n_state) + + x = self.linear1(x) + x = self.relu(x) + x = self.linear2(x) + + return x + + +# ============================================================================ +# Feature Length Calculation Utility +# ============================================================================ + + +def calculate_audio_feature_length(mel_length: int) -> int: + """ + Calculate output feature length after encoder + adapter processing. + + Processing chain: + 1. Mel spectrogram input: length = mel_length + 2. Encoder (AudioEncoder): + - conv1: kernel=3, padding=1, stride=1 → no change + - conv2: kernel=3, padding=1, stride=2 → T // 2 + - avg_pool: kernel=2, stride=2 → T // 4 + - Output length: (mel_length + 1) // 2 // 2 ≈ mel_length // 4 + 3. Adapter (Adaptor): + - conv: kernel=3, padding=1, stride=2 → T // 2 + - Output length: (encoder_output - 1) // 2 + 1 + + Args: + mel_length: Length of mel spectrogram (time steps) + + Returns: + Final feature length after encoder + adapter + + Example: + >>> calculate_audio_feature_length(1000) + 125 + """ + encoder_output_len = (mel_length + 1) // 4 + + adapter_output_len = (encoder_output_len - 1) // 2 + 1 + + return max(1, adapter_output_len) + + +class StepAudio2Processor(ProcessorMixin): + """Processor for Step-Audio2 that handles text tokenization and audio mel-spectrograms.""" + + attributes = ["tokenizer"] + attribute_class = {"tokenizer": "PreTrainedTokenizerBase"} + + def __init__(self, tokenizer): + self.tokenizer = tokenizer + self.audio_token = "" + + def __call__(self, text=None, audio=None, **kwargs): + """Process text and audio inputs.""" + if text is None: + text = "" + + allowed = { + "padding", + "truncation", + "max_length", + "return_tensors", + "add_special_tokens", + "pad_to_multiple_of", + "stride", + "return_attention_mask", + "return_token_type_ids", + "return_overflowing_tokens", + "return_offsets_mapping", + "return_special_tokens_mask", + "verbose", + "is_split_into_words", + } + tok_kwargs = {k: v for k, v in kwargs.items() if k in allowed} + if "return_tensors" not in tok_kwargs: + tok_kwargs["return_tensors"] = "pt" + + encoded = self.tokenizer(text, **tok_kwargs) + + if audio is None: + encoded["audio_mels"] = torch.empty((0, 128, 0)) + encoded["audio_lens"] = torch.tensor([], dtype=torch.int32) + return encoded + + audio_list = list(audio) if isinstance(audio, (list, tuple)) else [audio] + + mels = [log_mel_spectrogram(a) for a in audio_list] + audio_mels, audio_lens = padding_mels(mels) + + encoded["audio_mels"] = audio_mels + encoded["audio_lens"] = audio_lens + return encoded + + +class StepAudio2AudioInputs(TypedDict): + """Audio inputs for Step-Audio2""" + + audio_mels: torch.Tensor + """Shape: (batch, num_mel_bins, time_steps)""" + + audio_lens: list[int] + """Shape: (batch,) - length of each audio in time steps""" + + +class StepAudio2ProcessingInfo(BaseProcessingInfo): + """Processing info for Step-Audio2""" + + def get_tokenizer(self, **kwargs): + return AutoTokenizer.from_pretrained( + self.ctx.model_config.model, + trust_remote_code=True, + fix_mistral_regex=True, + **kwargs, + ) + + def get_hf_processor(self, **kwargs: object): + return StepAudio2Processor(self.get_tokenizer()) + + def build_data_parser(self) -> MultiModalDataParser: + return MultiModalDataParser(target_sr=16000) + + # Backward/branch compatibility: + # some code paths still call get_data_parser() on ProcessingInfo. + def get_data_parser(self) -> MultiModalDataParser: + return MultiModalDataParser(target_sr=16000) + + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + return {"audio": None} + + def get_mm_max_tokens_per_item( + self, + seq_len: int, + mm_counts: Mapping[str, int], + ) -> Mapping[str, int]: + max_audio_tokens = 250 + return {"audio": max_audio_tokens} + + +class StepAudio2DummyInputsBuilder(BaseDummyInputsBuilder[StepAudio2ProcessingInfo]): + """Dummy inputs builder for Step-Audio2""" + + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: + num_audios = mm_counts.get("audio", 0) + return "" * num_audios + + def get_dummy_mm_data( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: Mapping[str, object] | None = None, + ) -> MultiModalDataDict: + audio_len = 16000 * 25 + num_audios = mm_counts.get("audio", 0) + return {"audio": self._get_dummy_audios(length=audio_len, num_audios=num_audios)} + + +class StepAudio2MultiModalProcessor(BaseMultiModalProcessor[StepAudio2ProcessingInfo]): + """Multi-modal processor for Step-Audio2""" + + def _get_mm_fields_config( + self, + hf_inputs, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + """Map processor outputs to vLLM multimodal fields. + + audio_mels is a padded 3D tensor per audio item with matching audio_lens. + Use flat_from_sizes to keep one mm item per audio, so downstream mm_kwargs + alignment does not collapse to a single item. + """ + audio_lens = hf_inputs.get("audio_lens", torch.empty(0)) + if isinstance(audio_lens, torch.Tensor): + lens_tensor = audio_lens.flatten() + elif audio_lens is None: + lens_tensor = torch.empty(0, dtype=torch.int32) + else: + lens_tensor = torch.tensor( + list(audio_lens) if hasattr(audio_lens, "__iter__") else [int(audio_lens)], dtype=torch.int32 + ) + + return dict( + audio_mels=MultiModalFieldConfig.flat_from_sizes("audio", lens_tensor), + audio_lens=MultiModalFieldConfig.batched("audio"), + ) + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptUpdate]: + processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) + + audio_token = getattr(processor, "audio_token", "") + + audio_items = mm_items.get("audio", []) + num_audio_items = len(audio_items) if audio_items else 0 + + out_mm_data = out_mm_kwargs.get_data() + audio_lens = out_mm_data.get("audio_lens") + + feature_lens = [] + if audio_lens is not None: + if isinstance(audio_lens, torch.Tensor): + audio_lens_list = audio_lens.tolist() + else: + audio_lens_list = list(audio_lens) if hasattr(audio_lens, "__iter__") else [audio_lens] + + for length in audio_lens_list: + if length > 0: + feature_len = calculate_audio_feature_length(int(length)) + feature_lens.append(feature_len) + + # CRITICAL: Align feature_lens with mm_items['audio'] count + if num_audio_items > 0: + if len(feature_lens) < num_audio_items: + default_feature_len = 250 + pad_count = num_audio_items - len(feature_lens) + feature_lens.extend([default_feature_len] * pad_count) + elif len(feature_lens) > num_audio_items: + feature_lens = feature_lens[:num_audio_items] + elif not feature_lens: + feature_lens = [250] + + def get_replacement_audio(item_idx: int): + """Generate replacement tokens for audio placeholder. + + Following Qwen2.5-Omni pattern: returns list[int] of token IDs. + """ + if item_idx >= len(feature_lens): + num_features = 1 + else: + num_features = feature_lens[item_idx] + + return [STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID] * num_features + + return [ + PromptReplacement( + modality="audio", + target=audio_token, + replacement=get_replacement_audio, + ) + ] + + def _call_hf_processor( + self, + prompt: str, + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object] | None = None, + tok_kwargs: Mapping[str, object] | None = None, + **_: object, + ): + """Call HF processor and post-process outputs""" + mm_data = dict(mm_data) + audios = mm_data.pop("audios", []) + + mm_kwargs = mm_kwargs or {} + tok_kwargs = tok_kwargs or {} + + # CRITICAL: Ensure audios is ALWAYS a list to prevent string iteration + if audios: + if isinstance(audios, str): + mm_data["audio"] = [audios] + elif isinstance(audios, (list, tuple)): + mm_data["audio"] = audios + else: + mm_data["audio"] = [audios] + + hf_inputs = super()._call_hf_processor( + prompt=prompt, + mm_data=mm_data, + mm_kwargs=mm_kwargs, + tok_kwargs=tok_kwargs, + ) + + if "audio_mels" not in hf_inputs: + hf_inputs["audio_mels"] = torch.empty((0, 128, 0)) + if "audio_lens" not in hf_inputs: + hf_inputs["audio_lens"] = torch.tensor([], dtype=torch.int32) + + return hf_inputs + + +# ============================================================================ +# Main Thinker Model +# ============================================================================ + + +@MULTIMODAL_REGISTRY.register_processor( + StepAudio2MultiModalProcessor, + info=StepAudio2ProcessingInfo, + dummy_inputs=StepAudio2DummyInputsBuilder, +) +class StepAudio2ThinkerForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP): + """ + Step-Audio2 Thinker - Stage 1 LLM + + Architecture: + AudioEncoder (6 layers, 512 hidden) → + Adaptor (512 → LLM dim) → + Qwen2 LLM (vocab=64012) + """ + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.have_multimodal_outputs = True + + config = vllm_config.model_config.hf_config + multimodal_config = vllm_config.model_config.multimodal_config + + self.config = config + self.multimodal_config = multimodal_config + + audio_enc_cfg = getattr(config, "audio_encoder_config", {}) + + def _get_config_value(cfg, attr_name: str, default_value): + """Get value from config object or dict""" + if isinstance(cfg, dict): + return cfg.get(attr_name, default_value) + else: + return getattr(cfg, attr_name, default_value) + + n_mels = _get_config_value(audio_enc_cfg, "n_mels", DEFAULT_MODEL_CONFIG.n_mels) + n_audio_ctx = _get_config_value(audio_enc_cfg, "n_audio_ctx", DEFAULT_MODEL_CONFIG.n_audio_ctx) + n_audio_state = _get_config_value(audio_enc_cfg, "n_audio_state", DEFAULT_MODEL_CONFIG.n_audio_state) + n_audio_head = _get_config_value(audio_enc_cfg, "n_audio_head", DEFAULT_MODEL_CONFIG.n_audio_head) + n_audio_layer = _get_config_value(audio_enc_cfg, "n_audio_layer", DEFAULT_MODEL_CONFIG.n_audio_layer) + kernel_size = _get_config_value(audio_enc_cfg, "kernel_size", DEFAULT_MODEL_CONFIG.kernel_size) + adapter_stride = _get_config_value(audio_enc_cfg, "adapter_stride", DEFAULT_MODEL_CONFIG.adapter_stride) + + n_hidden = None + + if audio_enc_cfg: + n_hidden = _get_config_value(audio_enc_cfg, "llm_dim", None) + + if n_hidden is None: + text_config = getattr(config, "text_config", None) + if text_config is not None: + n_hidden = _get_config_value(text_config, "hidden_size", None) + + if n_hidden is None: + n_hidden = _get_config_value(config, "hidden_size", None) + + if n_hidden is None: + n_hidden = DEFAULT_MODEL_CONFIG.hidden_size + + self.encoder = AudioEncoder( + n_mels=n_mels, + n_ctx=n_audio_ctx, + n_state=n_audio_state, + n_head=n_audio_head, + n_layer=n_audio_layer, + ) + + self.adapter = Adaptor( + n_state=n_audio_state, + n_hidden=n_hidden, + kernel_size=kernel_size, + stride=adapter_stride, + ) + + from transformers import PretrainedConfig + + text_config = getattr(config, "text_config", None) + lm_config = text_config if isinstance(text_config, PretrainedConfig) else config + architectures = getattr(lm_config, "architectures", None) or getattr(config, "architectures", None) + if not architectures: + model_type = getattr(lm_config, "model_type", None) or getattr(config, "model_type", None) + model_type_to_arch = { + "qwen2": "Qwen2ForCausalLM", + "qwen2_5": "Qwen2_5ForCausalLM", + } + arch = model_type_to_arch.get(model_type) + architectures = [arch] if arch else ["Qwen2ForCausalLM"] + lm_vllm_config = vllm_config.with_hf_config(lm_config, architectures=architectures) + self.language_model = init_vllm_registered_model( + vllm_config=lm_vllm_config, hf_config=lm_config, prefix=maybe_prefix(prefix, "language_model") + ) + + self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + if modality.startswith("audio"): + return "" + raise ValueError("Only audio modality is supported") + + @property + def device(self): + return next(self.parameters()).device + + @property + def dtype(self): + return next(self.parameters()).dtype + + def _parse_and_validate_audio_input(self, **kwargs: object) -> StepAudio2AudioInputs | None: + """Parse audio inputs from kwargs""" + audio_mels = kwargs.get("audio_mels", None) + audio_lens = kwargs.get("audio_lens", None) + + if audio_mels is None: + return None + + if isinstance(audio_mels, torch.Tensor): + if audio_mels.dim() >= 4: + audio_mels = audio_mels.flatten(0, 1) + elif audio_mels.dim() == 3: + # Already (B, n_mels, T) + pass + elif audio_mels.dim() == 2: + audio_mels = audio_mels.unsqueeze(0) + else: + return None + else: + audio_mels = flatten_bn(audio_mels, concat=True) + + if isinstance(audio_lens, torch.Tensor): + if audio_lens.dim() >= 2: + audio_lens = audio_lens.flatten(0, 1) + elif audio_lens.dim() == 0: + audio_lens = audio_lens.unsqueeze(0) + # dim == 1 is already fine + else: + audio_lens = flatten_bn(audio_lens, concat=True) + + expected_n_mels = 128 + if audio_mels.ndim == 3 and audio_mels.size(1) != expected_n_mels: + return None + + return StepAudio2AudioInputs( + audio_mels=audio_mels.to(self.dtype).to(self.device), + audio_lens=audio_lens, + ) + + def _process_audio_input(self, audio_input: StepAudio2AudioInputs) -> tuple[torch.Tensor, ...]: + """Process audio mels through encoder and adapter""" + audio_mels = audio_input["audio_mels"] + audio_lens = audio_input["audio_lens"] + if not isinstance(audio_lens, torch.Tensor): + audio_lens = torch.tensor(audio_lens, device=self.device) + else: + audio_lens = audio_lens.to(self.device) + + audio_features, audio_lens = self.encoder(audio_mels, audio_lens) + + audio_features = self.adapter(audio_features) + + audio_feature_lens = (audio_lens - 1) // 2 + 1 + + audio_feature_list = [audio_features[i, : audio_feature_lens[i]] for i in range(audio_features.size(0))] + + return audio_feature_list + + def get_multimodal_embeddings(self, **kwargs) -> NestedTensors | None: + """Get multimodal embeddings""" + audio_input = self._parse_and_validate_audio_input(**kwargs) + if audio_input is None: + return None + else: + audio_embeddings = self._process_audio_input(audio_input) + return audio_embeddings + + def embed_multimodal(self, **kwargs: object) -> NestedTensors: + """vLLM multimodal encoder entrypoint.""" + audio_embeddings = self.get_multimodal_embeddings(**kwargs) + if audio_embeddings is None: + return [] + return audio_embeddings + + def get_input_embeddings( + self, + input_ids: torch.Tensor, + multimodal_embeddings: NestedTensors | None = None, + ) -> torch.Tensor: + """Get input embeddings with multimodal fusion""" + # Under newer vLLM versions the registered LM wrapper exposes + # get_input_embeddings(), while the inner `.model` may be a bare + # Qwen2Model without that compatibility method. + inputs_embeds = self.language_model.embed_input_ids(input_ids) + if multimodal_embeddings is not None: + inputs_embeds = _merge_multimodal_embeddings( + input_ids, inputs_embeds, multimodal_embeddings, STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID + ) + return inputs_embeds + + def get_language_model(self) -> nn.Module: + return self.language_model + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ): + """Forward pass - returns OmniOutput for multi-stage pipeline""" + if intermediate_tensors is not None: + inputs_embeds = None + elif inputs_embeds is None: + audio_embeddings = self.get_multimodal_embeddings(**kwargs) + inputs_embeds = self.get_input_embeddings(input_ids, audio_embeddings) + input_ids = None + + hidden_states = self.language_model( + input_ids, + positions, + intermediate_tensors, + inputs_embeds=inputs_embeds, + ) + + return OmniOutput( + text_hidden_states=hidden_states, + multimodal_outputs={}, + ) + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: + return self.language_model.compute_logits(hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load weights with name mapping from HuggingFace to vLLM structure""" + from vllm.model_executor.models.utils import AutoWeightsLoader + + # Map HuggingFace weight names to our model structure + # HF: model.*, lm_head.* → Our: language_model.model.*, language_model.lm_head.* + # HF: encoder.*, adapter.* → Our: encoder.*, adapter.* (no change) + def weight_mapper(name: str) -> str: + if name.startswith("lm_head."): + return name.replace("lm_head.", "language_model.lm_head.") + elif name.startswith("model."): + return name.replace("model.", "language_model.model.") + else: + # encoder.* and adapter.* remain unchanged + return name + + # Apply mapping to weights + mapped_weights = [(weight_mapper(name), tensor) for name, tensor in weights] + + loader = AutoWeightsLoader(self) + return loader.load_weights(mapped_weights) + + @staticmethod + def separate_tokens( + token_ids: list[int], + text_max: int = DEFAULT_TOKEN_CONFIG.text_max, + audio_start: int = DEFAULT_TOKEN_CONFIG.audio_start, + ) -> tuple[list[int], list[int]]: + """Separate generated tokens into text and audio tokens""" + text_tokens = [tid for tid in token_ids if tid < text_max] + audio_tokens = [tid - audio_start for tid in token_ids if tid >= audio_start] + return text_tokens, audio_tokens + + @staticmethod + def has_audio_output(token_ids: list[int], audio_start: int = DEFAULT_TOKEN_CONFIG.audio_start) -> bool: + """Check if generated tokens contain audio tokens""" + return any(tid >= audio_start for tid in token_ids) + + +class StepAudio2OutputProcessor: + """Helper class to process Step-Audio2 outputs""" + + @staticmethod + def process_output(output_ids: torch.Tensor, tokenizer, remove_audio_padding: bool = True) -> dict: + """Process model output and separate text/audio tokens""" + if isinstance(output_ids, torch.Tensor): + output_ids = output_ids.squeeze().tolist() + + # Separate tokens + text_tokens, audio_tokens = StepAudio2ThinkerForConditionalGeneration.separate_tokens(output_ids) + + # Remove audio padding if requested + if remove_audio_padding and audio_tokens: + audio_tokens = [t for t in audio_tokens if t < DEFAULT_TOKEN_CONFIG.audio_eos] + + # Decode text + text = tokenizer.decode(text_tokens, skip_special_tokens=False) + + return { + "text": text, + "text_tokens": text_tokens, + "audio_tokens": audio_tokens, + "has_audio": len(audio_tokens) > 0, + "all_tokens": output_ids, + } diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py b/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py new file mode 100644 index 000000000..a59ba3fef --- /dev/null +++ b/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py @@ -0,0 +1,759 @@ +import io +import os +from collections.abc import Iterable +from dataclasses import dataclass, field +from typing import Any + +import librosa # noqa: TID251 +import numpy as np +import onnxruntime +import s3tokenizer +import soundfile as sf +import torch +import torch.nn as nn +import torchaudio +import torchaudio.compliance.kaldi as kaldi +from flashcosyvoice.modules.hifigan import HiFTGenerator +from flashcosyvoice.utils.audio import mel_spectrogram +from hyperpyyaml import load_hyperpyyaml +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.models.interfaces import SupportsPP +from vllm.model_executor.models.utils import WeightsMapper +from vllm.sequence import IntermediateTensors +from vllm.v1.outputs import SamplerOutput +from vllm.v1.sample.metadata import SamplingMetadata +from vllm.v1.sample.sampler import Sampler + +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.model_executor.models.step_audio2.step_audio2_constants import ( + DEFAULT_STREAM_CONFIG, + DEFAULT_TOKEN2WAV_CONFIG, + STEP_AUDIO2_DEFAULT_PROMPT_WAV, + STREAM_SOURCE_CACHE_LEN, +) + +logger = init_logger(__name__) + + +def fade_in_out( + fade_in_mel: torch.Tensor, + fade_out_mel: torch.Tensor, + window: torch.Tensor, +) -> torch.Tensor: + """Cross-fade two overlapping waveform segments using a Hamming window. + + The window is split in half: the first half ramps *up* (fade-in) and + the second half ramps *down* (fade-out). The overlap region of + ``fade_in_mel`` is blended with the tail of ``fade_out_mel``. + """ + mel_overlap_len = int(window.shape[0] / 2) + fade_in_mel = fade_in_mel.clone() + fade_in_mel[..., :mel_overlap_len] = ( + fade_in_mel[..., :mel_overlap_len] * window[:mel_overlap_len] + + fade_out_mel[..., -mel_overlap_len:] * window[mel_overlap_len:] + ) + return fade_in_mel + + +@dataclass +class _StreamState: + """Per-request streaming state for concurrent Token2Wav sessions.""" + + setup_done: bool = False + stream_cache: dict | None = None + hift_cache_dict: dict = field(default_factory=dict) + finished: bool = False + + +class StepAudio2Token2WavCore(nn.Module): + """Core Token2Wav model - handles token to waveform conversion""" + + def __init__( + self, + model_path: str, + float16: bool = False, + device: str = "cuda", + n_timesteps: int = DEFAULT_TOKEN2WAV_CONFIG.n_timesteps, + ): + super().__init__() + self.model_path = model_path + self.float16 = float16 + self.device = torch.device(device) + self.n_timesteps = n_timesteps + + self._models_loaded = False + self._audio_tokenizer = None + self._spk_model = None + self._flow = None + self._hift = None + + self.cache = {} + + # Streaming state (constants from centralised config) + self.mel_cache_len = DEFAULT_STREAM_CONFIG.mel_cache_len + self.source_cache_len = STREAM_SOURCE_CACHE_LEN + self.speech_window: torch.Tensor | None = None # created lazily on device + self.stream_cache: dict | None = None # flow model causal cache + self.hift_cache_dict: dict = {} # HiFT vocoder overlap cache + + def _ensure_models_loaded(self): + """Lazy load models on first use""" + if self._models_loaded: + return + + from vllm.logger import init_logger + + logger = init_logger(__name__) + logger.info(f"Loading Token2Wav models from: {self.model_path}") + + self._audio_tokenizer = ( + s3tokenizer.load_model(f"{self.model_path}/speech_tokenizer_v2_25hz.onnx").to(self.device).eval() + ) + + option = onnxruntime.SessionOptions() + option.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL + option.intra_op_num_threads = 1 + self._spk_model = onnxruntime.InferenceSession( + f"{self.model_path}/campplus.onnx", sess_options=option, providers=["CPUExecutionProvider"] + ) + + with open(f"{self.model_path}/flow.yaml") as f: + configs = load_hyperpyyaml(f) + self._flow = configs["flow"] + if self.float16: + self._flow.half() + self._flow.load_state_dict( + torch.load(f"{self.model_path}/flow.pt", map_location="cpu", weights_only=True), strict=True + ) + self._flow.to(self.device).eval() + + self._hift = HiFTGenerator() + hift_state_dict = { + k.replace("generator.", ""): v + for k, v in torch.load(f"{self.model_path}/hift.pt", map_location="cpu", weights_only=True).items() + } + self._hift.load_state_dict(hift_state_dict, strict=True) + self._hift.to(self.device).eval() + + self._models_loaded = True + logger.info("Token2Wav models loaded successfully") + + @property + def audio_tokenizer(self): + """Lazy-loaded audio tokenizer""" + self._ensure_models_loaded() + return self._audio_tokenizer + + @property + def spk_model(self): + """Lazy-loaded speaker model""" + self._ensure_models_loaded() + return self._spk_model + + @property + def flow(self): + """Lazy-loaded flow model""" + self._ensure_models_loaded() + return self._flow + + @property + def hift(self): + """Lazy-loaded HiFT vocoder""" + self._ensure_models_loaded() + return self._hift + + def _prepare_prompt(self, prompt_wav: str): + """Prepare prompt audio for conditioning""" + # Prefer soundfile/librosa path to avoid torchaudio->torchcodec runtime coupling. + try: + audio_np_raw, sample_rate = sf.read(prompt_wav, dtype="float32", always_2d=False) + except Exception: + audio_np_raw, sample_rate = librosa.load(prompt_wav, sr=None, mono=True) + if isinstance(audio_np_raw, np.ndarray) and audio_np_raw.ndim > 1: + audio_np_raw = np.mean(audio_np_raw, axis=-1) + sample_rate = int(sample_rate) + + # 16 kHz branch: s3tokenizer + speaker embedding + audio_16k = audio_np_raw.astype(np.float32) + if sample_rate != 16000: + audio_16k = librosa.resample(y=audio_16k, orig_sr=sample_rate, target_sr=16000) + audio = torch.from_numpy(audio_16k) + mels = s3tokenizer.log_mel_spectrogram(audio) + mels, mels_lens = s3tokenizer.padding([mels]) + prompt_speech_tokens, prompt_speech_tokens_lens = self.audio_tokenizer.quantize( + mels.to(self.device), mels_lens.to(self.device) + ) + + spk_feat = kaldi.fbank(audio.unsqueeze(0), num_mel_bins=80, dither=0, sample_frequency=16000) + spk_feat = spk_feat - spk_feat.mean(dim=0, keepdim=True) + spk_emb = torch.tensor( + self.spk_model.run(None, {self.spk_model.get_inputs()[0].name: spk_feat.unsqueeze(dim=0).cpu().numpy()})[0], + device=self.device, + ) + + # 24 kHz branch: mel spectrogram for flow model conditioning + # Must resample from the ORIGINAL audio, not the 16 kHz version. + audio_24k = audio_np_raw.astype(np.float32) + if sample_rate != 24000: + audio_24k = librosa.resample(y=audio_24k, orig_sr=sample_rate, target_sr=24000) + audio_24k_t = torch.from_numpy(audio_24k).unsqueeze(0) # [1, T] + prompt_mel = mel_spectrogram(audio_24k_t).transpose(1, 2).squeeze(0) # [T, num_mels] + prompt_mels = prompt_mel.unsqueeze(0).to(self.device) + prompt_mels_lens = torch.tensor([prompt_mels.shape[1]], dtype=torch.int32, device=self.device) + prompt_mels = torch.nn.functional.pad( + prompt_mels, + (0, 0, 0, prompt_speech_tokens.shape[1] * self.flow.up_rate - prompt_mels.shape[1]), + mode="replicate", + ) + + return prompt_speech_tokens, prompt_speech_tokens_lens, spk_emb, prompt_mels, prompt_mels_lens + + def forward( + self, generated_speech_tokens: list, prompt_wav: str, return_bytes: bool = True + ) -> torch.Tensor | bytes: + """ + Convert audio tokens to waveform + + Args: + generated_speech_tokens: List of audio token IDs (0-6560) + prompt_wav: Path to prompt wav file for speaker conditioning + return_bytes: If True, return WAV bytes; else return tensor + + Returns: + Audio waveform as bytes or tensor + """ + if prompt_wav not in self.cache: + self.cache[prompt_wav] = self._prepare_prompt(prompt_wav) + prompt_speech_tokens, prompt_speech_tokens_lens, spk_emb, prompt_mels, prompt_mels_lens = self.cache[prompt_wav] + + generated_speech_tokens = torch.tensor([generated_speech_tokens], dtype=torch.int32, device=self.device) + generated_speech_tokens_lens = torch.tensor( + [generated_speech_tokens.shape[1]], dtype=torch.int32, device=self.device + ) + + with torch.amp.autocast(str(self.device.type), dtype=torch.float16 if self.float16 else torch.float32): + mel = self.flow.inference( + generated_speech_tokens, + generated_speech_tokens_lens, + prompt_speech_tokens, + prompt_speech_tokens_lens, + prompt_mels, + prompt_mels_lens, + spk_emb, + self.n_timesteps, + ) + + wav, _ = self.hift(speech_feat=mel) + + if return_bytes: + output = io.BytesIO() + torchaudio.save(output, wav.cpu(), sample_rate=DEFAULT_TOKEN2WAV_CONFIG.sample_rate, format="wav") + return output.getvalue() + else: + return wav + + # ------------------------------------------------------------------ + # Streaming (chunked) inference + # ------------------------------------------------------------------ + + def setup_stream(self, prompt_wav: str) -> None: + """Initialize flow + HiFT caches for chunked streaming inference. + + Must be called once before the first ``stream_chunk()`` call for a + given utterance. + """ + if prompt_wav not in self.cache: + self.cache[prompt_wav] = self._prepare_prompt(prompt_wav) + ( + prompt_speech_tokens, + _, + spk_emb, + prompt_mels, + _, + ) = self.cache[prompt_wav] + + # Lazily create the Hamming window on the correct device + if self.speech_window is None or self.speech_window.device != self.device: + self.speech_window = torch.from_numpy(np.hamming(2 * self.source_cache_len)).to( + device=self.device, dtype=torch.float32 + ) + + # Flow model: setup causal cache with prompt tokens + lookahead + pre_lookahead = DEFAULT_STREAM_CONFIG.pre_lookahead_len + self.stream_cache = self.flow.setup_cache( + torch.cat([prompt_speech_tokens, prompt_speech_tokens[:, :pre_lookahead]], dim=1), + prompt_mels, + spk_emb, + n_timesteps=self.n_timesteps, + ) + + # HiFT vocoder: empty overlap buffers + self.hift_cache_dict = { + "mel": torch.zeros(1, prompt_mels.shape[2], 0, device=self.device), + "source": torch.zeros(1, 1, 0, device=self.device), + "speech": torch.zeros(1, 0, device=self.device), + } + + def stream_chunk( + self, + audio_tokens: list[int], + prompt_wav: str, + last_chunk: bool = False, + ) -> torch.Tensor: + """Process one chunk of audio tokens and return a waveform segment. + + Args: + audio_tokens: Token IDs for this chunk (includes lookahead for + non-last chunks). + prompt_wav: Path to the prompt wav (for cached speaker embeddings). + last_chunk: Whether this is the final chunk of the utterance. + + Returns: + 1-D float tensor of audio samples (24 kHz). + """ + if self.stream_cache is None: + raise ValueError("stream_cache not initialised – call setup_stream() first") + + if prompt_wav not in self.cache: + self.cache[prompt_wav] = self._prepare_prompt(prompt_wav) + _, _, spk_emb, prompt_mels, _ = self.cache[prompt_wav] + + token_tensor = torch.tensor([audio_tokens], dtype=torch.int32, device=self.device) + + with torch.amp.autocast( + str(self.device.type), + dtype=torch.float16 if self.float16 else torch.float32, + ): + chunk_mel, self.stream_cache = self.flow.inference_chunk( + token=token_tensor, + spk=spk_emb, + cache=self.stream_cache, + last_chunk=last_chunk, + n_timesteps=self.n_timesteps, + ) + + # Trim estimator attention cache to avoid unbounded growth + keep = DEFAULT_STREAM_CONFIG.estimator_cache_keep + est_att = self.stream_cache.get("estimator_att_cache") + if est_att is not None and est_att.shape[4] > (prompt_mels.shape[1] + keep): + self.stream_cache["estimator_att_cache"] = torch.cat( + [est_att[:, :, :, :, : prompt_mels.shape[1]], est_att[:, :, :, :, -keep:]], + dim=4, + ) + + # ---- HiFT vocoder with overlap-add ---- + hift_cache_mel = self.hift_cache_dict["mel"] + hift_cache_source = self.hift_cache_dict["source"] + hift_cache_speech = self.hift_cache_dict["speech"] + + mel = torch.cat([hift_cache_mel, chunk_mel], dim=2) + speech, source = self.hift(mel, hift_cache_source) + + # Cross-fade overlap region for smooth concatenation + if hift_cache_speech.shape[-1] > 0: + speech = fade_in_out(speech, hift_cache_speech, self.speech_window) + + # Update vocoder cache for the next chunk + self.hift_cache_dict = { + "mel": mel[..., -self.mel_cache_len :].clone().detach(), + "source": source[:, :, -self.source_cache_len :].clone().detach(), + "speech": speech[:, -self.source_cache_len :].clone().detach(), + } + + # For non-last chunks, trim the tail that will be re-synthesised + if not last_chunk: + speech = speech[:, : -self.source_cache_len] + + return speech.squeeze(0) # [samples] + + def reset_stream(self) -> None: + """Clear all streaming state after an utterance is complete.""" + self.stream_cache = None + self.hift_cache_dict = {} + + # ------------------------------------------------------------------ + # Per-request streaming (does NOT mutate self – uses external state) + # ------------------------------------------------------------------ + + def setup_stream_for(self, prompt_wav: str, state: _StreamState) -> None: + """Initialise flow + HiFT caches into *state* (no self mutation).""" + if prompt_wav not in self.cache: + self.cache[prompt_wav] = self._prepare_prompt(prompt_wav) + prompt_speech_tokens, _, spk_emb, prompt_mels, _ = self.cache[prompt_wav] + + if self.speech_window is None or self.speech_window.device != self.device: + self.speech_window = torch.from_numpy(np.hamming(2 * self.source_cache_len)).to( + device=self.device, dtype=torch.float32 + ) + + pre_lookahead = DEFAULT_STREAM_CONFIG.pre_lookahead_len + state.stream_cache = self.flow.setup_cache( + torch.cat( + [prompt_speech_tokens, prompt_speech_tokens[:, :pre_lookahead]], + dim=1, + ), + prompt_mels, + spk_emb, + n_timesteps=self.n_timesteps, + ) + state.hift_cache_dict = { + "mel": torch.zeros(1, prompt_mels.shape[2], 0, device=self.device), + "source": torch.zeros(1, 1, 0, device=self.device), + "speech": torch.zeros(1, 0, device=self.device), + } + state.setup_done = True + + def stream_chunk_for( + self, + audio_tokens: list[int], + prompt_wav: str, + last_chunk: bool, + state: _StreamState, + ) -> torch.Tensor: + """Process one chunk using *state* (no self mutation except speech_window).""" + if state.stream_cache is None: + raise ValueError("stream_cache not initialised – call setup_stream_for() first") + + if prompt_wav not in self.cache: + self.cache[prompt_wav] = self._prepare_prompt(prompt_wav) + _, _, spk_emb, prompt_mels, _ = self.cache[prompt_wav] + + token_tensor = torch.tensor([audio_tokens], dtype=torch.int32, device=self.device) + + with torch.amp.autocast( + str(self.device.type), + dtype=torch.float16 if self.float16 else torch.float32, + ): + chunk_mel, state.stream_cache = self.flow.inference_chunk( + token=token_tensor, + spk=spk_emb, + cache=state.stream_cache, + last_chunk=last_chunk, + n_timesteps=self.n_timesteps, + ) + + # Trim estimator attention cache to avoid unbounded growth + keep = DEFAULT_STREAM_CONFIG.estimator_cache_keep + est_att = state.stream_cache.get("estimator_att_cache") + if est_att is not None and est_att.shape[4] > (prompt_mels.shape[1] + keep): + state.stream_cache["estimator_att_cache"] = torch.cat( + [ + est_att[:, :, :, :, : prompt_mels.shape[1]], + est_att[:, :, :, :, -keep:], + ], + dim=4, + ) + + # ---- HiFT vocoder with overlap-add ---- + hift_cache_mel = state.hift_cache_dict["mel"] + hift_cache_source = state.hift_cache_dict["source"] + hift_cache_speech = state.hift_cache_dict["speech"] + + mel = torch.cat([hift_cache_mel, chunk_mel], dim=2) + speech, source = self.hift(mel, hift_cache_source) + + if hift_cache_speech.shape[-1] > 0: + speech = fade_in_out(speech, hift_cache_speech, self.speech_window) + + state.hift_cache_dict = { + "mel": mel[..., -self.mel_cache_len :].clone().detach(), + "source": source[:, :, -self.source_cache_len :].clone().detach(), + "speech": speech[:, -self.source_cache_len :].clone().detach(), + } + + if not last_chunk: + speech = speech[:, : -self.source_cache_len] + + return speech.squeeze(0) + + @staticmethod + def reset_stream_for(state: _StreamState) -> None: + """Clear streaming caches for the given state.""" + state.stream_cache = None + state.hift_cache_dict = {} + state.setup_done = False + state.finished = True + + +class StepAudio2Token2WavForConditionalGeneration(nn.Module, SupportsPP): + """vLLM-compatible wrapper for Step-Audio2 Token2Wav""" + + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "model.": "token2wav.", + } + ) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.vllm_config = vllm_config + self.config = vllm_config.model_config.hf_config + + model_path = getattr(self.config, "token2wav_path", None) + model_name_or_path = vllm_config.model_config.model + if model_path is None: + # Resolve HF repo names to local cache path + if not os.path.isdir(model_name_or_path): + from huggingface_hub import snapshot_download + + model_name_or_path = snapshot_download(model_name_or_path) + model_path = f"{model_name_or_path}/token2wav" + + float16 = getattr(self.config, "token2wav_float16", False) + + device = "cuda" + if hasattr(vllm_config, "device_config") and vllm_config.device_config: + device = str(vllm_config.device_config.device) + + self.token2wav = StepAudio2Token2WavCore(model_path=model_path, float16=float16, device=device) + + self.have_multimodal_outputs = True + # Required for the runner to pass async_chunk info via + # runtime_additional_information on every forward step. + self.enable_update_additional_information = True + + # Per-request streaming states (ordered list matching batch order). + self._stream_states: list[_StreamState] = [] + + # Resolve default prompt wav path from model directory. + # Priority: env var > {model_dir}/assets/default_female.wav > bare filename (cwd) + self._default_prompt_wav = self._resolve_prompt_wav(model_name_or_path) + + self.make_empty_intermediate_tensors = lambda: None + + @staticmethod + def _resolve_prompt_wav(model_root: str) -> str: + """Find the default prompt wav, searching model assets directory.""" + env_val = os.environ.get("STEP_AUDIO2_DEFAULT_PROMPT_WAV") + if env_val and os.path.exists(env_val): + return env_val + # Search common locations: model dir, then bundled assets in this package + pkg_dir = os.path.dirname(__file__) + candidates = [ + os.path.join(model_root, "assets", STEP_AUDIO2_DEFAULT_PROMPT_WAV), + os.path.join(model_root, STEP_AUDIO2_DEFAULT_PROMPT_WAV), + os.path.join(pkg_dir, "assets", STEP_AUDIO2_DEFAULT_PROMPT_WAV), + os.path.join(pkg_dir, STEP_AUDIO2_DEFAULT_PROMPT_WAV), + STEP_AUDIO2_DEFAULT_PROMPT_WAV, + ] + for path in candidates: + if os.path.exists(path): + return path + # Return the env var (even if not found) or the bare filename as fallback + return env_val or STEP_AUDIO2_DEFAULT_PROMPT_WAV + + def get_language_model(self) -> torch.nn.Module: + # Return the wrapper so vLLM interface helpers resolve + # embed_input_ids() on this adapter instead of the bare core model. + return self + + @property + def sampler(self): + return Sampler() + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings=None, + is_multimodal=None, + ) -> torch.Tensor: + hidden_size = self.vllm_config.model_config.get_hidden_size() + return torch.zeros_like(input_ids).reshape(-1, 1).repeat(1, hidden_size) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + additional: dict | None = None, + runtime_additional_information: list[dict[str, Any]] | None = None, + **kwargs: Any, + ) -> torch.Tensor | bytes: + """ + Forward pass for Token2Wav + + Args: + input_ids: Audio token IDs (from stage input processor) + positions: Position IDs (unused) + intermediate_tensors: Intermediate tensors (unused) + inputs_embeds: Input embeddings (unused) + additional: Additional information dict (deprecated) + runtime_additional_information: Per-request dicts from the + async_chunk processor. Each dict contains ``audio_tokens`` + (list[int]) and ``last_chunk`` (bool). + **kwargs: Other kwargs + + Returns: + OmniOutput with waveform tensor + """ + # ----- async_chunk streaming path ----- + # Only enter when at least one entry carries ``left_context_size`` + # (the field the framework forwards from the async_chunk payload). + # The runner always passes runtime_additional_information (possibly + # [{}]), so a bare truthiness check would incorrectly enter this + # branch in synchronous mode. + if runtime_additional_information and any( + "left_context_size" in info for info in runtime_additional_information + ): + return self._forward_async_chunk(input_ids, runtime_additional_information, **kwargs) + + # ----- profiling / warmup guard ----- + # Profile/warmup passes use all-zero token IDs. Real audio tokens + # are non-zero, so check sum == 0 rather than a length threshold + # (short utterances may have fewer than 50 tokens). + if isinstance(input_ids, torch.Tensor): + is_profile_mode = input_ids.numel() == 0 or input_ids.sum().item() == 0 + else: + is_profile_mode = not input_ids if not hasattr(input_ids, "__len__") else len(input_ids) == 0 + + if is_profile_mode: + device = torch.device("cpu") + if isinstance(inputs_embeds, torch.Tensor) and inputs_embeds.numel() > 0: + device = inputs_embeds.device + elif isinstance(input_ids, torch.Tensor) and input_ids.numel() > 0: + device = input_ids.device + + dummy_audio = torch.zeros(1000, device=device, dtype=torch.float32) + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": [dummy_audio]}, + ) + + # ----- synchronous (full-sequence) path ----- + if isinstance(input_ids, torch.Tensor): + if input_ids.dim() == 1: + input_ids_list = input_ids.cpu().tolist() + else: + input_ids_list = input_ids[0].cpu().tolist() + else: + input_ids_list = input_ids if isinstance(input_ids, list) else [input_ids] + + prompt_wav = self._default_prompt_wav + + if not os.path.exists(prompt_wav): + raise FileNotFoundError(f"Token2Wav: prompt_wav file not found: {prompt_wav}") + + waveform_tensor = self.token2wav(input_ids_list, prompt_wav=prompt_wav, return_bytes=False) + + if waveform_tensor.dim() == 2 and waveform_tensor.shape[0] == 1: + waveform_tensor = waveform_tensor.squeeze(0) + + waveform_tensor_cpu = waveform_tensor.detach().cpu().contiguous() + + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": [waveform_tensor_cpu]}, + ) + + # ------------------------------------------------------------------ + # Async-chunk streaming forward + # ------------------------------------------------------------------ + + def _forward_async_chunk( + self, + input_ids: torch.Tensor, + runtime_additional_information: list[dict[str, Any]], + **kwargs: Any, + ) -> OmniOutput: + """Handle one async_chunk forward step (batch=1). + + The end-of-stream flag is encoded as ``left_context_size == 1`` + (the only field the framework transparently forwards from the + async_chunk payload to ``request.additional_information``). + """ + # --- batch=1 guard --- + batch_size = sum(1 for info in runtime_additional_information if "left_context_size" in info) + if batch_size != 1: + raise RuntimeError( + f"Token2Wav async_chunk only supports batch=1, got {batch_size}. " + "Batch>1 requires framework support for batch_req_ids." + ) + + info = next(info for info in runtime_additional_information if "left_context_size" in info) + last_chunk = info.get("left_context_size", 0) == 1 + + # --- Manage single stream state --- + # If the previous request was preempted (setup_done but not finished), + # its stale caches must be discarded before we start a new session. + if self._stream_states and not self._stream_states[0].finished: + state = self._stream_states[0] + if state.setup_done and not last_chunk: + # Heuristic: a non-last chunk arriving into an already-active + # state that we did NOT set up this call means a new request + # reused the slot. For a genuinely continuing request the + # state is fine; we cannot distinguish without request IDs, + # so we only forcibly reset on an explicit new-session signal + # (setup_done=True but stream_cache already consumed all + # previous data -- detected below when stream_chunk fails). + pass + if not self._stream_states or self._stream_states[0].finished: + self._stream_states = [_StreamState()] + state = self._stream_states[0] + + # --- Extract audio tokens --- + audio_tokens = input_ids.flatten().cpu().tolist() + + # Empty chunk (e.g. EOF with no audio tokens) + if not audio_tokens: + if state.setup_done: + self.token2wav.reset_stream_for(state) + else: + state.finished = True + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": [torch.zeros(1, dtype=torch.float32)]}, + ) + + prompt_wav = self._default_prompt_wav + if not os.path.exists(prompt_wav): + raise FileNotFoundError(f"Token2Wav: prompt_wav file not found: {prompt_wav}") + + try: + if not state.setup_done: + logger.info("Token2Wav: initialising streaming caches") + self.token2wav.setup_stream_for(prompt_wav, state) + + waveform = self.token2wav.stream_chunk_for( + audio_tokens, + prompt_wav=prompt_wav, + last_chunk=last_chunk, + state=state, + ) + except Exception: + logger.exception("Token2Wav: stream error, resetting state") + self.token2wav.reset_stream_for(state) + raise + + if last_chunk: + self.token2wav.reset_stream_for(state) + + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": [waveform.detach().cpu().contiguous()]}, + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + return hidden_states + + def sample( + self, + logits: torch.Tensor, + sampling_metadata: SamplingMetadata, + ) -> SamplerOutput | None: + return None + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loaded_params = set() + for name, _ in self.token2wav.named_parameters(): + loaded_params.add(f"token2wav.{name}") + for name, _ in self.token2wav.named_buffers(): + loaded_params.add(f"token2wav.{name}") + return loaded_params + + def load_weights_without_buffers(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loaded_params = set() + for name, _ in self.token2wav.named_parameters(): + loaded_params.add(f"token2wav.{name}") + return loaded_params diff --git a/vllm_omni/model_executor/stage_configs/step_audio_2.yaml b/vllm_omni/model_executor/stage_configs/step_audio_2.yaml new file mode 100644 index 000000000..b560fea9b --- /dev/null +++ b/vllm_omni/model_executor/stage_configs/step_audio_2.yaml @@ -0,0 +1,122 @@ +# 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 new file mode 100644 index 000000000..1cecfddc5 --- /dev/null +++ b/vllm_omni/model_executor/stage_configs/step_audio_2_asr.yaml @@ -0,0 +1,35 @@ +# 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 new file mode 100644 index 000000000..98c484072 --- /dev/null +++ b/vllm_omni/model_executor/stage_configs/step_audio_2_async_chunk.yaml @@ -0,0 +1,78 @@ +# 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_input_processors/step_audio2.py b/vllm_omni/model_executor/stage_input_processors/step_audio2.py new file mode 100644 index 000000000..5e015aa7c --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/step_audio2.py @@ -0,0 +1,204 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Stage input processor for Step-Audio2: Thinker → Token2Wav transition.""" + +from typing import Any + +import torch +from vllm.inputs import TextPrompt +from vllm.logger import init_logger + +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 ( + DEFAULT_STREAM_CONFIG, + DEFAULT_TOKEN_CONFIG, +) + +logger = init_logger(__name__) + + +def _ensure_list(x): + """Convert ConstantList / tensor / iterable to Python list.""" + if hasattr(x, "_x"): + return list(x._x) + elif isinstance(x, torch.Tensor): + return x.tolist() + elif not isinstance(x, list): + return list(x) if hasattr(x, "__iter__") else [x] + return list(x) + + +# ========================= +# Async chunk processor +# ========================= + + +def thinker2token2wav_async_chunk( + transfer_manager: Any, + pooling_output: dict[str, Any], + request: OmniEngineCoreRequest, + is_finished: bool = False, +) -> dict[str, Any] | None: + """ + Async chunk processor: stream audio tokens from Thinker to Token2Wav. + + Unlike Qwen3-Omni which passes hidden states via pooling_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. + + The flow model's conformer encoder requires ``pre_lookahead_len`` extra + tokens beyond the chunk boundary. Each chunk therefore contains + ``chunk_size + pre_lookahead_len`` tokens, but the consumed pointer + only advances by ``chunk_size`` so the lookahead tokens are re-used as + the start of the next chunk's context (cached inside the encoder). + + Payload adapts to the framework's ``_poll_single_request`` non-AR path: + - ``code_predictor_codes`` → ``request.prompt_token_ids`` → ``input_ids`` + - ``left_context_size`` → ``request.additional_information`` + → ``runtime_additional_information`` (0 = not last, 1 = last chunk) + - ``finished`` → marks request as finished in adapter + + Args: + transfer_manager: OmniChunkTransferAdapter instance (has + ``code_prompt_token_ids`` defaultdict for state tracking). + pooling_output: Pooler 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. + """ + audio_start = DEFAULT_TOKEN_CONFIG.audio_start + audio_eos = DEFAULT_TOKEN_CONFIG.audio_eos + finished = bool(is_finished or request.is_finished()) + + # Only look at decode (generated) tokens — the prompt may contain + # historical audio tokens from prior conversation turns. + all_token_ids = _ensure_list(request.all_token_ids) + prompt_len = len(_ensure_list(request.prompt_token_ids)) + generated_ids = all_token_ids[prompt_len:] + + # Extract audio tokens and convert to 0-based IDs for Token2Wav + audio_tokens = [tid - audio_start for tid in generated_ids if tid >= audio_start] + # Remove padding / EOS tokens + audio_tokens = [t for t in audio_tokens if t < audio_eos] + + # Flow model streaming parameters (from centralised config) + chunk_size = DEFAULT_STREAM_CONFIG.chunk_size + pre_lookahead_len = DEFAULT_STREAM_CONFIG.pre_lookahead_len + + # consumed = number of tokens whose mel output has been produced. + # We track this via transfer_manager.code_prompt_token_ids[request_id]. + request_id = request.external_req_id + consumed = len(transfer_manager.code_prompt_token_ids[request_id]) + available = len(audio_tokens) - consumed + + if finished: + # 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), + } + 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), + } + else: + # Non-last chunk: need chunk_size + pre_lookahead_len tokens + required = chunk_size + pre_lookahead_len + if available < required: + return None + 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), + } + + +def thinker2token2wav( + source_outputs: list[Any], + prompt: OmniTokensPrompt | TextPrompt | None = None, + requires_multimodal_data: bool = False, +) -> list[OmniTokensPrompt]: + """ + Process thinker outputs to create token2wav inputs. + + Workflow: + 1. Extract generated token IDs from thinker output + 2. Separate audio tokens from text tokens + 3. Package audio tokens for token2wav stage + + Step-Audio2 token ranges: + - Text tokens: 0 - 151688 + - Audio tokens: 151696 - 158257 (vocab size 6562) + + Args: + source_outputs: Resolved list of upstream (thinker) engine outputs + prompt: Original prompt data (unused for Token2Wav in default-voice mode) + requires_multimodal_data: Whether multimodal data is required + + Returns: + List of OmniTokensPrompt for token2wav stage + + """ + if not source_outputs: + raise ValueError("source_outputs cannot be empty") + + thinker_outputs = source_outputs + token2wav_inputs = [] + + # Token configuration uses the fixed Step-Audio2 model constants. + audio_start = DEFAULT_TOKEN_CONFIG.audio_start + audio_eos = DEFAULT_TOKEN_CONFIG.audio_eos # Relative to audio start + + # Process each thinker output + for i, thinker_output in enumerate(thinker_outputs): + output = thinker_output.outputs[0] + + # Only look at decode (generated) tokens — the prompt may contain + # historical audio tokens from prior conversation turns. + gen_ids = output.token_ids + + # Convert to CPU list if needed (for multi-process/multi-GPU setup) + if isinstance(gen_ids, torch.Tensor): + gen_ids = gen_ids.cpu().tolist() + + # Separate audio tokens from text tokens + # Audio tokens are >= audio_start + audio_tokens = [ + tid - audio_start # Convert to 0-based for Token2Wav + for tid in gen_ids + if tid >= audio_start + ] + + # Remove padding tokens (anything >= EOS) + audio_tokens = [t for t in audio_tokens if t < audio_eos] + + if not audio_tokens: + # No audio tokens generated, skip Token2Wav for this request + logger.info(f"Request {i}: No audio tokens generated, skipping Token2Wav stage") + continue + + logger.debug("Creating Token2Wav input with %d audio tokens", len(audio_tokens)) + + token2wav_inputs.append( + OmniTokensPrompt( + prompt_token_ids=audio_tokens, # Pass original tokens (no encoding) + multi_modal_data=None, + additional_information=None, # Not used due to framework limitation + ) + ) + + return token2wav_inputs