[Feat] support layerwise offload for Bagel (#2734)

Signed-off-by: Nick Cao <ncao@redhat.com>
Signed-off-by: siyuan.lei <siyuanlei37@gmail.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: Nick Cao <ncao@redhat.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>

.buildkite/test-merge.yml

@@ -235,28 +235,6 @@ 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

benchmarks/build_dataset/seed_tts_design/en/meta.lst

@@ -0,0 +1,20 @@
+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.

benchmarks/build_dataset/seed_tts_smoke/en/meta.lst

@@ -0,0 +1,20 @@
+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.

benchmarks/qwen3-tts/README.md

@@ -1,107 +0,0 @@
-# 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

benchmarks/qwen3-tts/plot_results.py

@@ -1,254 +0,0 @@
-"""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)

benchmarks/qwen3-tts/results/.gitignore

@@ -1,3 +0,0 @@
-# Benchmark results are machine-specific - do not commit
-*
-!.gitignore

benchmarks/qwen3-tts/run_benchmark.sh

@@ -1,285 +0,0 @@
-#!/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 "============================================================"

benchmarks/qwen3-tts/transformers/bench_tts_hf.py

@@ -1,301 +0,0 @@
-"""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)

benchmarks/qwen3-tts/vllm_omni/bench_async_chunk.py

@@ -1,301 +0,0 @@
-"""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))

benchmarks/qwen3-tts/vllm_omni/bench_tts_serve.py

@@ -1,371 +0,0 @@
-"""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))

benchmarks/qwen3-tts/vllm_omni/plot_async_chunk.py

@@ -1,249 +0,0 @@
-"""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)

benchmarks/qwen3-tts/vllm_omni/run_async_chunk_benchmark.sh

@@ -1,169 +0,0 @@
-#!/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 "============================================================"

benchmarks/tts/README.md

@@ -0,0 +1,227 @@
+# 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

benchmarks/tts/bench_tts.py

@@ -0,0 +1,308 @@
+#!/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()

benchmarks/tts/bench_voxcpm_offline.py

@@ -6,6 +6,27 @@ 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
@@ -26,7 +47,21 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 from vllm_omni import AsyncOmni, Omni
 
-REPO_ROOT = Path(__file__).resolve().parents[3]
+
+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())
 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"
 
@@ -473,9 +508,6 @@ 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
 
 

benchmarks/tts/model_configs.yaml

@@ -0,0 +1,39 @@
+# 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: {}

benchmarks/tts/plot_results.py

@@ -0,0 +1,324 @@
+"""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()

benchmarks/voxcpm/README.md

@@ -1,119 +0,0 @@
-# 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.

benchmarks/voxcpm/vllm_omni/bench_tts_serve.py

@@ -1,283 +0,0 @@
-"""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()))

benchmarks/voxcpm/vllm_omni/run_offline_matrix.py

@@ -1,303 +0,0 @@
-"""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()))

docs/configuration/stage_configs.md

@@ -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 | optional | `"mp"` | **Pipeline-wide.** Executor backend (`"mp"` or `"ray"`). |
+| `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. |
 | `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,6 +27,9 @@ 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.

docs/user_guide/diffusion/cpu_offload_diffusion.md

@@ -195,6 +195,7 @@ 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

docs/user_guide/diffusion/quantization/msmodelslim.md

@@ -0,0 +1,56 @@
+# 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.

docs/user_guide/diffusion_features.md

@@ -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** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |

examples/offline_inference/moss_tts_nano/README.md

@@ -1,101 +0,0 @@
-# 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).

examples/offline_inference/moss_tts_nano/end2end.py

@@ -1,180 +0,0 @@
-# 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())

examples/offline_inference/text_to_image/text_to_image.py

@@ -300,6 +300,7 @@ 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)

examples/online_serving/moss_tts_nano/README.md

@@ -1,158 +0,0 @@
-# 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.

examples/online_serving/moss_tts_nano/gradio_demo.py

@@ -1,690 +0,0 @@
-"""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()

examples/online_serving/moss_tts_nano/run_gradio_demo.sh

@@ -1,43 +0,0 @@
-#!/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"

examples/online_serving/moss_tts_nano/run_server.sh

@@ -1,19 +0,0 @@
-#!/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

tests/benchmarks/conftest.py

@@ -0,0 +1,103 @@
+# 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()

tests/benchmarks/test_bench_tts_cli.py

@@ -0,0 +1,139 @@
+"""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()

tests/benchmarks/test_seed_tts_dataset_variants.py

@@ -0,0 +1,168 @@
+"""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'"

tests/dfx/conftest.py

@@ -120,7 +120,11 @@ def create_test_parameter_mapping(configs: list[dict[str, Any]]) -> dict[str, di
                 "test_name": test_name,
                 "benchmark_params": [],
             }
-        mapping[test_name]["benchmark_params"].extend(config["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)
     return mapping
 
 

tests/dfx/perf/scripts/run_benchmark.py

@@ -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"}
+    exclude_keys = {"request_rate", "baseline", "num_prompts", "max_concurrency", "task", "enabled", "eval_phase"}
 
     for key, value in params.items():
         if key in exclude_keys or value is None:

tests/dfx/perf/tests/test_runner_metadata.py

@@ -0,0 +1,79 @@
+"""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"

tests/dfx/perf/tests/test_tts.json

@@ -1,34 +1,155 @@
 [
-    {
-        "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]
-                }
-            }
-        ]
-    }
+  {
+    "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]
+        }
+      }
+    ]
+  }
 ]

tests/e2e/accuracy/qwen3_omni/run_qwen_omni_acc_benchmark.py

@@ -340,7 +340,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
     p.add_argument(
         "--max-seed-tts-mean-wer",
         type=float,
-        default=0.02,
+        default=0.5,
         help="If set, fail when seed_tts_content_error_mean is strictly above this value.",
     )
     p.add_argument(

tests/e2e/offline_inference/test_moss_tts_nano.py

@@ -1,131 +0,0 @@
-# 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

tests/e2e/online_serving/test_bagel_expansion.py

@@ -9,6 +9,7 @@ Coverage:
 - Tensor-Parallel
 - Ulysses-SP
 - Ring-Attention
+- Layerwise Offloading
 
 assert_diffusion_response validates successful generation and the expected
 512x512 resolution.
@@ -31,7 +32,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.
+    Ulysses-SP, Ring-Attention, Layerwise Offloading.
     """
 
     return [
@@ -111,6 +112,15 @@ 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,
+        ),
     ]
 
 
@@ -132,6 +142,7 @@ 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.

tests/e2e/online_serving/test_moss_tts_nano.py

@@ -1,117 +0,0 @@
-# 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)

vllm_omni/benchmarks/data_modules/seed_tts_dataset.py

@@ -259,6 +259,146 @@ 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,
@@ -277,3 +417,65 @@ 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

vllm_omni/benchmarks/patch/patch.py

@@ -40,7 +40,9 @@ 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
@@ -85,23 +87,16 @@ 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 (ref_audio, ref_text, task_type, …) into ``extra_body``.
+    """Merge Seed-TTS per-row TTS fields into ``extra_body`` and mark for PCM capture.
 
-    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.
+    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.
     """
     if not isinstance(sample, SeedTTSSampleRequest):
         return
-    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.
+    # Mark for PCM capture (WER / UTMOS eval) regardless of extra body presence.
     setattr(rfi, "seed_tts_row", True)
     sys_prompt = (sample.seed_tts_system_prompt or "").strip() or SEED_TTS_DEFAULT_OMNI_SYSTEM_PROMPT
     setattr(
@@ -112,6 +107,12 @@ 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:
@@ -136,7 +137,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 == "seed-tts"
+    is_seed_tts = args.dataset_name in ("seed-tts", "seed-tts-text", "seed-tts-design")
 
     # Check if we need to handle omni-related backends/datasets
     is_omni_backend = args.backend in ["openai-chat-omni", "openai-audio-speech", "daily-omni"]
@@ -249,7 +250,13 @@ def get_samples(args, tokenizer):
                 "--hf-name for the Hub dataset id."
             )
 
-        dataset = SeedTTSDataset(
+        _cls_map = {
+            "seed-tts": SeedTTSDataset,
+            "seed-tts-text": SeedTTSTextDataset,
+            "seed-tts-design": SeedTTSDesignDataset,
+        }
+        DatasetCls = _cls_map[args.dataset_name]
+        dataset = DatasetCls(
             dataset_path=repo_id,
             random_seed=args.seed,
             locale=getattr(args, "seed_tts_locale", "en"),

vllm_omni/config/pipeline_registry.py

@@ -79,10 +79,6 @@ _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) ---

vllm_omni/config/stage_config.py

@@ -439,7 +439,7 @@ class DeployConfig:
 
     # === Pipeline-wide engine settings (applied uniformly to every stage) ===
     trust_remote_code: bool = True
-    distributed_executor_backend: str = "mp"
+    distributed_executor_backend: str | None = None
     dtype: str | None = None
     quantization: str | None = None
     enable_prefix_caching: bool = False
@@ -601,23 +601,6 @@ 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.
 
@@ -636,7 +619,9 @@ def _apply_platform_overrides(
 ) -> DeployConfig:
     """Merge platform-specific stage overrides into deploy config."""
     if platform is None:
-        platform = _detect_platform()
+        from vllm_omni.platforms import current_omni_platform
+
+        platform = current_omni_platform.device_name.lower()
     if platform is None or deploy.platforms is None:
         return deploy
     platform_section = deploy.platforms.get(platform)

vllm_omni/deploy/moss_tts_nano.yaml

@@ -1,31 +0,0 @@
-# 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

vllm_omni/diffusion/data.py

@@ -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") is None:
+        if "quantization" in kwargs and kwargs.get("quantization_config", None) is None:
             kwargs["quantization_config"] = kwargs.pop("quantization")
         else:
             kwargs.pop("quantization", None)

vllm_omni/diffusion/executor/abstract.py

@@ -22,6 +22,10 @@ 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):

vllm_omni/diffusion/models/bagel/bagel_transformer.py

@@ -738,6 +738,8 @@ class Qwen2MoTDecoderLayer(nn.Module):
 
 
 class Qwen2MoTModel(Qwen2PreTrainedModel):
+    _layerwise_offload_blocks_attrs = ["layers"]
+
     def __init__(
         self,
         config,

vllm_omni/diffusion/models/bagel/pipeline_bagel.py

@@ -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:
+        if quant_config is None and not od_config.enable_layerwise_offload:
             self.to(self.device)
         self.setup_diffusion_pipeline_profiler(
             enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -1484,7 +1484,7 @@ class HunYuanSparseMoeBlock(nn.Module):
             config.hidden_size,
             config.num_experts,
             bias=False,
-            quant_config=None,
+            quant_config=quant_config,
             prefix=f"{prefix}.gate",
         )
         if config.use_mixed_mlp_moe > 0:
@@ -1658,8 +1658,10 @@ class HunYuanAttention(nn.Module):
         custom_pos_emb: tuple[torch.FloatTensor] | None = None,
         **kwargs,
     ) -> torch.Tensor:
-        bsz, q_len, _ = hidden_states.size()
+        bsz, q_len, hidden_size = hidden_states.size()
+        hidden_states = hidden_states.reshape(-1, hidden_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)
@@ -1723,7 +1725,7 @@ class HunyuanImage3DecoderLayer(nn.Module):
                 rope_theta=rope_theta,
                 rope_scaling=rope_scaling,
                 max_position_embeddings=max_position_embeddings,
-                quant_config=None,
+                quant_config=quant_config,
                 bias=attention_bias,
                 cache_config=None,
                 prefix=f"{prefix}.self_attn",
@@ -1933,7 +1935,7 @@ class HunyuanImage3Model(nn.Module):
                 layer_idx=int(prefix.split(".")[-1]),
                 prefix=prefix,
             ),
-            prefix=f"{prefix}.layers",
+            prefix=f"{prefix}.layers" if prefix else "layers",
         )
         if get_pp_group().is_last_rank:
             self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -1948,7 +1950,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 = self.config.hidden_size
+        hidden_size = qkv.shape[1]
 
         if hasattr(self.config, "head_dim"):
             attention_head_dim = self.config.head_dim
@@ -2001,8 +2003,15 @@ class HunyuanImage3Model(nn.Module):
         split_params_mapping = [
             (".gate_up_proj", ".gate_and_up_proj", 2, [(1, 1), (0, 1)], None),
             (
-                ".qkv_proj",
-                ".qkv_proj",
+                ".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",
                 num_attention_heads + num_kv_heads * 2,
                 [("q", num_attention_heads), ("k", num_kv_heads), ("v", num_kv_heads)],
                 self._split_qkv_weight,
@@ -2101,6 +2110,8 @@ 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:

vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py

@@ -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
+from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
 from vllm.transformers_utils.config import get_config
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
@@ -64,6 +64,15 @@ 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",

vllm_omni/diffusion/registry.py

@@ -5,6 +5,7 @@ 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
@@ -13,6 +14,7 @@ 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__)
 
@@ -242,6 +244,22 @@ _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:
@@ -264,6 +282,7 @@ 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

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -13,6 +13,7 @@ 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
@@ -21,6 +22,7 @@ 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
@@ -120,6 +122,20 @@ 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

vllm_omni/engine/async_omni_engine.py

@@ -1451,16 +1451,17 @@ 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)
 

vllm_omni/entrypoints/cli/benchmark/serve.py

@@ -147,7 +147,9 @@ 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") if c not in action.choices]
+                extra = [
+                    c for c in ("daily-omni", "seed-tts", "seed-tts-text", "seed-tts-design") if c not in action.choices
+                ]
                 if extra:
                     action.choices = list(action.choices) + extra
 

vllm_omni/entrypoints/omni_base.py

@@ -7,7 +7,6 @@ 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
@@ -250,8 +249,6 @@ 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",

vllm_omni/entrypoints/openai/api_server.py

@@ -1572,13 +1572,23 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
 
         logger.info(f"Successfully generated {len(images)} image(s)")
 
-        # Encode images to base64
-        image_data = [ImageData(b64_json=encode_image_base64(img), revised_prompt=None) for img in images]
+        # Determine output format (default to png)
+        output_format = _choose_output_format(request.output_format or "png", None)
 
-        return ImageGenerationResponse(
-            created=int(time.time()),
-            data=image_data,
-        )
+        # 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)
 
     except (EngineGenerateError, EngineDeadError) as exc:
         return _create_engine_error_json_response(raw_request, exc)

vllm_omni/entrypoints/openai/protocol/images.py

@@ -139,6 +139,12 @@ 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"""

vllm_omni/entrypoints/openai/serving_speech.py

@@ -63,7 +63,6 @@ _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
@@ -73,7 +72,6 @@ _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] = {
@@ -317,10 +315,6 @@ 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"
@@ -858,8 +852,6 @@ 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)
@@ -1117,42 +1109,6 @@ 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.
 
@@ -1724,9 +1680,6 @@ 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 = {}

vllm_omni/model_executor/models/moss_tts_nano/__init__.py

@@ -1,5 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from .modeling_moss_tts_nano import MossTTSNanoForGeneration
-
-__all__ = ["MossTTSNanoForGeneration"]

vllm_omni/model_executor/models/moss_tts_nano/configuration_moss_tts_nano.py

@@ -1,60 +0,0 @@
-# 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

vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py

@@ -1,515 +0,0 @@
-# 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,
-        )

vllm_omni/model_executor/models/moss_tts_nano/pipeline.py

@@ -1,40 +0,0 @@
-# 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],
-            },
-        ),
-    ),
-)

vllm_omni/model_executor/models/registry.py

@@ -174,12 +174,6 @@ _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",

vllm_omni/platforms/interface.py

@@ -6,6 +6,7 @@ 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
 
@@ -71,6 +72,13 @@ 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,

vllm_omni/platforms/musa/platform.py

@@ -25,11 +25,11 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
 
     @classmethod
     def get_omni_ar_worker_cls(cls) -> str:
-        return "vllm_omni.platforms.musa.worker.musa_ar_worker.MUSAARWorker"
+        return "vllm_omni.worker.gpu_ar_worker.GPUARWorker"
 
     @classmethod
     def get_omni_generation_worker_cls(cls) -> str:
-        return "vllm_omni.platforms.musa.worker.musa_generation_worker.MUSAGenerationWorker"
+        return "vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker"
 
     @classmethod
     def get_default_stage_config_path(cls) -> str:

vllm_omni/platforms/musa/worker/__init__.py

@@ -1,9 +0,0 @@
-# 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"]

vllm_omni/platforms/musa/worker/musa_ar_worker.py

@@ -1,103 +0,0 @@
-# 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)

vllm_omni/platforms/musa/worker/musa_generation_worker.py

@@ -1,106 +0,0 @@
-# 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)

vllm_omni/platforms/npu/models/hunyuan_fused_moe.py

@@ -107,12 +107,6 @@ 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])

vllm_omni/platforms/npu/platform.py

@@ -5,6 +5,7 @@ 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
 
@@ -13,6 +14,12 @@ 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.
@@ -53,6 +60,13 @@ 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,

vllm_omni/platforms/xpu/platform.py

@@ -45,8 +45,7 @@ class XPUOmniPlatform(OmniPlatform, XPUPlatform):
 
     @classmethod
     def supports_torch_inductor(cls) -> bool:
-        # TODO: Enable this when torch compile bugs are resolved
-        return False
+        return True
 
     @classmethod
     def get_default_stage_config_path(cls) -> str:

vllm_omni/worker/gpu_ar_worker.py

@@ -29,7 +29,7 @@ class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase):
 
     @instrument(span_name="Init device")
     def init_device(self):
-        if self.device_config.device_type == "cuda":
+        if self.device_config.device_type in ("cuda", "musa"):
             # 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() if torch.cuda.is_available() else 0
+                visible_device_count = torch.accelerator.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 "

vllm_omni/worker/gpu_generation_worker.py

@@ -28,7 +28,7 @@ class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase):
 
     @instrument(span_name="Init device")
     def init_device(self):
-        if self.device_config.device_type == "cuda":
+        if self.device_config.device_type in ("cuda", "musa"):
             # 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() if torch.cuda.is_available() else 0
+                visible_device_count = torch.accelerator.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 "