[MR V2] Migrate to Model Runner V2
Squashed migration rebased onto current main. Collapses the 46-commit dev/migrate-MR-v2 history (10 merge commits + superseded pre-rebase work) into a single commit, so the diff reflects only the genuine MR-V2 feature (worker_v2/ runners + model_states, omni schedulers, stage engine pipeline, pooling_output serialization bridge). Content already merged to main separately (step_audio2, tts_adapters, npu/_310p, etc.) drops out. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Taichang ZhouandClaude Opus 4.8 committed
1 parent
cead615288
commit
9678ae58cf
74 files changed
+8244
-95
No files matched your search
@@ -257,8 +257,11 @@ async def run_single_request(
|
||||
new_chunks = audio_data[audio_list_consumed:]
|
||||
audio_list_consumed = len(audio_data)
|
||||
elif isinstance(audio_data, torch.Tensor):
|
||||
new_chunks = [audio_data]
|
||||
audio_last_tensor = audio_data
|
||||
if audio_data is not audio_last_tensor:
|
||||
new_chunks = [audio_data]
|
||||
audio_last_tensor = audio_data
|
||||
else:
|
||||
new_chunks = []
|
||||
else:
|
||||
new_chunks = []
|
||||
|
||||
|
||||
+140
@@ -0,0 +1,140 @@
|
||||
#!/usr/bin/env bash
|
||||
# =============================================================================
|
||||
# Qwen3-Omni end2end functional test script
|
||||
# Uses default built-in assets, iterates over query-type x output-modalities combinations
|
||||
# =============================================================================
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
export VLLM_OMNI_USE_V2_RUNNER=1
|
||||
|
||||
MODEL_PATH="/models/Qwen3-Omni-30B-A3B-Instruct"
|
||||
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
|
||||
END2END="$SCRIPT_DIR/end2end.py"
|
||||
LOG_DIR="$SCRIPT_DIR/test_logs"
|
||||
SUMMARY_FILE="$LOG_DIR/summary.log"
|
||||
|
||||
rm -rf "$LOG_DIR"
|
||||
mkdir -p "$LOG_DIR"
|
||||
> "$SUMMARY_FILE"
|
||||
|
||||
PASS=0
|
||||
FAIL=0
|
||||
TOTAL=0
|
||||
|
||||
run_test() {
|
||||
local name="$1"
|
||||
shift
|
||||
local log_file="$LOG_DIR/${name}.log"
|
||||
local output_dir="$LOG_DIR/${name}_output"
|
||||
|
||||
TOTAL=$((TOTAL + 1))
|
||||
echo "========================================"
|
||||
echo "[${TOTAL}] Running: ${name}"
|
||||
echo " Command: python $END2END $*"
|
||||
echo "========================================"
|
||||
|
||||
local start_time
|
||||
start_time=$(date +%s)
|
||||
|
||||
if python "$END2END" "$@" --output-dir "$output_dir" > "$log_file" 2>&1; then
|
||||
local end_time
|
||||
end_time=$(date +%s)
|
||||
local elapsed=$((end_time - start_time))
|
||||
echo " => PASS (${elapsed}s)"
|
||||
echo "[PASS] ${name} (${elapsed}s)" >> "$SUMMARY_FILE"
|
||||
PASS=$((PASS + 1))
|
||||
else
|
||||
local exit_code=$?
|
||||
local end_time
|
||||
end_time=$(date +%s)
|
||||
local elapsed=$((end_time - start_time))
|
||||
echo " => FAIL (exit code: ${exit_code}, ${elapsed}s)"
|
||||
echo "[FAIL] ${name} (exit code: ${exit_code}, ${elapsed}s)" >> "$SUMMARY_FILE"
|
||||
# Append last 30 lines of log to summary
|
||||
echo " --- last 30 lines of log ---" >> "$SUMMARY_FILE"
|
||||
tail -n 30 "$log_file" >> "$SUMMARY_FILE"
|
||||
echo " --- end ---" >> "$SUMMARY_FILE"
|
||||
FAIL=$((FAIL + 1))
|
||||
fi
|
||||
echo ""
|
||||
}
|
||||
|
||||
# =============================================================================
|
||||
# Test cases
|
||||
# =============================================================================
|
||||
|
||||
# --- 1. Text-only input ---
|
||||
run_test "text_to_text" \
|
||||
--model "$MODEL_PATH" --query-type text --modalities text
|
||||
|
||||
run_test "text_to_text_audio" \
|
||||
--model "$MODEL_PATH" --query-type text --modalities text,audio
|
||||
|
||||
run_test "text_to_audio" \
|
||||
--model "$MODEL_PATH" --query-type text --modalities audio
|
||||
|
||||
# --- 2. Image input ---
|
||||
run_test "image_to_text" \
|
||||
--model "$MODEL_PATH" --query-type use_image --modalities text
|
||||
|
||||
run_test "image_to_text_audio" \
|
||||
--model "$MODEL_PATH" --query-type use_image --modalities text,audio
|
||||
|
||||
# --- 3. Audio input ---
|
||||
run_test "audio_to_text" \
|
||||
--model "$MODEL_PATH" --query-type use_audio --modalities text
|
||||
|
||||
run_test "audio_to_text_audio" \
|
||||
--model "$MODEL_PATH" --query-type use_audio --modalities text,audio
|
||||
|
||||
# --- 4. Video input ---
|
||||
run_test "video_to_text" \
|
||||
--model "$MODEL_PATH" --query-type use_video --modalities text
|
||||
|
||||
run_test "video_to_text_audio" \
|
||||
--model "$MODEL_PATH" --query-type use_video --modalities text,audio
|
||||
|
||||
# --- 5. Multi-audio input ---
|
||||
run_test "multi_audios_to_text" \
|
||||
--model "$MODEL_PATH" --query-type use_multi_audios --modalities text
|
||||
|
||||
run_test "multi_audios_to_text_audio" \
|
||||
--model "$MODEL_PATH" --query-type use_multi_audios --modalities text,audio
|
||||
|
||||
# --- 6. Mixed modality input (audio+image+video) ---
|
||||
run_test "mixed_to_text" \
|
||||
--model "$MODEL_PATH" --query-type use_mixed_modalities --modalities text
|
||||
|
||||
run_test "mixed_to_text_audio" \
|
||||
--model "$MODEL_PATH" --query-type use_mixed_modalities --modalities text,audio
|
||||
|
||||
# --- 7. Audio from video ---
|
||||
run_test "audio_in_video_to_text" \
|
||||
--model "$MODEL_PATH" --query-type use_audio_in_video --modalities text
|
||||
|
||||
run_test "audio_in_video_to_text_audio" \
|
||||
--model "$MODEL_PATH" --query-type use_audio_in_video --modalities text,audio
|
||||
|
||||
# --- 8. py-generator mode ---
|
||||
run_test "text_to_text_generator" \
|
||||
--model "$MODEL_PATH" --query-type text --modalities text --py-generator
|
||||
|
||||
run_test "audio_to_text_audio_generator" \
|
||||
--model "$MODEL_PATH" --query-type use_audio --modalities text,audio --py-generator
|
||||
|
||||
# --- 9. Multi-prompt batch inference ---
|
||||
run_test "text_to_text_batch3" \
|
||||
--model "$MODEL_PATH" --query-type text --modalities text --num-prompts 3
|
||||
|
||||
# =============================================================================
|
||||
# Summary
|
||||
# =============================================================================
|
||||
echo "========================================" | tee -a "$SUMMARY_FILE"
|
||||
echo "Tests done: total ${TOTAL}, passed ${PASS}, failed ${FAIL}" | tee -a "$SUMMARY_FILE"
|
||||
echo "========================================" | tee -a "$SUMMARY_FILE"
|
||||
echo ""
|
||||
echo "Detailed logs: $LOG_DIR"
|
||||
echo "Summary report: $SUMMARY_FILE"
|
||||
|
||||
exit $FAIL
|
||||
@@ -0,0 +1,143 @@
|
||||
#!/usr/bin/env bash
|
||||
# =============================================================================
|
||||
# Qwen3-Omni end2end_async_chunk functional test script
|
||||
# Uses default built-in assets, iterates over AsyncOmni query-type x output-modalities combinations
|
||||
# =============================================================================
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
export VLLM_OMNI_USE_V2_RUNNER=1
|
||||
|
||||
MODEL_PATH="/models/Qwen3-Omni-30B-A3B-Instruct"
|
||||
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
|
||||
END2END="$SCRIPT_DIR/end2end_async_chunk.py"
|
||||
STAGE_CONFIGS="$SCRIPT_DIR/../../../vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml"
|
||||
LOG_DIR="$SCRIPT_DIR/test_logs_async_chunk"
|
||||
SUMMARY_FILE="$LOG_DIR/summary.log"
|
||||
|
||||
rm -rf "$LOG_DIR"
|
||||
mkdir -p "$LOG_DIR"
|
||||
> "$SUMMARY_FILE"
|
||||
|
||||
PASS=0
|
||||
FAIL=0
|
||||
TOTAL=0
|
||||
|
||||
run_test() {
|
||||
local name="$1"
|
||||
shift
|
||||
local log_file="$LOG_DIR/${name}.log"
|
||||
local output_dir="$LOG_DIR/${name}_output"
|
||||
|
||||
TOTAL=$((TOTAL + 1))
|
||||
echo "========================================"
|
||||
echo "[${TOTAL}] Running: ${name}"
|
||||
echo " Command: python $END2END $*"
|
||||
echo "========================================"
|
||||
|
||||
local start_time
|
||||
start_time=$(date +%s)
|
||||
|
||||
if python "$END2END" "$@" --output-dir "$output_dir" > "$log_file" 2>&1; then
|
||||
local end_time
|
||||
end_time=$(date +%s)
|
||||
local elapsed=$((end_time - start_time))
|
||||
echo " => PASS (${elapsed}s)"
|
||||
echo "[PASS] ${name} (${elapsed}s)" >> "$SUMMARY_FILE"
|
||||
PASS=$((PASS + 1))
|
||||
else
|
||||
local exit_code=$?
|
||||
local end_time
|
||||
end_time=$(date +%s)
|
||||
local elapsed=$((end_time - start_time))
|
||||
echo " => FAIL (exit code: ${exit_code}, ${elapsed}s)"
|
||||
echo "[FAIL] ${name} (exit code: ${exit_code}, ${elapsed}s)" >> "$SUMMARY_FILE"
|
||||
# Append last 30 lines of log to summary
|
||||
echo " --- last 30 lines of log ---" >> "$SUMMARY_FILE"
|
||||
tail -n 30 "$log_file" >> "$SUMMARY_FILE"
|
||||
echo " --- end ---" >> "$SUMMARY_FILE"
|
||||
FAIL=$((FAIL + 1))
|
||||
fi
|
||||
echo ""
|
||||
}
|
||||
|
||||
# =============================================================================
|
||||
# Test cases
|
||||
# =============================================================================
|
||||
|
||||
# --- 1. Text-only input ---
|
||||
run_test "text_to_text" \
|
||||
--model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \
|
||||
--query-type text --modalities text
|
||||
|
||||
run_test "text_to_text_audio" \
|
||||
--model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \
|
||||
--query-type text --modalities text,audio
|
||||
|
||||
run_test "text_to_audio" \
|
||||
--model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \
|
||||
--query-type text --modalities audio
|
||||
|
||||
# --- 2. Image input ---
|
||||
run_test "image_to_text" \
|
||||
--model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \
|
||||
--query-type use_image --modalities text
|
||||
|
||||
run_test "image_to_text_audio" \
|
||||
--model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \
|
||||
--query-type use_image --modalities text,audio
|
||||
|
||||
# --- 3. Audio input ---
|
||||
run_test "audio_to_text" \
|
||||
--model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \
|
||||
--query-type use_audio --modalities text
|
||||
|
||||
run_test "audio_to_text_audio" \
|
||||
--model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \
|
||||
--query-type use_audio --modalities text,audio
|
||||
|
||||
# --- 4. Video input ---
|
||||
run_test "video_to_text" \
|
||||
--model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \
|
||||
--query-type use_video --modalities text
|
||||
|
||||
run_test "video_to_text_audio" \
|
||||
--model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \
|
||||
--query-type use_video --modalities text,audio
|
||||
|
||||
# --- 5. Multi-prompt batch inference ---
|
||||
run_test "text_to_text_batch3" \
|
||||
--model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \
|
||||
--query-type text --modalities text --num-prompts 3
|
||||
|
||||
# --- 6. Concurrent requests (max-in-flight) ---
|
||||
run_test "text_to_text_audio_concurrent2" \
|
||||
--model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \
|
||||
--query-type text --modalities text,audio --num-prompts 2 --max-in-flight 2
|
||||
|
||||
# --- 7. stream-audio-to-disk mode ---
|
||||
run_test "audio_to_text_audio_stream" \
|
||||
--model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \
|
||||
--query-type use_audio --modalities text,audio --stream-audio-to-disk
|
||||
|
||||
# --- 8. request-timeout test (long enough timeout, should PASS) ---
|
||||
run_test "text_to_text_with_timeout" \
|
||||
--model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \
|
||||
--query-type text --modalities text --request-timeout-s 600
|
||||
|
||||
# --- 9. batch-timeout test (long enough timeout, should PASS) ---
|
||||
run_test "text_to_text_audio_batch_timeout" \
|
||||
--model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \
|
||||
--query-type text --modalities text,audio --batch-timeout-s 600
|
||||
|
||||
# =============================================================================
|
||||
# Summary
|
||||
# =============================================================================
|
||||
echo "========================================" | tee -a "$SUMMARY_FILE"
|
||||
echo "Tests done: total ${TOTAL}, passed ${PASS}, failed ${FAIL}" | tee -a "$SUMMARY_FILE"
|
||||
echo "========================================" | tee -a "$SUMMARY_FILE"
|
||||
echo ""
|
||||
echo "Detailed logs: $LOG_DIR"
|
||||
echo "Summary report: $SUMMARY_FILE"
|
||||
|
||||
exit $FAIL
|
||||
@@ -135,17 +135,18 @@ def _estimate_prompt_len(
|
||||
return 2048
|
||||
|
||||
|
||||
def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult:
|
||||
def get_custom_voice_query(use_batch_sample: bool = False, model_name: str | None = None) -> QueryResult:
|
||||
"""Build CustomVoice sample inputs.
|
||||
|
||||
Args:
|
||||
use_batch_sample: When True, return a batch of prompts; otherwise a single prompt.
|
||||
model_name: Override the default HF model path (e.g. with a local path).
|
||||
|
||||
Returns:
|
||||
QueryResult with Omni inputs and the CustomVoice model path.
|
||||
"""
|
||||
task_type = "CustomVoice"
|
||||
model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
|
||||
model_name = model_name or "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
|
||||
if use_batch_sample:
|
||||
texts = [
|
||||
"其实我真的有发现，我是一个特别善于观察别人情绪的人。",
|
||||
@@ -196,17 +197,18 @@ def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult:
|
||||
)
|
||||
|
||||
|
||||
def get_voice_design_query(use_batch_sample: bool = False) -> QueryResult:
|
||||
def get_voice_design_query(use_batch_sample: bool = False, model_name: str | None = None) -> QueryResult:
|
||||
"""Build VoiceDesign sample inputs.
|
||||
|
||||
Args:
|
||||
use_batch_sample: When True, return a batch of prompts; otherwise a single prompt.
|
||||
model_name: Override the default HF model path (e.g. with a local path).
|
||||
|
||||
Returns:
|
||||
QueryResult with Omni inputs and the VoiceDesign model path.
|
||||
"""
|
||||
task_type = "VoiceDesign"
|
||||
model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign"
|
||||
model_name = model_name or "Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign"
|
||||
if use_batch_sample:
|
||||
texts = [
|
||||
"哥哥，你回来啦，人家等了你好久好久了，要抱抱！",
|
||||
@@ -255,18 +257,19 @@ def get_voice_design_query(use_batch_sample: bool = False) -> QueryResult:
|
||||
)
|
||||
|
||||
|
||||
def get_base_query(use_batch_sample: bool = False, mode_tag: str = "icl") -> QueryResult:
|
||||
def get_base_query(use_batch_sample: bool = False, mode_tag: str = "icl", model_name: str | None = None) -> QueryResult:
|
||||
"""Build Base (voice clone) sample inputs.
|
||||
|
||||
Args:
|
||||
use_batch_sample: When True, return a batch of prompts (Case 2).
|
||||
mode_tag: "icl" or "xvec_only" to control x_vector_only_mode behavior.
|
||||
model_name: Override the default HF model path (e.g. with a local path).
|
||||
|
||||
Returns:
|
||||
QueryResult with Omni inputs and the Base model path.
|
||||
"""
|
||||
task_type = "Base"
|
||||
model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-Base"
|
||||
model_name = model_name or "Qwen/Qwen3-TTS-12Hz-1.7B-Base"
|
||||
ref_audio_path_1 = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav"
|
||||
ref_audio_single = ref_audio_path_1
|
||||
ref_text_single = (
|
||||
@@ -335,14 +338,17 @@ def _build_inputs(args) -> tuple[str, list]:
|
||||
)
|
||||
|
||||
query_func = query_map[args.query_type]
|
||||
model_override = args.model if args.model else None
|
||||
if args.query_type in {"CustomVoice", "VoiceDesign"}:
|
||||
query_result = query_func(use_batch_sample=args.use_batch_sample)
|
||||
query_result = query_func(use_batch_sample=args.use_batch_sample, model_name=model_override)
|
||||
elif args.query_type == "Base":
|
||||
query_result = query_func(use_batch_sample=args.use_batch_sample, mode_tag=args.mode_tag)
|
||||
query_result = query_func(
|
||||
use_batch_sample=args.use_batch_sample, mode_tag=args.mode_tag, model_name=model_override
|
||||
)
|
||||
else:
|
||||
query_result = query_func()
|
||||
|
||||
model_name = query_result.model_name
|
||||
model_name = args.model if args.model else query_result.model_name
|
||||
|
||||
if args.txt_prompts:
|
||||
with open(args.txt_prompts) as f:
|
||||
@@ -359,7 +365,10 @@ def _build_inputs(args) -> tuple[str, list]:
|
||||
for t in lines
|
||||
]
|
||||
else:
|
||||
inputs = query_result.inputs if isinstance(query_result.inputs, list) else [query_result.inputs]
|
||||
if isinstance(query_result.inputs, list):
|
||||
inputs = query_result.inputs * args.num_prompts
|
||||
else:
|
||||
inputs = [query_result.inputs] * args.num_prompts
|
||||
|
||||
return model_name, inputs
|
||||
|
||||
@@ -434,6 +443,12 @@ async def main_streaming(args):
|
||||
|
||||
def parse_args():
|
||||
parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Override the default model path for the selected query type.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--query-type",
|
||||
"-q",
|
||||
|
||||
@@ -0,0 +1,640 @@
|
||||
"""Offline benchmark script for Qwen3 TTS via vLLM Omni.
|
||||
|
||||
Supports warmup + test rounds with a single Omni instance to avoid
|
||||
repeated model loading. Prints per-round wall time summary at the end.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Any, NamedTuple
|
||||
|
||||
import soundfile as sf
|
||||
import torch
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni import AsyncOmni, Omni
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class QueryResult(NamedTuple):
|
||||
"""Container for a prepared Omni request."""
|
||||
|
||||
inputs: dict
|
||||
model_name: str
|
||||
|
||||
|
||||
_model_cache: dict[str, Any] = {}
|
||||
|
||||
|
||||
def _estimate_prompt_len(
|
||||
additional_information: dict[str, Any],
|
||||
model_name: str,
|
||||
) -> int:
|
||||
"""Estimate prompt_token_ids placeholder length for the Talker stage."""
|
||||
try:
|
||||
from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSConfig
|
||||
from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
|
||||
Qwen3TTSTalkerForConditionalGeneration,
|
||||
)
|
||||
|
||||
if model_name not in _model_cache:
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
tok = AutoTokenizer.from_pretrained(
|
||||
model_name, trust_remote_code=True, padding_side="left", fix_mistral_regex=True
|
||||
)
|
||||
cfg = Qwen3TTSConfig.from_pretrained(model_name, trust_remote_code=True)
|
||||
|
||||
speech_tok = None
|
||||
try:
|
||||
from transformers.utils import cached_file
|
||||
|
||||
from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_tokenizer import Qwen3TTSTokenizer
|
||||
|
||||
st_cfg_path = cached_file(model_name, "speech_tokenizer/config.json")
|
||||
if st_cfg_path:
|
||||
speech_tok = Qwen3TTSTokenizer.from_pretrained(
|
||||
# os.path.dirname(st_cfg_path), torch_dtype=torch.bfloat16
|
||||
os.path.dirname(st_cfg_path),
|
||||
torch_dtype=torch.bfloat16,
|
||||
fix_mistral_regex=True,
|
||||
)
|
||||
logger.info("Loaded speech tokenizer for exact ref_code_len estimation")
|
||||
except Exception as e:
|
||||
logger.debug("Could not load speech tokenizer: %s", e)
|
||||
|
||||
_model_cache[model_name] = (tok, getattr(cfg, "talker_config", None), speech_tok)
|
||||
|
||||
tok, tcfg, speech_tok = _model_cache[model_name]
|
||||
task_type = (additional_information.get("task_type") or ["CustomVoice"])[0]
|
||||
|
||||
def _estimate_ref_code_len(ref_audio: object) -> int | None:
|
||||
if not isinstance(ref_audio, (str, list)):
|
||||
return None
|
||||
audio_path = ref_audio[0] if isinstance(ref_audio, list) else ref_audio
|
||||
if not isinstance(audio_path, str) or not audio_path.strip():
|
||||
return None
|
||||
try:
|
||||
from vllm.multimodal.media import MediaConnector
|
||||
|
||||
connector = MediaConnector(allowed_local_media_path="/")
|
||||
audio, sr = connector.fetch_audio(audio_path)
|
||||
import numpy as np
|
||||
|
||||
wav_np = np.asarray(audio, dtype=np.float32)
|
||||
|
||||
if speech_tok is not None:
|
||||
enc = speech_tok.encode(wav_np, sr=int(sr), return_dict=True)
|
||||
ref_code = getattr(enc, "audio_codes", None)
|
||||
if isinstance(ref_code, list):
|
||||
ref_code = ref_code[0] if ref_code else None
|
||||
if ref_code is not None and hasattr(ref_code, "shape"):
|
||||
shape = ref_code.shape
|
||||
return int(shape[0]) if len(shape) == 2 else int(shape[1]) if len(shape) == 3 else None
|
||||
|
||||
codec_hz = getattr(tcfg, "codec_frame_rate", None) or 12
|
||||
return int(len(audio) / sr * codec_hz)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information(
|
||||
additional_information=additional_information,
|
||||
task_type=task_type,
|
||||
tokenize_prompt=lambda t: tok(t, padding=False)["input_ids"],
|
||||
codec_language_id=getattr(tcfg, "codec_language_id", None),
|
||||
spk_is_dialect=getattr(tcfg, "spk_is_dialect", None),
|
||||
estimate_ref_code_len=_estimate_ref_code_len,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to estimate prompt length, using fallback 2048: %s", exc)
|
||||
return 2048
|
||||
|
||||
|
||||
def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult:
|
||||
task_type = "CustomVoice"
|
||||
model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
|
||||
if use_batch_sample:
|
||||
texts = [
|
||||
"其实我真的有发现，我是一个特别善于观察别人情绪的人。",
|
||||
"She said she would be here by noon.",
|
||||
"I like you very much.",
|
||||
"Really, you do?",
|
||||
"Yes, absolutely.",
|
||||
]
|
||||
instructs = ["", "Very happy.", "Very happy.", "Very happy.", "Very happy."]
|
||||
languages = ["Chinese", "English", "English", "English", "English"]
|
||||
speakers = ["Vivian", "Ryan", "Ryan", "Ryan", "Ryan"]
|
||||
inputs = []
|
||||
for text, instruct, language, speaker in zip(texts, instructs, languages, speakers):
|
||||
additional_information = {
|
||||
"task_type": [task_type],
|
||||
"text": [text],
|
||||
"instruct": [instruct],
|
||||
"language": [language],
|
||||
"speaker": [speaker],
|
||||
"max_new_tokens": [2048],
|
||||
}
|
||||
inputs.append(
|
||||
{
|
||||
"prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
|
||||
"additional_information": additional_information,
|
||||
}
|
||||
)
|
||||
else:
|
||||
text = "其实我真的有发现，我是一个特别善于观察别人情绪的人。"
|
||||
language = "Chinese"
|
||||
speaker = "Vivian"
|
||||
instruct = "用特别愤怒的语气说"
|
||||
additional_information = {
|
||||
"task_type": [task_type],
|
||||
"text": [text],
|
||||
"language": [language],
|
||||
"speaker": [speaker],
|
||||
"instruct": [instruct],
|
||||
"max_new_tokens": [2048],
|
||||
}
|
||||
inputs = {
|
||||
"prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
|
||||
"additional_information": additional_information,
|
||||
}
|
||||
return QueryResult(inputs=inputs, model_name=model_name)
|
||||
|
||||
|
||||
def get_voice_design_query(use_batch_sample: bool = False) -> QueryResult:
|
||||
task_type = "VoiceDesign"
|
||||
model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign"
|
||||
if use_batch_sample:
|
||||
texts = [
|
||||
"哥哥，你回来啦，人家等了你好久好久了，要抱抱！",
|
||||
"It's in the top drawer... wait, it's empty? No way, that's impossible! I'm sure I put it there!",
|
||||
]
|
||||
instructs = [
|
||||
"体现撒娇稚嫩的萝莉女声，音调偏高且起伏明显，营造出黏人、做作又刻意卖萌的听觉效果。",
|
||||
"Speak in an incredulous tone, but with a hint of panic beginning to creep into your voice.",
|
||||
]
|
||||
languages = ["Chinese", "English"]
|
||||
inputs = []
|
||||
for text, instruct, language in zip(texts, instructs, languages):
|
||||
additional_information = {
|
||||
"task_type": [task_type],
|
||||
"text": [text],
|
||||
"language": [language],
|
||||
"instruct": [instruct],
|
||||
"max_new_tokens": [2048],
|
||||
"non_streaming_mode": [True],
|
||||
}
|
||||
inputs.append(
|
||||
{
|
||||
"prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
|
||||
"additional_information": additional_information,
|
||||
}
|
||||
)
|
||||
else:
|
||||
text = "哥哥，你回来啦，人家等了你好久好久了，要抱抱！"
|
||||
instruct = "体现撒娇稚嫩的萝莉女声，音调偏高且起伏明显，营造出黏人、做作又刻意卖萌的听觉效果。"
|
||||
language = "Chinese"
|
||||
additional_information = {
|
||||
"task_type": [task_type],
|
||||
"text": [text],
|
||||
"language": [language],
|
||||
"instruct": [instruct],
|
||||
"max_new_tokens": [2048],
|
||||
"non_streaming_mode": [True],
|
||||
}
|
||||
inputs = {
|
||||
"prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
|
||||
"additional_information": additional_information,
|
||||
}
|
||||
return QueryResult(inputs=inputs, model_name=model_name)
|
||||
|
||||
|
||||
def get_base_query(use_batch_sample: bool = False, mode_tag: str = "icl") -> QueryResult:
|
||||
task_type = "Base"
|
||||
model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-Base"
|
||||
ref_audio_path_1 = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav"
|
||||
ref_audio_single = ref_audio_path_1
|
||||
ref_text_single = (
|
||||
"Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
|
||||
)
|
||||
syn_text_single = "Good one. Okay, fine, I'm just gonna leave this sock monkey here. Goodbye."
|
||||
syn_lang_single = "Auto"
|
||||
x_vector_only_mode = mode_tag == "xvec_only"
|
||||
if use_batch_sample:
|
||||
syn_text_batch = [
|
||||
"Good one. Okay, fine, I'm just gonna leave this sock monkey here. Goodbye.",
|
||||
"其实我真的有发现，我是一个特别善于观察别人情绪的人。",
|
||||
]
|
||||
syn_lang_batch = ["Chinese", "English"]
|
||||
inputs = []
|
||||
for text, language in zip(syn_text_batch, syn_lang_batch):
|
||||
additional_information = {
|
||||
"task_type": [task_type],
|
||||
"ref_audio": [ref_audio_single],
|
||||
"ref_text": [ref_text_single],
|
||||
"text": [text],
|
||||
"language": [language],
|
||||
"x_vector_only_mode": [x_vector_only_mode],
|
||||
"max_new_tokens": [2048],
|
||||
}
|
||||
inputs.append(
|
||||
{
|
||||
"prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
|
||||
"additional_information": additional_information,
|
||||
}
|
||||
)
|
||||
else:
|
||||
additional_information = {
|
||||
"task_type": [task_type],
|
||||
"ref_audio": [ref_audio_single],
|
||||
"ref_text": [ref_text_single],
|
||||
"text": [syn_text_single],
|
||||
"language": [syn_lang_single],
|
||||
"x_vector_only_mode": [x_vector_only_mode],
|
||||
"max_new_tokens": [2048],
|
||||
}
|
||||
inputs = {
|
||||
"prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
|
||||
"additional_information": additional_information,
|
||||
}
|
||||
return QueryResult(inputs=inputs, model_name=model_name)
|
||||
|
||||
|
||||
query_map = {
|
||||
"CustomVoice": get_custom_voice_query,
|
||||
"VoiceDesign": get_voice_design_query,
|
||||
"Base": get_base_query,
|
||||
}
|
||||
|
||||
|
||||
def _build_inputs(args) -> tuple[str, list]:
|
||||
"""Resolve model name and inputs list from CLI args."""
|
||||
if args.batch_size < 1 or (args.batch_size & (args.batch_size - 1)) != 0:
|
||||
raise ValueError(
|
||||
f"--batch-size must be a power of two (got {args.batch_size}); "
|
||||
"non-power-of-two values do not align with CUDA graph capture sizes "
|
||||
"of Code2Wav."
|
||||
)
|
||||
|
||||
query_func = query_map[args.query_type]
|
||||
if args.query_type in {"CustomVoice", "VoiceDesign"}:
|
||||
query_result = query_func(use_batch_sample=args.use_batch_sample)
|
||||
elif args.query_type == "Base":
|
||||
query_result = query_func(use_batch_sample=args.use_batch_sample, mode_tag=args.mode_tag)
|
||||
else:
|
||||
query_result = query_func()
|
||||
|
||||
model_name = query_result.model_name
|
||||
|
||||
if args.txt_prompts:
|
||||
with open(args.txt_prompts) as f:
|
||||
lines = [line.strip() for line in f if line.strip()]
|
||||
if not lines:
|
||||
raise ValueError(f"No valid prompts found in {args.txt_prompts}")
|
||||
template = query_result.inputs if not isinstance(query_result.inputs, list) else query_result.inputs[0]
|
||||
template_info = template["additional_information"]
|
||||
inputs = [
|
||||
{
|
||||
"prompt_token_ids": [0] * _estimate_prompt_len({**template_info, "text": [t]}, model_name),
|
||||
"additional_information": {**template_info, "text": [t]},
|
||||
}
|
||||
for t in lines
|
||||
]
|
||||
else:
|
||||
inputs = query_result.inputs if isinstance(query_result.inputs, list) else [query_result.inputs]
|
||||
|
||||
return model_name, inputs
|
||||
|
||||
|
||||
def _save_wav(output_dir: str, request_id: str, mm: dict) -> None:
|
||||
"""Concatenate audio chunks and write to a wav file."""
|
||||
audio_data = mm["audio"]
|
||||
sr_raw = mm["sr"]
|
||||
sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
|
||||
sr = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
|
||||
audio_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data
|
||||
out_wav = os.path.join(output_dir, f"output_{request_id}.wav")
|
||||
sf.write(out_wav, audio_tensor.float().cpu().numpy().flatten(), samplerate=sr, format="WAV")
|
||||
logger.info("Saved audio to %s", out_wav)
|
||||
|
||||
|
||||
class _SummaryCollector:
|
||||
"""Monkey-patches omni._log_summary_and_cleanup to capture summary dicts.
|
||||
|
||||
NOTE: This relies on a private method ``_log_summary_and_cleanup`` which may
|
||||
break if the upstream Omni implementation is refactored. If Omni exposes a
|
||||
public hook for summary collection in the future, prefer that instead.
|
||||
"""
|
||||
|
||||
def __init__(self, omni):
|
||||
self.omni = omni
|
||||
self.summaries: list[dict[str, Any]] = []
|
||||
self.collecting = False
|
||||
self._orig = omni._log_summary_and_cleanup
|
||||
|
||||
def install(self):
|
||||
from pprint import pformat as _pformat
|
||||
|
||||
collector = self
|
||||
|
||||
def _patched(request_id: str) -> None:
|
||||
req_state = collector.omni.request_states.get(request_id)
|
||||
try:
|
||||
if req_state is None or req_state.metrics is None:
|
||||
return
|
||||
summary = req_state.metrics.build_and_log_summary()
|
||||
logger.info("[Summary] %s", _pformat(summary, sort_dicts=False))
|
||||
if collector.collecting and summary:
|
||||
collector.summaries.append(summary)
|
||||
except Exception:
|
||||
logger.exception("Failed to build/log summary for req=%s", request_id)
|
||||
finally:
|
||||
collector.omni.request_states.pop(request_id, None)
|
||||
|
||||
collector.omni._log_summary_and_cleanup = _patched
|
||||
|
||||
|
||||
def _print_summary(round_times: list[float], summaries: list[dict[str, Any]] | None = None) -> None:
|
||||
"""Print per-round wall time and optional detailed pipeline summary."""
|
||||
print("\n" + "=" * 60)
|
||||
print("BENCHMARK SUMMARY")
|
||||
print("=" * 60)
|
||||
for i, t in enumerate(round_times):
|
||||
print(f" Test round {i}: {t:.3f}s")
|
||||
print("-" * 60)
|
||||
avg = sum(round_times) / len(round_times)
|
||||
print(f" Rounds: {len(round_times)}")
|
||||
print(f" Total: {sum(round_times):.3f}s")
|
||||
print(f" Average: {avg:.3f}s")
|
||||
print(f" Min: {min(round_times):.3f}s")
|
||||
print(f" Max: {max(round_times):.3f}s")
|
||||
print("=" * 60)
|
||||
|
||||
if not summaries:
|
||||
print()
|
||||
return
|
||||
|
||||
# --- Aggregate from captured summaries ---
|
||||
overall_list = [s["overall_summary"] for s in summaries if "overall_summary" in s]
|
||||
if not overall_list:
|
||||
print()
|
||||
return
|
||||
|
||||
n = len(overall_list)
|
||||
|
||||
# Average Wall time
|
||||
avg_wall = sum(d["e2e_wall_time_ms"] for d in overall_list) / n
|
||||
print(f"\n{'=' * 60}")
|
||||
print("DETAILED PIPELINE SUMMARY (test rounds average)")
|
||||
print("=" * 60)
|
||||
print(f" Avg Wall Time: {avg_wall:.3f} ms")
|
||||
|
||||
# Per-stage timing
|
||||
stage_keys = sorted(k for k in overall_list[0] if k.startswith("e2e_stage_") and k.endswith("_wall_time_ms"))
|
||||
if stage_keys:
|
||||
print("-" * 60)
|
||||
print(" Per-stage timing (avg):")
|
||||
for sk in stage_keys:
|
||||
stage_avg = sum(d.get(sk, 0.0) for d in overall_list) / n
|
||||
# extract stage index from key like e2e_stage_0_wall_time_ms
|
||||
stage_idx = sk.replace("e2e_stage_", "").replace("_wall_time_ms", "")
|
||||
print(f" Stage {stage_idx} wall time: {stage_avg:.3f} ms")
|
||||
|
||||
# Per-stage gen time from stage_table
|
||||
all_stage_gen: dict[int, list[float]] = {}
|
||||
for s in summaries:
|
||||
for req_entry in s.get("stage_table", []):
|
||||
for stage_info in req_entry.get("stages", []):
|
||||
sid = stage_info.get("stage_id", -1)
|
||||
gen_ms = stage_info.get("stage_gen_time_ms", 0.0)
|
||||
all_stage_gen.setdefault(sid, []).append(gen_ms)
|
||||
if all_stage_gen:
|
||||
print(" Per-stage gen time (avg):")
|
||||
for sid in sorted(all_stage_gen):
|
||||
vals = all_stage_gen[sid]
|
||||
print(f" Stage {sid} gen time: {sum(vals) / len(vals):.3f} ms")
|
||||
|
||||
# Per-prompt throughput
|
||||
avg_tokens_per_s = sum(d.get("e2e_avg_tokens_per_s", 0.0) for d in overall_list) / n
|
||||
avg_time_per_req = sum(d.get("e2e_avg_time_per_request_ms", 0.0) for d in overall_list) / n
|
||||
total_tokens_list = [d.get("e2e_total_tokens", 0) for d in overall_list]
|
||||
avg_total_tokens = sum(total_tokens_list) / n
|
||||
print("-" * 60)
|
||||
print(" Per-prompt throughput (avg):")
|
||||
print(f" Tokens per request: {avg_total_tokens:.1f}")
|
||||
print(f" Time per request: {avg_time_per_req:.3f} ms")
|
||||
print(f" Tokens/s: {avg_tokens_per_s:.3f}")
|
||||
print("=" * 60 + "\n")
|
||||
|
||||
|
||||
def main(args):
|
||||
"""Run sync benchmark: warmup rounds then test rounds."""
|
||||
model_name, inputs = _build_inputs(args)
|
||||
output_dir = args.output_dir
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
omni = Omni(
|
||||
model=model_name,
|
||||
stage_configs_path=args.stage_configs_path,
|
||||
log_stats=args.log_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
|
||||
)
|
||||
|
||||
collector = _SummaryCollector(omni)
|
||||
collector.install()
|
||||
|
||||
batch_size = args.batch_size
|
||||
|
||||
# --- Warmup ---
|
||||
for r in range(args.warmup_rounds):
|
||||
print(f"[Warmup {r}/{args.warmup_rounds}] running...")
|
||||
for batch_start in range(0, len(inputs), batch_size):
|
||||
batch = inputs[batch_start : batch_start + batch_size]
|
||||
for _ in omni.generate(batch):
|
||||
pass
|
||||
print(f"[Warmup {r}/{args.warmup_rounds}] done")
|
||||
|
||||
# --- Test ---
|
||||
collector.collecting = True
|
||||
round_times = []
|
||||
for r in range(args.test_rounds):
|
||||
print(f"[Test {r}/{args.test_rounds}] running...")
|
||||
t_start = time.perf_counter()
|
||||
for batch_start in range(0, len(inputs), batch_size):
|
||||
batch = inputs[batch_start : batch_start + batch_size]
|
||||
for stage_outputs in omni.generate(batch):
|
||||
output = stage_outputs.request_output
|
||||
_save_wav(
|
||||
output_dir,
|
||||
f"round{r}_{output.request_id}",
|
||||
output.outputs[0].multimodal_output,
|
||||
)
|
||||
elapsed = time.perf_counter() - t_start
|
||||
round_times.append(elapsed)
|
||||
print(f"[Test {r}/{args.test_rounds}] done in {elapsed:.3f}s")
|
||||
collector.collecting = False
|
||||
|
||||
_print_summary(round_times, collector.summaries if args.log_stats else None)
|
||||
|
||||
|
||||
async def main_streaming(args):
|
||||
"""Run async streaming benchmark: warmup rounds then test rounds."""
|
||||
model_name, inputs = _build_inputs(args)
|
||||
output_dir = args.output_dir
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
omni = AsyncOmni(
|
||||
model=model_name,
|
||||
stage_configs_path=args.stage_configs_path,
|
||||
log_stats=args.log_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
|
||||
)
|
||||
|
||||
collector = _SummaryCollector(omni)
|
||||
collector.install()
|
||||
|
||||
# --- Warmup ---
|
||||
for r in range(args.warmup_rounds):
|
||||
print(f"[Warmup {r}/{args.warmup_rounds}] running...")
|
||||
for i, prompt in enumerate(inputs):
|
||||
async for _ in omni.generate(prompt, request_id=f"warmup_{r}_{i}"):
|
||||
pass
|
||||
print(f"[Warmup {r}/{args.warmup_rounds}] done")
|
||||
|
||||
# --- Test ---
|
||||
collector.collecting = True
|
||||
round_times = []
|
||||
for r in range(args.test_rounds):
|
||||
print(f"[Test {r}/{args.test_rounds}] running...")
|
||||
t_start = time.perf_counter()
|
||||
for i, prompt in enumerate(inputs):
|
||||
request_id = f"round{r}_{i}"
|
||||
t_req_start = time.perf_counter()
|
||||
t_prev = t_req_start
|
||||
chunk_idx = 0
|
||||
async for stage_output in omni.generate(prompt, request_id=request_id):
|
||||
mm = stage_output.request_output.outputs[0].multimodal_output
|
||||
if not stage_output.finished:
|
||||
t_now = time.perf_counter()
|
||||
audio = mm.get("audio")
|
||||
n = len(audio) if isinstance(audio, list) else (0 if audio is None else 1)
|
||||
if chunk_idx == 0:
|
||||
ttfa_ms = (t_now - t_req_start) * 1000
|
||||
logger.info("Request %s: chunk %d samples=%d TTFA=%.1fms", request_id, chunk_idx, n, ttfa_ms)
|
||||
else:
|
||||
dt_ms = (t_now - t_prev) * 1000
|
||||
logger.info(
|
||||
"Request %s: chunk %d samples=%d inter_chunk=%.1fms", request_id, chunk_idx, n, dt_ms
|
||||
)
|
||||
t_prev = t_now
|
||||
chunk_idx += 1
|
||||
else:
|
||||
t_end = time.perf_counter()
|
||||
total_ms = (t_end - t_req_start) * 1000
|
||||
logger.info("Request %s: done total=%.1fms chunks=%d", request_id, total_ms, chunk_idx)
|
||||
_save_wav(output_dir, request_id, mm)
|
||||
elapsed = time.perf_counter() - t_start
|
||||
round_times.append(elapsed)
|
||||
print(f"[Test {r}/{args.test_rounds}] done in {elapsed:.3f}s")
|
||||
collector.collecting = False
|
||||
|
||||
_print_summary(round_times, collector.summaries if args.log_stats else None)
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(description="Benchmark script for Qwen3 TTS via vLLM Omni")
|
||||
parser.add_argument(
|
||||
"--query-type",
|
||||
"-q",
|
||||
type=str,
|
||||
default="CustomVoice",
|
||||
choices=query_map.keys(),
|
||||
help="Query type.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--warmup-rounds",
|
||||
type=int,
|
||||
default=3,
|
||||
help="Number of warmup rounds (results discarded, default: 3).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--test-rounds",
|
||||
type=int,
|
||||
default=7,
|
||||
help="Number of test rounds (results saved and timed, default: 7).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-stats",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Enable writing detailed statistics (default: disabled).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-init-timeout",
|
||||
type=int,
|
||||
default=300,
|
||||
help="Timeout for initializing a single stage in seconds (default: 300).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
default="output_audio",
|
||||
help="Output directory for generated wav files (default: output_audio).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--txt-prompts",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to a .txt file with one prompt per line.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to a stage configs file.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use-batch-sample",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Use batch input sample for CustomVoice/VoiceDesign/Base query.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mode-tag",
|
||||
type=str,
|
||||
default="icl",
|
||||
choices=["icl", "xvec_only"],
|
||||
help="Mode tag for Base query x_vector_only_mode (default: icl).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--streaming",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Stream audio chunks via AsyncOmni.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batch-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of prompts per batch (default: 1).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-diffusion-pipeline-profiler",
|
||||
action="store_true",
|
||||
help="Enable diffusion pipeline profiler to display stage durations.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
if args.streaming:
|
||||
asyncio.run(main_streaming(args))
|
||||
else:
|
||||
main(args)
|
||||
@@ -0,0 +1,173 @@
|
||||
#!/usr/bin/env bash
|
||||
# =============================================================================
|
||||
# Qwen3-TTS end2end functional test script
|
||||
# Iterates over query-type and mode combinations for qwen3_tts/end2end.py
|
||||
#
|
||||
# Usage:
|
||||
# ./run_tests.sh --custom-voice-model /path/a --voice-design-model /path/b --base-model /path/c
|
||||
# ./run_tests.sh --custom-voice-model /path/a # only run CustomVoice tests
|
||||
# ./run_tests.sh --base-model /path/c # only run Base tests
|
||||
# =============================================================================
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
export VLLM_OMNI_USE_V2_RUNNER=1
|
||||
# export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
|
||||
|
||||
# ---- Parse named arguments ----
|
||||
CUSTOM_VOICE_MODEL=""
|
||||
VOICE_DESIGN_MODEL=""
|
||||
BASE_MODEL=""
|
||||
|
||||
usage() {
|
||||
echo "Usage: $0 [--custom-voice-model PATH] [--voice-design-model PATH] [--base-model PATH]"
|
||||
echo ""
|
||||
echo "Pass one or more model paths. Only the query types with a model provided will be tested."
|
||||
exit 1
|
||||
}
|
||||
|
||||
if [[ $# -eq 0 ]]; then
|
||||
usage
|
||||
fi
|
||||
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case "$1" in
|
||||
--custom-voice-model) CUSTOM_VOICE_MODEL="$2"; shift 2 ;;
|
||||
--voice-design-model) VOICE_DESIGN_MODEL="$2"; shift 2 ;;
|
||||
--base-model) BASE_MODEL="$2"; shift 2 ;;
|
||||
-h|--help) usage ;;
|
||||
*) echo "Unknown option: $1"; usage ;;
|
||||
esac
|
||||
done
|
||||
|
||||
if [[ -z "$CUSTOM_VOICE_MODEL" && -z "$VOICE_DESIGN_MODEL" && -z "$BASE_MODEL" ]]; then
|
||||
echo "Error: at least one model path must be provided."
|
||||
usage
|
||||
fi
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
|
||||
END2END="$SCRIPT_DIR/end2end.py"
|
||||
LOG_DIR="$SCRIPT_DIR/test_logs"
|
||||
SUMMARY_FILE="$LOG_DIR/summary.log"
|
||||
|
||||
rm -rf "$LOG_DIR"
|
||||
mkdir -p "$LOG_DIR"
|
||||
> "$SUMMARY_FILE"
|
||||
|
||||
PASS=0
|
||||
FAIL=0
|
||||
TOTAL=0
|
||||
|
||||
run_test() {
|
||||
local name="$1"
|
||||
shift
|
||||
local log_file="$LOG_DIR/${name}.log"
|
||||
local output_dir="$LOG_DIR/${name}_output"
|
||||
|
||||
TOTAL=$((TOTAL + 1))
|
||||
echo "========================================"
|
||||
echo "[${TOTAL}] Running: ${name}"
|
||||
echo " Command: python $END2END $*"
|
||||
echo "========================================"
|
||||
|
||||
local start_time
|
||||
start_time=$(date +%s)
|
||||
|
||||
if python "$END2END" "$@" --output-dir "$output_dir" > "$log_file" 2>&1; then
|
||||
local end_time
|
||||
end_time=$(date +%s)
|
||||
local elapsed=$((end_time - start_time))
|
||||
echo " => PASS (${elapsed}s)"
|
||||
echo "[PASS] ${name} (${elapsed}s)" >> "$SUMMARY_FILE"
|
||||
PASS=$((PASS + 1))
|
||||
else
|
||||
local exit_code=$?
|
||||
local end_time
|
||||
end_time=$(date +%s)
|
||||
local elapsed=$((end_time - start_time))
|
||||
echo " => FAIL (exit code: ${exit_code}, ${elapsed}s)"
|
||||
echo "[FAIL] ${name} (exit code: ${exit_code}, ${elapsed}s)" >> "$SUMMARY_FILE"
|
||||
echo " --- last 30 lines of log ---" >> "$SUMMARY_FILE"
|
||||
tail -n 30 "$log_file" >> "$SUMMARY_FILE"
|
||||
echo " --- end ---" >> "$SUMMARY_FILE"
|
||||
FAIL=$((FAIL + 1))
|
||||
fi
|
||||
echo ""
|
||||
}
|
||||
|
||||
# =============================================================================
|
||||
# Test cases — only run sections whose model path was provided
|
||||
# =============================================================================
|
||||
|
||||
# --- 1. CustomVoice ---
|
||||
if [[ -n "$CUSTOM_VOICE_MODEL" ]]; then
|
||||
run_test "custom_voice" \
|
||||
--model "$CUSTOM_VOICE_MODEL" --query-type CustomVoice
|
||||
|
||||
run_test "custom_voice_batch_sample" \
|
||||
--model "$CUSTOM_VOICE_MODEL" --query-type CustomVoice --use-batch-sample
|
||||
|
||||
run_test "custom_voice_py_generator" \
|
||||
--model "$CUSTOM_VOICE_MODEL" --query-type CustomVoice --py-generator
|
||||
|
||||
run_test "custom_voice_streaming" \
|
||||
--model "$CUSTOM_VOICE_MODEL" --query-type CustomVoice --streaming
|
||||
|
||||
run_test "custom_voice_batch4" \
|
||||
--model "$CUSTOM_VOICE_MODEL" --query-type CustomVoice --batch-size 4
|
||||
|
||||
run_test "custom_voice_batch4_batch_sample" \
|
||||
--model "$CUSTOM_VOICE_MODEL" --query-type CustomVoice --batch-size 4 --use-batch-sample --stage-configs-path "$SCRIPT_DIR/../../../vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml"
|
||||
|
||||
run_test "custom_voice_num_prompts3" \
|
||||
--model "$CUSTOM_VOICE_MODEL" --query-type CustomVoice --num-prompts 3
|
||||
fi
|
||||
|
||||
# --- 2. VoiceDesign ---
|
||||
if [[ -n "$VOICE_DESIGN_MODEL" ]]; then
|
||||
run_test "voice_design" \
|
||||
--model "$VOICE_DESIGN_MODEL" --query-type VoiceDesign
|
||||
|
||||
run_test "voice_design_batch_sample" \
|
||||
--model "$VOICE_DESIGN_MODEL" --query-type VoiceDesign --use-batch-sample
|
||||
|
||||
run_test "voice_design_py_generator" \
|
||||
--model "$VOICE_DESIGN_MODEL" --query-type VoiceDesign --py-generator
|
||||
|
||||
run_test "voice_design_streaming" \
|
||||
--model "$VOICE_DESIGN_MODEL" --query-type VoiceDesign --streaming
|
||||
fi
|
||||
|
||||
# --- 3. Base (icl mode) ---
|
||||
if [[ -n "$BASE_MODEL" ]]; then
|
||||
run_test "base_icl" \
|
||||
--model "$BASE_MODEL" --query-type Base --mode-tag icl
|
||||
|
||||
run_test "base_icl_batch_sample" \
|
||||
--model "$BASE_MODEL" --query-type Base --mode-tag icl --use-batch-sample
|
||||
|
||||
run_test "base_icl_py_generator" \
|
||||
--model "$BASE_MODEL" --query-type Base --mode-tag icl --py-generator
|
||||
|
||||
run_test "base_icl_streaming" \
|
||||
--model "$BASE_MODEL" --query-type Base --mode-tag icl --streaming
|
||||
|
||||
# --- 4. Base (xvec_only mode) ---
|
||||
run_test "base_xvec_only" \
|
||||
--model "$BASE_MODEL" --query-type Base --mode-tag xvec_only
|
||||
|
||||
run_test "base_xvec_only_batch_sample" \
|
||||
--model "$BASE_MODEL" --query-type Base --mode-tag xvec_only --use-batch-sample
|
||||
fi
|
||||
|
||||
# =============================================================================
|
||||
# Summary
|
||||
# =============================================================================
|
||||
echo "========================================" | tee -a "$SUMMARY_FILE"
|
||||
echo "Tests done: total ${TOTAL}, passed ${PASS}, failed ${FAIL}" | tee -a "$SUMMARY_FILE"
|
||||
echo "========================================" | tee -a "$SUMMARY_FILE"
|
||||
echo ""
|
||||
echo "Detailed logs: $LOG_DIR"
|
||||
echo "Summary report: $SUMMARY_FILE"
|
||||
|
||||
exit $FAIL
|
||||
@@ -1,6 +1,9 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
# Default model name
|
||||
MODEL="${MODEL:-Qwen/Qwen2.5-Omni-7B}"
|
||||
|
||||
# Default query type
|
||||
QUERY_TYPE="${1:-mixed_modalities}"
|
||||
|
||||
@@ -162,7 +165,7 @@ echo ""
|
||||
|
||||
request_body=$(cat <<EOF
|
||||
{
|
||||
"model": "Qwen/Qwen2.5-Omni-7B",
|
||||
"model": "$MODEL",
|
||||
"sampling_params_list": $sampling_params_list,
|
||||
"mm_processor_kwargs": $mm_processor_kwargs,
|
||||
"modalities": $MODALITIES,
|
||||
|
||||
Reference in new issue
Block a user