Compare commits

..
Author SHA1 Message Date
Claude 05883ff470 [CI] Fix moss_tts_nano test import paths to unblock collection
The test files imported from `tests.utils` (non-existent) and
`tests.conftest.OmniServerParams` (not lazy-exported), causing pytest
collection to fail for the whole suite. This broke unrelated CI jobs
(Diffusion Cache Backend Test, Simple Unit Test) because collection
runs across all test files even when they are deselected by marker.

Align with the pattern used by sibling TTS tests
(e.g. test_cosyvoice3_tts.py, test_voxcpm.py):
- hardware_test -> tests.helpers.mark
- OmniServerParams -> tests.helpers.runtime

https://claude.ai/code/session_01JYESRH7Ua1mc1wqKCQQs5q
2026-04-24 15:44:03 +00:00
Yueqian Lin 74522ca3ec Merge branch 'main' into feat/moss-tts-nano
Signed-off-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
2026-04-24 11:00:13 -04:00
Yueqian Lin 0a43fb94ac style(moss_tts_nano): ruff-format on_requests_finished log call
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-04-22 10:12:07 -04:00
Yueqian Lin 123dfc0772 fix(moss_tts_nano): release stream generators on abnormal termination
forward() only popped _stream_gens on normal completion, so cancelled,
timed-out, or preempted requests leaked their generator and skipped the
finally block that unlinks the temp WAV files. Implement
on_requests_finished to close each finished generator, which raises
GeneratorExit inside it and triggers the existing cleanup.

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-04-22 10:09:58 -04:00
Yueqian Lin b03a41a7f9 refactor(moss_tts_nano): migrate to Pipeline + Deploy schema
Follow-up to #2958. MOSS-TTS-Nano was authored before the schema refactor
and still shipped as a legacy stage_configs yaml; this aligns it with the
same layout the 5 migrated TTS models now use.

- vllm_omni/model_executor/models/moss_tts_nano/pipeline.py declares
  MOSS_TTS_NANO_PIPELINE (single LLM_AR stage, owns_tokenizer, audio
  output, stop_token_ids=[2] as a hard EOS backstop).
- vllm_omni/deploy/moss_tts_nano.yaml holds runtime knobs (max_num_seqs,
  gpu_memory_utilization, enforce_eager, default_sampling_params,
  skip_mm_profiling); trust_remote_code stays at deploy top-level.
- vllm_omni/config/pipeline_registry.py registers the entry so the lazy
  registry can resolve it.
- moss_tts_nano/__init__.py exports MossTTSNanoForGeneration (VoxCPM2 pattern).
- Removed vllm_omni/model_executor/stage_configs/moss_tts_nano.yaml.

Examples, shell scripts, READMEs, and buildkite-invoked tests are
updated to use `vllm serve <model> --omni` / `--deploy-config`
(auto-load kicks in; no --stage-configs-path or --trust-remote-code).

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-04-21 18:07:50 -04:00
Yueqian Lin 26bbe7b527 fix(moss_tts_nano): address PR review feedback
- Wire ref_audio through /v1/audio/speech: _build_moss_tts_params is now
  async, resolves ref_audio via MediaConnector and passes the waveform as
  prompt_audio_array so the model can materialise a temp WAV for upstream
  inference_stream (previously voice cloning via the REST endpoint was a
  no-op). Serving now also requires ref_text when ref_audio is provided.
- Fix per-request EOS in batched decode: replace the shared
  _ar_emit_stop_token bool with a _ar_last_chunk_flags list so
  compute_logits emits EOS per row; finished requests no longer wait for
  the slowest peer in a max_num_seqs=4 batch.
- Snapshot and restore CPU+CUDA RNG state around torch.manual_seed to
  limit global-state bleed; add comment noting deterministic output
  under concurrent batching is best-effort (upstream inference_stream
  uses the global RNG).
- Align _DEFAULT_MODE with the offline example and tests ("voice_clone").
- Clean up outdated "buffer first" comment in _create_stream_gen; document
  single-threaded AR-worker assumption for _stream_gens; add one-liner
  explaining the load_weights iterator drain.

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-04-21 17:58:50 -04:00
Yueqian Lin 3bd77d3f8f Merge branch 'main' into feat/moss-tts-nano
Signed-off-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
2026-04-21 17:50:21 -04:00
Yueqian Lin 033e4041be fix(moss-tts-nano): remove unused voice variable (ruff F841)
The upstream inference_stream() API has no voice/spk parameter;
voice preset selection is not yet wired into the call. Remove the
dead assignment to silence ruff F841.

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-04-14 13:11:43 -04:00
Yueqian Lin 38aaaebeda feat(moss-tts-nano): enable progressive streaming via AR runner
Switch from GPUGenerationWorker to AR worker with OmniARScheduler,
using the VoxCPM-style generator pattern for streaming:

- inference_stream() stored per-request in _stream_gens dict
- Each forward() call yields one audio chunk via next(generator)
- compute_logits() emits EOS only when last chunk is yielded
- AR scheduler loops until EOS, enabling progressive audio output

TTFP reduced from ~3.1s to ~0.11s (30x improvement).

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-04-14 13:06:44 -04:00
Yueqian Lin 71db0c5ef5 fix(moss-tts-nano): add buildkite pre-merge test, remove misplaced docs
- Add MOSS-TTS-Nano E2E test to .buildkite/test-merge.yml
- Remove docs from docs/ (will be synced from examples/ READMEs)

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-04-14 13:06:44 -04:00
Yueqian Lin 45e8829f07 docs(moss-tts-nano): add e2e online serving test and user guide docs
- Add online serving e2e test (non-streaming, streaming, Chinese)
- Add online serving user guide doc (API, voices, Gradio, curl/Python)
- Add offline inference user guide doc

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-04-14 13:06:44 -04:00
Yueqian Lin fc81b8504a feat(moss-tts-nano): add online serving, Gradio demo, and fix /v1/audio/speech integration
- Register moss_tts_nano as TTS model stage in serving_speech.py
  (model detection, validation, prompt building)
- Fix registry mod_relname: moss_tts_nano -> modeling_moss_tts_nano
- Fix stage config: is_comprehension=true (required for generate task)
- Fix default mode: voice_clone -> continuation (built-in presets)
- Add compute_logits stub for VllmModelForTextGeneration protocol
- Remove unused _sentinel nn.Parameter
- Add Gradio demo with AudioWorklet streaming player (48kHz stereo)
- Add online/offline serving docs and launch scripts

TODO: Single-stage generation models don't support true streaming
(progressive audio chunks). Current TTFP = full generation time.
Needs multi-step scheduling support in GPUGenerationModelRunner.

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-04-14 13:06:44 -04:00
Yueqian Lin 571c8d4235 fix(moss-tts-nano): also patch torchaudio.save for soundfile fallback
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-04-14 13:06:12 -04:00
Yueqian Lin 0b8190d07f fix(moss-tts-nano): patch torchaudio.load to use soundfile when torchcodec unavailable
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-04-14 13:06:12 -04:00
Yueqian Lin da6b1b5741 feat(moss-tts-nano): add single-stage MOSS-TTS-Nano TTS support
Integrates OpenMOSS-Team/MOSS-TTS-Nano (0.1B) as a single
GPUGenerationWorker stage.  Both the AR LM and MOSS-Audio-Tokenizer-Nano
codec run inside MossTTSNanoForGeneration.forward(), removing the need for
an inter-stage connector.

Key design choices:
- Weights loaded in load_weights() not __init__ (avoids pre-CUDA alloc)
- trust_remote_code delegates to upstream HF model classes
- codec path read from config.audio_tokenizer_pretrained_name_or_path
- inference_stream() collects progressive audio chunks for low latency
- 48 kHz stereo output; voice clone + continuation modes supported


Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-04-14 13:06:12 -04:00
76 changed files with 5245 additions and 2055 deletions

No files matched your search

+22
View File
@@ -235,6 +235,28 @@ steps:
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "MOSS-TTS-Nano E2E Test"
depends_on: upload-merge-pipeline
commands:
- |
timeout 20m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
pytest -s -v tests/e2e/online_serving/test_moss_tts_nano.py tests/e2e/offline_inference/test_moss_tts_nano.py -m "advanced_model" --run-level "advanced_model"
'
agents:
queue: "gpu_1_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: "Omni Model Test with H100"
timeout_in_minutes: 30
depends_on: upload-merge-pipeline
@@ -1,20 +0,0 @@
vd001|||The quick brown fox jumps over the lazy dog.|A warm, friendly female voice with a slight American Midwest accent, speaking at a moderate pace with natural inflection.
vd002|||Welcome to the future of text-to-speech synthesis.|A deep, authoritative male news anchor voice, clear and professional with a measured cadence.
vd003|||The sunset painted the sky in brilliant shades of orange and pink.|A gentle elderly female voice, soft and wise, with a slight Southern American accent.
vd004|||Scientists have discovered a new species of deep-sea creature.|A young male voice with an Australian accent, curious and enthusiastic.
vd005|||Breaking news: a major climate summit opens today in Geneva.|A crisp female newsreader voice, neutral accent, confident and precise.
vd006|||In the beginning, there was darkness and silence across the universe.|A rich, dramatic bass male narrator voice, slow and deeply resonant.
vd007|||Come closer, I have something important to tell you.|A soft, intimate female voice, slightly whispery, warm and gentle.
vd008|||And they're off! The horses race toward the first turn at incredible speed.|An energetic male sports commentator, fast-paced and excited.
vd009|||Once upon a time, in a land far away, lived a very clever fox.|A light, playful voice with childlike enthusiasm, bright and clear.
vd010|||The ancient manuscript reveals secrets hidden for a thousand years.|A wise, measured elderly male voice, slow and deliberate, British English accent.
vd011|||Good evening, ladies and gentlemen, and welcome to our show.|A sophisticated female voice with a slight French accent speaking English, elegant and refined.
vd012|||System initialized. Running diagnostics. All systems nominal.|A clear, precise robotic-sounding voice, neutral and monotone with slight synthetic quality.
vd013|||I hear what you are saying, and it is completely understandable to feel that way.|A warm, empathetic female therapist voice, calm and reassuring, unhurried pace.
vd014|||Attention all units: proceed to grid reference seven-seven-alpha.|A firm, authoritative military male voice, clipped and commanding.
vd015|||Oh my goodness, you have to try this amazing new recipe I just found!|An enthusiastic, bubbly female voice, high energy and friendly.
vd016|||Dude, the waves were totally amazing out there today. Super happy about it!|A relaxed male voice with a California accent, casual and laid-back.
vd017|||The quarterly results exceed expectations across all major metrics.|A sharp, businesslike female voice, confident and efficient, fast-paced delivery.
vd018|||Chapter one. The morning sun filtered gently through the forest canopy.|A smooth, rich male audiobook narrator voice, expressive and engaging.
vd019|||To be or not to be, that is the question.|A theatrical female voice, dramatic and expressive, stage projection quality.
vd020|||And that is all for tonight. Stay well out there, everyone.|A warm, velvety male late-night radio DJ voice, smooth and intimate.
@@ -1,20 +0,0 @@
smoke001|||The quick brown fox jumps over the lazy dog near the riverbank at sunset.
smoke002|||Welcome to the future of text-to-speech synthesis in production systems.
smoke003|||Yesterday the team finished rolling out the new authentication flow.
smoke004|||She walked carefully across the wet cobblestones, careful not to slip.
smoke005|||The conference call is scheduled for nine in the morning, Pacific time.
smoke006|||Please remember to save your work before closing the editor.
smoke007|||Two plus two equals four, but five hundred and forty three digits is long.
smoke008|||I would like a coffee with oat milk and a chocolate croissant please.
smoke009|||The library closes at eight on weekdays and six on Saturdays.
smoke010|||During the Renaissance, art and science flourished in European cities.
smoke011|||He whispered the secret word so quietly that no one else could hear.
smoke012|||Our flight departs from gate twenty three at eleven fifteen.
smoke013|||The storm knocked out power for six hours, but the backup generator kicked in.
smoke014|||Reading a good book on a rainy afternoon is one of life's great pleasures.
smoke015|||When the kettle whistled, she poured the hot water over the fresh tea leaves.
smoke016|||The algorithm runs in linear time, which is a big improvement over the previous approach.
smoke017|||In the distance, the mountains were shrouded in thick morning fog.
smoke018|||Our company reported record revenue for the fourth quarter of the fiscal year.
smoke019|||She explained the new policy in detail during the staff meeting this morning.
smoke020|||The children laughed and played in the garden until the sun began to set.
+107
View File
@@ -0,0 +1,107 @@
# Qwen3-TTS Benchmark
Benchmarks for Qwen3-TTS text-to-speech models, comparing vLLM-Omni streaming serving against HuggingFace Transformers offline inference.
## Prerequisites
```bash
pip install matplotlib aiohttp soundfile numpy tqdm
pip install qwen_tts # for HF baseline
```
## Quick Start
Run the full benchmark (vllm-omni + HF baseline) with a single command:
```bash
cd benchmarks/qwen3-tts
bash run_benchmark.sh
```
Results (JSON + PNG plots) are saved to `results/`.
### Common options
```bash
# Only vllm-omni (skip HF baseline)
bash run_benchmark.sh --async-only
# Only HF baseline
bash run_benchmark.sh --hf-only
# Use a different model (e.g. 1.7B)
MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
# Use a Voice Clone model
MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-Base TASK_TYPE=Base bash run_benchmark.sh --async-only
# Use batch size 16 for higher throughput
BATCH_SIZE=16 bash run_benchmark.sh --async-only
# Custom GPU, prompt count, concurrency levels
GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
```
## Manual Steps
### 1) Start the vLLM-Omni server
```bash
CUDA_VISIBLE_DEVICES=0 python -m vllm_omni.entrypoints.cli.main serve \
"Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
--omni --host 127.0.0.1 --port 8000 \
--deploy-config vllm_omni/deploy/qwen3_tts.yaml \
--stage-overrides '{"0":{"max_num_seqs":1,"gpu_memory_utilization":0.3,"max_num_batched_tokens":512},"1":{"max_num_seqs":1,"gpu_memory_utilization":0.3,"max_num_batched_tokens":8192}}' \
--trust-remote-code
```
### 2) Run online serving benchmark
```bash
python benchmarks/qwen3-tts/vllm_omni/bench_tts_serve.py \
--port 8000 \
--num-prompts 50 \
--max-concurrency 1 4 10 \
--config-name "async_chunk" \
--result-dir results/
```
### 3) Run HuggingFace baseline
```bash
python benchmarks/qwen3-tts/transformers/bench_tts_hf.py \
--model "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
--num-prompts 50 \
--gpu-device 0 \
--result-dir results/
```
### 4) Generate comparison plots
```bash
python benchmarks/qwen3-tts/plot_results.py \
--results results/bench_async_chunk_*.json results/bench_hf_transformers_*.json \
--labels "vllm-omni" "hf_transformers" \
--output results/comparison.png
```
## Batch-size presets
The bench script loads the bundled production deploy (`vllm_omni/deploy/qwen3_tts.yaml`) and layers per-stage budgets on top via `--stage-overrides`, driven by the `BATCH_SIZE` env var. Each batch size picks compatible per-stage `max_num_seqs`, `max_num_batched_tokens`, and `gpu_memory_utilization` defaults:
| `BATCH_SIZE` | Description |
|:--:|-------------|
| `1` (default) | Single-request processing (lowest latency) |
| `4` | Moderate-throughput concurrent processing |
| `16` | High-throughput concurrent processing |
The 2-stage pipeline (Talker -> Code2Wav) runs with `async_chunk` streaming enabled via the prod deploy; the `SharedMemoryConnector` streams codec frames (25-frame chunks with 25-frame context overlap) between stages.
The model is specified via the CLI `--model` flag (or `MODEL` env var), so the same bench script works for both the 0.6B and 1.7B model variants.
## Metrics
- **TTFP (Time to First Audio Packet)**: Time from request to first audio chunk (streaming latency)
- **E2E (End-to-End Latency)**: Total time from request to complete audio response
- **RTF (Real-Time Factor)**: E2E latency / audio duration. RTF < 1.0 means faster-than-real-time synthesis
- **Throughput**: Total audio seconds generated per wall-clock second
+254
View File
@@ -0,0 +1,254 @@
"""Plot Qwen3-TTS benchmark results.
Generates comparison bar charts similar to the async_chunk design doc:
- TTFP (Time-to-First-Packet) across concurrency levels
- E2E latency across concurrency levels
- RTF (Real-Time Factor) across concurrency levels
Usage:
# Compare two configs (async_chunk vs no_async_chunk):
python plot_results.py \
--results results/bench_async_chunk_*.json results/bench_no_async_chunk_*.json \
--labels "async_chunk" "no_async_chunk" \
--output results/qwen3_tts_benchmark.png
# Single config:
python plot_results.py \
--results results/bench_async_chunk_*.json \
--labels "async_chunk" \
--output results/qwen3_tts_benchmark.png
"""
import argparse
import json
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
def load_results(result_files: list[str]) -> list[list[dict]]:
"""Load benchmark results from JSON files."""
all_results = []
for f in result_files:
with open(f) as fh:
data = json.load(fh)
all_results.append(data)
return all_results
def plot_comparison(
all_results: list[list[dict]],
labels: list[str],
output_path: str,
title_prefix: str = "Qwen3-TTS",
):
"""Generate comparison bar charts."""
n_configs = len(all_results)
# Collect concurrency levels present in ALL configs (skip missing data)
all_concurrencies = [set(r["concurrency"] for r in results) for results in all_results]
concurrencies = sorted(set.union(*all_concurrencies))
# Build data arrays, using None for missing concurrency levels
ttfp_data = {label: [] for label in labels}
e2e_data = {label: [] for label in labels}
rtf_data = {label: [] for label in labels}
throughput_data = {label: [] for label in labels}
for results, label in zip(all_results, labels):
conc_map = {r["concurrency"]: r for r in results}
for c in concurrencies:
r = conc_map.get(c)
ttfp_data[label].append(r["mean_ttfp_ms"] if r else None)
e2e_data[label].append(r["mean_e2e_ms"] if r else None)
rtf_data[label].append(r["mean_rtf"] if r else None)
throughput_data[label].append(r["audio_throughput"] if r else None)
fig, axes = plt.subplots(2, 2, figsize=(14, 10))
fig.suptitle(f"{title_prefix} Performance Benchmark", fontsize=16, fontweight="bold")
x = np.arange(len(concurrencies))
width = 0.35 if n_configs == 2 else 0.5
if n_configs > 1:
offsets = np.linspace(-width / 2 * (n_configs - 1), width / 2 * (n_configs - 1), n_configs)
else:
offsets = [0]
colors = ["#2196F3", "#FF5722", "#4CAF50", "#FFC107"]
def plot_metric(ax, data_dict, ylabel, title, fmt=".1f"):
bars = []
for i, (label, values) in enumerate(data_dict.items()):
# Replace None with 0 for plotting, but track which are missing
plot_values = [v if v is not None else 0 for v in values]
color = colors[i % len(colors)]
bar = ax.bar(x + offsets[i], plot_values, width, label=label, color=color, alpha=0.85)
bars.append(bar)
# Add value labels on bars (skip None/missing data)
max_val = max((v for v in values if v is not None), default=1)
for rect, val in zip(bar, values):
if val is not None and val > 0:
ax.text(
rect.get_x() + rect.get_width() / 2,
rect.get_height() + max_val * 0.02,
f"{val:{fmt}}",
ha="center",
va="bottom",
fontsize=9,
fontweight="bold",
)
ax.set_xlabel("Concurrency", fontsize=12)
ax.set_ylabel(ylabel, fontsize=12)
ax.set_title(title, fontsize=13, fontweight="bold")
ax.set_xticks(x)
ax.set_xticklabels([str(c) for c in concurrencies])
ax.legend(fontsize=10)
ax.grid(axis="y", alpha=0.3)
ax.set_axisbelow(True)
plot_metric(axes[0, 0], ttfp_data, "TTFP (ms)", "Time to First Audio Packet (TTFP)")
plot_metric(axes[0, 1], e2e_data, "E2E Latency (ms)", "End-to-End Latency (E2E)")
plot_metric(axes[1, 0], rtf_data, "RTF", "Real-Time Factor (RTF)", fmt=".3f")
plot_metric(axes[1, 1], throughput_data, "Audio-sec / Wall-sec", "Audio Throughput", fmt=".2f")
plt.tight_layout()
plt.savefig(output_path, dpi=150, bbox_inches="tight")
print(f"Plot saved to {output_path}")
plt.close()
def plot_single_summary(results: list[dict], label: str, output_path: str):
"""Generate a single-config summary with percentile breakdown."""
concurrencies = [r["concurrency"] for r in results]
fig, axes = plt.subplots(1, 3, figsize=(16, 5))
fig.suptitle(f"Qwen3-TTS Benchmark - {label}", fontsize=15, fontweight="bold")
# TTFP breakdown
ax = axes[0]
means = [r["mean_ttfp_ms"] for r in results]
medians = [r["median_ttfp_ms"] for r in results]
p90s = [r["p90_ttfp_ms"] for r in results]
p99s = [r["p99_ttfp_ms"] for r in results]
x = np.arange(len(concurrencies))
w = 0.2
ax.bar(x - 1.5 * w, means, w, label="mean", color="#2196F3")
ax.bar(x - 0.5 * w, medians, w, label="median", color="#4CAF50")
ax.bar(x + 0.5 * w, p90s, w, label="p90", color="#FF9800")
ax.bar(x + 1.5 * w, p99s, w, label="p99", color="#F44336")
ax.set_xticks(x)
ax.set_xticklabels([str(c) for c in concurrencies])
ax.set_xlabel("Concurrency")
ax.set_ylabel("TTFP (ms)")
ax.set_title("Time to First Audio Packet")
ax.legend(fontsize=9)
ax.grid(axis="y", alpha=0.3)
# E2E breakdown
ax = axes[1]
means = [r["mean_e2e_ms"] for r in results]
medians = [r["median_e2e_ms"] for r in results]
p90s = [r["p90_e2e_ms"] for r in results]
p99s = [r["p99_e2e_ms"] for r in results]
ax.bar(x - 1.5 * w, means, w, label="mean", color="#2196F3")
ax.bar(x - 0.5 * w, medians, w, label="median", color="#4CAF50")
ax.bar(x + 0.5 * w, p90s, w, label="p90", color="#FF9800")
ax.bar(x + 1.5 * w, p99s, w, label="p99", color="#F44336")
ax.set_xticks(x)
ax.set_xticklabels([str(c) for c in concurrencies])
ax.set_xlabel("Concurrency")
ax.set_ylabel("E2E Latency (ms)")
ax.set_title("End-to-End Latency")
ax.legend(fontsize=9)
ax.grid(axis="y", alpha=0.3)
# RTF
ax = axes[2]
means = [r["mean_rtf"] for r in results]
medians = [r["median_rtf"] for r in results]
ax.bar(x - 0.15, means, 0.3, label="mean", color="#2196F3")
ax.bar(x + 0.15, medians, 0.3, label="median", color="#4CAF50")
ax.set_xticks(x)
ax.set_xticklabels([str(c) for c in concurrencies])
ax.set_xlabel("Concurrency")
ax.set_ylabel("RTF")
ax.set_title("Real-Time Factor")
ax.legend(fontsize=9)
ax.grid(axis="y", alpha=0.3)
plt.tight_layout()
plt.savefig(output_path, dpi=150, bbox_inches="tight")
print(f"Plot saved to {output_path}")
plt.close()
def print_comparison_table(all_results: list[list[dict]], labels: list[str]):
"""Print a markdown-formatted comparison table."""
concurrencies = sorted(set(r["concurrency"] for r in all_results[0]))
print("\n## Benchmark Results\n")
header = "| Metric | Concurrency |"
sep = "| --- | --- |"
for label in labels:
header += f" {label} |"
sep += " --- |"
print(header)
print(sep)
for metric, key, fmt in [
("TTFP (ms)", "mean_ttfp_ms", ".1f"),
("E2E (ms)", "mean_e2e_ms", ".1f"),
("RTF", "mean_rtf", ".3f"),
("Throughput (audio-s/s)", "audio_throughput", ".2f"),
]:
for c in concurrencies:
row = f"| {metric} | {c} |"
for results in all_results:
conc_map = {r["concurrency"]: r for r in results}
val = conc_map.get(c, {}).get(key, 0)
row += f" {val:{fmt}} |"
print(row)
# Improvement calculation (only if 2 configs)
if len(all_results) == 2:
print(f"\n## Improvement ({labels[0]} vs {labels[1]})\n")
print("| Metric | Concurrency | Improvement |")
print("| --- | --- | --- |")
for metric, key in [("TTFP", "mean_ttfp_ms"), ("E2E", "mean_e2e_ms"), ("RTF", "mean_rtf")]:
for c in concurrencies:
m0 = {r["concurrency"]: r for r in all_results[0]}
m1 = {r["concurrency"]: r for r in all_results[1]}
v0 = m0.get(c, {}).get(key, 0)
v1 = m1.get(c, {}).get(key, 0)
if v1 > 0:
pct = (v1 - v0) / v1 * 100
print(f"| {metric} | {c} | {pct:+.1f}% |")
def parse_args():
parser = argparse.ArgumentParser(description="Plot Qwen3-TTS benchmark results")
parser.add_argument(
"--results", type=str, nargs="+", required=True, help="Path(s) to result JSON files (one per config)"
)
parser.add_argument(
"--labels", type=str, nargs="+", required=True, help="Labels for each config (must match --results count)"
)
parser.add_argument("--output", type=str, default="results/qwen3_tts_benchmark.png", help="Output image path")
parser.add_argument("--title", type=str, default="Qwen3-TTS", help="Title prefix for the plot")
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
assert len(args.results) == len(args.labels), "--results and --labels must have the same count"
all_results = load_results(args.results)
print_comparison_table(all_results, args.labels)
Path(args.output).parent.mkdir(parents=True, exist_ok=True)
if len(all_results) == 1:
plot_single_summary(all_results[0], args.labels[0], args.output)
else:
plot_comparison(all_results, args.labels, args.output, title_prefix=args.title)
+3
View File
@@ -0,0 +1,3 @@
# Benchmark results are machine-specific - do not commit
*
!.gitignore
+285
View File
@@ -0,0 +1,285 @@
#!/bin/bash
# Qwen3-TTS Benchmark Runner
#
# Compares vllm-omni streaming serving vs HuggingFace transformers offline inference.
# Produces JSON results and comparison plots.
#
# Usage:
# # Full comparison (vllm-omni + HF):
# bash run_benchmark.sh
#
# # Only vllm-omni async_chunk config:
# bash run_benchmark.sh --async-only
#
# # Only HuggingFace baseline:
# bash run_benchmark.sh --hf-only
#
# # vllm-omni only (skip HF):
# bash run_benchmark.sh --skip-hf
#
# # Custom settings:
# GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
#
# # Use 1.7B model:
# MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
#
# # Use Voice Clone model
# MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-Base TASK_TYPE=Base bash run_benchmark.sh --async-only
#
# # Use batch_size=4:
# BATCH_SIZE=4 bash run_benchmark.sh --async-only
#
# Environment variables:
# GPU_DEVICE - GPU index to use (default: 0)
# NUM_PROMPTS - Number of prompts per concurrency level (default: 50)
# CONCURRENCY - Space-separated concurrency levels (default: "1 4 10")
# MODEL - Model name (default: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice)
# PORT - Server port (default: 8000)
# BATCH_SIZE - Per-stage ``max_num_seqs`` for both talker and code2wav (default: 1)
# GPU_MEM_TALKER - gpu_memory_utilization for talker stage (default: 0.3 at bs=1, else 0.2)
# GPU_MEM_CODE2WAV - gpu_memory_utilization for code2wav stage (default: 0.3 at bs=1, else 0.2)
# TASK_TYPE - Task type: CustomVoice, VoiceDesign, Base (default: CustomVoice)
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
# Defaults
GPU_DEVICE="${GPU_DEVICE:-0}"
NUM_PROMPTS="${NUM_PROMPTS:-50}"
CONCURRENCY="${CONCURRENCY:-1 4 10}"
MODEL="${MODEL:-Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice}"
PORT="${PORT:-8000}"
BATCH_SIZE="${BATCH_SIZE:-1}"
DEFAULT_MEM=$([ "${BATCH_SIZE}" = "1" ] && echo "0.3" || echo "0.2")
GPU_MEM_TALKER="${GPU_MEM_TALKER:-${DEFAULT_MEM}}"
GPU_MEM_CODE2WAV="${GPU_MEM_CODE2WAV:-${DEFAULT_MEM}}"
NUM_WARMUPS="${NUM_WARMUPS:-3}"
DEPLOY_CONFIG="vllm_omni/deploy/qwen3_tts.yaml"
RESULT_DIR="${SCRIPT_DIR}/results"
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
TASK_TYPE="${TASK_TYPE:-CustomVoice}"
# Build --stage-overrides JSON from BATCH_SIZE + GPU_MEM_*.
STAGE_OVERRIDES=$(
BATCH_SIZE="${BATCH_SIZE}" \
GPU_MEM_TALKER="${GPU_MEM_TALKER}" \
GPU_MEM_CODE2WAV="${GPU_MEM_CODE2WAV}" \
python - <<'PYEOF'
import json, os
bs = int(os.environ["BATCH_SIZE"])
mem_t = float(os.environ["GPU_MEM_TALKER"])
mem_c = float(os.environ["GPU_MEM_CODE2WAV"])
# Prefill budget grows with batch size on both stages.
talker_batched = 512 if bs <= 4 else 4096
code2wav_batched = 8192 if bs <= 4 else 32768
print(json.dumps({
"0": {"max_num_seqs": bs, "gpu_memory_utilization": mem_t, "max_num_batched_tokens": talker_batched},
"1": {"max_num_seqs": bs, "gpu_memory_utilization": mem_c, "max_num_batched_tokens": code2wav_batched},
}))
PYEOF
)
# Parse args
RUN_ASYNC=true
RUN_HF=true
for arg in "$@"; do
case "$arg" in
--async-only) RUN_HF=false ;;
--hf-only) RUN_ASYNC=false ;;
--skip-hf) RUN_HF=false ;;
esac
done
mkdir -p "${RESULT_DIR}"
echo "============================================================"
echo " Qwen3-TTS Benchmark"
echo "============================================================"
echo " GPU: ${GPU_DEVICE}"
echo " Model: ${MODEL}"
echo " Prompts: ${NUM_PROMPTS}"
echo " Concurrency: ${CONCURRENCY}"
echo " Port: ${PORT}"
echo " Deploy config: ${DEPLOY_CONFIG}"
echo " Batch size: ${BATCH_SIZE}"
echo " GPU mem T/C: ${GPU_MEM_TALKER} / ${GPU_MEM_CODE2WAV}"
echo " Results: ${RESULT_DIR}"
echo " Task type: ${TASK_TYPE}"
echo "============================================================"
# Start server and wait for it to be ready
start_server() {
local config_name="$1"
local log_file="${RESULT_DIR}/server_${config_name}_${TIMESTAMP}.log"
echo ""
echo "Starting server with config: ${config_name}"
echo " Deploy config: ${DEPLOY_CONFIG}"
echo " Stage overrides: ${STAGE_OVERRIDES}"
echo " Log file: ${log_file}"
VLLM_WORKER_MULTIPROC_METHOD=spawn \
CUDA_VISIBLE_DEVICES="${GPU_DEVICE}" \
python -m vllm_omni.entrypoints.cli.main serve "${MODEL}" \
--omni \
--host 127.0.0.1 \
--port "${PORT}" \
--deploy-config "${DEPLOY_CONFIG}" \
--stage-overrides "${STAGE_OVERRIDES}" \
--stage-init-timeout 120 \
--trust-remote-code \
--disable-log-stats \
> "${log_file}" 2>&1 &
SERVER_PID=$!
echo " Server PID: ${SERVER_PID}"
# Wait for server to be ready
echo " Waiting for server to be ready..."
local max_wait=300
local waited=0
while [ ${waited} -lt ${max_wait} ]; do
if curl -sf "http://127.0.0.1:${PORT}/v1/models" > /dev/null 2>&1; then
echo " Server is ready! (waited ${waited}s)"
return 0
fi
# Check if process is still alive
if ! kill -0 ${SERVER_PID} 2>/dev/null; then
echo " ERROR: Server process died. Check log: ${log_file}"
tail -20 "${log_file}"
return 1
fi
sleep 2
waited=$((waited + 2))
done
echo " ERROR: Server did not start within ${max_wait}s. Check log: ${log_file}"
kill ${SERVER_PID} 2>/dev/null || true
return 1
}
# Stop the server
stop_server() {
if [ -n "${SERVER_PID:-}" ]; then
echo " Stopping server (PID: ${SERVER_PID})..."
kill ${SERVER_PID} 2>/dev/null || true
wait ${SERVER_PID} 2>/dev/null || true
# Kill any remaining child processes on the port
local pids
pids=$(lsof -ti:${PORT} 2>/dev/null || true)
if [ -n "${pids}" ]; then
echo " Cleaning up remaining processes on port ${PORT}..."
echo "${pids}" | xargs kill -9 2>/dev/null || true
fi
echo " Server stopped."
SERVER_PID=""
fi
}
# Cleanup on exit
trap 'stop_server' EXIT
# Run benchmark for a given config
run_bench() {
local config_name="$1"
echo ""
echo "============================================================"
echo " Benchmarking: ${config_name}"
echo "============================================================"
start_server "${config_name}"
# Convert concurrency string to args
local conc_args=""
for c in ${CONCURRENCY}; do
conc_args="${conc_args} ${c}"
done
cd "${PROJECT_ROOT}"
python "${SCRIPT_DIR}/vllm_omni/bench_tts_serve.py" \
--host 127.0.0.1 \
--port "${PORT}" \
--num-prompts "${NUM_PROMPTS}" \
--max-concurrency ${conc_args} \
--num-warmups "${NUM_WARMUPS}" \
--config-name "${config_name}" \
--result-dir "${RESULT_DIR}" \
--task-type "${TASK_TYPE}"
stop_server
# Allow GPU memory to settle
sleep 5
}
# Run vllm-omni benchmark
if [ "${RUN_ASYNC}" = true ]; then
run_bench "async_chunk"
fi
# Run HuggingFace baseline benchmark
if [ "${RUN_HF}" = true ]; then
echo ""
echo "============================================================"
echo " Benchmarking: HuggingFace transformers (offline)"
echo "============================================================"
cd "${PROJECT_ROOT}"
python "${SCRIPT_DIR}/transformers/bench_tts_hf.py" \
--model "${MODEL}" \
--num-prompts "${NUM_PROMPTS}" \
--num-warmups "${NUM_WARMUPS}" \
--gpu-device "${GPU_DEVICE}" \
--config-name "hf_transformers" \
--result-dir "${RESULT_DIR}" \
--task-type "${TASK_TYPE}"
# Allow GPU memory to settle
sleep 5
fi
# Plot results
echo ""
echo "============================================================"
echo " Generating plots..."
echo "============================================================"
RESULT_FILES=""
LABELS=""
if [ "${RUN_ASYNC}" = true ]; then
ASYNC_FILE=$(ls -t "${RESULT_DIR}"/bench_async_chunk_*.json 2>/dev/null | head -1)
if [ -n "${ASYNC_FILE}" ]; then
RESULT_FILES="${ASYNC_FILE}"
LABELS="async_chunk"
fi
fi
if [ "${RUN_HF}" = true ]; then
HF_FILE=$(ls -t "${RESULT_DIR}"/bench_hf_transformers_*.json 2>/dev/null | head -1)
if [ -n "${HF_FILE}" ]; then
if [ -n "${RESULT_FILES}" ]; then
RESULT_FILES="${RESULT_FILES} ${HF_FILE}"
LABELS="${LABELS} hf_transformers"
else
RESULT_FILES="${HF_FILE}"
LABELS="hf_transformers"
fi
fi
fi
if [ -n "${RESULT_FILES}" ]; then
python "${SCRIPT_DIR}/plot_results.py" \
--results ${RESULT_FILES} \
--labels ${LABELS} \
--output "${RESULT_DIR}/qwen3_tts_benchmark_${TIMESTAMP}.png"
fi
echo ""
echo "============================================================"
echo " Benchmark complete!"
echo " Results: ${RESULT_DIR}"
echo "============================================================"
@@ -0,0 +1,301 @@
"""Benchmark Qwen3-TTS using HuggingFace transformers (qwen_tts library).
Measures E2E latency, RTF, and audio duration for offline (non-serving) inference.
Results are saved in the same JSON format as bench_tts_serve.py for unified plotting.
Usage:
python bench_tts_hf.py \
--model Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice \
--num-prompts 50 \
--num-warmups 3 \
--gpu-device 0 \
--result-dir results/
"""
import argparse
import json
import time
from dataclasses import asdict, dataclass, field
from datetime import datetime
from pathlib import Path
import numpy as np
import soundfile as sf
import torch
PROMPTS = [
"Hello, welcome to the voice synthesis benchmark test.",
"She said she would be here by noon, but nobody showed up.",
"The quick brown fox jumps over the lazy dog near the riverbank.",
"I can't believe how beautiful the sunset looks from up here on the mountain.",
"Please remember to bring your identification documents to the appointment tomorrow morning.",
"Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
"The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
"After the meeting, we should discuss the quarterly results and plan for the next phase.",
"Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
"The train leaves at half past seven, so we need to arrive at the station before then.",
"Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
"It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
]
REF_AUDIO = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav"
REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
INSTRUCT = "Speak in an incredulous tone, but with a hint of panic beginning to creep into your voice."
@dataclass
class BenchmarkResult:
config_name: str = ""
concurrency: int = 1 # always 1 for offline
num_prompts: int = 0
completed: int = 0
failed: int = 0
duration_s: float = 0.0
# TTFP stats - not applicable for HF offline, set to E2E for compatibility
mean_ttfp_ms: float = 0.0
median_ttfp_ms: float = 0.0
std_ttfp_ms: float = 0.0
p90_ttfp_ms: float = 0.0
p95_ttfp_ms: float = 0.0
p99_ttfp_ms: float = 0.0
# E2E stats (ms)
mean_e2e_ms: float = 0.0
median_e2e_ms: float = 0.0
std_e2e_ms: float = 0.0
p90_e2e_ms: float = 0.0
p95_e2e_ms: float = 0.0
p99_e2e_ms: float = 0.0
# RTF stats
mean_rtf: float = 0.0
median_rtf: float = 0.0
std_rtf: float = 0.0
p99_rtf: float = 0.0
# Audio stats
mean_audio_duration_s: float = 0.0
total_audio_duration_s: float = 0.0
audio_throughput: float = 0.0
request_throughput: float = 0.0
# Per-request details
per_request: list = field(default_factory=list)
def generate_audio(model, prompt: str, args):
if args.task_type == "Base":
return model.generate_voice_clone(
text=prompt,
language=args.language,
ref_audio=REF_AUDIO,
ref_text=REF_TEXT,
)
if args.task_type == "VoiceDesign":
return model.generate_voice_design(
text=prompt,
language=args.language,
instruct=INSTRUCT,
)
return model.generate_custom_voice(
text=prompt,
language=args.language,
speaker=args.voice,
)
def run_benchmark(args):
from qwen_tts import Qwen3TTSModel
device = f"cuda:{args.gpu_device}"
print(f"Loading model: {args.model} on {device}")
model = Qwen3TTSModel.from_pretrained(
args.model,
device_map=device,
dtype=torch.bfloat16,
)
print("Model loaded.")
# Build prompt list
prompts = [PROMPTS[i % len(PROMPTS)] for i in range(args.num_prompts)]
# Warmup
if args.num_warmups > 0:
print(f"Warming up with {args.num_warmups} requests...")
for i in range(args.num_warmups):
p = PROMPTS[i % len(PROMPTS)]
wavs, sr = generate_audio(model, p, args)
# Sync GPU
torch.cuda.synchronize(device)
print("Warmup done.")
# Benchmark
print(f"Running {args.num_prompts} requests sequentially...")
e2e_times = []
rtfs = []
audio_durations = []
per_request = []
failed = 0
audio_dir = None
if args.save_audio:
audio_dir = Path(args.result_dir) / "audio_hf"
audio_dir.mkdir(parents=True, exist_ok=True)
total_start = time.perf_counter()
for i, prompt in enumerate(prompts):
try:
torch.cuda.synchronize(device)
st = time.perf_counter()
wavs, sr = generate_audio(model, prompt, args)
torch.cuda.synchronize(device)
elapsed = time.perf_counter() - st
# Compute audio duration
audio_samples = wavs[0]
if isinstance(audio_samples, torch.Tensor):
audio_samples = audio_samples.cpu().numpy()
audio_dur = len(audio_samples) / sr
rtf = elapsed / audio_dur if audio_dur > 0 else 0.0
e2e_times.append(elapsed)
rtfs.append(rtf)
audio_durations.append(audio_dur)
per_request.append(
{
"e2e_ms": elapsed * 1000,
"ttfp_ms": elapsed * 1000, # no streaming, TTFP = E2E
"rtf": rtf,
"audio_duration_s": audio_dur,
"prompt": prompt,
}
)
if audio_dir:
sf.write(str(audio_dir / f"output_{i:04d}.wav"), audio_samples, sr)
if (i + 1) % 10 == 0 or i == 0:
print(
f" [{i + 1}/{args.num_prompts}] e2e={elapsed * 1000:.0f}ms rtf={rtf:.3f} audio={audio_dur:.2f}s"
)
except Exception as e:
print(f" [{i + 1}/{args.num_prompts}] FAILED: {e}")
failed += 1
total_duration = time.perf_counter() - total_start
completed = len(e2e_times)
# Compute stats
result = BenchmarkResult(
config_name=args.config_name,
concurrency=1,
num_prompts=args.num_prompts,
completed=completed,
failed=failed,
duration_s=total_duration,
)
if e2e_times:
e2e_ms = [t * 1000 for t in e2e_times]
result.mean_e2e_ms = float(np.mean(e2e_ms))
result.median_e2e_ms = float(np.median(e2e_ms))
result.std_e2e_ms = float(np.std(e2e_ms))
result.p90_e2e_ms = float(np.percentile(e2e_ms, 90))
result.p95_e2e_ms = float(np.percentile(e2e_ms, 95))
result.p99_e2e_ms = float(np.percentile(e2e_ms, 99))
# For HF offline, TTFP = E2E (no streaming)
result.mean_ttfp_ms = result.mean_e2e_ms
result.median_ttfp_ms = result.median_e2e_ms
result.std_ttfp_ms = result.std_e2e_ms
result.p90_ttfp_ms = result.p90_e2e_ms
result.p95_ttfp_ms = result.p95_e2e_ms
result.p99_ttfp_ms = result.p99_e2e_ms
result.mean_rtf = float(np.mean(rtfs))
result.median_rtf = float(np.median(rtfs))
result.std_rtf = float(np.std(rtfs))
result.p99_rtf = float(np.percentile(rtfs, 99))
result.mean_audio_duration_s = float(np.mean(audio_durations))
result.total_audio_duration_s = float(np.sum(audio_durations))
result.audio_throughput = result.total_audio_duration_s / total_duration
result.request_throughput = completed / total_duration
result.per_request = per_request
# Print summary in standardized performance template
W = 50
print("")
print(f"{'=' * W}")
print(f"{'Serving Benchmark Result':^{W}}")
print(f"{'=' * W}")
print(f"{'Successful requests:':<40}{completed:<10}")
print(f"{'Failed requests:':<40}{failed:<10}")
print(f"{'Maximum request concurrency:':<40}{1:<10}")
print(f"{'Benchmark duration (s):':<40}{total_duration:<10.2f}")
print(f"{'Request throughput (req/s):':<40}{result.request_throughput:<10.2f}")
print(f"{'-' * W}")
print(f"{'End-to-end Latency':^{W}}")
print(f"{'-' * W}")
print(f"{'Mean E2EL (ms):':<40}{result.mean_e2e_ms:<10.2f}")
print(f"{'Median E2EL (ms):':<40}{result.median_e2e_ms:<10.2f}")
print(f"{'P99 E2EL (ms):':<40}{result.p99_e2e_ms:<10.2f}")
print(f"{'=' * W}")
print(f"{'Audio Result':^{W}}")
print(f"{'=' * W}")
print(f"{'Total audio duration generated (s):':<40}{result.total_audio_duration_s:<10.2f}")
print(f"{'Audio throughput (audio duration/s):':<40}{result.audio_throughput:<10.2f}")
print(f"{'-' * W}")
print(f"{'Time to First Packet':^{W}}")
print(f"{'-' * W}")
print(f"{'Mean AUDIO_TTFP (ms):':<40}{result.mean_ttfp_ms:<10.2f}")
print(f"{'Median AUDIO_TTFP (ms):':<40}{result.median_ttfp_ms:<10.2f}")
print(f"{'P99 AUDIO_TTFP (ms):':<40}{result.p99_ttfp_ms:<10.2f}")
print(f"{'-' * W}")
print(f"{'Real Time Factor':^{W}}")
print(f"{'-' * W}")
print(f"{'Mean AUDIO_RTF:':<40}{result.mean_rtf:<10.3f}")
print(f"{'Median AUDIO_RTF:':<40}{result.median_rtf:<10.3f}")
print(f"{'P99 AUDIO_RTF:':<40}{result.p99_rtf:<10.3f}")
print(f"{'=' * W}")
print("")
# Save results (as a list with single concurrency=1 entry, matching serve format)
result_dir = Path(args.result_dir)
result_dir.mkdir(parents=True, exist_ok=True)
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
with open(result_file, "w") as f:
json.dump([asdict(result)], f, indent=2)
print(f"Results saved to {result_file}")
return result
def parse_args():
parser = argparse.ArgumentParser(description="Qwen3-TTS HuggingFace Benchmark")
parser.add_argument(
"--model", type=str, default="Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice", help="HuggingFace model name or path"
)
parser.add_argument("--num-prompts", type=int, default=50)
parser.add_argument("--num-warmups", type=int, default=3)
parser.add_argument("--gpu-device", type=int, default=0)
parser.add_argument("--voice", type=str, default="Vivian")
parser.add_argument("--language", type=str, default="English")
parser.add_argument("--task-type", type=str, default="CustomVoice", choices=["CustomVoice", "VoiceDesign", "Base"])
parser.add_argument(
"--config-name", type=str, default="hf_transformers", help="Label for this config (used in filenames)"
)
parser.add_argument("--result-dir", type=str, default="results")
parser.add_argument("--save-audio", action="store_true", help="Save generated audio files")
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
run_benchmark(args)
@@ -0,0 +1,301 @@
"""Benchmark comparing async_chunk on vs off for Qwen3-TTS.
Measures TTFP (Time-to-First-Packet), E2E latency, and RTF across
concurrency levels for both async_chunk modes. Saves results as JSON.
Usage:
# Run against a server already serving with a given config:
python bench_async_chunk.py \
--host 127.0.0.1 --port 8000 \
--config-name async_chunk_on \
--num-prompts 50 \
--max-concurrency 1 10 \
--result-dir results/
"""
import argparse
import asyncio
import json
import time
from dataclasses import asdict, dataclass, field
from datetime import datetime
from pathlib import Path
import aiohttp
import numpy as np
from tqdm.asyncio import tqdm
PROMPTS = [
"Hello, welcome to the voice synthesis benchmark test.",
"She said she would be here by noon, but nobody showed up.",
"The quick brown fox jumps over the lazy dog near the riverbank.",
"I can't believe how beautiful the sunset looks from up here on the mountain.",
"Please remember to bring your identification documents to the appointment tomorrow morning.",
"Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
"The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
"After the meeting, we should discuss the quarterly results and plan for the next phase.",
"Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
"The train leaves at half past seven, so we need to arrive at the station before then.",
"Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
"It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
]
@dataclass
class RequestResult:
success: bool = False
ttfp: float = 0.0
e2e: float = 0.0
audio_bytes: int = 0
audio_duration: float = 0.0
rtf: float = 0.0
prompt: str = ""
error: str = ""
@dataclass
class BenchmarkResult:
config_name: str = ""
concurrency: int = 0
num_prompts: int = 0
completed: int = 0
failed: int = 0
duration_s: float = 0.0
mean_ttfp_ms: float = 0.0
median_ttfp_ms: float = 0.0
std_ttfp_ms: float = 0.0
p90_ttfp_ms: float = 0.0
p95_ttfp_ms: float = 0.0
p99_ttfp_ms: float = 0.0
mean_e2e_ms: float = 0.0
median_e2e_ms: float = 0.0
std_e2e_ms: float = 0.0
p90_e2e_ms: float = 0.0
p95_e2e_ms: float = 0.0
p99_e2e_ms: float = 0.0
mean_rtf: float = 0.0
median_rtf: float = 0.0
std_rtf: float = 0.0
mean_audio_duration_s: float = 0.0
total_audio_duration_s: float = 0.0
audio_throughput: float = 0.0
request_throughput: float = 0.0
per_request: list = field(default_factory=list)
def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width: int = 2) -> float:
return num_bytes / sample_width / sample_rate
async def send_tts_request(
session: aiohttp.ClientSession,
api_url: str,
prompt: str,
voice: str = "vivian",
language: str = "English",
stream: bool = True,
pbar: tqdm | None = None,
) -> RequestResult:
payload = {
"input": prompt,
"voice": voice,
"language": language,
"stream": stream,
"response_format": "pcm",
}
result = RequestResult(prompt=prompt)
st = time.perf_counter()
try:
async with session.post(api_url, json=payload) as response:
if response.status != 200:
result.error = f"HTTP {response.status}: {await response.text()}"
return result
first_chunk = True
total_bytes = 0
async for chunk in response.content.iter_any():
if first_chunk and len(chunk) > 0:
result.ttfp = time.perf_counter() - st
first_chunk = False
total_bytes += len(chunk)
result.e2e = time.perf_counter() - st
result.audio_bytes = total_bytes
result.audio_duration = pcm_bytes_to_duration(total_bytes)
if result.audio_duration > 0:
result.rtf = result.e2e / result.audio_duration
result.success = True
except Exception as e:
result.error = str(e)
result.e2e = time.perf_counter() - st
if pbar:
pbar.update(1)
return result
async def run_benchmark(
host: str,
port: int,
num_prompts: int,
max_concurrency: int,
num_warmups: int = 3,
voice: str = "vivian",
language: str = "English",
stream: bool = True,
) -> BenchmarkResult:
api_url = f"http://{host}:{port}/v1/audio/speech"
connector = aiohttp.TCPConnector(limit=max_concurrency, limit_per_host=max_concurrency, keepalive_timeout=60)
session = aiohttp.ClientSession(connector=connector, timeout=aiohttp.ClientTimeout(total=600))
if num_warmups > 0:
print(f" Warming up with {num_warmups} requests...")
warmup_tasks = [
send_tts_request(session, api_url, PROMPTS[i % len(PROMPTS)], voice, language, stream)
for i in range(num_warmups)
]
await asyncio.gather(*warmup_tasks)
print(" Warmup done.")
request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
print(f" Running {num_prompts} requests with concurrency={max_concurrency}...")
semaphore = asyncio.Semaphore(max_concurrency)
pbar = tqdm(total=num_prompts, desc=f" concurrency={max_concurrency}")
async def limited_request(prompt):
async with semaphore:
return await send_tts_request(session, api_url, prompt, voice, language, stream, pbar)
start_time = time.perf_counter()
tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
results: list[RequestResult] = await asyncio.gather(*tasks)
duration = time.perf_counter() - start_time
pbar.close()
await session.close()
successful = [r for r in results if r.success]
failed = [r for r in results if not r.success]
bench = BenchmarkResult(
concurrency=max_concurrency,
num_prompts=num_prompts,
completed=len(successful),
failed=len(failed),
duration_s=duration,
)
if successful:
ttfps = [r.ttfp * 1000 for r in successful]
e2es = [r.e2e * 1000 for r in successful]
rtfs = [r.rtf for r in successful]
audio_durs = [r.audio_duration for r in successful]
bench.mean_ttfp_ms = float(np.mean(ttfps))
bench.median_ttfp_ms = float(np.median(ttfps))
bench.std_ttfp_ms = float(np.std(ttfps))
bench.p90_ttfp_ms = float(np.percentile(ttfps, 90))
bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
bench.p99_ttfp_ms = float(np.percentile(ttfps, 99))
bench.mean_e2e_ms = float(np.mean(e2es))
bench.median_e2e_ms = float(np.median(e2es))
bench.std_e2e_ms = float(np.std(e2es))
bench.p90_e2e_ms = float(np.percentile(e2es, 90))
bench.p95_e2e_ms = float(np.percentile(e2es, 95))
bench.p99_e2e_ms = float(np.percentile(e2es, 99))
bench.mean_rtf = float(np.mean(rtfs))
bench.median_rtf = float(np.median(rtfs))
bench.std_rtf = float(np.std(rtfs))
bench.mean_audio_duration_s = float(np.mean(audio_durs))
bench.total_audio_duration_s = float(np.sum(audio_durs))
bench.audio_throughput = bench.total_audio_duration_s / duration
bench.request_throughput = len(successful) / duration
bench.per_request = [
{
"ttfp_ms": r.ttfp * 1000,
"e2e_ms": r.e2e * 1000,
"rtf": r.rtf,
"audio_duration_s": r.audio_duration,
"prompt": r.prompt,
}
for r in successful
]
print(f"\n{'=' * 60}")
print(f" Concurrency: {max_concurrency} | Completed: {bench.completed} | Failed: {bench.failed}")
print(f" Duration: {duration:.2f}s | Throughput: {bench.request_throughput:.2f} req/s")
print(
f" TTFP (ms): mean={bench.mean_ttfp_ms:.1f} median={bench.median_ttfp_ms:.1f}"
f" p90={bench.p90_ttfp_ms:.1f} p99={bench.p99_ttfp_ms:.1f}"
)
print(
f" E2E (ms): mean={bench.mean_e2e_ms:.1f} median={bench.median_e2e_ms:.1f}"
f" p90={bench.p90_e2e_ms:.1f} p99={bench.p99_e2e_ms:.1f}"
)
print(f" RTF: mean={bench.mean_rtf:.3f} median={bench.median_rtf:.3f}")
print(f" Throughput: {bench.audio_throughput:.2f} audio-sec/wall-sec")
print(f"{'=' * 60}\n")
if failed:
for r in failed[:3]:
print(f" [ERROR] {r.error[:200]}")
return bench
async def main(args):
all_results = []
for concurrency in args.max_concurrency:
result = await run_benchmark(
host=args.host,
port=args.port,
num_prompts=args.num_prompts,
max_concurrency=concurrency,
num_warmups=args.num_warmups,
voice=args.voice,
language=args.language,
stream=args.stream,
)
result.config_name = args.config_name
all_results.append(asdict(result))
result_dir = Path(args.result_dir)
result_dir.mkdir(parents=True, exist_ok=True)
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
with open(result_file, "w") as f:
json.dump(all_results, f, indent=2)
print(f"Results saved to {result_file}")
def parse_args():
parser = argparse.ArgumentParser(description="Qwen3-TTS async_chunk benchmark client")
parser.add_argument("--host", type=str, default="127.0.0.1")
parser.add_argument("--port", type=int, default=8000)
parser.add_argument("--num-prompts", type=int, default=50)
parser.add_argument("--max-concurrency", type=int, nargs="+", default=[1, 10])
parser.add_argument("--num-warmups", type=int, default=3)
parser.add_argument("--voice", type=str, default="vivian")
parser.add_argument("--language", type=str, default="English")
parser.add_argument("--stream", action="store_true", default=True)
parser.add_argument("--no-stream", dest="stream", action="store_false")
parser.add_argument("--config-name", type=str, default="async_chunk_on")
parser.add_argument("--result-dir", type=str, default="results")
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
asyncio.run(main(args))
@@ -0,0 +1,371 @@
"""Benchmark client for Qwen3-TTS via /v1/audio/speech endpoint.
Measures TTFP (Time-to-First-Packet), E2E latency, and RTF (Real-Time Factor)
across configurable concurrency levels. Saves results as JSON for plotting.
Usage:
python bench_tts_serve.py \
--host 127.0.0.1 --port 8000 \
--num-prompts 50 \
--max-concurrency 1 4 10 \
--result-dir results/
"""
import argparse
import asyncio
import json
import time
from dataclasses import asdict, dataclass, field
from datetime import datetime
from pathlib import Path
import aiohttp
import numpy as np
from tqdm.asyncio import tqdm
PROMPTS = [
"Hello, welcome to the voice synthesis benchmark test.",
"She said she would be here by noon, but nobody showed up.",
"The quick brown fox jumps over the lazy dog near the riverbank.",
"I can't believe how beautiful the sunset looks from up here on the mountain.",
"Please remember to bring your identification documents to the appointment tomorrow morning.",
"Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
"The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
"After the meeting, we should discuss the quarterly results and plan for the next phase.",
"Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
"The train leaves at half past seven, so we need to arrive at the station before then.",
"Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
"It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
]
REF_AUDIO = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav"
REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
INSTRUCT = "Speak in an incredulous tone, but with a hint of panic beginning to creep into your voice."
@dataclass
class RequestResult:
success: bool = False
ttfp: float = 0.0 # Time to first audio packet (seconds)
e2e: float = 0.0 # End-to-end latency (seconds)
audio_bytes: int = 0 # Total audio bytes received
audio_duration: float = 0.0 # Audio duration in seconds (estimated from PCM)
rtf: float = 0.0 # Real-time factor = e2e / audio_duration
prompt: str = ""
error: str = ""
@dataclass
class BenchmarkResult:
config_name: str = ""
concurrency: int = 0
num_prompts: int = 0
completed: int = 0
failed: int = 0
duration_s: float = 0.0
# TTFP stats (ms)
mean_ttfp_ms: float = 0.0
median_ttfp_ms: float = 0.0
std_ttfp_ms: float = 0.0
p90_ttfp_ms: float = 0.0
p95_ttfp_ms: float = 0.0
p99_ttfp_ms: float = 0.0
# E2E stats (ms)
mean_e2e_ms: float = 0.0
median_e2e_ms: float = 0.0
std_e2e_ms: float = 0.0
p90_e2e_ms: float = 0.0
p95_e2e_ms: float = 0.0
p99_e2e_ms: float = 0.0
# RTF stats
mean_rtf: float = 0.0
median_rtf: float = 0.0
std_rtf: float = 0.0
p99_rtf: float = 0.0
# Audio stats
mean_audio_duration_s: float = 0.0
total_audio_duration_s: float = 0.0
audio_throughput: float = 0.0 # audio_duration / wall_time
request_throughput: float = 0.0 # requests / second
# Per-request details
per_request: list = field(default_factory=list)
def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width: int = 2) -> float:
"""Convert raw PCM byte count to duration in seconds."""
num_samples = num_bytes / sample_width
return num_samples / sample_rate
def create_payload(
prompt: str, task_type: str = "CustomVoice", voice: str = "vivian", language: str = "English"
) -> dict:
payload = {
"input": prompt,
"language": language,
"stream": True,
"response_format": "pcm",
"task_type": task_type,
}
if task_type == "Base":
payload["ref_audio"] = REF_AUDIO
payload["ref_text"] = REF_TEXT
elif task_type == "CustomVoice":
payload["voice"] = voice
elif task_type == "VoiceDesign":
payload["instructions"] = INSTRUCT
return payload
async def send_tts_request(
session: aiohttp.ClientSession,
api_url: str,
prompt: str,
task_type: str = "CustomVoice",
voice: str = "vivian",
language: str = "English",
pbar: tqdm | None = None,
) -> RequestResult:
"""Send a streaming TTS request and measure latency metrics."""
payload = create_payload(prompt, task_type, voice, language)
result = RequestResult(prompt=prompt)
st = time.perf_counter()
try:
async with session.post(api_url, json=payload) as response:
if response.status != 200:
result.error = f"HTTP {response.status}: {await response.text()}"
result.success = False
return result
first_chunk = True
total_bytes = 0
async for chunk in response.content.iter_any():
if first_chunk and len(chunk) > 0:
result.ttfp = time.perf_counter() - st
first_chunk = False
total_bytes += len(chunk)
result.e2e = time.perf_counter() - st
result.audio_bytes = total_bytes
result.audio_duration = pcm_bytes_to_duration(total_bytes)
if result.audio_duration > 0:
result.rtf = result.e2e / result.audio_duration
result.success = True
except Exception as e:
result.error = str(e)
result.success = False
result.e2e = time.perf_counter() - st
if pbar:
pbar.update(1)
return result
async def run_benchmark(
host: str,
port: int,
num_prompts: int,
max_concurrency: int,
num_warmups: int = 3,
task_type: str = "CustomVoice",
voice: str = "vivian",
language: str = "English",
) -> BenchmarkResult:
"""Run benchmark at a given concurrency level."""
api_url = f"http://{host}:{port}/v1/audio/speech"
connector = aiohttp.TCPConnector(
limit=max_concurrency,
limit_per_host=max_concurrency,
keepalive_timeout=60,
)
session = aiohttp.ClientSession(
connector=connector,
timeout=aiohttp.ClientTimeout(total=600),
)
# Warmup
if num_warmups > 0:
print(f" Warming up with {num_warmups} requests...")
warmup_tasks = []
for i in range(num_warmups):
prompt = PROMPTS[i % len(PROMPTS)]
warmup_tasks.append(send_tts_request(session, api_url, prompt, task_type, voice, language))
await asyncio.gather(*warmup_tasks)
print(" Warmup done.")
# Build request list
request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
# Run benchmark
print(f" Running {num_prompts} requests with concurrency={max_concurrency}...")
semaphore = asyncio.Semaphore(max_concurrency)
pbar = tqdm(total=num_prompts, desc=f" concurrency={max_concurrency}")
async def limited_request(prompt):
async with semaphore:
return await send_tts_request(session, api_url, prompt, task_type, voice, language, pbar)
start_time = time.perf_counter()
tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
results: list[RequestResult] = await asyncio.gather(*tasks)
duration = time.perf_counter() - start_time
pbar.close()
await session.close()
# Compute stats
successful = [r for r in results if r.success]
failed = [r for r in results if not r.success]
bench = BenchmarkResult(
concurrency=max_concurrency,
num_prompts=num_prompts,
completed=len(successful),
failed=len(failed),
duration_s=duration,
)
if successful:
ttfps = [r.ttfp * 1000 for r in successful] # convert to ms
e2es = [r.e2e * 1000 for r in successful]
rtfs = [r.rtf for r in successful]
audio_durs = [r.audio_duration for r in successful]
bench.mean_ttfp_ms = float(np.mean(ttfps))
bench.median_ttfp_ms = float(np.median(ttfps))
bench.std_ttfp_ms = float(np.std(ttfps))
bench.p90_ttfp_ms = float(np.percentile(ttfps, 90))
bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
bench.p99_ttfp_ms = float(np.percentile(ttfps, 99))
bench.mean_e2e_ms = float(np.mean(e2es))
bench.median_e2e_ms = float(np.median(e2es))
bench.std_e2e_ms = float(np.std(e2es))
bench.p90_e2e_ms = float(np.percentile(e2es, 90))
bench.p95_e2e_ms = float(np.percentile(e2es, 95))
bench.p99_e2e_ms = float(np.percentile(e2es, 99))
bench.mean_rtf = float(np.mean(rtfs))
bench.median_rtf = float(np.median(rtfs))
bench.std_rtf = float(np.std(rtfs))
bench.p99_rtf = float(np.percentile(rtfs, 99))
bench.mean_audio_duration_s = float(np.mean(audio_durs))
bench.total_audio_duration_s = float(np.sum(audio_durs))
bench.audio_throughput = bench.total_audio_duration_s / duration
bench.request_throughput = len(successful) / duration
bench.per_request = [
{
"ttfp_ms": r.ttfp * 1000,
"e2e_ms": r.e2e * 1000,
"rtf": r.rtf,
"audio_duration_s": r.audio_duration,
"prompt": r.prompt,
}
for r in successful
]
# Print summary in standardized performance template
W = 50
print("")
print(f"{'=' * W}")
print(f"{'Serving Benchmark Result':^{W}}")
print(f"{'=' * W}")
print(f"{'Successful requests:':<40}{bench.completed:<10}")
print(f"{'Failed requests:':<40}{bench.failed:<10}")
print(f"{'Maximum request concurrency:':<40}{max_concurrency:<10}")
print(f"{'Benchmark duration (s):':<40}{duration:<10.2f}")
print(f"{'Request throughput (req/s):':<40}{bench.request_throughput:<10.2f}")
print(f"{'-' * W}")
print(f"{'End-to-end Latency':^{W}}")
print(f"{'-' * W}")
print(f"{'Mean E2EL (ms):':<40}{bench.mean_e2e_ms:<10.2f}")
print(f"{'Median E2EL (ms):':<40}{bench.median_e2e_ms:<10.2f}")
print(f"{'P99 E2EL (ms):':<40}{bench.p99_e2e_ms:<10.2f}")
print(f"{'=' * W}")
print(f"{'Audio Result':^{W}}")
print(f"{'=' * W}")
print(f"{'Total audio duration generated (s):':<40}{bench.total_audio_duration_s:<10.2f}")
print(f"{'Audio throughput (audio duration/s):':<40}{bench.audio_throughput:<10.2f}")
print(f"{'-' * W}")
print(f"{'Time to First Packet':^{W}}")
print(f"{'-' * W}")
print(f"{'Mean AUDIO_TTFP (ms):':<40}{bench.mean_ttfp_ms:<10.2f}")
print(f"{'Median AUDIO_TTFP (ms):':<40}{bench.median_ttfp_ms:<10.2f}")
print(f"{'P99 AUDIO_TTFP (ms):':<40}{bench.p99_ttfp_ms:<10.2f}")
print(f"{'-' * W}")
print(f"{'Real Time Factor':^{W}}")
print(f"{'-' * W}")
print(f"{'Mean AUDIO_RTF:':<40}{bench.mean_rtf:<10.3f}")
print(f"{'Median AUDIO_RTF:':<40}{bench.median_rtf:<10.3f}")
print(f"{'P99 AUDIO_RTF:':<40}{bench.p99_rtf:<10.3f}")
print(f"{'=' * W}")
print("")
if failed:
for r in failed[:3]:
print(f" [ERROR] {r.error[:200]}")
return bench
async def main(args):
all_results = []
for concurrency in args.max_concurrency:
result = await run_benchmark(
host=args.host,
port=args.port,
num_prompts=args.num_prompts,
max_concurrency=concurrency,
num_warmups=args.num_warmups,
task_type=args.task_type,
voice=args.voice,
language=args.language,
)
result.config_name = args.config_name
all_results.append(asdict(result))
# Save results
result_dir = Path(args.result_dir)
result_dir.mkdir(parents=True, exist_ok=True)
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
with open(result_file, "w") as f:
json.dump(all_results, f, indent=2)
print(f"Results saved to {result_file}")
return all_results
def parse_args():
parser = argparse.ArgumentParser(description="Qwen3-TTS Benchmark Client")
parser.add_argument("--host", type=str, default="127.0.0.1")
parser.add_argument("--port", type=int, default=8000)
parser.add_argument("--num-prompts", type=int, default=50, help="Number of prompts per concurrency level")
parser.add_argument( # noqa: E501
"--max-concurrency", type=int, nargs="+", default=[1, 4, 10], help="Concurrency levels to test"
)
parser.add_argument("--num-warmups", type=int, default=3)
parser.add_argument("--task-type", type=str, default="CustomVoice", choices=["CustomVoice", "VoiceDesign", "Base"])
parser.add_argument("--voice", type=str, default="vivian")
parser.add_argument("--language", type=str, default="English")
parser.add_argument(
"--config-name", type=str, default="async_chunk", help="Label for this config (used in filenames)"
)
parser.add_argument("--result-dir", type=str, default="results")
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
asyncio.run(main(args))
@@ -0,0 +1,249 @@
"""Plot TTFP comparison: async_chunk off vs on.
Generates a bar chart with improvement arrows, matching the Qwen3-Omni
async_chunk benchmark figure style.
Usage:
python plot_async_chunk.py \
--off results/bench_async_chunk_off_*.json \
--on results/bench_async_chunk_on_*.json \
--output results/qwen3_tts_async_chunk_ttfp.png
# Also supports E2E and RTF metrics:
python plot_async_chunk.py \
--off results/bench_async_chunk_off_*.json \
--on results/bench_async_chunk_on_*.json \
--metric e2e \
--output results/qwen3_tts_async_chunk_e2e.png
"""
import argparse
import json
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
METRIC_CONFIG = {
"ttfp": {
"key": "mean_ttfp_ms",
"ylabel": "TTFP (s)",
"title": "TTFP (Time to First Audio Packet) - Qwen3-TTS, by concurrency",
"to_seconds": True,
},
"e2e": {
"key": "mean_e2e_ms",
"ylabel": "E2E (s)",
"title": "E2E Latency - Qwen3-TTS, by concurrency",
"to_seconds": True,
},
"rtf": {
"key": "mean_rtf",
"ylabel": "RTF",
"title": "Real-Time Factor - Qwen3-TTS, by concurrency",
"to_seconds": False,
},
}
def load_results(path: str) -> list[dict]:
with open(path) as f:
return json.load(f)
def plot_ttfp_comparison(
off_results: list[dict],
on_results: list[dict],
metric: str,
output_path: str,
title_override: str | None = None,
):
cfg = METRIC_CONFIG[metric]
key = cfg["key"]
to_seconds = cfg["to_seconds"]
off_map = {r["concurrency"]: r for r in off_results}
on_map = {r["concurrency"]: r for r in on_results}
concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
off_vals = []
on_vals = []
for c in concurrencies:
v_off = off_map[c][key]
v_on = on_map[c][key]
if to_seconds:
v_off /= 1000.0
v_on /= 1000.0
off_vals.append(v_off)
on_vals.append(v_on)
fig, ax = plt.subplots(figsize=(8, 6))
x = np.arange(len(concurrencies))
width = 0.3
ax.bar(x - width / 2, off_vals, width, label="async_chunk off", color="#87CEEB", edgecolor="none")
ax.bar(x + width / 2, on_vals, width, label="async_chunk on", color="#FFF8DC", edgecolor="#DDD8B8")
# Draw improvement arrows and labels
for i in range(len(concurrencies)):
v_off = off_vals[i]
v_on = on_vals[i]
if v_on > 0:
improvement = v_off / v_on
else:
improvement = float("inf")
# Arrow from top of off-bar to top of on-bar
arrow_start_x = x[i] - width / 2
arrow_start_y = v_off * 0.95
arrow_end_x = x[i] + width / 2
arrow_end_y = v_on * 1.05
ax.annotate(
"",
xy=(arrow_end_x, arrow_end_y),
xytext=(arrow_start_x, arrow_start_y),
arrowprops=dict(arrowstyle="->", color="red", lw=1.5),
)
# Improvement label
label_x = (arrow_start_x + arrow_end_x) / 2
label_y = arrow_start_y + (v_off - v_on) * 0.15
ax.text(
label_x,
label_y,
f"{improvement:.1f}x improvement",
ha="center",
va="bottom",
fontsize=10,
color="red",
fontweight="bold",
)
title = title_override or cfg["title"]
ax.set_title(title, fontsize=13, fontweight="bold")
ax.set_ylabel(cfg["ylabel"], fontsize=12)
ax.set_xlabel("Max concurrency", fontsize=12)
ax.set_xticks(x)
ax.set_xticklabels([str(c) for c in concurrencies])
ax.set_yscale("log")
ax.legend(loc="upper left", fontsize=11)
ax.grid(axis="y", alpha=0.3, linestyle="--")
ax.set_axisbelow(True)
plt.tight_layout()
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
plt.savefig(output_path, dpi=150, bbox_inches="tight")
print(f"Plot saved to {output_path}")
plt.close()
def plot_all_metrics(off_results: list[dict], on_results: list[dict], output_path: str):
"""Generate a 1x3 subplot with TTFP, E2E, and RTF comparisons."""
off_map = {r["concurrency"]: r for r in off_results}
on_map = {r["concurrency"]: r for r in on_results}
concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
fig, axes = plt.subplots(1, 3, figsize=(18, 6))
fig.suptitle("Qwen3-TTS: async_chunk on vs off", fontsize=15, fontweight="bold")
for ax, metric in zip(axes, ["ttfp", "e2e", "rtf"]):
cfg = METRIC_CONFIG[metric]
key = cfg["key"]
to_seconds = cfg["to_seconds"]
off_vals = []
on_vals = []
for c in concurrencies:
v_off = off_map[c][key]
v_on = on_map[c][key]
if to_seconds:
v_off /= 1000.0
v_on /= 1000.0
off_vals.append(v_off)
on_vals.append(v_on)
x = np.arange(len(concurrencies))
width = 0.3
ax.bar(x - width / 2, off_vals, width, label="async_chunk off", color="#87CEEB")
ax.bar(x + width / 2, on_vals, width, label="async_chunk on", color="#FFF8DC", edgecolor="#DDD8B8")
for i in range(len(concurrencies)):
if on_vals[i] > 0:
improvement = off_vals[i] / on_vals[i]
ax.annotate(
"",
xy=(x[i] + width / 2, on_vals[i] * 1.05),
xytext=(x[i] - width / 2, off_vals[i] * 0.95),
arrowprops=dict(arrowstyle="->", color="red", lw=1.5),
)
label_y = off_vals[i] * 0.85
ax.text(x[i], label_y, f"{improvement:.1f}x", ha="center", fontsize=10, color="red", fontweight="bold")
ax.set_title(cfg["title"].split(" - ")[0], fontsize=12, fontweight="bold")
ax.set_ylabel(cfg["ylabel"], fontsize=11)
ax.set_xlabel("Max concurrency", fontsize=11)
ax.set_xticks(x)
ax.set_xticklabels([str(c) for c in concurrencies])
if metric != "rtf":
ax.set_yscale("log")
ax.legend(fontsize=9)
ax.grid(axis="y", alpha=0.3, linestyle="--")
ax.set_axisbelow(True)
plt.tight_layout()
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
plt.savefig(output_path, dpi=150, bbox_inches="tight")
print(f"Plot saved to {output_path}")
plt.close()
def print_table(off_results: list[dict], on_results: list[dict]):
off_map = {r["concurrency"]: r for r in off_results}
on_map = {r["concurrency"]: r for r in on_results}
concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
print("\n## Benchmark Results: async_chunk off vs on\n")
print("| Metric | Concurrency | async_chunk off | async_chunk on | Improvement |")
print("| --- | --- | --- | --- | --- |")
for name, key, fmt in [
("TTFP (ms)", "mean_ttfp_ms", ".1f"),
("E2E (ms)", "mean_e2e_ms", ".1f"),
("RTF", "mean_rtf", ".3f"),
("Throughput", "audio_throughput", ".2f"),
]:
for c in concurrencies:
v_off = off_map[c].get(key, 0)
v_on = on_map[c].get(key, 0)
if v_on > 0 and key != "audio_throughput":
ratio = f"{v_off / v_on:.1f}x"
elif v_off > 0 and key == "audio_throughput":
ratio = f"{v_on / v_off:.1f}x"
else:
ratio = "N/A"
print(f"| {name} | {c} | {v_off:{fmt}} | {v_on:{fmt}} | {ratio} |")
def parse_args():
parser = argparse.ArgumentParser(description="Plot async_chunk comparison for Qwen3-TTS")
parser.add_argument("--off", type=str, required=True, help="JSON results for async_chunk off")
parser.add_argument("--on", type=str, required=True, help="JSON results for async_chunk on")
parser.add_argument("--metric", type=str, default="ttfp", choices=["ttfp", "e2e", "rtf", "all"])
parser.add_argument("--output", type=str, default="results/qwen3_tts_async_chunk.png")
parser.add_argument("--title", type=str, default=None, help="Custom title override")
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
off_results = load_results(args.off)
on_results = load_results(args.on)
print_table(off_results, on_results)
if args.metric == "all":
plot_all_metrics(off_results, on_results, args.output)
else:
plot_ttfp_comparison(off_results, on_results, args.metric, args.output, args.title)
+169
View File
@@ -0,0 +1,169 @@
#!/bin/bash
# Qwen3-TTS async_chunk on vs off Benchmark
#
# Starts two servers (async_chunk on and off), benchmarks both,
# and generates comparison plots.
#
# Usage:
# bash run_async_chunk_benchmark.sh
#
# Environment variables:
# GPU_DEVICE - GPU index (default: 0)
# MODEL - Model path (default: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice)
# NUM_PROMPTS - Prompts per concurrency level (default: 50)
# CONCURRENCY - Space-separated concurrency levels (default: "1 10")
# PORT_ON - Port for async_chunk on server (default: 8000)
# PORT_OFF - Port for async_chunk off server (default: 8001)
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
cd "$PROJECT_ROOT"
GPU_DEVICE="${GPU_DEVICE:-0}"
MODEL="${MODEL:-Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice}"
NUM_PROMPTS="${NUM_PROMPTS:-50}"
CONCURRENCY="${CONCURRENCY:-1 10}"
NUM_WARMUPS="${NUM_WARMUPS:-3}"
PORT_ON="${PORT_ON:-8000}"
PORT_OFF="${PORT_OFF:-8001}"
RESULT_DIR="${SCRIPT_DIR}/results"
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
# The bundled ``vllm_omni/deploy/qwen3_tts.yaml`` is auto-loaded by the model
# registry; no ``--deploy-config`` flag needed on the default (ON) path.
# async_chunk OFF is selected by the ``--no-async-chunk`` CLI flag —
# the single ``qwen3_tts`` pipeline dispatches to the end-to-end codec
# processor when ``deploy.async_chunk`` is false.
mkdir -p "${RESULT_DIR}"
echo "============================================================"
echo " Qwen3-TTS async_chunk Benchmark"
echo "============================================================"
echo " GPU: ${GPU_DEVICE}"
echo " Model: ${MODEL}"
echo " Prompts: ${NUM_PROMPTS}"
echo " Concurrency: ${CONCURRENCY}"
echo " Port (on/off): ${PORT_ON} / ${PORT_OFF}"
echo " Results: ${RESULT_DIR}"
echo "============================================================"
cleanup() {
echo "Cleaning up servers..."
kill "$PID_ON" 2>/dev/null || true
kill "$PID_OFF" 2>/dev/null || true
wait "$PID_ON" 2>/dev/null || true
wait "$PID_OFF" 2>/dev/null || true
}
trap cleanup EXIT
wait_for_server() {
local port=$1
local name=$2
local max_wait=300
local elapsed=0
echo "Waiting for ${name} server on port ${port}..."
while ! curl -s "http://localhost:${port}/health" >/dev/null 2>&1; do
sleep 5
elapsed=$((elapsed + 5))
if [ $elapsed -ge $max_wait ]; then
echo "ERROR: ${name} server failed to start within ${max_wait}s"
exit 1
fi
done
echo "${name} server ready (${elapsed}s)"
}
# ---- Phase 1: Start async_chunk ON server ----
echo ""
echo "[Phase 1] Starting async_chunk ON server on port ${PORT_ON}..."
CUDA_VISIBLE_DEVICES=${GPU_DEVICE} vllm-omni serve "${MODEL}" \
--host 0.0.0.0 --port "${PORT_ON}" \
--trust-remote-code --enforce-eager --omni \
> "${RESULT_DIR}/server_on_${TIMESTAMP}.log" 2>&1 &
PID_ON=$!
wait_for_server "${PORT_ON}" "async_chunk_on"
echo "[Phase 1] Benchmarking async_chunk ON..."
# shellcheck disable=SC2086
python "${SCRIPT_DIR}/bench_async_chunk.py" \
--host 127.0.0.1 --port "${PORT_ON}" \
--config-name "async_chunk_on" \
--num-prompts "${NUM_PROMPTS}" \
--max-concurrency ${CONCURRENCY} \
--num-warmups "${NUM_WARMUPS}" \
--result-dir "${RESULT_DIR}"
echo "[Phase 1] Stopping async_chunk ON server..."
kill "$PID_ON" 2>/dev/null || true
wait "$PID_ON" 2>/dev/null || true
sleep 5
# ---- Phase 2: Start async_chunk OFF server ----
echo ""
echo "[Phase 2] Starting async_chunk OFF server on port ${PORT_OFF}..."
CUDA_VISIBLE_DEVICES=${GPU_DEVICE} vllm-omni serve "${MODEL}" \
--no-async-chunk \
--host 0.0.0.0 --port "${PORT_OFF}" \
--trust-remote-code --enforce-eager --omni \
> "${RESULT_DIR}/server_off_${TIMESTAMP}.log" 2>&1 &
PID_OFF=$!
wait_for_server "${PORT_OFF}" "async_chunk_off"
echo "[Phase 2] Benchmarking async_chunk OFF (non-streaming)..."
# shellcheck disable=SC2086
python "${SCRIPT_DIR}/bench_async_chunk.py" \
--host 127.0.0.1 --port "${PORT_OFF}" \
--config-name "async_chunk_off" \
--num-prompts "${NUM_PROMPTS}" \
--max-concurrency ${CONCURRENCY} \
--num-warmups "${NUM_WARMUPS}" \
--no-stream \
--result-dir "${RESULT_DIR}"
echo "[Phase 2] Stopping async_chunk OFF server..."
kill "$PID_OFF" 2>/dev/null || true
wait "$PID_OFF" 2>/dev/null || true
# ---- Phase 3: Plot results ----
echo ""
echo "[Phase 3] Generating plots..."
# Find the latest result files
RESULT_ON=$(ls -t "${RESULT_DIR}"/bench_async_chunk_on_*.json 2>/dev/null | head -1)
RESULT_OFF=$(ls -t "${RESULT_DIR}"/bench_async_chunk_off_*.json 2>/dev/null | head -1)
if [ -z "$RESULT_ON" ] || [ -z "$RESULT_OFF" ]; then
echo "ERROR: Could not find result files. Check logs in ${RESULT_DIR}/"
exit 1
fi
echo " ON results: ${RESULT_ON}"
echo " OFF results: ${RESULT_OFF}"
# TTFP comparison (main figure)
python "${SCRIPT_DIR}/plot_async_chunk.py" \
--off "${RESULT_OFF}" \
--on "${RESULT_ON}" \
--metric ttfp \
--output "${RESULT_DIR}/qwen3_tts_async_chunk_ttfp.png"
# All metrics comparison
python "${SCRIPT_DIR}/plot_async_chunk.py" \
--off "${RESULT_OFF}" \
--on "${RESULT_ON}" \
--metric all \
--output "${RESULT_DIR}/qwen3_tts_async_chunk_all.png"
echo ""
echo "============================================================"
echo " Benchmark complete!"
echo " Results: ${RESULT_DIR}/"
echo " Plots:"
echo " - ${RESULT_DIR}/qwen3_tts_async_chunk_ttfp.png"
echo " - ${RESULT_DIR}/qwen3_tts_async_chunk_all.png"
echo "============================================================"
-227
View File
@@ -1,227 +0,0 @@
# TTS Universal Benchmark
A model-agnostic serving benchmark for TTS models in vllm-omni. One CLI
(`bench_tts.py`) + one YAML registry (`model_configs.yaml`) drive perf and
quality runs for every registered checkpoint: **Qwen3-TTS** (Base / CustomVoice)
and **VoxCPM2** today, more to come.
The same three task types — `voice_clone`, `default_voice`, `voice_design` —
are wired into both the manual CLI and the DFX nightly CI matrix
(`tests/dfx/perf/tests/test_tts.json`).
## Quick start
### 1. Start the server
```bash
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni --port 8000
```
The server auto-loads its Deploy YAML from `vllm_omni/deploy/qwen3_tts.yaml`
(Pipeline + Deploy schema introduced in #2383). No `--stage-configs-path` or
`--deploy-config` flag is needed for any registered model.
### 2. Run the benchmark (`vllm bench serve --omni`)
The primary, directly-controllable path. Copy-paste one of these and tweak
any bench flag (sampling params, endpoint, extra body, warmups, etc.):
#### voice_clone (Qwen3-TTS-Base, seed-tts dataset)
```bash
vllm bench serve --omni \
--host 127.0.0.1 --port 8000 \
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
--backend openai-audio-speech \
--endpoint /v1/audio/speech \
--dataset-name seed-tts \
--dataset-path /path/to/seed-tts-eval \
--seed-tts-locale en \
--num-prompts 20 --num-warmups 2 \
--extra-body '{"task_type":"Base"}' \
--max-concurrency 1 --request-rate inf \
--percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration \
--save-result --result-dir ./results
```
#### default_voice (Qwen3-TTS-CustomVoice, bundled seed_tts_smoke)
```bash
vllm bench serve --omni \
--host 127.0.0.1 --port 8000 \
--model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
--backend openai-audio-speech \
--endpoint /v1/audio/speech \
--dataset-name seed-tts-text \
--dataset-path benchmarks/build_dataset/seed_tts_smoke \
--seed-tts-locale en \
--num-prompts 20 --num-warmups 2 \
--extra-body '{"voice":"Vivian","language":"English","task_type":"CustomVoice"}' \
--max-concurrency 1 --request-rate inf \
--percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration \
--save-result --result-dir ./results
```
#### voice_design (Qwen3-TTS-CustomVoice, bundled seed_tts_design)
```bash
vllm bench serve --omni \
--host 127.0.0.1 --port 8000 \
--model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
--backend openai-audio-speech \
--endpoint /v1/audio/speech \
--dataset-name seed-tts-design \
--dataset-path benchmarks/build_dataset/seed_tts_design \
--seed-tts-locale en \
--num-prompts 20 --num-warmups 2 \
--extra-body '{"task_type":"VoiceDesign","language":"English"}' \
--max-concurrency 1 --request-rate inf \
--percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration \
--save-result --result-dir ./results
```
#### Add WER / SIM / UTMOS to any of the above
Append `--seed-tts-wer-eval` (and optionally `SEED_TTS_EVAL_DEVICE=cuda:0`
in the env, per PR #2558). This triggers the seed-tts-eval protocol:
Whisper-large-v3 ASR → WER, WavLM embeddings → SIM, balacoon/utmos → UTMOS.
### 3. Convenience wrapper (`bench_tts.py`)
If you're running the **canonical** configuration for a registered model,
`bench_tts.py` loads the right defaults from `model_configs.yaml` and
emits the exact `vllm bench serve --omni` command above — useful for
concurrency sweeps and multi-task runs:
```bash
# Smallest smoke — 5 prompts, concurrency=1
python benchmarks/tts/bench_tts.py \
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
--task voice_clone \
--dataset-path /path/to/seed-tts-eval \
--concurrency 1 --num-prompts 5 \
--output-dir ./results
# Full concurrency sweep
python benchmarks/tts/bench_tts.py \
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
--task voice_clone \
--dataset-path /path/to/seed-tts-eval \
--concurrency 1 2 4 8 16 32 \
--num-prompts 20 \
--output-dir ./results
# With WER / SIM / UTMOS quality eval (adds ASR + embedding compute)
python benchmarks/tts/bench_tts.py \
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
--task voice_clone \
--dataset-path /path/to/seed-tts-eval \
--wer-eval \
--concurrency 4 --num-prompts 200 \
--output-dir ./results
```
### 4. Plot a sweep
```bash
python benchmarks/tts/plot_results.py \
--results ./results/*.json \
--output ./results/curve.png
```
Outputs TTFP / RTF / throughput curves (and a markdown table) for every
`(task, concurrency)` combination in the result set.
## Task types
| Task | Dataset | Request body | Checkpoints that support it |
|-----------------|-------------------|-----------------------------------------------------|------------------------------------------|
| `voice_clone` | `seed-tts` | `ref_audio` + `ref_text` + `task_type=Base` | `Qwen3-TTS-*-Base`, `VoxCPM2` |
| `default_voice` | `seed-tts-text` | `voice=Vivian` + `task_type=CustomVoice` | `Qwen3-TTS-*-CustomVoice` |
| `voice_design` | `seed-tts-design` | `instructions=<natural-language description>` + `task_type=VoiceDesign` | `Qwen3-TTS-*-CustomVoice` |
**`-CustomVoice` checkpoints do NOT ship `speaker_encoder` weights**, so
voice_clone requests raise `ValueError` at model runtime. Use `-Base` for
voice_clone.
## Adding a new TTS model
Drop an entry into `model_configs.yaml` — no Python changes required:
```yaml
models:
<org>/<model-id>:
supported_tasks: [voice_clone] # or default_voice / voice_design
backend: openai-audio-speech # vllm bench serve backend
endpoint: /v1/audio/speech # OpenAI-compatible endpoint
task_extra_body: # merged into every request's body
voice_clone:
task_type: Base
```
Then add the model's Deploy YAML under `vllm_omni/deploy/<model>.yaml`
(Pipeline + Deploy schema) and it's immediately benchable.
## Datasets
| Dataset | Bundled? | Format | Source |
|--------------------|----------|-------------------|----------------------------------------------------------------|
| `seed-tts-design` | ✅ | 5-field meta.lst | `benchmarks/build_dataset/seed_tts_design/en/meta.lst` (20 prompts) |
| `seed_tts_smoke` | ✅ | 4-field meta.lst | `benchmarks/build_dataset/seed_tts_smoke/en/meta.lst` (20 text-only) |
| `seed-tts` | ❌ | 4-field meta.lst + WAVs | Google-Drive: [BytedanceSpeech/seed-tts-eval][seedtts] (~1.2 GB) |
| `seed-tts-text` | ❌ | 4-field meta.lst | Same archive as `seed-tts` (wav column unused) |
[seedtts]: https://github.com/BytedanceSpeech/seed-tts-eval
For manual voice_clone / default_voice runs against the full corpus, follow
`benchmarks/build_dataset/download_process_data_seedtts.md` and point
`--dataset-path` at the extracted `seedtts_testset` directory.
## DFX nightly CI
`tests/dfx/perf/tests/test_tts.json` wires three perf regimes plus quality:
| eval_phase | concurrency | purpose | Baseline metrics |
|---------------|-------------|---------------------------------------------------------|-----------------------------------------|
| `latency` | 1 | Single-request TTFP / RTF SLO | `median_audio_ttfp_ms`, `median_audio_rtf` |
| `throughput` | 8 | Codec-batching cliff sentinel (PDF #272 concurrency≥8) | `median_audio_ttfp_ms`, `median_audio_rtf` |
| `quality` | 4 | WER / SIM / UTMOS regression (disabled in CI by default)| `mean_audio_rtf` |
Why `median_*` for latency/throughput and `mean_*` for quality: latency
distributions have cold-start tails that drag the mean; quality aggregates
over 200 prompts so single-request outliers don't matter.
Quality entries are `enabled: false` in CI because seed-tts-eval is not
staged in the Buildkite container (matches the precedent in
PR #2558 — quality runs are manual / release-validation, not nightly).
## Concurrency cliff regression sentinel
Observed on H20-3e, Qwen3-TTS-1.7B (measured pre-merge on this branch):
| Task | Model | c=1 | c=4 | **c=8** | c=16 | c=32 |
|---------------|---------------|--------|--------|------------|--------|--------|
| voice_clone | 1.7B-Base | RTF 0.15 / TTFP 165ms | 0.28 / 412ms | **0.49 / 1701ms** | 0.72 / 3355ms | 0.77 / 3772ms |
| voice_design | 1.7B-CustomVoice | RTF 0.08 / TTFP 53ms | 0.11 / 154ms | **0.21 / 872ms** | 0.33 / 1801ms | 0.38 / 1989ms |
Both models show a **4–6× TTFP jump from c=4 to c=8** while audio throughput
saturates around c=4–8 — the codec-bs=1 bottleneck documented in
vllm-project/vllm-omni#272. The `throughput` CI regime at c=8 is the
sentinel for regressions in this area.
## File layout
```
benchmarks/tts/
├── README.md (this file)
├── bench_tts.py CLI — serve-mode benchmark driver
├── bench_voxcpm_offline.py CLI — offline VoxCPM benchmark (sync + streaming)
├── plot_results.py Generate per-task / per-concurrency curves
└── model_configs.yaml Model registry (supported tasks + extra body)
```
## Related
- Upstream seed-tts-eval integration: vllm-project/vllm-omni#2558
- Pipeline + Deploy schema: vllm-project/vllm-omni#2383
- Concurrency cliff RFC: vllm-project/vllm-omni#272
-308
View File
@@ -1,308 +0,0 @@
#!/usr/bin/env python3
"""Universal TTS benchmark CLI for vllm-omni.
Runs ``vllm bench serve --omni`` with model-aware defaults loaded from
``model_configs.yaml``. Supports Qwen3-TTS, VoxCPM2, and any future TTS
model registered in the config file -- no code changes needed to add models.
Usage::
python benchmarks/tts/bench_tts.py \\
--model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \\
--task voice_clone \\
--locale en \\
--concurrency 1 4 \\
--num-prompts 20 \\
--dataset-path /path/to/seed-tts-eval \\
--host localhost --port 8000
See ``--help`` for full option list.
"""
from __future__ import annotations
import argparse
import json
import math
import os
import subprocess
import sys
from datetime import datetime
from pathlib import Path
from typing import Any
import yaml
def _vllm_omni_bin() -> str:
"""Return the vllm-omni (or vllm) binary co-located with the current Python."""
bin_dir = Path(sys.executable).parent
for candidate in ("vllm-omni", "vllm"):
p = bin_dir / candidate
if p.is_file():
return str(p)
return "vllm-omni" # fall back and let the shell resolve it
_REPO_ROOT = Path(__file__).resolve().parent.parent.parent
_SCRIPT_DIR = Path(__file__).resolve().parent
_DEFAULT_MODEL_CONFIGS = _SCRIPT_DIR / "model_configs.yaml"
# Maps task name to the dataset_name used with vllm bench serve
_TASK_TO_DATASET: dict[str, str] = {
"voice_clone": "seed-tts",
"default_voice": "seed-tts-text",
"voice_design": "seed-tts-design",
}
# Default design dataset path (bundled with the repo)
_DEFAULT_DESIGN_DATASET_PATH = str(_REPO_ROOT / "benchmarks" / "build_dataset" / "seed_tts_design")
def load_model_configs(path: Path) -> dict[str, Any]:
"""Load model registry from YAML file."""
with open(path, encoding="utf-8") as f:
data = yaml.safe_load(f)
return data.get("models", {})
def build_bench_args(
*,
host: str,
port: int,
model: str,
task: str,
model_cfg: dict[str, Any],
locale: str,
num_prompts: int,
concurrency: int | None,
dataset_path: str | None,
wer_eval: bool,
output_dir: str | None,
result_filename: str | None,
extra_cli_args: list[str],
) -> list[str]:
"""Build the ``vllm bench serve --omni`` command for one (task, concurrency) run."""
dataset_name = _TASK_TO_DATASET[task]
backend: str = model_cfg["backend"]
endpoint: str = model_cfg["endpoint"]
task_extra_body: dict[str, Any] = (model_cfg.get("task_extra_body") or {}).get(task) or {}
# Resolve dataset path
if dataset_path:
resolved_dataset_path = dataset_path
elif task == "voice_design":
resolved_dataset_path = _DEFAULT_DESIGN_DATASET_PATH
else:
resolved_dataset_path = None
cmd = [
_vllm_omni_bin(),
"bench",
"serve",
"--omni",
"--host",
host,
"--port",
str(port),
"--model",
model,
"--backend",
backend,
"--endpoint",
endpoint,
"--dataset-name",
dataset_name,
"--num-prompts",
str(num_prompts),
"--num-warmups",
"2",
"--percentile-metrics",
"ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
]
if resolved_dataset_path:
cmd += ["--dataset-path", resolved_dataset_path]
if locale:
cmd += ["--seed-tts-locale", locale]
if task_extra_body:
cmd += ["--extra-body", json.dumps(task_extra_body, separators=(",", ":"))]
if concurrency is not None:
cmd += ["--max-concurrency", str(concurrency), "--request-rate", "inf"]
if wer_eval:
cmd.append("--seed-tts-wer-eval")
if output_dir or result_filename:
out_dir = output_dir or "."
os.makedirs(out_dir, exist_ok=True)
cmd += ["--save-result", "--result-dir", out_dir]
if result_filename:
cmd += ["--result-filename", result_filename]
cmd += extra_cli_args
return cmd
def run_one_benchmark(cmd: list[str]) -> dict[str, Any] | None:
"""Run a single benchmark subprocess and return parsed JSON result if available."""
print(f"\n{'=' * 60}")
print("Running:", " ".join(cmd))
print("=" * 60)
result = subprocess.run(cmd, check=False)
if result.returncode != 0:
print(f"[bench_tts] WARNING: benchmark exited with code {result.returncode}")
return None
# If --save-result was used, find the result file
try:
result_dir_idx = cmd.index("--result-dir")
result_dir = Path(cmd[result_dir_idx + 1])
if "--result-filename" in cmd:
fname_idx = cmd.index("--result-filename")
result_file = result_dir / cmd[fname_idx + 1]
else:
# find most recently modified json
jsons = sorted(result_dir.glob("result_*.json"), key=lambda p: p.stat().st_mtime)
result_file = jsons[-1] if jsons else None
if result_file and result_file.is_file():
return json.loads(result_file.read_text(encoding="utf-8"))
except (ValueError, IndexError, OSError):
pass
return None
def print_summary_table(results: list[dict[str, Any]]) -> None:
"""Print a unified metrics table across all (task, concurrency) runs."""
if not results:
return
header = (
f"{'Task':<16} {'Concurrency':>11} {'RTF mean':>10} "
f"{'TTFP (ms)':>10} {'Throughput':>12} {'WER':>7} {'SIM':>7} {'UTMOS':>7}"
)
print(f"\n{'=' * len(header)}")
print("BENCHMARK SUMMARY")
print("=" * len(header))
print(header)
print("-" * len(header))
for r in results:
task = r.get("_task", "?")
conc = r.get("_concurrency", "?")
rtf = r.get("mean_audio_rtf", float("nan"))
ttfp = r.get("mean_audio_ttfp_ms", float("nan"))
throughput = r.get("audio_throughput", float("nan"))
wer = r.get("seed_tts_mean_wer", float("nan"))
sim = r.get("seed_tts_mean_sim", float("nan"))
utmos = r.get("seed_tts_mean_utmos", float("nan"))
def fmt(v: float, digits: int = 3) -> str:
return f"{v:.{digits}f}" if not math.isnan(v) else " n/a"
print(
f"{task:<16} {str(conc):>11} {fmt(rtf):>10} {fmt(ttfp, 0):>10} "
f"{fmt(throughput):>12} {fmt(wer):>7} {fmt(sim):>7} {fmt(utmos):>7}"
)
print("=" * len(header))
def main() -> None:
"""Entry point for the universal TTS benchmark CLI."""
parser = argparse.ArgumentParser(
description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter,
)
parser.add_argument(
"--model", required=True, help="HuggingFace model ID (e.g. Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice)"
)
parser.add_argument("--task", default="all", help="Task type: voice_clone | default_voice | voice_design | all")
parser.add_argument("--locale", default="en", choices=["en", "zh"])
parser.add_argument("--concurrency", type=int, nargs="+", default=[1, 4], metavar="N")
parser.add_argument(
"--num-prompts",
type=int,
nargs="+",
default=[20],
metavar="N",
help="Number of prompts per run. If one value, applied to all concurrency levels.",
)
parser.add_argument(
"--dataset-path", default=None, help="Root of seed-tts-eval dataset (required for voice_clone/default_voice)"
)
parser.add_argument("--wer-eval", action="store_true", help="Enable WER/SIM/UTMOS quality eval")
parser.add_argument("--output-dir", default=None, help="Directory to save result JSON files")
parser.add_argument("--host", default="localhost")
parser.add_argument("--port", type=int, default=8000)
parser.add_argument("--model-configs", default=str(_DEFAULT_MODEL_CONFIGS), help="Path to model_configs.yaml")
parser.add_argument("extra", nargs=argparse.REMAINDER, help="Extra args passed directly to vllm bench serve")
args = parser.parse_args()
model_configs = load_model_configs(Path(args.model_configs))
if args.model not in model_configs:
known = "\n ".join(model_configs.keys())
print(f"[bench_tts] ERROR: model '{args.model}' not in model_configs.yaml.\nKnown models:\n {known}")
sys.exit(1)
model_cfg = model_configs[args.model]
supported_tasks: list[str] = model_cfg.get("supported_tasks", [])
tasks_to_run: list[str]
if args.task == "all":
tasks_to_run = supported_tasks
elif args.task in supported_tasks:
tasks_to_run = [args.task]
else:
print(
f"[bench_tts] ERROR: task '{args.task}' not supported by {args.model}.\nSupported tasks: {supported_tasks}"
)
sys.exit(1)
# Align num_prompts list with concurrency list
num_prompts_list: list[int] = args.num_prompts
if len(num_prompts_list) == 1:
num_prompts_list = num_prompts_list * len(args.concurrency)
elif len(num_prompts_list) != len(args.concurrency):
print(
f"[bench_tts] ERROR: --num-prompts ({len(num_prompts_list)} values) must be "
f"length 1 or match --concurrency ({len(args.concurrency)} values)."
)
sys.exit(1)
all_results: list[dict[str, Any]] = []
for task in tasks_to_run:
for concurrency, num_prompts in zip(args.concurrency, num_prompts_list):
ts = datetime.now().strftime("%Y%m%d-%H%M%S")
result_filename = f"bench_tts_{args.model.replace('/', '_')}_{task}_c{concurrency}_{ts}.json"
cmd = build_bench_args(
host=args.host,
port=args.port,
model=args.model,
task=task,
model_cfg=model_cfg,
locale=args.locale,
num_prompts=num_prompts,
concurrency=concurrency,
dataset_path=args.dataset_path,
wer_eval=args.wer_eval,
output_dir=args.output_dir,
result_filename=result_filename,
extra_cli_args=args.extra or [],
)
result = run_one_benchmark(cmd)
if result is not None:
result["_task"] = task
result["_concurrency"] = concurrency
all_results.append(result)
# Persist the metadata so plot_results.py can pick it up.
if args.output_dir and result_filename:
result_path = Path(args.output_dir) / result_filename
if result_path.is_file():
result_path.write_text(json.dumps(result, indent=2), encoding="utf-8")
print_summary_table(all_results)
if __name__ == "__main__":
main()
-39
View File
@@ -1,39 +0,0 @@
# Universal TTS benchmark model registry.
# Maps HuggingFace model ID → supported tasks + per-task extra body fields.
# To add a new TTS model: add an entry here. No code changes required.
#
# The server auto-loads its Deploy YAML from vllm_omni/deploy/<model>.yaml via
# the Pipeline + Deploy schema introduced in #2383, so no stage_config path
# is tracked here.
models:
# -CustomVoice checkpoints lack speaker_encoder weights, so voice_clone is
# NOT supported (an attempt raises ValueError from _extract_speaker_embedding
# at model runtime). Use -Base for voice_clone.
Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice:
supported_tasks: [default_voice, voice_design]
backend: openai-audio-speech
endpoint: /v1/audio/speech
task_extra_body:
default_voice:
voice: Vivian
language: English
task_type: CustomVoice
voice_design:
task_type: VoiceDesign
language: English
Qwen/Qwen3-TTS-12Hz-1.7B-Base:
supported_tasks: [voice_clone]
backend: openai-audio-speech
endpoint: /v1/audio/speech
task_extra_body:
voice_clone:
task_type: Base
openbmb/VoxCPM2:
supported_tasks: [voice_clone]
backend: openai-audio-speech
endpoint: /v1/audio/speech
task_extra_body:
voice_clone: {}
-324
View File
@@ -1,324 +0,0 @@
"""Plot universal TTS benchmark results.
Reads JSON files saved by ``bench_tts.py`` (via ``vllm bench serve --omni``)
and generates comparison bar charts grouped by task type.
Metrics plotted:
- AUDIO_TTFP (mean audio time-to-first-packet, ms)
- E2EL (mean end-to-end latency, ms)
- Audio RTF (mean real-time factor)
- Audio throughput (audio-seconds / wall-second)
Quality metrics (WER / SIM / UTMOS) are printed in a table when present.
Usage::
# Single run — one JSON per task, all in results/
python benchmarks/tts/plot_results.py \\
--results results/bench_tts_*.json \\
--output results/tts_benchmark.png
# Compare two runs (e.g. async_chunk on vs off)
python benchmarks/tts/plot_results.py \\
--results run_a/bench_tts_*.json \\
--results run_b/bench_tts_*.json \\
--labels "async_chunk_on" "async_chunk_off" \\
--output results/comparison.png
"""
from __future__ import annotations
import argparse
import json
import math
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
# ---------------------------------------------------------------------------
# JSON loading
# ---------------------------------------------------------------------------
def load_run(paths: list[str]) -> list[dict]:
"""Load and merge all JSON files for one run into a flat list of records.
Each record is expected to have at least ``_concurrency`` (int) and
``_task`` (str) keys injected by ``bench_tts.py``. Records that come
from a file that contains a list are flattened.
"""
records: list[dict] = []
for p in paths:
raw = json.loads(Path(p).read_text(encoding="utf-8"))
if isinstance(raw, list):
records.extend(raw)
elif isinstance(raw, dict):
records.append(raw)
return records
def _get(record: dict, key: str) -> float:
v = record.get(key, float("nan"))
if v is None or (isinstance(v, float) and math.isnan(v)):
return float("nan")
try:
return float(v)
except (TypeError, ValueError):
return float("nan")
# ---------------------------------------------------------------------------
# Plotting helpers
# ---------------------------------------------------------------------------
def _bar_group(
ax: plt.Axes,
x: np.ndarray,
data_per_label: dict[str, list[float]],
width: float,
colors: list[str],
ylabel: str,
title: str,
concurrency_labels: list[str],
fmt: str = ".1f",
) -> None:
n = len(data_per_label)
offsets = np.linspace(-(n - 1) * width / 2, (n - 1) * width / 2, n) if n > 1 else [0.0]
for i, (label, values) in enumerate(data_per_label.items()):
plot_vals = [0.0 if math.isnan(v) else v for v in values]
bar = ax.bar(x + offsets[i], plot_vals, width, label=label, color=colors[i % len(colors)], alpha=0.85)
max_val = max((v for v in values if not math.isnan(v)), default=1.0)
for rect, val in zip(bar, values):
if not math.isnan(val) and val > 0:
ax.text(
rect.get_x() + rect.get_width() / 2,
rect.get_height() + max_val * 0.02,
f"{val:{fmt}}",
ha="center",
va="bottom",
fontsize=8,
fontweight="bold",
)
ax.set_xlabel("Concurrency", fontsize=11)
ax.set_ylabel(ylabel, fontsize=11)
ax.set_title(title, fontsize=12, fontweight="bold")
ax.set_xticks(x)
ax.set_xticklabels(concurrency_labels)
ax.legend(fontsize=9)
ax.grid(axis="y", alpha=0.3)
ax.set_axisbelow(True)
COLORS = ["#2196F3", "#FF5722", "#4CAF50", "#FFC107", "#9C27B0"]
# ---------------------------------------------------------------------------
# Comparison plot (multiple labels / runs)
# ---------------------------------------------------------------------------
def plot_comparison(
all_runs: list[list[dict]],
labels: list[str],
output_path: str,
task_filter: str | None = None,
title_prefix: str = "TTS",
) -> None:
"""One 2×2 subplot per task found in the data."""
# Determine tasks to plot
tasks: list[str] = []
for run in all_runs:
for r in run:
t = r.get("_task", "unknown")
if t not in tasks:
tasks.append(t)
if task_filter:
tasks = [t for t in tasks if t == task_filter]
n_tasks = len(tasks)
if n_tasks == 0:
print("[plot_results] No tasks found in data.")
return
fig, axes_grid = plt.subplots(n_tasks, 4, figsize=(18, 4.5 * n_tasks))
fig.suptitle(f"{title_prefix} Benchmark", fontsize=15, fontweight="bold")
# Ensure axes_grid is always 2D
if n_tasks == 1:
axes_grid = [axes_grid]
for row_idx, task in enumerate(tasks):
# Collect concurrencies across all runs for this task
all_concs: set[int] = set()
for run in all_runs:
for r in run:
if r.get("_task") == task:
c = r.get("_concurrency")
if c is not None:
all_concs.add(int(c))
concurrencies = sorted(all_concs)
x = np.arange(len(concurrencies))
conc_labels = [str(c) for c in concurrencies]
def _series(run: list[dict], metric_key: str) -> list[float]:
conc_map = {int(r["_concurrency"]): r for r in run if r.get("_task") == task and "_concurrency" in r}
return [_get(conc_map.get(c, {}), metric_key) for c in concurrencies]
metrics = [
("mean_audio_ttfp_ms", "TTFP (ms)", "Time-to-First-Packet", ".0f"),
("mean_e2el_ms", "E2E Latency (ms)", "End-to-End Latency", ".0f"),
("mean_audio_rtf", "RTF", "Real-Time Factor (RTF)", ".3f"),
("audio_throughput", "audio-s / wall-s", "Audio Throughput", ".2f"),
]
axes_row = axes_grid[row_idx]
for col_idx, (key, ylabel, subtitle, fmt) in enumerate(metrics):
data_per_label = {lbl: _series(run, key) for lbl, run in zip(labels, all_runs)}
_bar_group(
axes_row[col_idx],
x,
data_per_label,
width=0.3 if len(labels) > 1 else 0.5,
colors=COLORS,
ylabel=ylabel,
title=f"{task} — {subtitle}",
concurrency_labels=conc_labels,
fmt=fmt,
)
plt.tight_layout()
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
plt.savefig(output_path, dpi=150, bbox_inches="tight")
print(f"Plot saved to {output_path}")
plt.close()
# ---------------------------------------------------------------------------
# Markdown comparison table
# ---------------------------------------------------------------------------
def print_comparison_table(all_runs: list[list[dict]], labels: list[str]) -> None:
tasks: list[str] = []
for run in all_runs:
for r in run:
t = r.get("_task", "unknown")
if t not in tasks:
tasks.append(t)
perf_metrics = [
("TTFP (ms)", "mean_audio_ttfp_ms", ".1f"),
("E2E (ms)", "mean_e2el_ms", ".1f"),
("RTF", "mean_audio_rtf", ".3f"),
("Throughput (a-s/s)", "audio_throughput", ".2f"),
]
quality_metrics = [
("WER (%)", "seed_tts_mean_wer", ".1f"),
("SIM", "seed_tts_mean_sim", ".3f"),
("UTMOS", "seed_tts_mean_utmos", ".2f"),
]
for task in tasks:
all_concs: set[int] = set()
for run in all_runs:
for r in run:
if r.get("_task") == task:
c = r.get("_concurrency")
if c is not None:
all_concs.add(int(c))
concurrencies = sorted(all_concs)
print(f"\n## {task}\n")
col_header = "| Metric | Concurrency |" + "".join(f" {lbl} |" for lbl in labels)
sep = "| --- | --- |" + " --- |" * len(labels)
print(col_header)
print(sep)
for metric, key, fmt in perf_metrics + quality_metrics:
for c in concurrencies:
row = f"| {metric} | {c} |"
for run in all_runs:
conc_map = {
int(r["_concurrency"]): r for r in run if r.get("_task") == task and "_concurrency" in r
}
val = _get(conc_map.get(c, {}), key)
row += f" {val:{fmt}} |" if not math.isnan(val) else " n/a |"
print(row)
# Improvement column (2-run comparison only)
if len(all_runs) == 2:
print(f"\n### Improvement ({labels[0]} vs {labels[1]})\n")
print("| Metric | Concurrency | Change |")
print("| --- | --- | --- |")
for metric, key, _ in perf_metrics:
for c in concurrencies:
conc_map0 = {
int(r["_concurrency"]): r for r in all_runs[0] if r.get("_task") == task and "_concurrency" in r
}
conc_map1 = {
int(r["_concurrency"]): r for r in all_runs[1] if r.get("_task") == task and "_concurrency" in r
}
v0 = _get(conc_map0.get(c, {}), key)
v1 = _get(conc_map1.get(c, {}), key)
if not math.isnan(v0) and not math.isnan(v1) and v1 > 0:
pct = (v1 - v0) / v1 * 100
print(f"| {metric} | {c} | {pct:+.1f}% |")
# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter,
)
parser.add_argument(
"--results",
type=str,
nargs="+",
action="append",
required=True,
metavar="FILE",
help="JSON result file(s) for one run. Repeat --results for multiple runs to compare.",
)
parser.add_argument(
"--labels",
type=str,
nargs="+",
default=None,
help="Label for each --results group (must match the number of --results groups).",
)
parser.add_argument("--output", type=str, default="results/tts_benchmark.png", help="Output image path.")
parser.add_argument("--title", type=str, default="TTS", help="Title prefix for the plot.")
parser.add_argument("--task", type=str, default=None, help="Filter to a single task (e.g. voice_clone).")
return parser.parse_args()
def main() -> None:
args = parse_args()
# args.results is a list-of-lists due to action="append"
all_runs: list[list[dict]] = [load_run(group) for group in args.results]
n_runs = len(all_runs)
labels: list[str]
if args.labels:
if len(args.labels) != n_runs:
raise SystemExit(f"--labels count ({len(args.labels)}) must match --results groups ({n_runs})")
labels = args.labels
else:
labels = [f"run{i + 1}" for i in range(n_runs)]
print_comparison_table(all_runs, labels)
plot_comparison(all_runs, labels, args.output, task_filter=args.task, title_prefix=args.title)
if __name__ == "__main__":
main()
+119
View File
@@ -0,0 +1,119 @@
# VoxCPM Benchmark
This directory contains both:
- online serving benchmark through the OpenAI-compatible `/v1/audio/speech` API
- offline benchmark for `Omni` / `AsyncOmni`
- full offline smoke-matrix orchestration
Both benchmark paths report:
- TTFP: time to first PCM packet
- E2E latency
- RTF: real-time factor (`e2e / audio_duration`)
## Offline Benchmark
Single offline benchmark run:
```bash
python benchmarks/voxcpm/vllm_omni/bench_tts_offline.py \
--model /path/to/voxcpm-model \
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm.yaml \
--text "This is a split-stage VoxCPM synthesis example running on vLLM Omni." \
--warmup-runs 1 \
--output-dir benchmarks/voxcpm/results/offline_single
```
Streaming offline benchmark:
```bash
python benchmarks/voxcpm/vllm_omni/bench_tts_offline.py \
--model /path/to/voxcpm-model \
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \
--text "This is a split-stage VoxCPM streaming example running on vLLM Omni." \
--warmup-runs 1 \
--output-dir benchmarks/voxcpm/results/offline_streaming
```
Full fixed offline matrix, equivalent to the old `examples/offline_inference/voxcpm/test.py`:
```bash
python benchmarks/voxcpm/vllm_omni/run_offline_matrix.py \
--model /path/to/voxcpm-model \
--ref-audio /path/to/reference.wav \
--ref-text "The exact transcript spoken in reference.wav." \
--output-root benchmarks/voxcpm/results/offline_matrix
```
The full matrix covers both routes:
- streaming: `voxcpm_async_chunk.yaml`
- sync: `voxcpm.yaml`
And these six scenarios under each route:
- warmup + single TTS
- warmup + single voice cloning
- warmup + batch TTS
- warmup + batch voice cloning
- cold single TTS
- cold single voice cloning
`bench_tts_offline.py` itself no longer writes `summary.json` / `results.json`; it prints TTFP / RTF inline and saves generated WAV files only. The matrix runner keeps only per-case `run.log`.
## Start the Server
Async-chunk:
```bash
vllm serve /path/to/voxcpm-model \
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \
--trust-remote-code \
--enforce-eager \
--omni \
--port 8091
```
Non-streaming:
```bash
vllm serve /path/to/voxcpm-model \
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm.yaml \
--trust-remote-code \
--enforce-eager \
--omni \
--port 8091
```
## Run the Benchmark
```bash
python benchmarks/voxcpm/vllm_omni/bench_tts_serve.py \
--host 127.0.0.1 \
--port 8091 \
--num-prompts 20 \
--max-concurrency 1 \
--result-dir /tmp/voxcpm_bench
```
Voice cloning benchmark:
```bash
python benchmarks/voxcpm/vllm_omni/bench_tts_serve.py \
--host 127.0.0.1 \
--port 8091 \
--num-prompts 10 \
--max-concurrency 1 \
--ref-audio https://example.com/reference.wav \
--ref-text "The exact transcript spoken in the reference audio." \
--result-dir /tmp/voxcpm_clone_bench
```
## Notes
- The benchmark uses `stream=true` and `response_format=pcm` so TTFP is measured from the first audio packet.
- `RTF < 1.0` means the server generates audio faster than real time.
- For `voxcpm_async_chunk.yaml`, keep concurrency at `1`. This matches native VoxCPM streaming more closely.
- Do not benchmark concurrent online streaming on `voxcpm_async_chunk.yaml`; use `voxcpm.yaml` for multi-request throughput runs.
- For the offline matrix mode, `--ref-audio` and `--ref-text` are required because clone cases are part of the fixed coverage set.
@@ -6,27 +6,6 @@ Supports both:
- text-only synthesis
- voice cloning
- text/clone batch inputs from txt or jsonl
Usage::
# Sync (default voice)
python benchmarks/tts/bench_voxcpm_offline.py \\
--model /path/to/VoxCPM \\
--text "Hello world" \\
--output-dir results/audio/
# Streaming (async_chunk)
python benchmarks/tts/bench_voxcpm_offline.py \\
--model /path/to/VoxCPM \\
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \\
--txt-prompts prompts.txt \\
--output-dir results/audio/
# Voice cloning batch via JSONL
python benchmarks/tts/bench_voxcpm_offline.py \\
--model /path/to/VoxCPM \\
--jsonl-prompts prompts.jsonl \\
--output-dir results/audio/
"""
from __future__ import annotations
@@ -47,21 +26,7 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni import AsyncOmni, Omni
def _find_repo_root(start: Path) -> Path:
"""Walk up from ``start`` until a repo marker is found.
Falls back to ``parents[2]`` for backwards compatibility if no marker hits
(which can only happen in unusual checkouts — the tree should always have
pyproject.toml + vllm_omni/ at the top level).
"""
for candidate in [start, *start.parents]:
if (candidate / "pyproject.toml").is_file() and (candidate / "vllm_omni").is_dir():
return candidate
return start.parents[2]
REPO_ROOT = _find_repo_root(Path(__file__).resolve())
REPO_ROOT = Path(__file__).resolve().parents[3]
DEFAULT_STAGE_ASYNC = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm_async_chunk.yaml"
DEFAULT_STAGE_SYNC = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm.yaml"
@@ -508,6 +473,9 @@ def parse_args():
def _is_streaming_stage_config(stage_configs_path: str) -> bool:
cfg_name = Path(stage_configs_path).name.lower()
# Keep routing purely config-path based:
# - voxcpm.yaml => sync
# - voxcpm_async_chunk.yaml => streaming
return "async_chunk" in cfg_name
@@ -0,0 +1,283 @@
"""Benchmark VoxCPM via /v1/audio/speech.
Reports TTFP (time to first packet), E2E latency, and RTF (real-time factor).
"""
from __future__ import annotations
import argparse
import asyncio
import json
import time
from dataclasses import asdict, dataclass, field
from datetime import datetime
from pathlib import Path
import aiohttp
import numpy as np
from tqdm.asyncio import tqdm
DEFAULT_MODEL = "OpenBMB/VoxCPM1.5"
DEFAULT_SAMPLE_RATE = 24000
PROMPTS = [
"Hello, welcome to the VoxCPM speech benchmark.",
"This is a short benchmark prompt for online text-to-speech generation.",
"The quick brown fox jumps over the lazy dog near the riverbank.",
"Please remember to bring your identification documents tomorrow morning.",
"Learning a new language takes patience, practice, and curiosity.",
"This benchmark reports TTFP and RTF for the VoxCPM online serving path.",
]
@dataclass
class RequestResult:
success: bool = False
ttfp: float = 0.0
e2e: float = 0.0
audio_bytes: int = 0
audio_duration: float = 0.0
rtf: float = 0.0
prompt: str = ""
error: str = ""
@dataclass
class BenchmarkResult:
concurrency: int = 0
num_prompts: int = 0
completed: int = 0
failed: int = 0
duration_s: float = 0.0
mean_ttfp_ms: float = 0.0
median_ttfp_ms: float = 0.0
p95_ttfp_ms: float = 0.0
mean_e2e_ms: float = 0.0
median_e2e_ms: float = 0.0
p95_e2e_ms: float = 0.0
mean_rtf: float = 0.0
median_rtf: float = 0.0
p95_rtf: float = 0.0
total_audio_duration_s: float = 0.0
request_throughput: float = 0.0
per_request: list[dict[str, float | str]] = field(default_factory=list)
def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = DEFAULT_SAMPLE_RATE, sample_width: int = 2) -> float:
num_samples = num_bytes / sample_width
return num_samples / sample_rate
async def send_tts_request(
session: aiohttp.ClientSession,
api_url: str,
*,
model: str,
prompt: str,
ref_audio: str | None,
ref_text: str | None,
pbar: tqdm | None = None,
) -> RequestResult:
payload: dict[str, object] = {
"model": model,
"input": prompt,
"stream": True,
"response_format": "pcm",
}
if ref_audio is not None:
payload["ref_audio"] = ref_audio
if ref_text is not None:
payload["ref_text"] = ref_text
result = RequestResult(prompt=prompt)
started_at = time.perf_counter()
try:
async with session.post(api_url, json=payload) as response:
if response.status != 200:
result.error = f"HTTP {response.status}: {await response.text()}"
return result
first_chunk = True
total_bytes = 0
async for chunk in response.content.iter_any():
if not chunk:
continue
if first_chunk:
result.ttfp = time.perf_counter() - started_at
first_chunk = False
total_bytes += len(chunk)
result.e2e = time.perf_counter() - started_at
result.audio_bytes = total_bytes
result.audio_duration = pcm_bytes_to_duration(total_bytes)
if result.audio_duration > 0:
result.rtf = result.e2e / result.audio_duration
result.success = True
except Exception as e:
result.error = str(e)
result.e2e = time.perf_counter() - started_at
if pbar is not None:
pbar.update(1)
return result
async def run_benchmark(
*,
host: str,
port: int,
model: str,
num_prompts: int,
max_concurrency: int,
num_warmups: int,
ref_audio: str | None,
ref_text: str | None,
) -> BenchmarkResult:
api_url = f"http://{host}:{port}/v1/audio/speech"
connector = aiohttp.TCPConnector(limit=max_concurrency, limit_per_host=max_concurrency, keepalive_timeout=60)
timeout = aiohttp.ClientTimeout(total=600)
async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
if num_warmups > 0:
print(f" Warming up with {num_warmups} requests...")
warmup_tasks = [
send_tts_request(
session,
api_url,
model=model,
prompt=PROMPTS[i % len(PROMPTS)],
ref_audio=ref_audio,
ref_text=ref_text,
)
for i in range(num_warmups)
]
await asyncio.gather(*warmup_tasks)
print(" Warmup done.")
request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
semaphore = asyncio.Semaphore(max_concurrency)
pbar = tqdm(total=num_prompts, desc=f" concurrency={max_concurrency}")
async def limited_request(prompt: str) -> RequestResult:
async with semaphore:
return await send_tts_request(
session,
api_url,
model=model,
prompt=prompt,
ref_audio=ref_audio,
ref_text=ref_text,
pbar=pbar,
)
started_at = time.perf_counter()
results = await asyncio.gather(*[asyncio.create_task(limited_request(prompt)) for prompt in request_prompts])
duration = time.perf_counter() - started_at
pbar.close()
succeeded = [result for result in results if result.success]
bench = BenchmarkResult(
concurrency=max_concurrency,
num_prompts=num_prompts,
completed=len(succeeded),
failed=len(results) - len(succeeded),
duration_s=duration,
)
if not succeeded:
return bench
ttfps = np.array([result.ttfp * 1000 for result in succeeded], dtype=np.float64)
e2es = np.array([result.e2e * 1000 for result in succeeded], dtype=np.float64)
rtfs = np.array([result.rtf for result in succeeded], dtype=np.float64)
audio_durations = np.array([result.audio_duration for result in succeeded], dtype=np.float64)
bench.mean_ttfp_ms = float(np.mean(ttfps))
bench.median_ttfp_ms = float(np.median(ttfps))
bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
bench.mean_e2e_ms = float(np.mean(e2es))
bench.median_e2e_ms = float(np.median(e2es))
bench.p95_e2e_ms = float(np.percentile(e2es, 95))
bench.mean_rtf = float(np.mean(rtfs))
bench.median_rtf = float(np.median(rtfs))
bench.p95_rtf = float(np.percentile(rtfs, 95))
bench.total_audio_duration_s = float(np.sum(audio_durations))
bench.request_throughput = len(succeeded) / duration if duration > 0 else 0.0
bench.per_request = [
{
"prompt": result.prompt,
"ttfp_ms": result.ttfp * 1000,
"e2e_ms": result.e2e * 1000,
"rtf": result.rtf,
"audio_duration_s": result.audio_duration,
}
for result in succeeded
]
return bench
def print_summary(result: BenchmarkResult) -> None:
width = 54
print("")
print("=" * width)
print(f"{'VoxCPM Serving Benchmark':^{width}}")
print("=" * width)
print(f"concurrency : {result.concurrency}")
print(f"requests : {result.completed}/{result.num_prompts} succeeded")
print(f"wall time (s) : {result.duration_s:.3f}")
print(f"mean TTFP (ms) : {result.mean_ttfp_ms:.2f}")
print(f"p95 TTFP (ms) : {result.p95_ttfp_ms:.2f}")
print(f"mean E2E (ms) : {result.mean_e2e_ms:.2f}")
print(f"p95 E2E (ms) : {result.p95_e2e_ms:.2f}")
print(f"mean RTF : {result.mean_rtf:.3f}")
print(f"p95 RTF : {result.p95_rtf:.3f}")
print(f"request throughput : {result.request_throughput:.2f} req/s")
print("=" * width)
async def main_async(args) -> None:
result_dir = Path(args.result_dir)
result_dir.mkdir(parents=True, exist_ok=True)
all_results: list[BenchmarkResult] = []
for concurrency in args.max_concurrency:
result = await run_benchmark(
host=args.host,
port=args.port,
model=args.model,
num_prompts=args.num_prompts,
max_concurrency=concurrency,
num_warmups=args.num_warmups,
ref_audio=args.ref_audio,
ref_text=args.ref_text,
)
print_summary(result)
all_results.append(result)
payload = {
"model": args.model,
"created_at": datetime.utcnow().isoformat() + "Z",
"results": [asdict(result) for result in all_results],
}
result_path = result_dir / "bench_tts_serve.json"
result_path.write_text(json.dumps(payload, indent=2), encoding="utf-8")
print(f"Saved results to: {result_path}")
def parse_args():
parser = argparse.ArgumentParser(description="Benchmark VoxCPM via /v1/audio/speech")
parser.add_argument("--host", default="127.0.0.1", help="Server host")
parser.add_argument("--port", type=int, default=8091, help="Server port")
parser.add_argument("--model", default=DEFAULT_MODEL, help="Model name or path")
parser.add_argument("--num-prompts", type=int, default=20, help="Number of prompts to send")
parser.add_argument("--max-concurrency", type=int, nargs="+", default=[1], help="Concurrency levels to benchmark")
parser.add_argument("--num-warmups", type=int, default=3, help="Warmup request count")
parser.add_argument("--ref-audio", default=None, help="Reference audio URL or data URL for voice cloning")
parser.add_argument("--ref-text", default=None, help="Reference audio transcript for voice cloning")
parser.add_argument("--result-dir", default="results", help="Directory to save benchmark JSON")
return parser.parse_args()
if __name__ == "__main__":
asyncio.run(main_async(parse_args()))
@@ -0,0 +1,303 @@
"""Run the full offline VoxCPM smoke matrix.
This script keeps the old `test.py` coverage, but delegates each case to
`bench_tts_offline.py` so the benchmark runner itself stays focused on a
single execution path.
"""
from __future__ import annotations
import shlex
import subprocess
import sys
import time
from dataclasses import dataclass
from pathlib import Path
from vllm.utils.argparse_utils import FlexibleArgumentParser
REPO_ROOT = Path(__file__).resolve().parents[3]
BENCH_SCRIPT = Path(__file__).with_name("bench_tts_offline.py")
DEFAULT_STAGE_ASYNC = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm_async_chunk.yaml"
DEFAULT_STAGE_SYNC = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm.yaml"
DEFAULT_OUTPUT_ROOT = BENCH_SCRIPT.parents[1] / "results" / "offline_matrix"
SINGLE_TTS_TEXT = "This is a single text-to-speech smoke test for VoxCPM on vLLM Omni."
SINGLE_CLONE_TEXT = "This sentence is synthesized with the cloned voice for validation."
BATCH_TTS_TEXTS = [
"The first batch text-to-speech sample validates sequential batch execution.",
"The second batch text-to-speech sample checks another prompt in the same file.",
"The third batch text-to-speech sample completes the sequential batch path.",
]
BATCH_CLONE_TEXTS = [
"The first cloned sample validates sequential batch voice cloning.",
"The second cloned sample checks the same reference voice on another prompt.",
"The third cloned sample finishes the shared-reference clone batch path.",
]
@dataclass(frozen=True, slots=True)
class ModeSpec:
name: str
stage_config: Path
@dataclass(frozen=True, slots=True)
class CaseSpec:
name: str
warmup_runs: int
prompt_kind: str
voice_clone: bool
@dataclass(frozen=True, slots=True)
class CaseResult:
mode: str
case: str
returncode: int
elapsed_s: float
output_dir: Path
log_path: Path
@property
def ok(self) -> bool:
return self.returncode == 0
MODE_SPECS = [
ModeSpec(name="streaming", stage_config=DEFAULT_STAGE_ASYNC),
ModeSpec(name="sync", stage_config=DEFAULT_STAGE_SYNC),
]
CASE_SPECS = [
CaseSpec(name="warmup_single_tts", warmup_runs=1, prompt_kind="single", voice_clone=False),
CaseSpec(name="warmup_single_clone", warmup_runs=1, prompt_kind="single", voice_clone=True),
CaseSpec(name="warmup_batch_tts", warmup_runs=1, prompt_kind="batch", voice_clone=False),
CaseSpec(name="warmup_batch_clone", warmup_runs=1, prompt_kind="batch", voice_clone=True),
CaseSpec(name="cold_single_tts", warmup_runs=0, prompt_kind="single", voice_clone=False),
CaseSpec(name="cold_single_clone", warmup_runs=0, prompt_kind="single", voice_clone=True),
]
def _write_lines(path: Path, lines: list[str]) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text("\n".join(lines) + "\n", encoding="utf-8")
def _prepare_batch_inputs(output_root: Path) -> tuple[Path, Path]:
input_dir = output_root / "inputs"
batch_tts_path = input_dir / "batch_tts_prompts.txt"
batch_clone_path = input_dir / "batch_clone_prompts.txt"
_write_lines(batch_tts_path, BATCH_TTS_TEXTS)
_write_lines(batch_clone_path, BATCH_CLONE_TEXTS)
return batch_tts_path, batch_clone_path
def _base_command(args, mode: ModeSpec, output_dir: Path) -> list[str]:
cmd = [
args.python,
str(BENCH_SCRIPT),
"--model",
args.model,
"--stage-configs-path",
str(mode.stage_config),
"--output-dir",
str(output_dir),
"--num-runs",
str(args.num_runs),
"--stage-init-timeout",
str(args.stage_init_timeout),
]
cmd.append("--log-stats" if args.log_stats else "--no-log-stats")
cmd.extend(["--cfg-value", str(args.cfg_value)])
cmd.extend(["--inference-timesteps", str(args.inference_timesteps)])
cmd.extend(["--min-len", str(args.min_len)])
cmd.extend(["--max-new-tokens", str(args.max_new_tokens)])
if args.streaming_prefix_len is not None:
cmd.extend(["--streaming-prefix-len", str(args.streaming_prefix_len)])
if args.enable_profiler:
profiler_dir = Path(args.profiler_dir) if args.profiler_dir is not None else (output_dir / "profiler")
cmd.append("--enable-profiler")
cmd.extend(["--profiler-dir", str(profiler_dir)])
cmd.extend(["--profiler-wait-seconds", str(args.profiler_wait_seconds)])
if args.profiler_stages is not None:
cmd.append("--profiler-stages")
cmd.extend(str(stage_id) for stage_id in args.profiler_stages)
return cmd
def _build_case_command(
args,
mode: ModeSpec,
case: CaseSpec,
*,
batch_tts_path: Path,
batch_clone_path: Path,
output_dir: Path,
) -> list[str]:
cmd = _base_command(args, mode, output_dir)
cmd.extend(["--warmup-runs", str(case.warmup_runs)])
if case.prompt_kind == "single":
cmd.extend(["--text", SINGLE_CLONE_TEXT if case.voice_clone else SINGLE_TTS_TEXT])
else:
cmd.extend(["--txt-prompts", str(batch_clone_path if case.voice_clone else batch_tts_path)])
if case.voice_clone:
cmd.extend(["--ref-audio", args.ref_audio, "--ref-text", args.ref_text])
return cmd
def _run_case(
args,
mode: ModeSpec,
case: CaseSpec,
*,
batch_tts_path: Path,
batch_clone_path: Path,
output_root: Path,
) -> CaseResult:
case_output_dir = output_root / mode.name / case.name
case_output_dir.mkdir(parents=True, exist_ok=True)
case_log_path = case_output_dir / "run.log"
cmd = _build_case_command(
args,
mode,
case,
batch_tts_path=batch_tts_path,
batch_clone_path=batch_clone_path,
output_dir=case_output_dir,
)
print()
print("=" * 80)
print(f"[{mode.name}] {case.name}")
print(f"Output directory: {case_output_dir}")
print(shlex.join(cmd))
start = time.perf_counter()
with case_log_path.open("w", encoding="utf-8") as log_fp:
process = subprocess.Popen(
cmd,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
bufsize=1,
)
assert process.stdout is not None
for line in process.stdout:
print(line, end="")
log_fp.write(line)
process.wait()
elapsed_s = time.perf_counter() - start
status = "PASS" if (process.returncode or 0) == 0 else f"FAIL({process.returncode})"
print(f"[{mode.name}] {case.name} -> {status} ({elapsed_s:.2f}s)")
return CaseResult(
mode=mode.name,
case=case.name,
returncode=int(process.returncode or 0),
elapsed_s=elapsed_s,
output_dir=case_output_dir,
log_path=case_log_path,
)
def parse_args():
parser = FlexibleArgumentParser(description="Run the full offline VoxCPM smoke matrix.")
parser.add_argument("--model", type=str, required=True, help="Local VoxCPM model directory.")
parser.add_argument("--ref-audio", type=str, required=True, help="Reference audio path for clone cases.")
parser.add_argument("--ref-text", type=str, required=True, help="Exact transcript spoken in --ref-audio.")
parser.add_argument("--output-root", type=str, default=str(DEFAULT_OUTPUT_ROOT), help="Root directory for outputs.")
parser.add_argument("--python", type=str, default=sys.executable, help="Python executable used to launch cases.")
parser.add_argument("--stage-init-timeout", type=int, default=600, help="Stage initialization timeout in seconds.")
parser.add_argument("--log-stats", dest="log_stats", action="store_true", help="Enable vLLM Omni stats logging.")
parser.add_argument(
"--no-log-stats",
dest="log_stats",
action="store_false",
help="Disable vLLM Omni stats logging.",
)
parser.set_defaults(log_stats=True)
parser.add_argument("--num-runs", type=int, default=1, help="Number of measured runs per case.")
parser.add_argument("--cfg-value", type=float, default=2.0, help="Classifier-free guidance value for VoxCPM.")
parser.add_argument("--inference-timesteps", type=int, default=10, help="Number of inference timesteps.")
parser.add_argument("--min-len", type=int, default=2, help="Minimum generated token length.")
parser.add_argument("--max-new-tokens", type=int, default=4096, help="Maximum generated token length.")
parser.add_argument(
"--streaming-prefix-len",
type=int,
default=None,
help="Optional VoxCPM streaming window passed to streaming cases.",
)
parser.add_argument("--enable-profiler", action="store_true", help="Enable torch profiler for each case.")
parser.add_argument(
"--profiler-dir",
type=str,
default=None,
help="Profiler output root. Defaults to <case-output-dir>/profiler.",
)
parser.add_argument(
"--profiler-stages",
type=int,
nargs="*",
default=None,
help="Optional stage ids to profile. Defaults to all configured stages.",
)
parser.add_argument(
"--profiler-wait-seconds",
type=float,
default=30.0,
help="Seconds to wait after stopping profiler for traces to flush.",
)
args = parser.parse_args()
if args.num_runs < 1:
parser.error("--num-runs must be >= 1")
return args
def main(args) -> int:
output_root = Path(args.output_root)
output_root.mkdir(parents=True, exist_ok=True)
batch_tts_path, batch_clone_path = _prepare_batch_inputs(output_root)
print(f"Model: {args.model}")
print(f"Reference audio: {args.ref_audio}")
print(f"Reference text: {args.ref_text}")
print(f"Python: {args.python}")
print(f"Output root: {output_root}")
print(f"Cases: {len(MODE_SPECS) * len(CASE_SPECS)}")
results: list[CaseResult] = []
for mode in MODE_SPECS:
for case in CASE_SPECS:
results.append(
_run_case(
args,
mode,
case,
batch_tts_path=batch_tts_path,
batch_clone_path=batch_clone_path,
output_root=output_root,
)
)
failed = [result for result in results if not result.ok]
print()
print("=" * 80)
print("Summary:")
for result in results:
status = "PASS" if result.ok else f"FAIL({result.returncode})"
print(f"- [{result.mode}] {result.case}: {status} ({result.elapsed_s:.2f}s)")
print(f" output_dir={result.output_dir}")
print(f" log={result.log_path}")
print(f"Passed: {len(results) - len(failed)}/{len(results)}")
if failed:
print("Failed cases:")
for result in failed:
print(f"- [{result.mode}] {result.case}: see {result.log_path}")
return 1
return 0
if __name__ == "__main__":
raise SystemExit(main(parse_args()))
+1 -4
View File
@@ -19,7 +19,7 @@ The new deploy schema lives under `vllm_omni/deploy/` and is paired with a froze
| `platforms` | dict | optional | `null` | Keyed by `npu` / `rocm` / `xpu`, each contains a `stages:` list with per-platform overrides applied on top of the CUDA defaults. |
| `pipeline` | str | optional | `null` | Override the auto-detected pipeline registry key (used for structural variants like `qwen2_5_omni_thinker_only`). |
| `trust_remote_code` | bool | optional | `true` | **Pipeline-wide.** Trust HF remote code on model load; applies to every stage. |
| `distributed_executor_backend` | str \| null | optional | `null` | **Pipeline-wide.** Distributed executor backend forwarded to vLLM (`"mp"`, `"ray"`, `"external_launcher"`). If omitted, vLLM auto-selects backend from runtime topology. |
| `distributed_executor_backend` | str | optional | `"mp"` | **Pipeline-wide.** Executor backend (`"mp"` or `"ray"`). |
| `dtype` | str \| null | optional | `null` | **Pipeline-wide.** Model dtype for every stage. |
| `quantization` | str \| null | optional | `null` | **Pipeline-wide.** Quantization method for every stage. |
| `enable_prefix_caching` | bool | optional | `false` | **Pipeline-wide.** Prefix cache toggle applied to every stage. |
@@ -27,9 +27,6 @@ The new deploy schema lives under `vllm_omni/deploy/` and is paired with a froze
| `data_parallel_size` | int | optional | `1` | **Pipeline-wide.** DP degree for every stage. |
| `pipeline_parallel_size` | int | optional | `1` | **Pipeline-wide.** PP degree for every stage. |
Note: for diffusion path, `distributed_executor_backend` currently defaults to
`mp`, and `ray` / `external_launcher` are not fully supported yet.
### Stage fields
Each entry under `stages:` accepts any `StageDeployConfig` field directly (no nested `engine_args:`). Only fields whose value legitimately varies across stages live here; pipeline-wide settings (trust_remote_code, distributed_executor_backend, dtype, quantization, prefix/chunked prefill, DP/PP sizes) are declared at the top level and applied to every stage. Unknown keys fall through to `engine_extras:` and are forwarded to the engine.
@@ -195,7 +195,6 @@ Factory function `get_offload_backend()` selects the appropriate backend based o
| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | `"transformer_blocks"` |
| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | - | ✓ | `"layers"`, `"customized modules"` |
**Notes:**
- Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally
@@ -1,56 +0,0 @@
# msModelSlim Quantization
## Overview
[msModelSlim](https://github.com/Ascend/msmodelslim) is an Ascend-friendly compression tool focused on acceleration, using compression techniques, and built for Ascend hardware. It includes a series of inference optimization technologies such as quantization and compression, aiming to accelerate large language dense models, MoE models, multimodal understanding models, multimodal generation models, etc.
Once you have a quantized model which is generated by **msModelSlim**, you can use vLLM Omni for inference by specifying the --quantization ascend parameter to enable quantization features.
### Supported Schemes
| Scheme | Bits | Status |
|--------|------|--------|
| W8A8 | 8 | ✅ Supported |
| W4A4 | 4 | Planned |
W8A8 is the first supported scheme. Additional schemes will be added in future releases.
## Model Quantization
The following example shows how to generate W8A8 quantized weights for the [Wan2_2 model](https://gitcode.com/Ascend/msmodelslim/blob/master/example/multimodal_sd/Wan2_2/README.md).
**Quantization Script:**
```bash
msmodelslim quant \
--model_path /path/to/wan2_2_t2v_float_weights \
--save_path /path/to/wan2_2_t2v_quantized_weights \
--device npu \
--model_type Wan2_2 \
--config_path /lab_practice/wan2_2/wan2_2_w8a8f8_mxfp_t2v.yaml \
--trust_remote_code True
```
After quantization completes, the output directory will contain the quantized model files.
For more examples, refer to the [official examples](https://gitcode.com/Ascend/msit/tree/master/msmodelslim/example).
## Configuration
1. **CLI**: pass `--quantization ascend`.
```bash
# Offline inference
python text_to_image.py --model <your-model> --quantization ascend
# Online serving
vllm serve <your-model> --omni --quantization ascend
```
## Supported Models
| Model | HF Models | Recommendation | `ignored_layers` |
|-------|-----------|---------------|------------------|
| HunyuanImage-3.0 | - | All layers | None |
Currently, quantized HunyuanImage-3.0 weights have not been uploaded to public model platforms such as Hugging Face. You can use a [HunyuanImage-3.0-adapted msModelSlim version](https://gitcode.com/betta18/msmodelslim/tree/hyimage3_mxfp8) to generate the quantized weights manually. We will upload the quantized weights as soon as possible.
+1 -1
View File
@@ -107,7 +107,7 @@ The following tables show which models support each feature:
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
| **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **FLUX.1-dev** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
| **FLUX.1-schnell** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
| **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
@@ -0,0 +1,101 @@
# MOSS-TTS-Nano Offline Inference
## Overview
Single-stage offline TTS pipeline using the 0.1B MOSS-TTS-Nano AR LM and MOSS-Audio-Tokenizer-Nano codec. Outputs 48 kHz stereo WAV.
## Quick Start
```bash
python end2end.py --text "Hello, this is MOSS-TTS-Nano."
```
The first run downloads `OpenMOSS-Team/MOSS-TTS-Nano` and `OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano` from Hugging Face.
## Usage
```
python end2end.py [OPTIONS]
Options:
--text TEXT Text to synthesize (default: "Hello, this is MOSS-TTS-Nano speaking.")
--voice VOICE Built-in voice preset (default: Junhao). See voice table below.
--mode MODE voice_clone (default) or continuation
--prompt-audio PATH Reference WAV/MP3 for custom voice cloning
--prompt-text TEXT Reference transcript (continuation mode)
--max-new-frames N Max AR frames, default 375 (~14 s audio)
--seed INT Random seed for reproducibility
--audio-temperature F Audio sampling temperature (default: 0.8)
--audio-top-k N Audio top-k sampling (default: 25)
--audio-top-p F Audio top-p sampling (default: 0.95)
--text-temperature F Text layer temperature (default: 1.0)
--batch Run a built-in batch of diverse samples (ZH/EN/FR)
--output-dir DIR Directory for WAV outputs (default: /tmp/moss_tts_nano_output)
--deploy-config PATH Override deploy YAML (defaults to vllm_omni/deploy/moss_tts_nano.yaml)
--stage-init-timeout INT Timeout in seconds for stage init (default: 120)
```
## Examples
```bash
# Built-in Chinese voice
python end2end.py --text "你好,这是MOSS-TTS-Nano的语音合成演示。" --voice Junhao
# Built-in English voice
python end2end.py --text "Hello from MOSS-TTS-Nano." --voice Ava
# Custom voice clone
python end2end.py \
--text "Hello, this is a cloned voice." \
--prompt-audio /path/to/reference.wav \
--prompt-text "Exact transcript of the reference audio."
# Batch synthesis (ZH/EN/FR)
python end2end.py --batch --output-dir /tmp/batch_output
# Reproducible output
python end2end.py --text "Deterministic test." --seed 42
```
## Built-in Voice Presets
| Voice | Language |
|-------|----------|
| `Junhao` | ZH |
| `Zhiming` | ZH |
| `Weiguo` | ZH |
| `Xiaoyu` | ZH |
| `Yuewen` | ZH |
| `Lingyu` | ZH |
| `Ava` | EN |
| `Bella` | EN |
| `Adam` | EN |
| `Nathan` | EN |
| `Sakura` | JA |
| `Yui` | JA |
| `Aoi` | JA |
| `Hina` | JA |
| `Mei` | JA |
## Deploy Config
Runtime knobs live in `vllm_omni/deploy/moss_tts_nano.yaml` (auto-loaded;
override with `--deploy-config PATH`). Key stage-level settings:
```yaml
stages:
- stage_id: 0
gpu_memory_utilization: 0.3 # ~2 GB VRAM; increase for faster init
max_num_seqs: 4 # concurrent requests
max_model_len: 4096
```
## Output Format
WAV files, 48 kHz, stereo (2-channel). The codec interleaves stereo as `[L, R, L, R, ...]` in the flat tensor returned by the model.
## Troubleshooting
- **`libnvrtc.so.13: cannot open shared object file`**: torchaudio 2.10+ torchcodec backend requires NVRTC. The model patches `torchaudio.load/save` automatically at load time to fall back to soundfile.
- **`flash_attn not installed`**: The model falls back to `sdpa` attention automatically.
- **Empty audio**: Check that `--text` is non-empty and the model loaded successfully (look for "MOSS-TTS-Nano LM loaded" in logs).
@@ -0,0 +1,180 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Offline inference example for MOSS-TTS-Nano via vLLM-Omni.
Single-stage pipeline: the 0.1B AR LM and MOSS-Audio-Tokenizer-Nano codec
both run inside one generation stage. Output is 48 kHz stereo WAV.
Supports:
- Built-in voice presets (--voice)
- Custom voice cloning via reference audio (--prompt-audio)
- Batch synthesis of multiple texts
- Reproducible output via --seed
Usage examples:
# Built-in voice
python end2end.py --text "Hello from MOSS-TTS-Nano."
# Voice clone with reference audio
python end2end.py --text "Hello!" --prompt-audio /path/to/ref.wav
# Batch with different voices
python end2end.py --batch --output-dir /tmp/moss_output
"""
from __future__ import annotations
import os
from pathlib import Path
import soundfile as sf
import torch
from vllm import SamplingParams
from vllm.utils.argparse_utils import FlexibleArgumentParser
# Prevent multiprocessing from re-importing CUDA in the wrong context.
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
from vllm_omni import Omni # noqa: E402
MODEL = "OpenMOSS-Team/MOSS-TTS-Nano"
BATCH_SAMPLES = [
{"text": "Hello, this is a test of MOSS-TTS-Nano.", "voice": "Ava", "label": "en_hello"},
{"text": "你好,这是 MOSS-TTS-Nano 的语音合成测试。", "voice": "Junhao", "label": "zh_hello"},
{"text": "Bonjour, ceci est un test de synthèse vocale.", "voice": "Bella", "label": "fr_hello"},
]
def build_request(
text: str,
voice: str = "Junhao",
mode: str = "voice_clone",
prompt_audio_path: str | None = None,
prompt_text: str | None = None,
max_new_frames: int = 375,
seed: int | None = None,
audio_temperature: float = 0.8,
audio_top_k: int = 25,
audio_top_p: float = 0.95,
text_temperature: float = 1.0,
) -> dict:
"""Build an Omni request payload for MOSS-TTS-Nano."""
additional: dict = {
"text": [text],
"voice": [voice],
"mode": [mode],
"max_new_frames": [max_new_frames],
"audio_temperature": [audio_temperature],
"audio_top_k": [audio_top_k],
"audio_top_p": [audio_top_p],
"text_temperature": [text_temperature],
}
if prompt_audio_path:
additional["prompt_audio_path"] = [str(prompt_audio_path)]
if prompt_text:
additional["prompt_text"] = [prompt_text]
if seed is not None:
additional["seed"] = [seed]
return {
"prompt": "<|im_start|>assistant\n", # minimal placeholder prompt
"additional_information": additional,
}
def save_audio(waveform: torch.Tensor, path: str, sample_rate: int = 48000) -> None:
audio_np = waveform.float().numpy()
# Reshape stereo: inference_stream yields interleaved [samples*2]; reshape to [samples, 2]
if audio_np.ndim == 1 and audio_np.shape[0] % 2 == 0:
audio_np = audio_np.reshape(-1, 2)
sf.write(path, audio_np, sample_rate)
print(f" Saved {path} ({audio_np.shape}, {sample_rate} Hz)")
def main(args) -> None:
omni = Omni(
model=MODEL,
deploy_config=args.deploy_config,
stage_init_timeout=args.stage_init_timeout,
)
sampling_params = SamplingParams(
temperature=1.0,
top_p=1.0,
top_k=50,
max_tokens=4096,
seed=args.seed if args.seed is not None else 42,
detokenize=False,
)
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
if args.batch:
print(f"Running batch synthesis ({len(BATCH_SAMPLES)} samples)...")
inputs = [build_request(s["text"], voice=s["voice"], seed=args.seed) for s in BATCH_SAMPLES]
params_list = [sampling_params] * len(inputs)
else:
print(f"Synthesizing: {args.text!r}")
inputs = build_request(
text=args.text,
voice=args.voice,
mode=args.mode,
prompt_audio_path=args.prompt_audio,
prompt_text=args.prompt_text,
max_new_frames=args.max_new_frames,
seed=args.seed,
audio_temperature=args.audio_temperature,
audio_top_k=args.audio_top_k,
audio_top_p=args.audio_top_p,
text_temperature=args.text_temperature,
)
params_list = sampling_params
for stage_outputs in omni.generate(inputs, params_list):
for i, req_output in enumerate(stage_outputs.request_output):
for j, out in enumerate(req_output.outputs):
mm = out.multimodal_output
if mm is None:
print(f" [req {i}] No audio output.")
continue
audio = mm.get("audio")
sr_tensor = mm.get("sr")
if audio is None:
print(f" [req {i}] No waveform in multimodal_output.")
continue
sr = int(sr_tensor.item()) if sr_tensor is not None else 48000
label = BATCH_SAMPLES[i]["label"] if args.batch else f"output_{i}_{j}"
out_path = str(output_dir / f"{label}.wav")
save_audio(audio.cpu(), out_path, sr)
print("Done.")
def parse_args():
parser = FlexibleArgumentParser(description="MOSS-TTS-Nano offline inference")
parser.add_argument("--text", default="Hello, this is MOSS-TTS-Nano speaking.", help="Text to synthesize.")
parser.add_argument("--voice", default="Junhao", help="Built-in voice preset name.")
parser.add_argument("--mode", default="voice_clone", choices=["voice_clone", "continuation"])
parser.add_argument("--prompt-audio", default=None, help="Path to reference audio for voice cloning.")
parser.add_argument("--prompt-text", default=None, help="Reference transcript (continuation mode).")
parser.add_argument("--max-new-frames", type=int, default=375, help="Max AR frames (~14s at default).")
parser.add_argument("--seed", type=int, default=None, help="Random seed.")
parser.add_argument("--audio-temperature", type=float, default=0.8)
parser.add_argument("--audio-top-k", type=int, default=25)
parser.add_argument("--audio-top-p", type=float, default=0.95)
parser.add_argument("--text-temperature", type=float, default=1.0)
parser.add_argument("--batch", action="store_true", help="Run built-in batch of diverse samples.")
parser.add_argument("--output-dir", default="/tmp/moss_tts_nano_output", help="Directory for WAV outputs.")
parser.add_argument(
"--deploy-config",
default=None,
help="Path to a deploy YAML; leave unset to auto-load vllm_omni/deploy/moss_tts_nano.yaml.",
)
parser.add_argument("--stage-init-timeout", type=int, default=120)
return parser.parse_args()
if __name__ == "__main__":
main(parse_args())
@@ -300,7 +300,6 @@ def parse_args() -> argparse.Namespace:
default=None,
help=("Custom system prompt. Used when --use-system-prompt is custom. "),
)
current_omni_platform.pre_register_and_update(parser)
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
nullify_stage_engine_defaults(parser)
@@ -0,0 +1,158 @@
# MOSS-TTS-Nano
## Model
| Model | Description |
|-------|-------------|
| `OpenMOSS-Team/MOSS-TTS-Nano` | 0.1B AR LM + MOSS-Audio-Tokenizer-Nano codec, 48 kHz stereo, ZH/EN/JA |
## Gradio Demo
An interactive Gradio demo is available with multilingual voice presets, custom voice cloning, and streaming support.
```bash
# Option 1: Launch server + Gradio together
./run_gradio_demo.sh
# Option 2: If server is already running
python gradio_demo.py --api-base http://localhost:8091
```
Then open http://localhost:7860 in your browser.
Features:
- 15 built-in voice presets (6 ZH · 4 EN · 5 JA)
- Custom voice cloning from uploaded reference audio
- Streaming mode (progressive PCM playback)
## Launch the Server
```bash
vllm serve OpenMOSS-Team/MOSS-TTS-Nano --omni --port 8091
```
The deploy config at `vllm_omni/deploy/moss_tts_nano.yaml` auto-loads; no
`--stage-configs-path`, `--trust-remote-code`, or `--enforce-eager` flags
are needed.
Or use the convenience script:
```bash
./run_server.sh
```
## Send TTS Request
### Using curl
```bash
# Built-in voice preset
curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"input": "你好,这是语音合成测试。",
"voice": "Junhao",
"response_format": "wav"
}' --output output.wav
# English preset
curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"input": "Hello, this is MOSS-TTS-Nano.",
"voice": "Ava",
"response_format": "wav"
}' --output output_en.wav
```
### Custom Voice Cloning
Provide a reference audio (base64 data URL) and its transcript:
```bash
REF_AUDIO=$(base64 -w 0 /path/to/reference.wav)
curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d "{
\"input\": \"Hello, this is a cloned voice.\",
\"voice\": \"Junhao\",
\"ref_audio\": \"data:audio/wav;base64,${REF_AUDIO}\",
\"ref_text\": \"Exact transcript of the reference audio.\"
}" --output cloned.wav
```
### Using Python
```python
import httpx
response = httpx.post(
"http://localhost:8091/v1/audio/speech",
json={
"input": "你好,这是语音合成测试。",
"voice": "Junhao",
"response_format": "wav",
},
timeout=300.0,
)
with open("output.wav", "wb") as f:
f.write(response.content)
```
### Streaming
```bash
curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"input": "Hello, streaming output from MOSS-TTS-Nano.",
"voice": "Ava",
"stream": true,
"response_format": "pcm"
}' --no-buffer | play -t raw -r 48000 -e signed -b 16 -c 2 -
```
**Note:** MOSS-TTS-Nano outputs at 48 kHz stereo (2-channel).
## API Parameters
MOSS-TTS-Nano uses the standard `/v1/audio/speech` endpoint.
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `input` | string | **required** | Text to synthesize (ZH / EN / JA) |
| `voice` | string | `"Junhao"` | Built-in voice name (see table below) |
| `response_format` | string | `"wav"` | Audio format: wav, mp3, flac, pcm |
| `ref_audio` | string | null | Base64 data URL for custom voice cloning |
| `ref_text` | string | null | Transcript of reference audio (required when `ref_audio` is set) |
| `stream` | bool | false | Stream raw PCM chunks |
| `max_new_tokens` | int | 4096 | Maximum tokens to generate |
## Built-in Voice Presets
| Voice | Language | Description |
|-------|----------|-------------|
| `Junhao` | ZH | Male, standard Mandarin |
| `Zhiming` | ZH | Male |
| `Weiguo` | ZH | Male |
| `Xiaoyu` | ZH | Female |
| `Yuewen` | ZH | Female |
| `Lingyu` | ZH | Female |
| `Ava` | EN | Female, American English |
| `Bella` | EN | Female |
| `Adam` | EN | Male |
| `Nathan` | EN | Male |
| `Sakura` | JA | Female |
| `Yui` | JA | Female |
| `Aoi` | JA | Female |
| `Hina` | JA | Female |
| `Mei` | JA | Female |
## Troubleshooting
1. **`libnvrtc.so.13: cannot open shared object file`**: torchaudio 2.10+ defaults to torchcodec which requires NVRTC. vLLM-Omni patches this automatically at model load time to use soundfile instead.
2. **Connection refused**: Ensure the server is running on the correct port.
3. **Flashinfer version mismatch**: Set `FLASHINFER_DISABLE_VERSION_CHECK=1` if you see version warnings.
4. **Out of memory**: The default `gpu_memory_utilization=0.3` is conservative. Increase it in the stage config if you have more VRAM available.
@@ -0,0 +1,690 @@
"""Gradio demo for MOSS-TTS-Nano with gapless streaming audio playback.
Uses a custom AudioWorklet-based player (adapted from the Qwen3-TTS demo)
for gap-free streaming. Audio is streamed from the vLLM-Omni server
through a same-origin proxy and played via the Web Audio API's AudioWorklet.
MOSS-TTS-Nano outputs 48 kHz stereo (2-channel interleaved L,R,L,R,...).
The AudioWorklet downmixes stereo to mono during playback.
Also supports non-streaming mode (full audio download) via gr.Audio.
Usage:
# Start the server first (see run_server.sh), then:
python gradio_demo.py --api-base http://localhost:8091
# Or use run_gradio_demo.sh to start both server and demo together.
"""
from __future__ import annotations
import argparse
import base64
import io
import json
import logging
try:
import gradio as gr
except ImportError:
raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
import httpx
import numpy as np
import soundfile as sf
from fastapi import FastAPI, Request
from fastapi.responses import Response, StreamingResponse
logger = logging.getLogger(__name__)
# MOSS-TTS-Nano outputs 48 kHz stereo (2-channel interleaved).
PCM_SAMPLE_RATE = 48000
PCM_CHANNELS = 2
DEFAULT_API_BASE = "http://localhost:8091"
# Built-in voice presets (matches the MOSS-TTS-Nano GitHub repo assets/)
VOICE_PRESETS = {
# Chinese
"Junhao (ZH)": "Junhao",
"Zhiming (ZH)": "Zhiming",
"Weiguo (ZH)": "Weiguo",
"Xiaoyu (ZH)": "Xiaoyu",
"Yuewen (ZH)": "Yuewen",
"Lingyu (ZH)": "Lingyu",
# English
"Ava (EN)": "Ava",
"Bella (EN)": "Bella",
"Adam (EN)": "Adam",
"Nathan (EN)": "Nathan",
# Japanese
"Sakura (JA)": "Sakura",
"Yui (JA)": "Yui",
"Aoi (JA)": "Aoi",
"Hina (JA)": "Hina",
"Mei (JA)": "Mei",
}
PRESET_LABELS = list(VOICE_PRESETS.keys())
DEFAULT_PRESET = "Junhao (ZH)"
# ── AudioWorklet processor (loaded in browser via Blob URL) ──────────
# Downmixes stereo-interleaved int16 PCM to mono float32 for playback.
WORKLET_JS = r"""
class TTSPlaybackProcessor extends AudioWorkletProcessor {
constructor() {
super();
this.queue = [];
this.buf = null;
this.pos = 0;
this.playing = false;
this.played = 0;
this.port.onmessage = (e) => {
if (e.data && e.data.type === 'clear') {
this.queue = []; this.buf = null; this.pos = 0; this.played = 0;
if (this.playing) { this.playing = false; this.port.postMessage({type:'stopped'}); }
return;
}
// Input: Int16Array of stereo-interleaved samples [L0,R0,L1,R1,...]
// Convert to mono float32 by averaging L+R channels.
const pcm = e.data;
const monoLen = Math.floor(pcm.length / 2);
const mono = new Float32Array(monoLen);
for (let i = 0; i < monoLen; i++) {
mono[i] = (pcm[i * 2] + pcm[i * 2 + 1]) / 65536.0;
}
this.queue.push(mono);
};
}
process(inputs, outputs) {
const out = outputs[0][0];
for (let i = 0; i < out.length; i++) {
if (!this.buf || this.pos >= this.buf.length) {
if (this.queue.length > 0) {
this.buf = this.queue.shift(); this.pos = 0;
} else {
for (let j = i; j < out.length; j++) out[j] = 0;
if (this.playing) { this.playing = false; this.port.postMessage({type:'stopped', played:this.played}); }
return true;
}
}
out[i] = this.buf[this.pos++];
this.played++;
}
if (!this.playing) { this.playing = true; this.port.postMessage({type:'started'}); }
return true;
}
}
registerProcessor('tts-playback-processor', TTSPlaybackProcessor);
"""
# ── Player HTML (container with metric cards) ────────────────────────
PLAYER_HTML = """
<div id="tts-player">
<div style="display:flex; align-items:center; gap:10px;">
<div id="tts-status-dot" style="width:10px;height:10px;border-radius:50%;background:#ccc;flex-shrink:0;"></div>
<span id="tts-status" style="font-weight:600;font-size:1.05em;">Ready</span>
<button id="tts-stop-btn" onclick="window.ttsStop()"
style="display:none; margin-left:auto; padding:5px 16px; border-radius:6px; border:1px solid #EF5552;
background:#fff; color:#EF5552; cursor:pointer; font-size:0.85em; transition:all 0.15s;">Stop</button>
</div>
<div id="tts-metrics" style="display:none; grid-template-columns:repeat(4,1fr); gap:10px; margin-top:12px;">
<div style="background:#f8f9fa;border-radius:6px;padding:8px 10px;text-align:center;">
<div style="font-size:0.7em;text-transform:uppercase;color:#888;letter-spacing:0.5px;margin-bottom:2px;">TTFP</div>
<div id="tts-m-ttfp" style="font-size:1.2em;font-weight:700;color:#333;">—</div>
</div>
<div style="background:#f8f9fa;border-radius:6px;padding:8px 10px;text-align:center;">
<div style="font-size:0.7em;text-transform:uppercase;color:#888;letter-spacing:0.5px;margin-bottom:2px;">RTF</div>
<div id="tts-m-rtf" style="font-size:1.2em;font-weight:700;color:#333;">—</div>
</div>
<div style="background:#f8f9fa;border-radius:6px;padding:8px 10px;text-align:center;">
<div style="font-size:0.7em;text-transform:uppercase;color:#888;letter-spacing:0.5px;margin-bottom:2px;">Audio</div>
<div id="tts-m-dur" style="font-size:1.2em;font-weight:700;color:#333;">—</div>
</div>
<div style="background:#f8f9fa;border-radius:6px;padding:8px 10px;text-align:center;">
<div style="font-size:0.7em;text-transform:uppercase;color:#888;letter-spacing:0.5px;margin-bottom:2px;">Speed</div>
<div id="tts-m-speed" style="font-size:1.2em;font-weight:700;color:#333;">—</div>
</div>
</div>
<div id="tts-rtf-bar-wrap" style="display:none; background:#e8ecf1; border-radius:4px; height:20px; overflow:hidden; position:relative; margin-top:10px;">
<div id="tts-rtf-bar" style="height:100%; border-radius:4px; transition:width 0.3s ease, background 0.3s ease; width:0%;"></div>
<span id="tts-rtf-label" style="position:absolute;top:50%;left:50%;transform:translate(-50%,-50%);font-size:0.75em;font-weight:600;color:#444;"></span>
</div>
<div id="tts-elapsed" style="display:none; margin-top:6px; font-size:0.8em; color:#999; text-align:right;"></div>
</div>
"""
def _build_player_js(sample_rate: int) -> str:
"""Build the JavaScript that powers the AudioWorklet player."""
return f"""
<script>
const SR = {sample_rate};
const WC = {json.dumps(WORKLET_JS)};
let ctx = null, node = null, abort = null, gen = false, st = {{}};
async function init() {{
if (ctx) return;
ctx = new AudioContext({{ sampleRate: SR }});
const b = new Blob([WC], {{ type: 'application/javascript' }});
const u = URL.createObjectURL(b);
await ctx.audioWorklet.addModule(u);
URL.revokeObjectURL(u);
node = new AudioWorkletNode(ctx, 'tts-playback-processor');
node.connect(ctx.destination);
node.port.onmessage = (e) => {{
if (e.data.type === 'started') setStatus('Playing...', '#2E7D32');
else if (e.data.type === 'stopped' && !gen) {{
setStatus('Done', '#2E7D32'); showStats(true);
const btn = document.getElementById('tts-stop-btn');
if (btn) btn.style.display = 'none';
}}
}};
}}
function setStatus(text, color) {{
const s = document.getElementById('tts-status');
const d = document.getElementById('tts-status-dot');
if (s) s.textContent = text;
if (d) d.style.background = color || '#ccc';
}}
function showStats(fin) {{
if (!st.t0) return;
const elapsed = (fin && st.streamEnd ? (st.streamEnd - st.t0) : (performance.now() - st.t0)) / 1000;
// st.samples counts stereo int16 samples; mono frames = samples / 2
const dur = (st.samples / 2) / SR;
const mTtfp = document.getElementById('tts-m-ttfp');
const mRtf = document.getElementById('tts-m-rtf');
const mDur = document.getElementById('tts-m-dur');
const mSpeed = document.getElementById('tts-m-speed');
const bar = document.getElementById('tts-rtf-bar');
const barLabel = document.getElementById('tts-rtf-label');
const elapsedEl = document.getElementById('tts-elapsed');
if (mTtfp && st.ttfp != null) mTtfp.textContent = st.ttfp.toFixed(0) + 'ms';
if (mDur) mDur.textContent = dur.toFixed(1) + 's';
if (dur > 0 && elapsed > 0) {{
const rtf = elapsed / dur;
const speed = 1 / rtf;
if (mRtf) {{
mRtf.textContent = rtf.toFixed(2) + 'x';
mRtf.style.color = rtf < 1 ? '#2E7D32' : rtf < 1.5 ? '#e8a317' : '#EF5552';
}}
if (mSpeed) {{
mSpeed.textContent = speed.toFixed(1) + 'x';
mSpeed.style.color = speed > 1 ? '#2E7D32' : speed > 0.7 ? '#e8a317' : '#EF5552';
}}
if (bar) {{
const pct = Math.min(speed / 3 * 100, 100);
bar.style.width = pct + '%';
bar.style.background = speed > 1 ? 'linear-gradient(90deg,#66BB6A,#2E7D32)' : speed > 0.7 ? 'linear-gradient(90deg,#e8a317,#f0b866)' : 'linear-gradient(90deg,#EF5552,#f87171)';
}}
if (barLabel) barLabel.textContent = speed.toFixed(1) + 'x realtime';
}}
if (elapsedEl) {{
elapsedEl.style.display = 'block';
elapsedEl.textContent = fin ? 'Completed in ' + elapsed.toFixed(1) + 's (' + st.chunks + ' chunks)' : elapsed.toFixed(1) + 's elapsed (' + st.chunks + ' chunks)';
}}
}}
window.ttsStop = function() {{
if (abort) abort.abort();
if (node) node.port.postMessage({{ type: 'clear' }});
gen = false;
setStatus('Stopped', '#999');
const btn = document.getElementById('tts-stop-btn');
if (btn) btn.style.display = 'none';
}};
window.ttsGenerate = async function(payload) {{
try {{ await init(); if (ctx.state === 'suspended') await ctx.resume(); }}
catch (e) {{ const s = document.getElementById('tts-status'); if (s) s.textContent = 'Audio init error: ' + e.message; return; }}
if (abort) abort.abort();
node.port.postMessage({{ type: 'clear' }});
await new Promise(r => setTimeout(r, 50));
node.port.postMessage({{ type: 'clear' }});
gen = true;
st = {{ t0: null, chunks: 0, samples: 0, ttfp: null }};
setStatus('Connecting...', '#2E7D32');
const bEl = document.getElementById('tts-stop-btn');
if (bEl) bEl.style.display = 'inline-block';
const mp = document.getElementById('tts-metrics');
if (mp) {{ mp.style.display = 'grid'; ['tts-m-ttfp','tts-m-rtf','tts-m-dur','tts-m-speed'].forEach(id => {{ const e = document.getElementById(id); if(e) {{ e.textContent = '\\u2014'; e.style.color = '#333'; }} }}); }}
const bw = document.getElementById('tts-rtf-bar-wrap');
if (bw) bw.style.display = 'block';
const bar = document.getElementById('tts-rtf-bar');
if (bar) bar.style.width = '0%';
const bl = document.getElementById('tts-rtf-label');
if (bl) bl.textContent = '';
const ee = document.getElementById('tts-elapsed');
if (ee) {{ ee.style.display = 'none'; ee.textContent = ''; }}
abort = new AbortController();
try {{
st.t0 = performance.now();
const r = await fetch('/proxy/v1/audio/speech', {{
method: 'POST',
headers: {{ 'Content-Type': 'application/json' }},
body: JSON.stringify(payload),
signal: abort.signal,
}});
if (!r.ok) {{ const t = await r.text(); throw new Error('Server ' + r.status + ': ' + t.slice(0, 200)); }}
setStatus('Streaming...', '#43A047');
const reader = r.body.getReader();
let left = new Uint8Array(0);
while (true) {{
const {{ done, value }} = await reader.read();
if (done) break;
let raw;
if (left.length > 0) {{
raw = new Uint8Array(left.length + value.length);
raw.set(left); raw.set(value, left.length);
}} else {{ raw = value; }}
const usable = raw.length - (raw.length % 2);
left = usable < raw.length ? raw.slice(usable) : new Uint8Array(0);
if (usable > 0) {{
const ab = new ArrayBuffer(usable);
new Uint8Array(ab).set(raw.subarray(0, usable));
const pcm = new Int16Array(ab);
node.port.postMessage(pcm);
st.chunks++;
st.samples += pcm.length;
if (st.ttfp == null) st.ttfp = performance.now() - st.t0;
showStats(false);
}}
}}
}} catch (e) {{
if (e.name !== 'AbortError') {{
setStatus('Error: ' + e.message, '#EF5552');
console.error('TTS error:', e);
}}
}} finally {{
st.streamEnd = performance.now();
showStats(true);
gen = false;
if (st.samples > 0) setStatus('Finishing playback...', '#2E7D32');
else {{
setStatus('No audio received', '#999');
if (bEl) bEl.style.display = 'none';
}}
}}
}};
</script>
"""
def encode_audio_to_base64(audio_data: tuple) -> str:
"""Encode Gradio audio input (sample_rate, numpy_array) to base64 data URL."""
sample_rate, audio_np = audio_data
if audio_np.dtype != np.int16:
if audio_np.dtype in (np.float32, np.float64):
audio_np = np.clip(audio_np, -1.0, 1.0)
audio_np = (audio_np * 32767).astype(np.int16)
else:
audio_np = audio_np.astype(np.int16)
buf = io.BytesIO()
sf.write(buf, audio_np, sample_rate, format="WAV")
wav_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
return f"data:audio/wav;base64,{wav_b64}"
def build_payload(
text: str,
voice_preset: str,
ref_audio: tuple | None,
ref_text: str,
response_format: str = "pcm",
stream: bool = True,
) -> dict:
"""Build the /v1/audio/speech request payload."""
if not text or not text.strip():
raise gr.Error("Please enter text to synthesize.")
voice_name = VOICE_PRESETS.get(voice_preset, "Junhao")
payload: dict = {
"input": text.strip(),
"voice": voice_name,
"response_format": "pcm" if stream else response_format,
"stream": stream,
}
has_ref = ref_audio is not None
if has_ref:
payload["ref_audio"] = encode_audio_to_base64(ref_audio)
ref_text_stripped = ref_text.strip() if ref_text else ""
if ref_text_stripped:
payload["ref_text"] = ref_text_stripped
else:
raise gr.Error("Voice cloning requires a transcript of the reference audio.")
return payload
def generate_speech(
api_base: str,
text: str,
voice_preset: str,
ref_audio: tuple | None,
ref_text: str,
response_format: str,
) -> tuple:
"""Non-streaming: call /v1/audio/speech and return full audio as (sr, np_array)."""
payload = build_payload(text, voice_preset, ref_audio, ref_text, response_format, stream=False)
try:
with httpx.Client(timeout=300.0) as client:
resp = client.post(
f"{api_base}/v1/audio/speech",
json=payload,
headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"},
)
except httpx.TimeoutException:
raise gr.Error("Request timed out. The server may be busy.")
except httpx.ConnectError:
raise gr.Error(f"Cannot connect to server at {api_base}. Is the server running?")
if resp.status_code != 200:
raise gr.Error(f"Server error ({resp.status_code}): {resp.text}")
try:
if response_format == "pcm":
# PCM is stereo interleaved; downmix to mono for Gradio.
samples = np.frombuffer(resp.content, dtype=np.int16)
n = len(samples) - (len(samples) % 2)
left = samples[:n:2].astype(np.float32)
right = samples[1:n:2].astype(np.float32)
mono = (left + right) / 65536.0
return (PCM_SAMPLE_RATE, mono)
# WAV/MP3/FLAC: soundfile handles header parsing.
audio_np, sample_rate = sf.read(io.BytesIO(resp.content))
if audio_np.ndim > 1:
audio_np = audio_np.mean(axis=1)
return (int(sample_rate), audio_np.astype(np.float32))
except Exception as e:
raise gr.Error(f"Failed to decode audio: {e}")
def create_app(api_base: str):
"""Create the FastAPI app with streaming proxy + Gradio UI."""
fastapi_app = FastAPI()
# ── Streaming proxy (same-origin, no CORS issues) ────────────
@fastapi_app.post("/proxy/v1/audio/speech")
async def proxy_speech(request: Request):
body = await request.json()
logger.info(
"Proxy request: input=%r voice=%s stream=%s",
body.get("input", "")[:50],
body.get("voice"),
body.get("stream"),
)
try:
client = httpx.AsyncClient(timeout=300)
resp = await client.send(
client.build_request(
"POST",
f"{api_base}/v1/audio/speech",
json=body,
headers={"Authorization": "Bearer EMPTY", "Content-Type": "application/json"},
),
stream=True,
)
except Exception as exc:
logger.exception("Proxy connection error")
await client.aclose()
return Response(content=str(exc), status_code=502)
if resp.status_code != 200:
content = await resp.aread()
logger.error("Proxy upstream error %d: %s", resp.status_code, content[:200])
await resp.aclose()
await client.aclose()
return Response(content=content, status_code=resp.status_code)
async def relay():
total = 0
try:
async for chunk in resp.aiter_bytes():
total += len(chunk)
yield chunk
except Exception:
logger.exception("Proxy relay error after %d bytes", total)
finally:
logger.info("Proxy relay done: %d bytes", total)
await resp.aclose()
await client.aclose()
return StreamingResponse(relay(), media_type="application/octet-stream")
# ── Gradio UI ────────────────────────────────────────────────
css = """
#generate-btn button { width: 100%; }
#streaming-player { border: 1px solid var(--border-color-primary) !important; border-radius: var(--block-radius) !important; padding: var(--block-padding) !important; }
"""
theme = gr.themes.Default(
primary_hue=gr.themes.Color(
c50="#E8F5E9",
c100="#C8E6C9",
c200="#A5D6A7",
c300="#81C784",
c400="#66BB6A",
c500="#4CAF50",
c600="#43A047",
c700="#388E3C",
c800="#2E7D32",
c900="#1B5E20",
c950="#0D3B10",
),
)
with gr.Blocks(title="MOSS-TTS-Nano Demo") as demo:
gr.HTML(f"""
<div style="display:flex; align-items:center; gap:16px; margin-bottom:8px;">
<img src="https://raw.githubusercontent.com/vllm-project/vllm-omni/main/docs/source/logos/vllm-omni-logo.png"
alt="vLLM-Omni" style="height:42px;">
<div>
<h1 style="margin:0; font-size:1.5em;">MOSS-TTS-Nano Streaming Demo</h1>
<span style="font-size:0.85em; color:#666;">
0.1B AR LM + MOSS-Audio-Tokenizer-Nano &middot; 48 kHz &middot; ZH / EN / JA
&nbsp;&middot;&nbsp; Served by
<a href="https://github.com/vllm-project/vllm-omni" target="_blank"
style="color:#43A047; text-decoration:none; font-weight:600;">vLLM-Omni</a>
&nbsp;&middot;&nbsp; <code style="background:#eef2f7; padding:2px 6px; border-radius:4px; font-size:0.9em;">{api_base}</code>
</span>
</div>
</div>
""")
with gr.Row():
with gr.Column(scale=3):
text_input = gr.Textbox(
label="Text to Synthesize",
placeholder="Enter text in Chinese, English, or Japanese...",
lines=4,
)
voice_preset = gr.Dropdown(
choices=PRESET_LABELS,
value=DEFAULT_PRESET,
label="Voice Preset",
info="Built-in multilingual voices. ZH = Chinese, EN = English, JA = Japanese.",
)
with gr.Accordion("Custom Voice Clone (optional)", open=False):
gr.Markdown(
"Upload a reference audio clip (10-30 s) and provide its transcript "
"to clone that voice instead of using the built-in preset."
)
ref_audio = gr.Audio(
label="Reference Audio",
type="numpy",
sources=["upload", "microphone"],
)
ref_text = gr.Textbox(
label="Reference Audio Transcript (required for cloning)",
placeholder="Exact transcript of the reference audio...",
lines=2,
)
with gr.Row():
response_format = gr.Dropdown(
choices=["wav", "mp3", "flac", "pcm"],
value="pcm",
label="Audio Format",
interactive=False,
scale=1,
)
stream_checkbox = gr.Checkbox(
label="Stream (gapless)",
value=True,
info="AudioWorklet streaming (recommended)",
scale=1,
)
with gr.Row():
generate_btn = gr.Button(
"Generate Speech",
variant="primary",
size="lg",
elem_id="generate-btn",
scale=3,
)
reset_btn = gr.Button(
"Reset",
variant="secondary",
size="lg",
scale=1,
)
with gr.Column(scale=2):
player_html = gr.HTML(
value=PLAYER_HTML,
visible=True,
label="streaming player",
elem_id="streaming-player",
)
audio_output = gr.Audio(
label="Generated Audio",
interactive=False,
autoplay=True,
visible=False,
)
gr.Examples(
examples=[
["Hello, this is MOSS-TTS-Nano speaking.", "Ava (EN)"],
["The quick brown fox jumps over the lazy dog.", "Adam (EN)"],
["MOSS-TTS-Nanoは、軽量なテキスト読み上げモデルです。", "Sakura (JA)"],
],
inputs=[text_input, voice_preset],
label="Examples",
)
gr.HTML("""
<div style="text-align:center; padding:8px 0; margin-top:4px;">
<a href="https://github.com/vllm-project/vllm-omni" target="_blank">
<img src="https://raw.githubusercontent.com/vllm-project/vllm-omni/main/docs/source/logos/vllm-omni-logo.png"
alt="vLLM-Omni" style="height:28px; opacity:0.7;">
</a>
</div>
""")
# Hidden textbox to pass payload from Python -> JavaScript
hidden_payload = gr.Textbox(visible=False, elem_id="tts-payload")
# ── Event wiring ─────────────────────────────────────────
def on_stream_change(stream: bool):
if stream:
return (
gr.update(value="pcm", interactive=False),
gr.update(visible=True), # player
gr.update(visible=False), # audio
)
return (
gr.update(value="wav", interactive=True),
gr.update(visible=False),
gr.update(visible=True),
)
stream_checkbox.change(
fn=on_stream_change,
inputs=[stream_checkbox],
outputs=[response_format, player_html, audio_output],
)
def on_reset():
return (
"", # text
None, # audio_output
"", # hidden_payload
PLAYER_HTML, # reset player
)
reset_btn.click(
fn=on_reset,
outputs=[text_input, audio_output, hidden_payload, player_html],
js="() => { if (window.ttsStop) window.ttsStop(); }",
)
all_inputs = [text_input, voice_preset, ref_audio, ref_text, response_format]
def on_generate(stream_enabled, *args):
import time as _time
text, voice, ref_a, ref_t, _fmt = args
if stream_enabled:
payload = build_payload(text, voice, ref_a, ref_t, stream=True)
payload["_nonce"] = int(_time.time() * 1000)
return json.dumps(payload), gr.update()
else:
audio = generate_speech(api_base, text, voice, ref_a, ref_t, _fmt)
return "", audio
generate_btn.click(
fn=on_generate,
inputs=[stream_checkbox] + all_inputs,
outputs=[hidden_payload, audio_output],
).then(
fn=lambda p: p,
inputs=[hidden_payload],
outputs=[hidden_payload],
js="(p) => { if (p && p.trim()) { const d = JSON.parse(p); delete d._nonce; window.ttsGenerate(d); } return p; }",
)
demo.queue()
return gr.mount_gradio_app(
fastapi_app,
demo,
path="/",
css=css,
theme=theme,
head=_build_player_js(PCM_SAMPLE_RATE),
)
def main():
parser = argparse.ArgumentParser(
description="Gradio demo for MOSS-TTS-Nano with gapless AudioWorklet streaming.",
)
parser.add_argument("--api-base", default=DEFAULT_API_BASE, help=f"API base URL (default: {DEFAULT_API_BASE})")
parser.add_argument("--host", default="0.0.0.0", help="Gradio host (default: 0.0.0.0)")
parser.add_argument("--port", type=int, default=7860, help="Gradio port (default: 7860)")
parser.add_argument("--share", action="store_true", help="Share publicly via Gradio tunnel")
args = parser.parse_args()
logging.basicConfig(level=logging.INFO)
print(f"Connecting to vLLM-Omni server at: {args.api_base}")
app = create_app(args.api_base)
import uvicorn
uvicorn.run(app, host=args.host, port=args.port)
if __name__ == "__main__":
main()
+43
View File
@@ -0,0 +1,43 @@
#!/bin/bash
# Launch MOSS-TTS-Nano server + Gradio demo together.
#
# Usage:
# ./run_gradio_demo.sh
# CUDA_VISIBLE_DEVICES=0 PORT=8091 GRADIO_PORT=7860 ./run_gradio_demo.sh
set -e
MODEL="${MODEL:-OpenMOSS-Team/MOSS-TTS-Nano}"
PORT="${PORT:-8091}"
GRADIO_PORT="${GRADIO_PORT:-7860}"
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
echo "Starting MOSS-TTS-Nano server (port $PORT)..."
FLASHINFER_DISABLE_VERSION_CHECK=1 \
vllm serve "$MODEL" \
--host 0.0.0.0 \
--port "$PORT" \
--omni &
SERVER_PID=$!
cleanup() {
echo "Stopping server (PID $SERVER_PID)..."
kill $SERVER_PID 2>/dev/null
wait $SERVER_PID 2>/dev/null
}
trap cleanup EXIT
# Wait for server to be ready.
echo "Waiting for server to start..."
for i in $(seq 1 120); do
if curl -s "http://localhost:$PORT/health" > /dev/null 2>&1; then
echo "Server ready."
break
fi
sleep 2
done
echo "Starting Gradio demo (port $GRADIO_PORT)..."
python "$SCRIPT_DIR/gradio_demo.py" \
--api-base "http://localhost:$PORT" \
--port "$GRADIO_PORT"
+19
View File
@@ -0,0 +1,19 @@
#!/bin/bash
# Launch vLLM-Omni server for MOSS-TTS-Nano
#
# Usage:
# ./run_server.sh
# CUDA_VISIBLE_DEVICES=0 PORT=8091 ./run_server.sh
set -e
MODEL="${MODEL:-OpenMOSS-Team/MOSS-TTS-Nano}"
PORT="${PORT:-8091}"
echo "Starting MOSS-TTS-Nano server with model: $MODEL"
FLASHINFER_DISABLE_VERSION_CHECK=1 \
vllm serve "$MODEL" \
--host 0.0.0.0 \
--port "$PORT" \
--omni
-103
View File
@@ -1,103 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""conftest.py for benchmarks unit tests.
Installs lightweight mock stubs for ``vllm`` (and sub-packages) so the
data-module unit tests can run without a full vLLM installation. Only the
symbols actually imported by
``vllm_omni.benchmarks.data_modules.seed_tts_dataset`` are emulated.
"""
from __future__ import annotations
import sys
import types
from dataclasses import dataclass
from typing import Any
def _install_vllm_stubs() -> None:
"""Register minimal vllm stubs in sys.modules.
Only installs when real vllm is unavailable. We actively probe the
import because an empty or partial vllm may not yet have imported
the submodules we rely on, and unconditionally registering stubs
would shadow the real package for sibling tests (e.g.
``tests/benchmarks/metrics/test_metrics.py`` needs the real
``vllm.benchmarks.serve``).
"""
try:
import vllm.benchmarks.datasets # noqa: F401
import vllm.tokenizers # noqa: F401
except ImportError:
pass
else:
return # real vllm available — do not shadow it
if "vllm.benchmarks.datasets" in sys.modules:
return
# ------------------------------------------------------------------ #
# vllm.benchmarks.datasets #
# ------------------------------------------------------------------ #
@dataclass
class SampleRequest:
prompt: str = ""
prompt_len: int = 0
expected_output_len: int = 0
multi_modal_data: Any = None
request_id: str = ""
class BenchmarkDataset:
def __init__(
self,
dataset_path: str = "",
random_seed: int = 0,
disable_shuffle: bool = False,
**kwargs: Any,
) -> None:
self.dataset_path = dataset_path
self.random_seed = random_seed
self.disable_shuffle = disable_shuffle
def maybe_oversample_requests(
self,
out: list,
num_requests: int,
request_id_prefix: str,
no_oversample: bool,
) -> None:
pass
# ------------------------------------------------------------------ #
# vllm.tokenizers / vllm.tokenizers.hf #
# ------------------------------------------------------------------ #
class TokenizerLike:
pass
def get_cached_tokenizer(t: Any) -> Any:
return t
# ------------------------------------------------------------------ #
# Wire up sys.modules #
# ------------------------------------------------------------------ #
vllm_mod = types.ModuleType("vllm")
vllm_benchmarks = types.ModuleType("vllm.benchmarks")
vllm_benchmarks_datasets = types.ModuleType("vllm.benchmarks.datasets")
vllm_tokenizers = types.ModuleType("vllm.tokenizers")
vllm_tokenizers_hf = types.ModuleType("vllm.tokenizers.hf")
vllm_benchmarks_datasets.BenchmarkDataset = BenchmarkDataset # type: ignore[attr-defined]
vllm_benchmarks_datasets.SampleRequest = SampleRequest # type: ignore[attr-defined]
vllm_tokenizers.TokenizerLike = TokenizerLike # type: ignore[attr-defined]
vllm_tokenizers_hf.get_cached_tokenizer = get_cached_tokenizer # type: ignore[attr-defined]
sys.modules["vllm"] = vllm_mod
sys.modules["vllm.benchmarks"] = vllm_benchmarks
sys.modules["vllm.benchmarks.datasets"] = vllm_benchmarks_datasets
sys.modules["vllm.tokenizers"] = vllm_tokenizers
sys.modules["vllm.tokenizers.hf"] = vllm_tokenizers_hf
# Install stubs immediately at collection time (before any test import).
_install_vllm_stubs()
-139
View File
@@ -1,139 +0,0 @@
"""Tests for the universal benchmarks/tts/bench_tts.py CLI."""
from __future__ import annotations
import json
import sys
from pathlib import Path
import pytest
import yaml
# Add benchmarks/tts to path for import
sys.path.insert(0, str(Path(__file__).parent.parent.parent / "benchmarks" / "tts"))
import bench_tts
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@pytest.fixture()
def model_configs_path(tmp_path: Path) -> Path:
cfg = {
"models": {
"test/ModelA": {
"stage_config": "model_a.yaml",
"supported_tasks": ["voice_clone", "default_voice"],
"backend": "openai-audio-speech",
"endpoint": "/v1/audio/speech",
"task_extra_body": {
"voice_clone": {"task_type": "Base"},
"default_voice": {"voice": "Vivian", "task_type": "CustomVoice"},
},
},
"test/ModelB": {
"stage_config": "model_b.yaml",
"supported_tasks": ["voice_clone"],
"backend": "openai-audio-speech",
"endpoint": "/v1/audio/speech",
"task_extra_body": {"voice_clone": {}},
},
}
}
p = tmp_path / "model_configs.yaml"
p.write_text(yaml.dump(cfg), encoding="utf-8")
return p
def test_load_model_configs(model_configs_path: Path) -> None:
configs = bench_tts.load_model_configs(model_configs_path)
assert "test/ModelA" in configs
assert "test/ModelB" in configs
assert configs["test/ModelA"]["supported_tasks"] == ["voice_clone", "default_voice"]
def test_build_bench_args_voice_clone(model_configs_path: Path) -> None:
configs = bench_tts.load_model_configs(model_configs_path)
cmd = bench_tts.build_bench_args(
host="localhost",
port=8000,
model="test/ModelA",
task="voice_clone",
model_cfg=configs["test/ModelA"],
locale="en",
num_prompts=10,
concurrency=1,
dataset_path="/data/seed-tts",
wer_eval=False,
output_dir=None,
result_filename=None,
extra_cli_args=[],
)
assert "--dataset-name" in cmd
idx = cmd.index("--dataset-name")
assert cmd[idx + 1] == "seed-tts"
assert "--max-concurrency" in cmd
assert "--extra-body" in cmd
extra_body = json.loads(cmd[cmd.index("--extra-body") + 1])
assert extra_body.get("task_type") == "Base"
def test_build_bench_args_default_voice_has_voice_param(model_configs_path: Path) -> None:
configs = bench_tts.load_model_configs(model_configs_path)
cmd = bench_tts.build_bench_args(
host="localhost",
port=8000,
model="test/ModelA",
task="default_voice",
model_cfg=configs["test/ModelA"],
locale="en",
num_prompts=10,
concurrency=1,
dataset_path="/data/seed-tts",
wer_eval=False,
output_dir=None,
result_filename=None,
extra_cli_args=[],
)
idx = cmd.index("--dataset-name")
assert cmd[idx + 1] == "seed-tts-text"
extra_body = json.loads(cmd[cmd.index("--extra-body") + 1])
assert extra_body.get("voice") == "Vivian"
def test_build_bench_args_wer_eval_adds_flag(model_configs_path: Path) -> None:
configs = bench_tts.load_model_configs(model_configs_path)
cmd = bench_tts.build_bench_args(
host="localhost",
port=8000,
model="test/ModelA",
task="voice_clone",
model_cfg=configs["test/ModelA"],
locale="en",
num_prompts=10,
concurrency=1,
dataset_path="/data/seed-tts",
wer_eval=True,
output_dir=None,
result_filename=None,
extra_cli_args=[],
)
assert "--seed-tts-wer-eval" in cmd
def test_unsupported_task_exits(model_configs_path: Path, capsys: pytest.CaptureFixture, mocker) -> None:
# ModelB does not support voice_design
mocker.patch.object(
sys,
"argv",
[
"bench_tts.py",
"--model",
"test/ModelB",
"--task",
"voice_design",
"--model-configs",
str(model_configs_path),
],
)
with pytest.raises(SystemExit):
bench_tts.main()
@@ -1,168 +0,0 @@
"""Tests for SeedTTSTextDataset, SeedTTSTextSampleRequest, SeedTTSDesignDataset,
and SeedTTSDesignSampleRequest.
vllm stubs are installed by tests/benchmarks/conftest.py before collection.
"""
from __future__ import annotations
import importlib.util
import sys
from pathlib import Path
import pytest
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
# Load the data module directly (bypasses vllm_omni.__init__ heavy imports).
_REPO_ROOT = Path(__file__).resolve().parents[2]
_MODULE_PATH = _REPO_ROOT / "vllm_omni" / "benchmarks" / "data_modules" / "seed_tts_dataset.py"
_MODULE_NAME = "vllm_omni.benchmarks.data_modules.seed_tts_dataset"
if _MODULE_NAME not in sys.modules:
_spec = importlib.util.spec_from_file_location(_MODULE_NAME, _MODULE_PATH)
_mod = importlib.util.module_from_spec(_spec)
sys.modules[_MODULE_NAME] = _mod
_spec.loader.exec_module(_mod)
from vllm_omni.benchmarks.data_modules.seed_tts_dataset import ( # noqa: E402
SeedTTSDesignDataset,
SeedTTSDesignSampleRequest,
SeedTTSTextDataset,
SeedTTSTextSampleRequest,
)
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
@pytest.fixture()
def seed_tts_root(tmp_path: Path) -> Path:
"""Minimal seed-tts-style directory with 5 entries."""
locale_dir = tmp_path / "en"
locale_dir.mkdir()
wav_dir = locale_dir / "prompt-wavs"
wav_dir.mkdir()
for i in range(5):
(wav_dir / f"utt{i:03d}.wav").write_bytes(b"RIFF\x00\x00\x00\x00WAVE")
meta = "\n".join(f"utt{i:03d}|ref text {i}|prompt-wavs/utt{i:03d}.wav|target text {i}" for i in range(5))
(locale_dir / "meta.lst").write_text(meta, encoding="utf-8")
return tmp_path
@pytest.fixture()
def mock_tokenizer(mocker):
tokenizer = mocker.MagicMock()
tokenizer.encode = lambda text, **kw: [0] * len(text.split())
tokenizer.get_vocab.return_value = {"<pad>": 0}
tokenizer.all_special_ids = []
tokenizer.all_special_tokens = []
tokenizer.vocab_size = 1
return tokenizer
# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------
def test_seed_tts_text_dataset_omits_ref_audio(seed_tts_root, mock_tokenizer):
ds = SeedTTSTextDataset(
dataset_path=str(seed_tts_root),
random_seed=0,
locale="en",
disable_shuffle=True,
)
requests = ds.sample(mock_tokenizer, num_requests=3)
assert len(requests) == 3
for req in requests:
assert isinstance(req, SeedTTSTextSampleRequest)
assert req.seed_tts_speech_extra is None or "ref_audio" not in (req.seed_tts_speech_extra or {})
assert req.seed_tts_ref_wav_path == ""
assert "target text" in req.prompt
# ---------------------------------------------------------------------------
# SeedTTSDesignDataset tests
# ---------------------------------------------------------------------------
@pytest.fixture()
def seed_tts_design_root(tmp_path: Path) -> Path:
"""seed-tts-design directory with 5-field meta.lst entries."""
locale_dir = tmp_path / "en"
locale_dir.mkdir()
meta = "\n".join(
f"des{i:03d}|||target text {i}|A warm {['female', 'male'][i % 2]} voice with neutral accent." for i in range(5)
)
(locale_dir / "meta.lst").write_text(meta, encoding="utf-8")
return tmp_path
def test_seed_tts_design_dataset_has_instructions(seed_tts_design_root, mock_tokenizer):
ds = SeedTTSDesignDataset(
dataset_path=str(seed_tts_design_root),
random_seed=0,
locale="en",
disable_shuffle=True,
)
requests = ds.sample(mock_tokenizer, num_requests=3)
assert len(requests) == 3
for req in requests:
assert isinstance(req, SeedTTSDesignSampleRequest)
extra = req.seed_tts_speech_extra or {}
assert "instructions" in extra
assert extra["instructions"], "instructions must be non-empty"
assert extra.get("task_type") == "VoiceDesign"
assert "ref_audio" not in extra
assert req.seed_tts_ref_wav_path == ""
def test_seed_tts_design_dataset_rejects_missing_description(seed_tts_design_root, mock_tokenizer):
"""Lines without a voice_description should be skipped."""
locale_dir = seed_tts_design_root / "en"
# The bad line has 4 fields, not 5, so will be filtered
meta = "bad|||target text without description\n" + "\n".join(
f"ok|||target text {i}|A clear female voice." for i in range(9)
)
(locale_dir / "meta.lst").write_text(meta, encoding="utf-8")
ds = SeedTTSDesignDataset(
dataset_path=str(seed_tts_design_root),
random_seed=0,
locale="en",
disable_shuffle=True,
)
requests = ds.sample(mock_tokenizer, num_requests=10, no_oversample=True)
assert len(requests) == 9 # since we filter the bad row out and don't oversample
for req in requests:
assert isinstance(req, SeedTTSDesignSampleRequest)
assert req.seed_tts_utterance_id == "ok"
def test_attach_sets_seed_tts_row_even_without_extra_body():
"""seed_tts_row=True must be set for SeedTTSTextSampleRequest (no extra body)."""
from vllm_omni.benchmarks.data_modules.seed_tts_dataset import SeedTTSTextSampleRequest
req = SeedTTSTextSampleRequest(
prompt="hello world",
prompt_len=2,
expected_output_len=100,
multi_modal_data=None,
request_id="test-0",
seed_tts_speech_extra=None,
seed_tts_ref_wav_path="",
)
assert req.seed_tts_speech_extra is None
assert req.seed_tts_ref_wav_path == ""
# The fix ensures that even with speech_extra=None, the function
# sets seed_tts_row=True. We verify the source code has the fix.
import inspect
import vllm_omni.benchmarks.patch.patch as patch_mod
src = inspect.getsource(patch_mod._attach_seed_tts_to_request_func_input)
# seed_tts_row must be set BEFORE the 'if not ex: return' check
row_pos = src.index("seed_tts_row")
not_ex_pos = src.index("if not ex:")
assert row_pos < not_ex_pos, "seed_tts_row must be set before 'if not ex: return'"
+1 -5
View File
@@ -120,11 +120,7 @@ def create_test_parameter_mapping(configs: list[dict[str, Any]]) -> dict[str, di
"test_name": test_name,
"benchmark_params": [],
}
for entry in config["benchmark_params"]:
# Skip disabled entries
if not entry.get("enabled", True):
continue
mapping[test_name]["benchmark_params"].append(entry)
mapping[test_name]["benchmark_params"].extend(config["benchmark_params"])
return mapping
+1 -1
View File
@@ -333,7 +333,7 @@ def test_performance_benchmark(omni_server, benchmark_params):
raise ValueError("The number of prompts does not match the QPS or max_concurrency")
args = ["--host", host, "--port", str(port)]
exclude_keys = {"request_rate", "baseline", "num_prompts", "max_concurrency", "task", "enabled", "eval_phase"}
exclude_keys = {"request_rate", "baseline", "num_prompts", "max_concurrency"}
for key, value in params.items():
if key in exclude_keys or value is None:
@@ -1,79 +0,0 @@
"""Tests for DFX runner metadata field exclusion."""
import json
import pytest
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_task_excluded_from_cli_args():
"""'task' field must not become --task CLI arg."""
params = {
"task": "voice_clone",
"dataset_name": "seed-tts",
"backend": "openai-audio-speech",
"endpoint": "/v1/audio/speech",
"percentile-metrics": "audio_rtf,audio_ttfp",
"baseline": {"mean_audio_rtf": [0.5]},
}
exclude_keys = {"request_rate", "baseline", "num_prompts", "max_concurrency", "task", "enabled", "eval_phase"}
args = []
for key, value in params.items():
if key in exclude_keys or value is None:
continue
arg_name = f"--{key.replace('_', '-')}"
if isinstance(value, bool) and value:
args.append(arg_name)
elif isinstance(value, dict):
args.extend([arg_name, json.dumps(value)])
elif not isinstance(value, bool):
args.extend([arg_name, str(value)])
assert "--task" not in args
assert "--enabled" not in args
assert "--dataset-name" in args
def test_enabled_false_entry_is_skipped():
"""benchmark_params entry with enabled=false should be skipped."""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent))
from tests.dfx.conftest import create_test_parameter_mapping
configs = [
{
"test_name": "test_model",
"server_params": {"model": "some/model"},
"benchmark_params": [
{
"task": "voice_clone",
"enabled": True,
"dataset_name": "seed-tts",
"backend": "openai-audio-speech",
"endpoint": "/v1/audio/speech",
"num_prompts": [10],
"max_concurrency": [1],
"percentile-metrics": "audio_rtf",
"baseline": {},
},
{
"task": "voice_design",
"enabled": False,
"dataset_name": "seed-tts-design",
"backend": "openai-audio-speech",
"endpoint": "/v1/audio/speech",
"num_prompts": [5],
"max_concurrency": [1],
"percentile-metrics": "audio_rtf",
"baseline": {},
},
],
}
]
mapping = create_test_parameter_mapping(configs)
params = mapping["test_model"]["benchmark_params"]
# Only the enabled=True entry should appear
assert len(params) == 1
assert params[0].get("task") == "voice_clone"
+32 -153
View File
@@ -1,155 +1,34 @@
[
{
"test_name": "test_qwen3_tts_base",
"server_params": {
"model": "Qwen/Qwen3-TTS-12Hz-1.7B-Base"
},
"benchmark_params": [
{
"task": "voice_clone",
"eval_phase": "latency",
"enabled": false,
"dataset_name": "seed-tts",
"backend": "openai-audio-speech",
"endpoint": "/v1/audio/speech",
"num_prompts": [20],
"max_concurrency": [1],
"seed_tts_locale": "en",
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"median_audio_ttfp_ms": [350],
"median_audio_rtf": [0.25]
}
},
{
"task": "voice_clone",
"eval_phase": "throughput",
"enabled": false,
"dataset_name": "seed-tts",
"backend": "openai-audio-speech",
"endpoint": "/v1/audio/speech",
"num_prompts": [80],
"max_concurrency": [8],
"seed_tts_locale": "en",
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"median_audio_ttfp_ms": [3500],
"median_audio_rtf": [0.75],
"audio_throughput": [10.0]
}
},
{
"task": "voice_clone",
"eval_phase": "quality",
"enabled": false,
"dataset_name": "seed-tts",
"backend": "openai-audio-speech",
"endpoint": "/v1/audio/speech",
"num_prompts": [200],
"max_concurrency": [4],
"seed_tts_locale": "en",
"seed_tts_wer_eval": true,
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"mean_audio_rtf": [0.45]
}
}
]
},
{
"test_name": "test_qwen3_tts_customvoice",
"server_params": {
"model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
},
"benchmark_params": [
{
"task": "default_voice",
"eval_phase": "latency",
"dataset_name": "seed-tts-text",
"backend": "openai-audio-speech",
"endpoint": "/v1/audio/speech",
"dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
"num_prompts": [20],
"max_concurrency": [1],
"seed_tts_locale": "en",
"extra_body": {"voice": "Vivian", "language": "English", "task_type": "CustomVoice"},
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"median_audio_ttfp_ms": [150],
"median_audio_rtf": [0.15]
}
},
{
"task": "default_voice",
"eval_phase": "throughput",
"dataset_name": "seed-tts-text",
"backend": "openai-audio-speech",
"endpoint": "/v1/audio/speech",
"dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
"num_prompts": [80],
"max_concurrency": [8],
"seed_tts_locale": "en",
"extra_body": {"voice": "Vivian", "language": "English", "task_type": "CustomVoice"},
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"median_audio_ttfp_ms": [1500],
"median_audio_rtf": [0.30],
"audio_throughput": [30.0]
}
},
{
"task": "default_voice",
"eval_phase": "quality",
"enabled": false,
"dataset_name": "seed-tts-text",
"backend": "openai-audio-speech",
"endpoint": "/v1/audio/speech",
"dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
"num_prompts": [200],
"max_concurrency": [4],
"seed_tts_locale": "en",
"extra_body": {"voice": "Vivian", "language": "English", "task_type": "CustomVoice"},
"seed_tts_wer_eval": true,
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"mean_audio_rtf": [0.35]
}
},
{
"task": "voice_design",
"eval_phase": "latency",
"dataset_name": "seed-tts-design",
"backend": "openai-audio-speech",
"endpoint": "/v1/audio/speech",
"dataset_path": "benchmarks/build_dataset/seed_tts_design",
"num_prompts": [20],
"max_concurrency": [1],
"seed_tts_locale": "en",
"extra_body": {"task_type": "VoiceDesign", "language": "English"},
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"median_audio_ttfp_ms": [150],
"median_audio_rtf": [0.15]
}
},
{
"task": "voice_design",
"eval_phase": "throughput",
"dataset_name": "seed-tts-design",
"backend": "openai-audio-speech",
"endpoint": "/v1/audio/speech",
"dataset_path": "benchmarks/build_dataset/seed_tts_design",
"num_prompts": [80],
"max_concurrency": [8],
"seed_tts_locale": "en",
"extra_body": {"task_type": "VoiceDesign", "language": "English"},
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"median_audio_ttfp_ms": [1500],
"median_audio_rtf": [0.35],
"audio_throughput": [25.0]
}
}
]
}
{
"test_name": "test_qwen3_tts",
"server_params": {
"model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
},
"benchmark_params": [
{
"dataset_name": "random",
"backend": "openai-audio-speech",
"endpoint": "/v1/audio/speech",
"num_prompts": [
10,
40
],
"max_concurrency": [
1,
4
],
"random_input_len": 100,
"random_output_len": 100,
"extra_body": {
"voice": "Vivian",
"language": "English"
},
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"mean_audio_ttfp_ms": [6000, 6000],
"mean_audio_rtf": [0.3, 0.3]
}
}
]
}
]
@@ -340,7 +340,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
p.add_argument(
"--max-seed-tts-mean-wer",
type=float,
default=0.5,
default=0.02,
help="If set, fail when seed_tts_content_error_mean is strictly above this value.",
)
p.add_argument(
@@ -0,0 +1,131 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""E2E offline inference tests for MOSS-TTS-Nano single-stage pipeline."""
from __future__ import annotations
import pytest
import torch
from vllm import SamplingParams
from tests.helpers.mark import hardware_test
from vllm_omni import Omni
MODEL_NAME = "OpenMOSS-Team/MOSS-TTS-Nano"
SAMPLE_RATE = 48000
DEFAULT_SAMPLING = SamplingParams(
temperature=1.0,
top_p=1.0,
top_k=50,
max_tokens=4096,
seed=42,
detokenize=False,
)
def _build_request(
text: str,
voice: str = "Junhao",
mode: str = "voice_clone",
max_new_frames: int = 100, # short for tests
seed: int = 42,
) -> dict:
return {
"prompt": "<|im_start|>assistant\n",
"additional_information": {
"text": [text],
"voice": [voice],
"mode": [mode],
"max_new_frames": [max_new_frames],
"seed": [seed],
},
}
def _collect_audio(omni: Omni, request: dict) -> tuple[torch.Tensor, int]:
"""Run a single request and return (waveform, sample_rate)."""
for stage_outputs in omni.generate(request, DEFAULT_SAMPLING):
for req_output in stage_outputs.request_output:
mm = req_output.outputs[0].multimodal_output
assert mm is not None, "Expected multimodal_output to be non-None"
audio = mm.get("audio")
sr = mm.get("sr")
assert audio is not None, "Expected 'audio' key in multimodal_output"
assert isinstance(audio, torch.Tensor), f"audio should be Tensor, got {type(audio)}"
return audio.cpu(), int(sr.item()) if sr is not None else SAMPLE_RATE
raise AssertionError("No stage outputs received")
@pytest.fixture(scope="module")
def omni_engine():
"""Module-scoped Omni engine to avoid re-loading the model for each test."""
return Omni(model=MODEL_NAME, stage_init_timeout=180)
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"})
def test_moss_tts_nano_english(omni_engine):
"""English TTS produces non-empty 48 kHz stereo audio."""
req = _build_request("Hello, this is a test of MOSS-TTS-Nano.", voice="Ava")
audio, sr = _collect_audio(omni_engine, req)
assert sr == SAMPLE_RATE, f"Expected sample_rate={SAMPLE_RATE}, got {sr}"
assert audio.numel() > 0, "Audio tensor should not be empty"
assert not torch.all(audio == 0), "Audio should not be all-zeros (silence)"
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"})
def test_moss_tts_nano_chinese(omni_engine):
"""Chinese TTS produces non-empty audio."""
req = _build_request("你好,这是语音合成测试。", voice="Junhao")
audio, sr = _collect_audio(omni_engine, req)
assert sr == SAMPLE_RATE
assert audio.numel() > 0
assert not torch.all(audio == 0)
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"})
def test_moss_tts_nano_deterministic(omni_engine):
"""Same seed produces identical waveforms."""
req = _build_request("Reproducible output test.", voice="Adam", seed=123)
audio1, _ = _collect_audio(omni_engine, req)
audio2, _ = _collect_audio(omni_engine, req)
assert audio1.shape == audio2.shape, "Waveform shapes should match with same seed"
assert torch.allclose(audio1, audio2, atol=1e-4), "Waveforms should match with same seed"
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"})
def test_moss_tts_nano_batch(omni_engine):
"""Batch of two requests returns audio for each."""
requests = [
_build_request("First request.", voice="Ava"),
_build_request("Second request.", voice="Bella"),
]
results = []
for stage_outputs in omni_engine.generate(requests, [DEFAULT_SAMPLING] * 2):
for req_output in stage_outputs.request_output:
mm = req_output.outputs[0].multimodal_output
assert mm is not None
results.append(mm["audio"].cpu())
assert len(results) == 2, f"Expected 2 outputs, got {len(results)}"
for i, audio in enumerate(results):
assert audio.numel() > 0, f"Audio {i} is empty"
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"})
def test_moss_tts_nano_voice_presets(omni_engine):
"""Multiple built-in voice presets all produce valid audio."""
voices = ["Junhao", "Ava", "Bella"]
for voice in voices:
req = _build_request(f"Testing voice {voice}.", voice=voice)
audio, sr = _collect_audio(omni_engine, req)
assert audio.numel() > 0, f"Voice {voice} produced empty audio"
assert sr == SAMPLE_RATE
@@ -9,7 +9,6 @@ Coverage:
- Tensor-Parallel
- Ulysses-SP
- Ring-Attention
- Layerwise Offloading
assert_diffusion_response validates successful generation and the expected
512x512 resolution.
@@ -32,7 +31,7 @@ PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
def _get_diffusion_feature_cases(model: str):
"""Return L4 diffusion feature cases for Bagel.
TeaCache, Cache-DiT, CFG-Parallel, Tensor-Parallel,
Ulysses-SP, Ring-Attention, Layerwise Offloading.
Ulysses-SP, Ring-Attention.
"""
return [
@@ -112,15 +111,6 @@ def _get_diffusion_feature_cases(model: str):
id="sp_ring_2",
marks=PARALLEL_FEATURE_MARKS,
),
# Layerwise Offloading (single-card)
pytest.param(
OmniServerParams(
model=model,
server_args=["--enable-layerwise-offload"],
),
id="single_card_layerwise_offload",
marks=SINGLE_CARD_FEATURE_MARKS,
),
]
@@ -142,7 +132,6 @@ def test_bagel(
- Tensor-Parallel (size=2)
- Ulysses-SP (degree=2)
- Ring-Attention (degree=2)
- Layerwise Offloading
Validation is delegated to assert_diffusion_response in tests/helpers/assertions.py,
which checks output dimensions and basic correctness.
@@ -0,0 +1,117 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
E2E Online tests for MOSS-TTS-Nano model with text input and audio output.
These tests verify the /v1/audio/speech endpoint works correctly with
actual model inference, not mocks.
"""
import os
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
import pytest
from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniServerParams
from tests.helpers.stage_config import get_deploy_config_path
MODEL = "OpenMOSS-Team/MOSS-TTS-Nano"
def get_prompt(prompt_type="text"):
"""Text prompt for text-to-audio tests."""
prompts = {
"text": "Hello, this is a test of MOSS-TTS-Nano text to speech synthesis.",
"chinese": "你好,这是MOSS-TTS-Nano的语音合成测试。",
}
return prompts.get(prompt_type, prompts["text"])
tts_server_params = [
pytest.param(
OmniServerParams(
model=MODEL,
stage_config_path=get_deploy_config_path("moss_tts_nano.yaml"),
server_args=["--disable-log-stats"],
),
id="moss_tts_nano",
)
]
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
def test_text_to_audio_001(omni_server, openai_client) -> None:
"""
Test text input processing and audio output via /v1/audio/speech.
Deploy Setting: default yaml
Input Modal: text
Output Modal: audio (48 kHz)
Input Setting: stream=False, voice=Junhao
Datasets: single request
"""
request_config = {
"model": omni_server.model,
"input": get_prompt(),
"stream": False,
"response_format": "wav",
"voice": "Junhao",
}
openai_client.send_audio_speech_request(request_config)
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
def test_text_to_audio_002(omni_server, openai_client) -> None:
"""
Test streaming text-to-audio via /v1/audio/speech.
Deploy Setting: default yaml
Input Modal: text
Output Modal: audio (48 kHz, PCM stream)
Input Setting: stream=True, voice=Ava
Datasets: single request
"""
request_config = {
"model": omni_server.model,
"input": get_prompt(),
"stream": True,
"response_format": "pcm",
"voice": "Ava",
}
openai_client.send_audio_speech_request(request_config)
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
def test_text_to_audio_003(omni_server, openai_client) -> None:
"""
Test Chinese text-to-audio via /v1/audio/speech.
Deploy Setting: default yaml
Input Modal: text (Chinese)
Output Modal: audio (48 kHz)
Input Setting: stream=False, voice=Junhao
Datasets: single request
"""
request_config = {
"model": omni_server.model,
"input": get_prompt("chinese"),
"stream": False,
"response_format": "wav",
"voice": "Junhao",
}
openai_client.send_audio_speech_request(request_config)
@@ -259,146 +259,6 @@ class SeedTTSDataset(BenchmarkDataset):
return out
@dataclass
class _SeedTTSDesignRow:
utterance_id: str
target_text: str
voice_description: str
def _parse_design_meta_line(line: str) -> _SeedTTSDesignRow | None:
"""Parse a 5-field design meta.lst line.
Format: ``utt_id|ref_text|wav_rel|target_text|voice_description``
Returns None (with a warning) if the line has fewer than 5 fields or if
voice_description is empty.
"""
line = line.strip()
if not line or line.startswith("#"):
return None
parts = line.split("|")
if len(parts) < 5:
logger.warning("Skipping malformed design meta.lst line (need 5 '|'-fields): %r", line[:120])
return None
utt_id = parts[0].strip()
target_text = parts[3].strip()
voice_description = parts[4].strip()
if not voice_description:
logger.warning("Skipping design meta.lst line with empty voice_description: %r", line[:120])
return None
return _SeedTTSDesignRow(
utterance_id=utt_id,
target_text=target_text,
voice_description=voice_description,
)
@dataclass
class SeedTTSDesignSampleRequest(SeedTTSSampleRequest):
"""SampleRequest for voice-design TTS (no ref_audio; voice described via natural language).
The ``seed_tts_speech_extra`` dict carries ``instructions`` (natural-language
voice description, forwarded as-is to the Qwen3-TTS VoiceDesign endpoint) and
``task_type="VoiceDesign"`` instead of ``ref_audio`` / ``ref_text``.
SIM is skipped (``seed_tts_ref_wav_path`` is empty).
"""
class SeedTTSDesignDataset(SeedTTSDataset):
"""Seed-TTS prompts for voice-design benchmarking (dataset name: ``seed-tts-design``).
Loads a 5-field ``meta.lst``::
utt_id|ref_text|wav_rel|target_text|voice_description
and builds requests with ``task_type="VoiceDesign"`` and the natural-language
``voice_description`` column forwarded via the ``instructions`` field
(the Qwen3-TTS VoiceDesign endpoint's expected key) instead of
``ref_audio`` / ``ref_text``. Speaker-similarity (SIM) is not computed.
"""
def load_data(self) -> None:
# Does NOT call super().load_data() — the format is different (5 fields,
# no wav file). self._rows is intentionally left empty; the parent
# sample() is fully overridden so an empty self._rows is safe.
meta = self._root / self.locale / "meta.lst"
if not meta.is_file():
raise FileNotFoundError(
f"Seed-TTS-Design meta not found: {meta}. Expected layout: {self._root}/{self.locale}/meta.lst"
)
text = meta.read_text(encoding="utf-8")
design_rows: list[_SeedTTSDesignRow] = []
for line in text.splitlines():
r = _parse_design_meta_line(line)
if r is not None:
design_rows.append(r)
if not design_rows:
raise ValueError(f"No valid rows in {meta}")
if not self.disable_shuffle:
rng = random.Random(self.random_seed)
rng.shuffle(design_rows)
self._design_rows = design_rows
# Keep self._rows empty — parent sample() is overridden.
self._rows = []
self.data = self._design_rows
logger.info(
"Loaded Seed-TTS-Design: root=%s locale=%s rows=%d",
self._root,
self.locale,
len(self._design_rows),
)
def sample(
self,
tokenizer: TokenizerLike,
num_requests: int,
output_len: int | None = None,
request_id_prefix: str = "",
no_oversample: bool = False,
**kwargs: Any,
) -> list[SampleRequest]:
if output_len is None:
output_len = self.DEFAULT_OUTPUT_LEN
tok = get_cached_tokenizer(tokenizer)
lang = "English" if self.locale == "en" else "Chinese"
out: list[SampleRequest] = []
for i, row in enumerate(self._design_rows):
if len(out) >= num_requests:
break
target = row.target_text
prompt_len = len(tok.encode(target))
speech_extra: dict[str, Any] = {
"instructions": row.voice_description,
"task_type": "VoiceDesign",
"language": lang,
"max_new_tokens": output_len,
}
out.append(
SeedTTSDesignSampleRequest(
prompt=target,
prompt_len=prompt_len,
expected_output_len=output_len,
multi_modal_data=None,
request_id=f"{request_id_prefix}{i}",
seed_tts_speech_extra=speech_extra,
seed_tts_utterance_id=row.utterance_id,
seed_tts_locale=self.locale,
seed_tts_system_prompt=self._system_prompt,
seed_tts_ref_wav_path="", # SIM skipped for voice-design
)
)
logger.info(
"Seed-TTS-Design: built %d requests (asked %d) — no ref_audio (voice design)",
len(out),
num_requests,
)
self.maybe_oversample_requests(out, num_requests, request_id_prefix, no_oversample)
return out
def load_seed_tts_dataset(
dataset_path: str,
random_seed: int = 0,
@@ -417,65 +277,3 @@ def load_seed_tts_dataset(
system_prompt=system_prompt,
**kwargs,
)
@dataclass
class SeedTTSTextSampleRequest(SeedTTSSampleRequest):
"""SampleRequest for default-voice TTS (no ref_audio, no ref_text).
The voice param (e.g. ``voice: "Vivian"``) is supplied at request time via
``--extra-body`` in the benchmark config. SIM is skipped (empty ref_wav_path).
WER and UTMOS are computed normally.
"""
class SeedTTSTextDataset(SeedTTSDataset):
"""Seed-TTS prompts for default-voice benchmarking (dataset name: ``seed-tts-text``).
Loads the same ``meta.lst`` as :class:`SeedTTSDataset` but builds requests
WITHOUT ``ref_audio`` / ``ref_text`` body fields. The named voice must be
supplied via ``--extra-body`` in the benchmark config.
Speaker-similarity (SIM) is not computed.
"""
def sample(
self,
tokenizer: TokenizerLike,
num_requests: int,
output_len: int | None = None,
request_id_prefix: str = "",
no_oversample: bool = False,
**kwargs: Any,
) -> list[SampleRequest]:
if output_len is None:
output_len = self.DEFAULT_OUTPUT_LEN
tok = get_cached_tokenizer(tokenizer)
out: list[SampleRequest] = []
for i, row in enumerate(self._rows):
if len(out) >= num_requests:
break
target = row.target_text
prompt_len = len(tok.encode(target))
out.append(
SeedTTSTextSampleRequest(
prompt=target,
prompt_len=prompt_len,
expected_output_len=output_len,
multi_modal_data=None,
request_id=f"{request_id_prefix}{i}",
seed_tts_speech_extra=None, # voice supplied via --extra-body in config
seed_tts_utterance_id=row.utterance_id,
seed_tts_locale=self.locale,
seed_tts_system_prompt=self._system_prompt,
seed_tts_ref_wav_path="", # empty → SIM skipped in seed_tts_eval
)
)
logger.info(
"Seed-TTS-Text: built %d requests (asked %d) — no ref_audio (default voice)",
len(out),
num_requests,
)
self.maybe_oversample_requests(out, num_requests, request_id_prefix, no_oversample)
return out
+15 -22
View File
@@ -40,9 +40,7 @@ from vllm_omni.benchmarks.data_modules.random_multi_modal_dataset import OmniRan
from vllm_omni.benchmarks.data_modules.seed_tts_dataset import (
SEED_TTS_DEFAULT_OMNI_SYSTEM_PROMPT,
SeedTTSDataset,
SeedTTSDesignDataset,
SeedTTSSampleRequest,
SeedTTSTextDataset,
)
get_samples_old = datasets.get_samples
@@ -87,16 +85,23 @@ def _attach_daily_omni_to_request_func_input(sample: SampleRequest, rfi: Request
def _attach_seed_tts_to_request_func_input(sample: SampleRequest, rfi: RequestFuncInput) -> None:
"""Merge Seed-TTS per-row TTS fields into ``extra_body`` and mark for PCM capture.
"""Merge Seed-TTS per-row TTS fields (ref_audio, ref_text, task_type, …) into ``extra_body``.
Always sets ``seed_tts_row=True`` on the RequestFuncInput for any
:class:`SeedTTSSampleRequest` subclass (including text-only and design
variants that carry no ``ref_audio``). This enables PCM capture for WER /
UTMOS evaluation even when there is no reference audio.
Used by both ``/v1/audio/speech`` and ``/v1/chat/completions`` (flattened into JSON body).
For ``openai-chat-omni``, also sets ``omni_chat_messages`` (system + user) so Qwen3-Omni
follows the same role layout as official TTS / multimodal demos. ``/v1/audio/speech`` ignores
``messages`` and only uses ``input`` + body fields.
Flags ``openai-chat-omni`` to request audio output and optionally export PCM for WER.
"""
if not isinstance(sample, SeedTTSSampleRequest):
return
# Mark for PCM capture (WER / UTMOS eval) regardless of extra body presence.
ex = sample.seed_tts_speech_extra
if not ex:
return
base = dict(rfi.extra_body) if rfi.extra_body else {}
base.update(ex)
rfi.extra_body = base
# Used by request funcs to force streaming TTS behavior and to export PCM when WER is on.
setattr(rfi, "seed_tts_row", True)
sys_prompt = (sample.seed_tts_system_prompt or "").strip() or SEED_TTS_DEFAULT_OMNI_SYSTEM_PROMPT
setattr(
@@ -107,12 +112,6 @@ def _attach_seed_tts_to_request_func_input(sample: SampleRequest, rfi: RequestFu
{"role": "user", "content": [{"type": "text", "text": sample.prompt}]},
],
)
ex = sample.seed_tts_speech_extra
if not ex:
return # voice comes from --extra-body in config; no ref_audio to merge
base = dict(rfi.extra_body) if rfi.extra_body else {}
base.update(ex)
rfi.extra_body = base
def _daily_omni_repo_from_args(args) -> str | None:
@@ -137,7 +136,7 @@ def get_samples(args, tokenizer):
is_daily_omni = args.dataset_name == "daily-omni" or (
args.dataset_name == "hf" and _daily_omni_repo_from_args(args) is not None
)
is_seed_tts = args.dataset_name in ("seed-tts", "seed-tts-text", "seed-tts-design")
is_seed_tts = args.dataset_name == "seed-tts"
# Check if we need to handle omni-related backends/datasets
is_omni_backend = args.backend in ["openai-chat-omni", "openai-audio-speech", "daily-omni"]
@@ -250,13 +249,7 @@ def get_samples(args, tokenizer):
"--hf-name for the Hub dataset id."
)
_cls_map = {
"seed-tts": SeedTTSDataset,
"seed-tts-text": SeedTTSTextDataset,
"seed-tts-design": SeedTTSDesignDataset,
}
DatasetCls = _cls_map[args.dataset_name]
dataset = DatasetCls(
dataset = SeedTTSDataset(
dataset_path=repo_id,
random_seed=args.seed,
locale=getattr(args, "seed_tts_locale", "en"),
+4
View File
@@ -79,6 +79,10 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = {
"vllm_omni.model_executor.models.fish_speech.pipeline",
"FISH_SPEECH_PIPELINE",
),
"moss_tts_nano": (
"vllm_omni.model_executor.models.moss_tts_nano.pipeline",
"MOSS_TTS_NANO_PIPELINE",
),
}
# --- Single-stage diffusion pipelines (populated in PR 3/N) ---
+19 -4
View File
@@ -439,7 +439,7 @@ class DeployConfig:
# === Pipeline-wide engine settings (applied uniformly to every stage) ===
trust_remote_code: bool = True
distributed_executor_backend: str | None = None
distributed_executor_backend: str = "mp"
dtype: str | None = None
quantization: str | None = None
enable_prefix_caching: bool = False
@@ -601,6 +601,23 @@ def load_deploy_config(path: str | Path) -> DeployConfig:
return DeployConfig(**kwargs)
def _detect_platform() -> str | None:
"""Return "npu", "rocm", "xpu", or None (CUDA default)."""
try:
from vllm.platforms import current_platform
name = current_platform.device_name.lower()
if "npu" in name:
return "npu"
if "rocm" in name or "amd" in name:
return "rocm"
if "xpu" in name:
return "xpu"
except Exception as e:
logger.debug("Platform auto-detect failed, falling back to CUDA: %s", e)
return None
def _extract_platform_overrides(ps: dict[str, Any]) -> tuple[dict[str, Any], str | None]:
"""Return ``(overrides, devices)`` from a platform stage entry.
@@ -619,9 +636,7 @@ def _apply_platform_overrides(
) -> DeployConfig:
"""Merge platform-specific stage overrides into deploy config."""
if platform is None:
from vllm_omni.platforms import current_omni_platform
platform = current_omni_platform.device_name.lower()
platform = _detect_platform()
if platform is None or deploy.platforms is None:
return deploy
platform_section = deploy.platforms.get(platform)
+31
View File
@@ -0,0 +1,31 @@
# MOSS-TTS-Nano deploy: single-stage AR TTS (0.1B).
# Verified on 1x L4 24GB (and sufficient for H20 / A100 / 3090).
#
# Fields omitted from a stage fall back to StageDeployConfig defaults (see
# vllm_omni/config/stage_config.py). Notable choices:
# * enforce_eager=true — the model uses upstream inference_stream() with
# trust_remote_code; CUDA graph capture is not wired up.
# * gpu_memory_utilization=0.3 — the AR LM + codec fit in ~2 GiB.
# * max_num_seqs=4 — mirrors the legacy stage_configs contract; per-request
# streaming generators are stored in self._stream_gens.
#
# MOSS-TTS-Nano has no multimodal input preprocessing, so skip_mm_profiling
# bypasses the dummy-MM pass at profile time.
async_chunk: false
trust_remote_code: true
stages:
- stage_id: 0
max_num_seqs: 4
gpu_memory_utilization: 0.3
enforce_eager: true
max_num_batched_tokens: 4096
max_model_len: 4096
devices: "0"
skip_mm_profiling: true
default_sampling_params:
temperature: 1.0
top_p: 1.0
top_k: 50
max_tokens: 4096
seed: 42
+1 -1
View File
@@ -759,7 +759,7 @@ class OmniDiffusionConfig:
# Backwards-compatibility: map "quantization" to "quantization_config"
# so callers using the old field name still work.
if "quantization" in kwargs and kwargs.get("quantization_config", None) is None:
if "quantization" in kwargs and kwargs.get("quantization_config") is None:
kwargs["quantization_config"] = kwargs.pop("quantization")
else:
kwargs.pop("quantization", None)
-4
View File
@@ -22,10 +22,6 @@ class DiffusionExecutor(ABC):
def get_class(od_config: OmniDiffusionConfig) -> type[DiffusionExecutor]:
executor_class: type[DiffusionExecutor]
distributed_executor_backend = od_config.distributed_executor_backend
# Keep backward-compatible behavior for callers/configs that omit this
# field and rely on the historical diffusion default backend.
if distributed_executor_backend is None:
distributed_executor_backend = "mp"
if isinstance(distributed_executor_backend, type):
if not issubclass(distributed_executor_backend, DiffusionExecutor):
@@ -738,8 +738,6 @@ class Qwen2MoTDecoderLayer(nn.Module):
class Qwen2MoTModel(Qwen2PreTrainedModel):
_layerwise_offload_blocks_attrs = ["layers"]
def __init__(
self,
config,
@@ -288,7 +288,7 @@ class BagelPipeline(nn.Module, SupportsModuleOffload, DiffusionPipelineProfilerM
# device until the weight loader materializes them. Calling
# .to(device) would fail on those meta tensors, so we skip it
# entirely and let the weight loader handle device placement.
if quant_config is None and not od_config.enable_layerwise_offload:
if quant_config is None:
self.to(self.device)
self.setup_diffusion_pipeline_profiler(
enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
@@ -1484,7 +1484,7 @@ class HunYuanSparseMoeBlock(nn.Module):
config.hidden_size,
config.num_experts,
bias=False,
quant_config=quant_config,
quant_config=None,
prefix=f"{prefix}.gate",
)
if config.use_mixed_mlp_moe > 0:
@@ -1658,10 +1658,8 @@ class HunYuanAttention(nn.Module):
custom_pos_emb: tuple[torch.FloatTensor] | None = None,
**kwargs,
) -> torch.Tensor:
bsz, q_len, hidden_size = hidden_states.size()
hidden_states = hidden_states.reshape(-1, hidden_size)
bsz, q_len, _ = hidden_states.size()
qkv, _ = self.qkv_proj(hidden_states)
qkv = qkv.reshape(bsz, q_len, -1)
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
past_key_value: Cache | None = kwargs.get("past_key_value", None)
@@ -1725,7 +1723,7 @@ class HunyuanImage3DecoderLayer(nn.Module):
rope_theta=rope_theta,
rope_scaling=rope_scaling,
max_position_embeddings=max_position_embeddings,
quant_config=quant_config,
quant_config=None,
bias=attention_bias,
cache_config=None,
prefix=f"{prefix}.self_attn",
@@ -1935,7 +1933,7 @@ class HunyuanImage3Model(nn.Module):
layer_idx=int(prefix.split(".")[-1]),
prefix=prefix,
),
prefix=f"{prefix}.layers" if prefix else "layers",
prefix=f"{prefix}.layers",
)
if get_pp_group().is_last_rank:
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -1950,7 +1948,7 @@ class HunyuanImage3Model(nn.Module):
num_attention_heads = self.config.num_attention_heads
num_kv_heads = getattr(self.config, "num_key_value_heads", self.config.num_attention_heads)
num_key_value_groups = num_attention_heads // num_kv_heads
hidden_size = qkv.shape[1]
hidden_size = self.config.hidden_size
if hasattr(self.config, "head_dim"):
attention_head_dim = self.config.head_dim
@@ -2003,15 +2001,8 @@ class HunyuanImage3Model(nn.Module):
split_params_mapping = [
(".gate_up_proj", ".gate_and_up_proj", 2, [(1, 1), (0, 1)], None),
(
".qkv_proj.weight",
".qkv_proj.weight",
num_attention_heads + num_kv_heads * 2,
[("q", num_attention_heads), ("k", num_kv_heads), ("v", num_kv_heads)],
self._split_qkv_weight,
),
(
".qkv_proj.weight_scale",
".qkv_proj.weight_scale",
".qkv_proj",
".qkv_proj",
num_attention_heads + num_kv_heads * 2,
[("q", num_attention_heads), ("k", num_kv_heads), ("v", num_kv_heads)],
self._split_qkv_weight,
@@ -2110,8 +2101,6 @@ class HunyuanImage3Model(nn.Module):
continue
if "mlp.experts" in name:
continue
if ".qkv_proj" in name and not name.endswith(weight_name):
continue
name = name.replace(weight_name, param_name)
# Skip loading extra bias for GPTQ models.
if name.endswith(".bias") and name not in params_dict:
@@ -14,7 +14,7 @@ from transformers.generation.utils import ALL_CACHE_NAMES, GenerationMixin
from transformers.models.siglip2 import Siglip2VisionConfig, Siglip2VisionModel
from transformers.utils.generic import ModelOutput
from vllm.config.vllm import get_current_vllm_config
from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
from vllm.model_executor.models.utils import AutoWeightsLoader
from vllm.transformers_utils.config import get_config
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
@@ -64,15 +64,6 @@ def to_device(data, device):
class HunyuanImage3Pipeline(HunyuanImage3PreTrainedModel, GenerationMixin, DiffusionPipelineProfilerMixin):
hf_to_vllm_mapper = WeightsMapper(
orig_to_new_prefix={
"model.": "",
},
orig_to_new_substr={
"mlp.gate.wg.": "mlp.gate.",
"gate_and_up_proj.": "gate_up_proj.",
},
)
_PROFILER_TARGETS = [
"model.forward",
"model.layers[0].forward",
-19
View File
@@ -5,7 +5,6 @@ import importlib
import torch.nn as nn
from vllm.logger import init_logger
from vllm.model_executor.model_loader.utils import configure_quant_config
from vllm.model_executor.models.registry import _LazyRegisteredModel, _ModelRegistry
from vllm_omni.diffusion.data import OmniDiffusionConfig
@@ -14,7 +13,6 @@ from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig, get_
from vllm_omni.diffusion.forward_context import get_forward_context
from vllm_omni.diffusion.hooks.sequence_parallel import apply_sequence_parallel
from vllm_omni.diffusion.utils.tf_utils import find_module_with_attr
from vllm_omni.platforms import current_omni_platform
logger = init_logger(__name__)
@@ -244,22 +242,6 @@ _NO_CACHE_ACCELERATION = {
}
def _prepare_diffusion_quant_config(
od_config: OmniDiffusionConfig,
model_class: type[nn.Module],
) -> None:
"""Prepare diffusion quant config using vLLM-style model bindings."""
quant_config = od_config.quantization_config
if quant_config is None:
return
if hasattr(quant_config, "maybe_update_config"):
quant_config.maybe_update_config(od_config.model)
diffusion_packed_modules_mapping = current_omni_platform.get_diffusion_packed_modules_mapping(model_class)
if diffusion_packed_modules_mapping is not None:
model_class.packed_modules_mapping = diffusion_packed_modules_mapping
configure_quant_config(quant_config, model_class)
def initialize_model(
od_config: OmniDiffusionConfig,
) -> nn.Module:
@@ -282,7 +264,6 @@ def initialize_model(
"""
model_class = DiffusionModelRegistry._try_load_model_cls(od_config.model_class_name)
if model_class is not None:
_prepare_diffusion_quant_config(od_config, model_class)
model = model_class(od_config=od_config)
vae_pp_size = od_config.parallel_config.vae_patch_parallel_size
@@ -13,7 +13,6 @@ import multiprocessing as mp
import os
from collections.abc import Iterable
from contextlib import AbstractContextManager, nullcontext
from types import SimpleNamespace
from typing import Any
import torch
@@ -22,7 +21,6 @@ from vllm.config import CompilationConfig, DeviceConfig, VllmConfig, set_current
from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
from vllm.logger import init_logger
from vllm.profiler.wrapper import CudaProfilerWrapper, WorkerProfiler
from vllm.transformers_utils.config import get_config, get_hf_text_config
from vllm.utils.import_utils import resolve_obj_by_qualname
from vllm.utils.mem_utils import GiB_bytes
from vllm.v1.worker.workspace import init_workspace_manager
@@ -122,20 +120,6 @@ class DiffusionWorker:
vllm_config.parallel_config.data_parallel_size = self.od_config.parallel_config.data_parallel_size
vllm_config.parallel_config.enable_expert_parallel = self.od_config.parallel_config.enable_expert_parallel
vllm_config.profiler_config = self.od_config.profiler_config
try:
hf_config = get_config(self.od_config.model, trust_remote_code=self.od_config.trust_remote_code)
except ValueError:
hf_config = None
logger.info("Skipping hf_config loading for diffusion model %r", self.od_config.model_class_name)
hf_text_config = get_hf_text_config(hf_config) if hf_config is not None else None
vllm_config.model_config = SimpleNamespace(
hf_config=hf_config,
hf_text_config=hf_text_config,
enforce_eager=self.od_config.enforce_eager,
dtype=self.od_config.dtype,
enable_return_routed_experts=False,
)
vllm_config.quant_config = self.od_config.quantization_config
self.vllm_config = vllm_config
# Initialize distributed environment
+3 -4
View File
@@ -1451,17 +1451,16 @@ class AsyncOmniEngine:
if lora_scale is not None:
if not hasattr(cfg.engine_args, "lora_scale") or cfg.engine_args.lora_scale is None:
cfg.engine_args.lora_scale = lora_scale
# Prefer explicit quantization_config; fallback to legacy --quantization.
quantization_config = kwargs.get("quantization_config")
if quantization_config is None:
quantization_config = kwargs.get("quantization")
if quantization_config is not None:
if (
not hasattr(cfg.engine_args, "quantization_config")
or cfg.engine_args.quantization_config is None
):
cfg.engine_args.quantization_config = quantization_config
quantization = kwargs.get("quantization")
if quantization is not None:
if not hasattr(cfg.engine_args, "quantization") or cfg.engine_args.quantization is None:
cfg.engine_args.quantization = quantization
except Exception as e:
logger.warning("Failed to inject LoRA config for stage: %s", e)
+1 -3
View File
@@ -147,9 +147,7 @@ class OmniBenchmarkServingSubcommand(OmniBenchmarkSubcommandBase):
for action in parser._actions:
if action.dest == "dataset_name" and action.choices is not None:
extra = [
c for c in ("daily-omni", "seed-tts", "seed-tts-text", "seed-tts-design") if c not in action.choices
]
extra = [c for c in ("daily-omni", "seed-tts") if c not in action.choices]
if extra:
action.choices = list(action.choices) + extra
+3
View File
@@ -7,6 +7,7 @@ import time
import types
import weakref
from collections.abc import Sequence
from pprint import pformat
from typing import TYPE_CHECKING, Any, Literal
import huggingface_hub
@@ -249,6 +250,8 @@ class OmniBase(PDDisaggregationMixin):
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))
except Exception:
logger.exception(
"[%s] Failed to build/log summary for req=%s",
+6 -16
View File
@@ -1572,23 +1572,13 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
logger.info(f"Successfully generated {len(images)} image(s)")
# Determine output format (default to png)
output_format = _choose_output_format(request.output_format or "png", None)
# Encode images to base64
image_data = [ImageData(b64_json=encode_image_base64(img), revised_prompt=None) for img in images]
# Encode images to base64 with the specified format
image_data = [
ImageData(b64_json=_encode_image_base64_with_compression(img, format=output_format), revised_prompt=None)
for img in images
]
response_kwargs = {
"created": int(time.time()),
"data": image_data,
"output_format": output_format,
}
if request.size:
response_kwargs["size"] = size_str
return ImageGenerationResponse(**response_kwargs)
return ImageGenerationResponse(
created=int(time.time()),
data=image_data,
)
except (EngineGenerateError, EngineDeadError) as exc:
return _create_engine_error_json_response(raw_request, exc)
@@ -139,12 +139,6 @@ class ImageGenerationRequest(BaseModel):
vae_use_slicing: bool | None = Field(default=False, description="Enable VAE slicing")
vae_use_tiling: bool | None = Field(default=False, description="Enable VAE tiling")
# Output format for generated images
output_format: str | None = Field(
default=None,
description="Output image format: 'png', 'jpeg', or 'webp'. Defaults to 'png'.",
)
class ImageData(BaseModel):
"""Single generated image data"""
@@ -63,6 +63,7 @@ _COSYVOICE3_TTS_MODEL_STAGES = {"cosyvoice3_talker"}
_OMNIVOICE_TTS_MODEL_STAGES = {"omnivoice_generator"}
_VOXCPM_TTS_MODEL_STAGES = {"latent_generator", "vae"}
_VOXCPM2_TTS_MODEL_STAGES = {"latent_generator"}
_MOSS_TTS_MODEL_STAGES = {"moss_tts_nano"}
_MING_TTS_MODEL_STAGES = {"ming_tts"}
_TTS_MODEL_STAGES: set[str] = (
_VOXTRAL_TTS_MODEL_STAGES
@@ -72,6 +73,7 @@ _TTS_MODEL_STAGES: set[str] = (
| _OMNIVOICE_TTS_MODEL_STAGES
| _VOXCPM_TTS_MODEL_STAGES
| _VOXCPM2_TTS_MODEL_STAGES
| _MOSS_TTS_MODEL_STAGES
| _MING_TTS_MODEL_STAGES
)
_TTS_LANGUAGES: set[str] = {
@@ -315,6 +317,10 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
return "cosyvoice3"
if model_stage in _OMNIVOICE_TTS_MODEL_STAGES:
return "omnivoice"
if model_stage in _VOXCPM2_TTS_MODEL_STAGES:
return "voxcpm2"
if model_stage in _MOSS_TTS_MODEL_STAGES:
return "moss_tts_nano"
if model_stage in (_VOXCPM_TTS_MODEL_STAGES | _VOXCPM2_TTS_MODEL_STAGES):
has_vae_stage = any(
getattr(getattr(stage, "engine_args", None), "model_stage", None) == "vae"
@@ -852,6 +858,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
return self._validate_voxcpm_request(request)
if self._tts_model_type == "voxcpm2":
return None # VoxCPM2 accepts any text input
if self._tts_model_type == "moss_tts_nano":
return self._validate_moss_tts_request(request)
if self._tts_model_type == "ming_flash_omni_tts":
return self._validate_ming_tts_request(request)
return self._validate_qwen_tts_request(request)
@@ -1109,6 +1117,42 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
return None
def _validate_moss_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None:
"""Validate MOSS-TTS-Nano request. Accepts optional ref_audio for voice cloning."""
if not request.input or not request.input.strip():
return "Input text cannot be empty"
if request.ref_audio is not None:
fmt_err = self._validate_ref_audio_format(request.ref_audio)
if fmt_err:
return fmt_err
if not request.ref_text or not request.ref_text.strip():
return "Voice cloning requires 'ref_text' (transcript of the reference audio)"
return None
async def _build_moss_tts_params(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]:
"""Build additional_information for MOSS-TTS-Nano.
Maps the standard /v1/audio/speech fields to MOSS-TTS-Nano's
additional_information keys (text, voice, mode, prompt_audio_array,
prompt_text, ...). Values are wrapped in lists to match the
vllm-omni convention. When ``request.ref_audio`` is provided, it is
resolved via MediaConnector and passed as a (wav_list, sample_rate)
tuple so the model owns temp-file lifecycle.
"""
params: dict[str, Any] = {
"text": [request.input],
}
if request.voice is not None:
params["voice"] = [request.voice]
if request.ref_text is not None:
params["prompt_text"] = [request.ref_text]
if request.max_new_tokens is not None:
params["max_new_frames"] = [request.max_new_tokens]
if request.ref_audio is not None:
wav_list, sr = await self._resolve_ref_audio(request.ref_audio)
params["prompt_audio_array"] = [[wav_list, sr]]
return params
def _validate_fish_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None:
"""Validate Fish Speech request parameters. Returns error message or None.
@@ -1680,6 +1724,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
elif self._tts_model_type == "cosyvoice3":
prompt = await self._build_cosyvoice3_prompt(request)
tts_params = {}
elif self._tts_model_type == "moss_tts_nano":
tts_params = await self._build_moss_tts_params(request)
prompt = {"prompt_token_ids": [1], "additional_information": tts_params}
elif self._tts_model_type == "ming_flash_omni_tts":
prompt = self._build_ming_prompt(request)
tts_params = {}
@@ -0,0 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from .modeling_moss_tts_nano import MossTTSNanoForGeneration
__all__ = ["MossTTSNanoForGeneration"]
@@ -0,0 +1,60 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Configuration for MOSS-TTS-Nano in vLLM-Omni single-stage pipeline."""
from transformers.configuration_utils import PretrainedConfig
class MossTTSNanoConfig(PretrainedConfig):
"""Config for MOSS-TTS-Nano (OpenMOSS-Team/MOSS-TTS-Nano).
The model is a Global-Local GPT2-based AR LM (0.1B) paired with the
MOSS-Audio-Tokenizer-Nano codec. Both components run in a single
vLLM-Omni generation stage.
Relevant fields from config.json:
gpt2_config – backbone config (n_layer=12, n_embd=768, n_head=12, ...)
n_vq – number of RVQ codebooks (16)
audio_vocab_size – per-codebook vocabulary size (1024)
audio_tokenizer_pretrained_name_or_path – HF hub path for codec model
"""
model_type = "moss_tts_nano"
def __init__(self, **kwargs):
gpt2_cfg = kwargs.pop("gpt2_config", None) or {}
if hasattr(gpt2_cfg, "to_dict"):
gpt2_cfg = gpt2_cfg.to_dict()
super().__init__(**kwargs)
# --- GPT2 backbone parameters (exposed at top level for vLLM) ---
self.hidden_size = gpt2_cfg.get("n_embd", 768)
self.num_hidden_layers = gpt2_cfg.get("n_layer", 12)
self.num_attention_heads = gpt2_cfg.get("n_head", 12)
self.num_key_value_heads = self.num_attention_heads # no GQA in GPT2
self.head_dim = self.hidden_size // self.num_attention_heads
self.vocab_size = gpt2_cfg.get("vocab_size", 16384)
self.max_position_embeddings = gpt2_cfg.get("n_positions", 32768)
self.intermediate_size = gpt2_cfg.get("n_inner", self.hidden_size * 4)
# --- Audio codec parameters ---
self.n_vq: int = getattr(self, "n_vq", 16)
self.audio_vocab_size: int = getattr(self, "audio_vocab_size", 1024)
self.audio_start_token_id: int = getattr(self, "audio_start_token_id", 6)
self.audio_end_token_id: int = getattr(self, "audio_end_token_id", 7)
self.audio_user_slot_token_id: int = getattr(self, "audio_user_slot_token_id", 8)
self.audio_assistant_slot_token_id: int = getattr(self, "audio_assistant_slot_token_id", 9)
self.audio_tokenizer_pretrained_name_or_path: str = getattr(
self,
"audio_tokenizer_pretrained_name_or_path",
"OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano",
)
self.audio_tokenizer_sample_rate: int = getattr(self, "audio_tokenizer_sample_rate", 48000)
# vLLM requires speculative_config to be absent or None
self.speculative_config = None
def get_text_config(self, **kwargs):
"""Return self so vLLM uses our top-level config."""
return self
@@ -0,0 +1,515 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""MOSS-TTS-Nano single-stage model for vLLM-Omni.
Runs in a single AR worker stage. The 0.1B AR LM and the
MOSS-Audio-Tokenizer-Nano codec are both loaded here.
Streaming is supported via the VoxCPM-style generator pattern:
- On first forward() for a request, inference_stream() is started as
a Python generator and stored in self._stream_gens[request_key].
- Each subsequent forward() call pops one audio chunk from the generator
and returns it as multimodal_outputs.
- compute_logits() emits EOS only when the last chunk has been yielded,
telling the AR scheduler to finish the request.
Weight loading deliberately happens inside load_weights() -- NOT __init__ --
so that vLLM initialises distributed state before any CUDA allocations occur.
"""
from __future__ import annotations
import tempfile
import threading
from collections.abc import Iterable
from pathlib import Path
from typing import Any
import torch
import torch.nn as nn
from vllm.config import VllmConfig
from vllm.logger import init_logger
from vllm_omni.model_executor.models.output_templates import OmniOutput
logger = init_logger(__name__)
def _patch_torchaudio_load() -> None:
"""Patch torchaudio.load to use soundfile if torchcodec is unavailable."""
try:
import torchaudio
torchaudio # noqa
import torchcodec # noqa: F401
return
except Exception:
pass
import soundfile as sf
def _soundfile_load(path, frame_offset=0, num_frames=-1, normalize=True, channels_first=True, format=None):
data, sr = sf.read(str(path), dtype="float32", always_2d=True)
if frame_offset > 0:
data = data[frame_offset:]
if num_frames > 0:
data = data[:num_frames]
waveform = torch.from_numpy(data)
if channels_first:
waveform = waveform.T
return waveform, sr
def _soundfile_save(path, src, sample_rate, channels_first=True, **kwargs):
wav = src.detach().cpu().float().numpy()
if channels_first and wav.ndim == 2:
wav = wav.T
sf.write(str(path), wav, sample_rate)
try:
import torchaudio
torchaudio.load = _soundfile_load
torchaudio.save = _soundfile_save
logger.info("Patched torchaudio.load/save to use soundfile (torchcodec unavailable)")
except Exception as e:
logger.warning("Could not patch torchaudio: %s", e)
# Default sampling parameters matching the upstream demo defaults.
_DEFAULT_TEXT_TEMPERATURE = 1.0
_DEFAULT_TEXT_TOP_P = 1.0
_DEFAULT_TEXT_TOP_K = 50
_DEFAULT_AUDIO_TEMPERATURE = 0.8
_DEFAULT_AUDIO_TOP_P = 0.95
_DEFAULT_AUDIO_TOP_K = 25
_DEFAULT_AUDIO_REPETITION_PENALTY = 1.2
_DEFAULT_MAX_NEW_FRAMES = 375
_DEFAULT_VOICE = "Junhao"
# "voice_clone" mirrors the mode used by offline examples and tests:
# built-in voices (e.g. "Junhao") are backed by preset reference audio
# upstream, so the default produces the expected voice timbre.
_DEFAULT_MODE = "voice_clone"
def _pick(info: dict, key: str, default):
"""Extract scalar from additional_information dict (list or plain value)."""
val = info.get(key, default)
if isinstance(val, (list, tuple)) and len(val) > 0:
return val[0]
return val if val is not None else default
class MossTTSNanoForGeneration(nn.Module):
"""Single-stage MOSS-TTS-Nano model with streaming audio output.
Uses the VoxCPM-style generator pattern: inference_stream() is stored
per-request and yields one audio chunk per forward() call. The AR
scheduler keeps the request alive until compute_logits() emits EOS.
"""
requires_raw_input_tokens = True
have_multimodal_outputs = True
has_preprocess = False
has_postprocess = False
enable_update_additional_information = True
inject_omni_request_id_into_runtime_info = True
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
super().__init__()
self.vllm_config = vllm_config
self.config = vllm_config.model_config.hf_config
self.model_path: str = vllm_config.model_config.model
self._lm: nn.Module | None = None
self._audio_tokenizer: nn.Module | None = None
self._device: torch.device | None = None
self._lock = threading.Lock()
# Per-request streaming generators (VoxCPM pattern).
# Single-threaded AR worker: mutations happen only inside forward(),
# so we rely on the worker's serial execution instead of self._lock.
self._stream_gens: dict[str, Any] = {}
# Per-row EOS mask aligned with the most recent forward() batch.
# compute_logits() emits EOS for rows flagged True, keeps others alive.
self._ar_last_chunk_flags: list[bool] = []
# ------------------------------------------------------------------
# Weight loading
# ------------------------------------------------------------------
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
with self._lock:
if self._lm is not None:
return set()
_patch_torchaudio_load()
try:
device = next(self.parameters()).device
except StopIteration:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self._device = device
if device.type == "cuda" and torch.cuda.is_bf16_supported():
tts_dtype = torch.bfloat16
elif device.type == "cuda":
tts_dtype = torch.float16
else:
tts_dtype = torch.float32
logger.info("Loading MOSS-TTS-Nano LM from %s (dtype=%s)", self.model_path, tts_dtype)
from transformers import AutoModelForCausalLM
lm = AutoModelForCausalLM.from_pretrained(
self.model_path,
trust_remote_code=True,
torch_dtype=tts_dtype,
)
if device.type == "cuda":
try:
import flash_attn # noqa: F401
lm._set_attention_implementation("flash_attention_2")
logger.info("MOSS-TTS-Nano using flash_attention_2")
except ImportError:
lm._set_attention_implementation("sdpa")
logger.info("MOSS-TTS-Nano using sdpa (flash_attn not installed)")
lm.to(device=device)
lm.eval()
self._lm = lm
logger.info("MOSS-TTS-Nano LM loaded on %s", device)
codec_path: str = getattr(
self.config,
"audio_tokenizer_pretrained_name_or_path",
"OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano",
)
logger.info("Loading MOSS-Audio-Tokenizer-Nano from %s", codec_path)
from transformers import AutoModel
audio_tokenizer = AutoModel.from_pretrained(
codec_path,
trust_remote_code=True,
torch_dtype=torch.float32,
)
audio_tokenizer.to(device=device)
audio_tokenizer.eval()
self._audio_tokenizer = audio_tokenizer
logger.info("MOSS-Audio-Tokenizer-Nano loaded on %s", device)
# MOSS-TTS-Nano loads weights inline via from_pretrained() above;
# vLLM's weight-loading protocol still requires us to exhaust the
# iterator so the underlying stream is closed cleanly.
for _ in weights:
pass
return set()
# ------------------------------------------------------------------
# Dummy run support
# ------------------------------------------------------------------
def get_dummy_runtime_additional_information(self, num_reqs: int) -> list[dict]:
return [{"text": "hello", "voice": _DEFAULT_VOICE, "_is_dummy": True}] * num_reqs
# ------------------------------------------------------------------
# Streaming generator management
# ------------------------------------------------------------------
def _create_stream_gen(self, info: dict[str, Any]):
"""Create an inference_stream() generator for a request.
Yields (waveform_tensor, is_last) tuples.
"""
text: str = str(_pick(info, "text", "") or "")
if not text.strip():
logger.warning("MOSS-TTS-Nano received empty text; yielding silence.")
sr = getattr(self.config, "audio_tokenizer_sample_rate", 48000)
yield torch.zeros((sr,), dtype=torch.float32), True
return
mode: str = str(_pick(info, "mode", _DEFAULT_MODE))
prompt_audio_path: str | None = _pick(info, "prompt_audio_path", None)
if prompt_audio_path is not None:
prompt_audio_path = str(prompt_audio_path)
# Voice-cloning path: serving layer passes the resolved waveform as
# (wav_list, sample_rate) so we can avoid re-decoding base64 and the
# model owns temp-file lifecycle.
prompt_audio_array = _pick(info, "prompt_audio_array", None)
prompt_text: str | None = _pick(info, "prompt_text", None)
if prompt_text is not None:
prompt_text = str(prompt_text)
max_new_frames: int = int(_pick(info, "max_new_frames", _DEFAULT_MAX_NEW_FRAMES))
seed: int | None = _pick(info, "seed", None)
if seed is not None:
# Upstream inference_stream relies on global RNG state. With
# max_num_seqs > 1, concurrent seeded requests will race to set
# the global seed; snapshot and restore on exit so we at least
# don't leak state back to other components of the engine.
seed = int(seed)
_cpu_rng_state = torch.get_rng_state()
_cuda_rng_state = torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
else:
_cpu_rng_state = None
_cuda_rng_state = None
sampling = {
"text_temperature": float(_pick(info, "text_temperature", _DEFAULT_TEXT_TEMPERATURE)),
"text_top_p": float(_pick(info, "text_top_p", _DEFAULT_TEXT_TOP_P)),
"text_top_k": int(_pick(info, "text_top_k", _DEFAULT_TEXT_TOP_K)),
"audio_temperature": float(_pick(info, "audio_temperature", _DEFAULT_AUDIO_TEMPERATURE)),
"audio_top_p": float(_pick(info, "audio_top_p", _DEFAULT_AUDIO_TOP_P)),
"audio_top_k": int(_pick(info, "audio_top_k", _DEFAULT_AUDIO_TOP_K)),
"audio_repetition_penalty": float(
_pick(info, "audio_repetition_penalty", _DEFAULT_AUDIO_REPETITION_PENALTY)
),
}
device = self._device or torch.device("cpu")
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
output_path = tmp.name
# Materialise resolved reference audio to a temp WAV so upstream's
# prompt_audio_path (which expects a filesystem path) can consume it.
prompt_audio_tmp: str | None = None
if prompt_audio_array is not None and prompt_audio_path is None:
try:
import numpy as np
import soundfile as sf
wav_list, sr_in = prompt_audio_array
wav_np = np.asarray(wav_list, dtype=np.float32)
if wav_np.ndim > 1:
wav_np = np.mean(wav_np, axis=-1)
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp_in:
prompt_audio_tmp = tmp_in.name
sf.write(prompt_audio_tmp, wav_np, int(sr_in))
prompt_audio_path = prompt_audio_tmp
except Exception:
logger.exception("MOSS-TTS-Nano failed to stage prompt_audio_array to temp file")
prompt_audio_tmp = None
# Stream audio chunks from inference_stream as they arrive. A final
# "result" event is only used as a fallback when no audio events
# were emitted (e.g. very short utterances on some backends).
audio_chunks: list[torch.Tensor] = []
try:
for event in self._lm.inference_stream(
text=text,
output_audio_path=output_path,
mode=mode,
prompt_text=prompt_text,
prompt_audio_path=prompt_audio_path,
text_tokenizer_path=self.model_path,
audio_tokenizer=self._audio_tokenizer,
device=device,
nq=None,
max_new_frames=max_new_frames,
do_sample=True,
use_kv_cache=True,
**sampling,
):
event_type = str(event.get("type", ""))
if event_type == "audio":
waveform = event.get("waveform")
if waveform is not None:
chunk = torch.as_tensor(waveform, dtype=torch.float32).cpu()
if chunk.ndim == 2:
chunk = chunk.T.reshape(-1)
else:
chunk = chunk.reshape(-1)
audio_chunks.append(chunk)
# Yield each chunk as it arrives (is_last=False).
yield chunk, False
elif event_type == "result":
if not audio_chunks:
waveform = event.get("waveform")
if waveform is not None:
chunk = torch.as_tensor(waveform, dtype=torch.float32).cpu().reshape(-1)
yield chunk, True
return
except Exception:
logger.exception("MOSS-TTS-Nano inference failed for text=%r", text[:80])
finally:
try:
Path(output_path).unlink(missing_ok=True)
except Exception:
pass
if prompt_audio_tmp is not None:
try:
Path(prompt_audio_tmp).unlink(missing_ok=True)
except Exception:
pass
if _cpu_rng_state is not None:
torch.set_rng_state(_cpu_rng_state)
if _cuda_rng_state is not None:
torch.cuda.set_rng_state_all(_cuda_rng_state)
# Signal completion. If we yielded audio chunks above, the last
# yield was is_last=False, so emit a final empty sentinel.
yield torch.zeros((0,), dtype=torch.float32), True
# ------------------------------------------------------------------
# Core forward pass (streaming, VoxCPM pattern)
# ------------------------------------------------------------------
def _make_dummy_hidden(self, input_ids: torch.Tensor | None) -> torch.Tensor:
"""Return a dummy hidden_states tensor for the AR runner."""
device = self._device or torch.device("cpu")
hidden = int(getattr(self.config, "hidden_size", 768))
n = 1 if input_ids is None else max(1, input_ids.shape[0])
return torch.zeros((n, hidden), device=device, dtype=torch.float32)
@torch.inference_mode()
def forward(
self,
input_ids: torch.Tensor | None = None,
positions: torch.Tensor | None = None,
intermediate_tensors: Any = None,
inputs_embeds: torch.Tensor | None = None,
runtime_additional_information: list[dict[str, Any]] | None = None,
**kwargs: Any,
) -> OmniOutput:
sr = getattr(self.config, "audio_tokenizer_sample_rate", 48000)
sr_tensor = torch.tensor(sr, dtype=torch.int32)
empty = torch.zeros((0,), dtype=torch.float32)
hidden = self._make_dummy_hidden(input_ids)
infos = runtime_additional_information or [{}]
if not runtime_additional_information or all(info.get("_is_dummy") for info in infos):
# Dummy/warmup path: finish immediately for every row.
self._ar_last_chunk_flags = [True] * len(infos)
return OmniOutput(
text_hidden_states=hidden,
multimodal_outputs={
"model_outputs": [empty] * len(infos),
"sr": [sr_tensor] * len(infos),
},
)
if self._lm is None or self._audio_tokenizer is None:
raise RuntimeError("MOSS-TTS-Nano model not loaded. Was load_weights() called?")
outputs: list[torch.Tensor] = []
srs: list[torch.Tensor] = []
last_chunk_flags: list[bool] = []
for info in infos:
if info.get("_is_dummy"):
outputs.append(empty)
srs.append(sr_tensor)
last_chunk_flags.append(True)
continue
request_key = str(info.get("_omni_req_id", "0"))
# Create generator on first call for this request.
if request_key not in self._stream_gens:
self._stream_gens[request_key] = self._create_stream_gen(info)
generator = self._stream_gens[request_key]
try:
chunk, is_last = next(generator)
except StopIteration:
self._stream_gens.pop(request_key, None)
outputs.append(empty)
last_chunk_flags.append(True)
else:
if is_last:
self._stream_gens.pop(request_key, None)
outputs.append(chunk)
last_chunk_flags.append(bool(is_last))
srs.append(sr_tensor)
# Per-row EOS mask: compute_logits() finishes each request the step
# its last chunk was yielded, without waiting for the slowest peer.
self._ar_last_chunk_flags = last_chunk_flags
return OmniOutput(
text_hidden_states=hidden,
multimodal_outputs={"model_outputs": outputs, "sr": srs},
)
def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None:
"""Release streaming generators for requests the scheduler finished.
forward() only pops generators on normal completion (is_last or
StopIteration). Abnormal termination (cancel, timeout, preempt) would
otherwise leak the generator and skip its ``finally`` block, stranding
temp WAV files. Closing here raises GeneratorExit inside the generator
so the cleanup block runs.
"""
for req_id in finished_req_ids:
gen = self._stream_gens.pop(str(req_id), None)
if gen is not None:
try:
gen.close()
except Exception:
logger.exception("MOSS-TTS-Nano failed to close stream gen for request %s", req_id)
# ------------------------------------------------------------------
# AR runner interface
# ------------------------------------------------------------------
def compute_logits(
self,
hidden_states: torch.Tensor | OmniOutput,
sampling_metadata: Any = None,
) -> torch.Tensor:
"""Emit per-row EOS / non-EOS logits to control AR scheduler lifetime.
Rows whose ``_ar_last_chunk_flags`` entry is True get EOS-dominant
logits so the scheduler finishes that request; other rows get a
non-EOS token so they stay alive for the next streaming chunk.
"""
if isinstance(hidden_states, OmniOutput):
hidden_states = hidden_states.text_hidden_states
if hidden_states is None:
device = self._device or torch.device("cpu")
hidden_states = torch.zeros((0, 1), device=device, dtype=torch.float32)
if hidden_states.ndim == 1:
hidden_states = hidden_states.unsqueeze(-1)
elif hidden_states.ndim > 2:
hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1])
vocab_size = int(getattr(self.config, "vocab_size", 32000))
num_rows = int(hidden_states.shape[0])
logits = torch.zeros(
(num_rows, vocab_size),
dtype=torch.float32,
device=hidden_states.device,
)
eos_id = 2 if vocab_size > 2 else 0
safe_id = 1 if vocab_size > 1 and 1 != eos_id else 0
flags = self._ar_last_chunk_flags
# If we lost alignment with the forward batch (e.g. scheduler dropped
# a row), conservatively treat missing rows as "finished" to avoid
# stranding requests.
for row in range(num_rows):
is_last = flags[row] if row < len(flags) else True
if is_last:
logits[row, eos_id] = 1.0e6
else:
logits[row, eos_id] = -1.0e9
logits[row, safe_id] = 1.0e6
return logits
def embed_input_ids(
self,
input_ids: torch.Tensor,
multimodal_embeddings=None,
is_multimodal=None,
) -> torch.Tensor:
hidden = int(getattr(self.config, "hidden_size", 768))
return torch.zeros(
(input_ids.shape[0], hidden),
device=input_ids.device,
dtype=torch.float32,
)
@@ -0,0 +1,40 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""MOSS-TTS-Nano pipeline topology (frozen).
Single-stage AR TTS: text -> speech waveform in one pass. The 0.1B AR LM
and the MOSS-Audio-Tokenizer-Nano codec both run inside
``MossTTSNanoForGeneration.forward()``, which uses the VoxCPM-style
generator pattern (``inference_stream()`` stored per-request; one audio
chunk yielded per forward) to drive progressive streaming through the
AR scheduler.
"""
from vllm_omni.config.stage_config import (
PipelineConfig,
StageExecutionType,
StagePipelineConfig,
)
MOSS_TTS_NANO_PIPELINE = PipelineConfig(
model_type="moss_tts_nano",
model_arch="MossTTSNanoForCausalLM",
stages=(
StagePipelineConfig(
stage_id=0,
model_stage="moss_tts_nano",
execution_type=StageExecutionType.LLM_AR,
input_sources=(),
final_output=True,
final_output_type="audio",
owns_tokenizer=True,
engine_output_type="audio",
sampling_constraints={
"detokenize": False,
# compute_logits() forces EOS (token id 2) when the last
# streaming chunk is yielded; keep a hard backstop here.
"stop_token_ids": [2],
},
),
),
)
@@ -174,6 +174,12 @@ _OMNI_MODELS = {
"VoxtralTTSAudioGenerationForConditionalGeneration",
),
"VoxtralTTSAudioTokenizer": ("voxtral_tts", "voxtral_tts_audio_tokenizer", "VoxtralTTSAudioTokenizer"),
## MOSS-TTS-Nano
"MossTTSNanoForCausalLM": (
"moss_tts_nano",
"modeling_moss_tts_nano",
"MossTTSNanoForGeneration",
),
"DyninOmniForConditionalGeneration": (
"dynin_omni",
"dynin_omni",
-8
View File
@@ -6,7 +6,6 @@ from enum import Enum
from typing import Any
import torch
import torch.nn as nn
from vllm.logger import init_logger
from vllm.platforms import Platform
@@ -72,13 +71,6 @@ class OmniPlatform(Platform):
def prepare_diffusion_op_runtime(cls, op_name: str, **kwargs: Any) -> None:
return None
@classmethod
def get_diffusion_packed_modules_mapping(
cls,
model_class: type[nn.Module],
) -> dict[str, list[str]] | None:
return None
@classmethod
def get_diffusion_attn_backend_cls(
cls,
+2 -2
View File
@@ -25,11 +25,11 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
@classmethod
def get_omni_ar_worker_cls(cls) -> str:
return "vllm_omni.worker.gpu_ar_worker.GPUARWorker"
return "vllm_omni.platforms.musa.worker.musa_ar_worker.MUSAARWorker"
@classmethod
def get_omni_generation_worker_cls(cls) -> str:
return "vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker"
return "vllm_omni.platforms.musa.worker.musa_generation_worker.MUSAGenerationWorker"
@classmethod
def get_default_stage_config_path(cls) -> str:
@@ -0,0 +1,9 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from vllm_omni.platforms.musa.worker.musa_ar_worker import MUSAARWorker
from vllm_omni.platforms.musa.worker.musa_generation_worker import (
MUSAGenerationWorker,
)
__all__ = ["MUSAARWorker", "MUSAGenerationWorker"]
@@ -0,0 +1,103 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""MUSA AR (Autoregressive) worker for vLLM-Omni.
This worker handles autoregressive model stages (thinker/talker) on MUSA devices.
"""
import gc
import os
import torch
from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.utils.mem_utils import MemorySnapshot, format_gib
from vllm.utils.torch_utils import set_random_seed
from vllm.v1.utils import report_usage_stats
from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
from vllm.v1.worker.utils import request_memory
from vllm.v1.worker.workspace import init_workspace_manager
from vllm_omni.worker.base import OmniGPUWorkerBase
from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
from vllm_omni.worker.mixins import OmniWorkerMixin
logger = init_logger(__name__)
class MUSAARWorker(OmniWorkerMixin, OmniGPUWorkerBase):
"""MUSA AR worker for thinker/talker stages in Omni model."""
def init_device(self):
"""Initialize the MUSA device for this worker."""
# This env var set by Ray causes exceptions with graph building.
os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
parallel_config = self.parallel_config
if (
parallel_config.distributed_executor_backend not in ("ray", "external_launcher")
and parallel_config.data_parallel_backend != "ray"
and parallel_config.nnodes_within_dp == 1
):
# Use local DP rank if available, otherwise use global DP rank.
dp_local_rank = self.parallel_config.data_parallel_rank_local
if dp_local_rank is None:
dp_local_rank = self.parallel_config.data_parallel_index
tp_pp_world_size = self.parallel_config.pipeline_parallel_size * self.parallel_config.tensor_parallel_size
# DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK
self.local_rank += dp_local_rank * tp_pp_world_size
assert self.local_rank < torch.musa.device_count(), (
f"DP adjusted local rank {self.local_rank} is out of bounds. "
)
visible_device_count = torch.musa.device_count()
assert self.parallel_config.local_world_size <= visible_device_count, (
f"local_world_size ({self.parallel_config.local_world_size}) must "
f"be less than or equal to the number of visible devices "
f"({visible_device_count})."
)
self.device = torch.device(f"musa:{self.local_rank}")
torch.musa.set_device(self.device)
current_platform.check_if_supports_dtype(self.model_config.dtype)
# Initialize the distributed environment BEFORE taking memory snapshot
# This ensures NCCL buffers are allocated before we measure available memory
init_worker_distributed_environment(
self.vllm_config,
self.rank,
self.distributed_init_method,
self.local_rank,
current_platform.dist_backend,
)
# Set random seed.
set_random_seed(self.model_config.seed)
# Now take memory snapshot after distributed environment is initialized
gc.collect()
torch.musa.empty_cache()
# Take current memory snapshot
self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device)
self.requested_memory = request_memory(init_snapshot, self.cache_config)
logger.debug("worker init memory snapshot: %r", self.init_snapshot)
logger.debug("worker requested memory: %sGiB", format_gib(self.requested_memory))
# Initialize workspace manager
num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1
init_workspace_manager(self.device, num_ubatches)
if self.use_v2_model_runner:
# OMNI: v2 model runner does not yet include omni hooks.
logger.warning("OMNI MUSAARWorker forces v1 model runner for omni hooks.")
self.use_v2_model_runner = False
# Construct the model runner
self.model_runner = GPUARModelRunner(self.vllm_config, self.device)
if self.rank == 0:
# If usage stat is enabled, collect relevant info.
report_usage_stats(self.vllm_config)
@@ -0,0 +1,106 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""MUSA generation worker for vLLM-Omni.
This worker handles non-autoregressive generation stages (e.g., code2wav waveform
generation) on MUSA devices.
"""
import gc
import os
import torch
from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.tracing import instrument
from vllm.utils.mem_utils import MemorySnapshot, format_gib
from vllm.utils.torch_utils import set_random_seed
from vllm.v1.utils import report_usage_stats
from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
from vllm.v1.worker.utils import request_memory
from vllm.v1.worker.workspace import init_workspace_manager
from vllm_omni.worker.base import OmniGPUWorkerBase
from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner
from vllm_omni.worker.mixins import OmniWorkerMixin
logger = init_logger(__name__)
class MUSAGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase):
"""MUSA generation worker for non-AR waveform generation stage."""
@instrument(span_name="Init device")
def init_device(self):
"""Initialize the MUSA device for this worker."""
# This env var set by Ray causes exceptions with graph building.
os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
parallel_config = self.parallel_config
if (
parallel_config.distributed_executor_backend not in ("ray", "external_launcher")
and parallel_config.data_parallel_backend != "ray"
and parallel_config.nnodes_within_dp == 1
):
# Use local DP rank if available, otherwise use global DP rank.
dp_local_rank = self.parallel_config.data_parallel_rank_local
if dp_local_rank is None:
dp_local_rank = self.parallel_config.data_parallel_index
tp_pp_world_size = self.parallel_config.pipeline_parallel_size * self.parallel_config.tensor_parallel_size
# DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK
self.local_rank += dp_local_rank * tp_pp_world_size
assert self.local_rank < torch.musa.device_count(), (
f"DP adjusted local rank {self.local_rank} is out of bounds. "
)
visible_device_count = torch.musa.device_count()
assert self.parallel_config.local_world_size <= visible_device_count, (
f"local_world_size ({self.parallel_config.local_world_size}) must "
f"be less than or equal to the number of visible devices "
f"({visible_device_count})."
)
self.device = torch.device(f"musa:{self.local_rank}")
torch.musa.set_device(self.device)
current_platform.check_if_supports_dtype(self.model_config.dtype)
# Initialize the distributed environment BEFORE taking memory snapshot
# This ensures NCCL buffers are allocated before we measure available memory
init_worker_distributed_environment(
self.vllm_config,
self.rank,
self.distributed_init_method,
self.local_rank,
current_platform.dist_backend,
)
# Set random seed.
set_random_seed(self.model_config.seed)
# Now take memory snapshot after distributed environment is initialized
gc.collect()
torch.musa.empty_cache()
# Take current memory snapshot
self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device)
self.requested_memory = request_memory(init_snapshot, self.cache_config)
logger.debug("worker init memory snapshot: %r", self.init_snapshot)
logger.debug("worker requested memory: %sGiB", format_gib(self.requested_memory))
# Initialize workspace manager
num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1
init_workspace_manager(self.device, num_ubatches)
if self.use_v2_model_runner:
# OMNI: v2 model runner does not yet include omni hooks.
logger.warning("OMNI MUSAGenerationWorker forces v1 model runner for omni hooks.")
self.use_v2_model_runner = False
# Construct the model runner
self.model_runner = GPUGenerationModelRunner(self.vllm_config, self.device)
if self.rank == 0:
# If usage stat is enabled, collect relevant info.
report_usage_stats(self.vllm_config)
@@ -107,6 +107,12 @@ class AscendHunyuanFusedMoE(AscendSharedFusedMoE):
def __init__(self, *, prefix: str = "", **kwargs: Any) -> None:
super().__init__(prefix=prefix, **kwargs)
self._prefix = prefix
self._init_hook_handle = self.register_forward_pre_hook(self._initialize_kernel_hook, with_kwargs=True)
def _initialize_kernel_hook(self, module: Any, args: Any, kwargs: Any) -> None:
if self.quant_method:
self.quant_method.process_weights_after_loading(self)
self._init_hook_handle.remove()
def forward(self, hidden_states: Any, router_logits: Any) -> Any:
_set_hunyuan_fused_moe_forward_context(hidden_states.shape[0])
-14
View File
@@ -5,7 +5,6 @@ from contextlib import nullcontext
from typing import Any
import torch
import torch.nn as nn
from vllm.logger import init_logger
from vllm_ascend.platform import NPUPlatform
@@ -14,12 +13,6 @@ from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum
logger = init_logger(__name__)
_DIFFUSION_PACKED_MODULES_MAPPING = {
"HunyuanImage3Pipeline": {
"experts": ["experts.0.gate_up_proj", "experts.0.down_proj"],
},
}
class NPUOmniPlatform(OmniPlatform, NPUPlatform):
"""NPU/Ascend implementation of OmniPlatform.
@@ -60,13 +53,6 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
prepare_hunyuan_fused_moe_runtime()
@classmethod
def get_diffusion_packed_modules_mapping(
cls,
model_class: type[nn.Module],
) -> dict[str, list[str]] | None:
return _DIFFUSION_PACKED_MODULES_MAPPING.get(model_class.__name__, None)
@classmethod
def get_diffusion_attn_backend_cls(
cls,
+2 -1
View File
@@ -45,7 +45,8 @@ class XPUOmniPlatform(OmniPlatform, XPUPlatform):
@classmethod
def supports_torch_inductor(cls) -> bool:
return True
# TODO: Enable this when torch compile bugs are resolved
return False
@classmethod
def get_default_stage_config_path(cls) -> str:
+2 -2
View File
@@ -29,7 +29,7 @@ class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase):
@instrument(span_name="Init device")
def init_device(self):
if self.device_config.device_type in ("cuda", "musa"):
if self.device_config.device_type == "cuda":
# This env var set by Ray causes exceptions with graph building.
os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
parallel_config = self.parallel_config
@@ -52,7 +52,7 @@ class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase):
assert self.local_rank < torch.accelerator.device_count(), (
f"DP adjusted local rank {self.local_rank} is out of bounds. "
)
visible_device_count = torch.accelerator.device_count()
visible_device_count = torch.accelerator.device_count() if torch.cuda.is_available() else 0
assert self.parallel_config.local_world_size <= visible_device_count, (
f"local_world_size ({self.parallel_config.local_world_size}) must "
f"be less than or equal to the number of visible devices "
+2 -2
View File
@@ -28,7 +28,7 @@ class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase):
@instrument(span_name="Init device")
def init_device(self):
if self.device_config.device_type in ("cuda", "musa"):
if self.device_config.device_type == "cuda":
# This env var set by Ray causes exceptions with graph building.
os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
parallel_config = self.parallel_config
@@ -51,7 +51,7 @@ class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase):
assert self.local_rank < torch.accelerator.device_count(), (
f"DP adjusted local rank {self.local_rank} is out of bounds. "
)
visible_device_count = torch.accelerator.device_count()
visible_device_count = torch.accelerator.device_count() if torch.cuda.is_available() else 0
assert self.parallel_config.local_world_size <= visible_device_count, (
f"local_world_size ({self.parallel_config.local_world_size}) must "
f"be less than or equal to the number of visible devices "