Merge branch 'main' into dev/rebase-v0170

Signed-off-by: Ricardo Noriega De Soto <rnoriega@redhat.com>

.buildkite/release-pipeline.yaml

@@ -0,0 +1,38 @@
+steps:
+  - label: ":docker: Build image"
+    key: image-build
+    commands:
+      - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
+      - "docker build --progress=plain --file docker/Dockerfile.ci -t vllm-omni-ci ."
+      - "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
+      - "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
+    agents:
+      queue: "cpu_queue_premerge"
+
+  # L2 Test
+  - label: "Upload Ready Pipeline"
+    depends_on: image-build
+    key: upload-ready-pipeline
+    commands:
+      - buildkite-agent pipeline upload .buildkite/test-ready.yml
+    agents:
+      queue: "cpu_queue_premerge"
+
+  # L3 Test
+  - label: "Upload Merge Pipeline"
+    depends_on: image-build
+    key: upload-merge-pipeline
+    commands:
+      - buildkite-agent pipeline upload .buildkite/test-merge.yml
+    agents:
+      queue: "cpu_queue_premerge"
+
+  # L4 Test
+  - label: "Upload Nightly Pipeline"
+    depends_on: image-build
+    key: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1"
+    commands:
+      - buildkite-agent pipeline upload .buildkite/test-nightly.yml
+    agents:
+      queue: "cpu_queue_premerge"

.buildkite/test-nightly.yml

@@ -1,5 +1,5 @@
 steps:
-  - label: "Omni Model Test with H100"
+  - label: ":full_moon: Omni Model Test with H100"
     timeout_in_minutes: 90
     depends_on: upload-nightly-pipeline
     if: build.env("NIGHTLY") == "1"
@@ -41,7 +41,7 @@ steps:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
 
-  - label: "Omni Model Test"
+  - label: ":full_moon: Omni Model Test"
     timeout_in_minutes: 60
     depends_on: upload-nightly-pipeline
     if: build.env("NIGHTLY") == "1"
@@ -56,18 +56,44 @@ steps:
           always-pull: true
           shm-size: "8gb"
           propagate-environment: true
+          shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "Omni Model Perf Test"
+  - label: ":full_moon: Qwen3-TTS Non-Async-Chunk E2E Test"
-    timeout_in_minutes: 120
+    timeout_in_minutes: 30
+    depends_on: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1"
+    commands:
+      - |
+        huggingface-cli download Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice
+        export VLLM_WORKER_MULTIPROC_METHOD=spawn
+        pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py::TestQwen3TTSNoAsyncChunk
+    agents:
+      queue: "gpu_4_queue"
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          shm-size: "8gb"
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: ":full_moon: Omni Model Perf Test with H100"
+    key: nightly-performance
+    timeout_in_minutes: 180
     depends_on: upload-nightly-pipeline
     if: build.env("NIGHTLY") == "1"
     commands:
       - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - export BENCHMARK_DIR=tests
       - pytest -s -v tests/perf/scripts/run_benchmark.py
+      - buildkite-agent artifact upload "tests/*.json"
     agents:
       queue: "mithril-h100-pool"
     plugins:
@@ -96,3 +122,18 @@ steps:
                 hostPath:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
+
+  - label: ":email: Nightly Perf Collection & Email"
+    key: nightly-perf-distribution
+    depends_on: nightly-performance
+    if: build.env("NIGHTLY") == "1"
+    commands:
+      - pip install openpyxl
+      - export DEFAULT_INPUT_DIR=tests
+      - export DEFAULT_OUTPUT_DIR=tests
+      - buildkite-agent artifact download "tests/*.json" . --step nightly-performance
+      - python tools/nightly/generate_nightly_perf_excel.py
+      - python tools/nightly/send_nightly_perf_email.py
+      - buildkite-agent artifact upload "tests/*.xlsx"
+    agents:
+      queue: "cpu_queue_premerge"

.buildkite/test-ready.yml

@@ -214,9 +214,10 @@ steps:
     commands:
       - |
         timeout 20m bash -c '
+          huggingface-cli download Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice
           export VLLM_LOGGING_LEVEL=DEBUG
           export VLLM_WORKER_MULTIPROC_METHOD=spawn
-          pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
+          pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py -k "not NoAsyncChunk"
         '
     agents:
       queue: "gpu_4_queue"

.pre-commit-config.yaml

@@ -8,6 +8,7 @@ repos:
       # list of supported hooks: https://pre-commit.com/hooks.html
       - id: check-yaml
         args: ["--unsafe"]
+        exclude: "examples/online_serving/chart-helm/templates/"
       - id: debug-statements
       - id: end-of-file-fixer
       - id: mixed-line-ending

README.md

@@ -1,6 +1,6 @@
 <p align="center">
   <picture>
-    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png">
+    <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png">
     <img alt="vllm-omni" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png" width=55%>
   </picture>
 </p>
@@ -9,7 +9,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 </h3>
 
 <p align="center">
-| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true"><b>Slides</b></a> |
+| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://drive.google.com/file/d/12k_DOo2xoEROvd4xQxYNuEi6C2GECoRx/view?usp=drive_link"><b>Slides</b></a> |
 </p>
 
 
@@ -17,6 +17,8 @@ Easy, fast, and cheap omni-modality model serving for everyone
 
 *Latest News* 🔥
 
+- [2026/03] **[vllm-omni-skills](https://github.com/hsliuustc0106/vllm-omni-skills)** is a community-driven collection of AI assistant skills that help developers work with vLLM-Omni more effectively. These skills can be used with popular agentic AI coding assistants like **Cursor IDE**, **Claude**, **Codex**, and more.
+- [2026/02] We released [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0) - A major alignment + capability release that rebases onto **upstream vLLM v0.16.0** and significantly expands performance, distributed execution, and production readiness across **Qwen3-Omni / Qwen3-TTS**, **Bagel**, **MiMo-Audio**, **GLM-Image** and the **Diffusion (DiT) image/video stack**—while also improving platform coverage (CUDA / ROCm / NPU / XPU), CI quality, and documentation.
 - [2026/02] We released [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0) - This is the first **stable release** of vLLM-Omni that expands Omni’s diffusion / image-video generation and audio / TTS stack, improves distributed execution and memory efficiency, and broadens platform/backend coverage (GPU/ROCm/NPU/XPU). It also brings meaningful upgrades to serving APIs, profiling & benchmarking, and overall stability. Please check our latest [paper](https://arxiv.org/abs/2602.02204) for architecture design and performance results.
 - [2026/01] We released [0.12.0rc1](https://github.com/vllm-project/vllm-omni/releases/tag/v0.12.0rc1) - a major RC milestone focused on maturing the diffusion stack, strengthening OpenAI-compatible serving, expanding omni-model coverage, and improving stability across platforms (GPU/NPU/ROCm).
 - [2025/11] vLLM community officially released [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni) in order to support omni-modality models serving.

benchmarks/qwen3-tts/vllm_omni/bench_async_chunk.py

@@ -0,0 +1,301 @@
+"""Benchmark comparing async_chunk on vs off for Qwen3-TTS.
+
+Measures TTFP (Time-to-First-Packet), E2E latency, and RTF across
+concurrency levels for both async_chunk modes. Saves results as JSON.
+
+Usage:
+    # Run against a server already serving with a given config:
+    python bench_async_chunk.py \
+        --host 127.0.0.1 --port 8000 \
+        --config-name async_chunk_on \
+        --num-prompts 50 \
+        --max-concurrency 1 10 \
+        --result-dir results/
+"""
+
+import argparse
+import asyncio
+import json
+import time
+from dataclasses import asdict, dataclass, field
+from datetime import datetime
+from pathlib import Path
+
+import aiohttp
+import numpy as np
+from tqdm.asyncio import tqdm
+
+PROMPTS = [
+    "Hello, welcome to the voice synthesis benchmark test.",
+    "She said she would be here by noon, but nobody showed up.",
+    "The quick brown fox jumps over the lazy dog near the riverbank.",
+    "I can't believe how beautiful the sunset looks from up here on the mountain.",
+    "Please remember to bring your identification documents to the appointment tomorrow morning.",
+    "Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
+    "The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
+    "After the meeting, we should discuss the quarterly results and plan for the next phase.",
+    "Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
+    "The train leaves at half past seven, so we need to arrive at the station before then.",
+    "Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
+    "It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
+]
+
+
+@dataclass
+class RequestResult:
+    success: bool = False
+    ttfp: float = 0.0
+    e2e: float = 0.0
+    audio_bytes: int = 0
+    audio_duration: float = 0.0
+    rtf: float = 0.0
+    prompt: str = ""
+    error: str = ""
+
+
+@dataclass
+class BenchmarkResult:
+    config_name: str = ""
+    concurrency: int = 0
+    num_prompts: int = 0
+    completed: int = 0
+    failed: int = 0
+    duration_s: float = 0.0
+    mean_ttfp_ms: float = 0.0
+    median_ttfp_ms: float = 0.0
+    std_ttfp_ms: float = 0.0
+    p90_ttfp_ms: float = 0.0
+    p95_ttfp_ms: float = 0.0
+    p99_ttfp_ms: float = 0.0
+    mean_e2e_ms: float = 0.0
+    median_e2e_ms: float = 0.0
+    std_e2e_ms: float = 0.0
+    p90_e2e_ms: float = 0.0
+    p95_e2e_ms: float = 0.0
+    p99_e2e_ms: float = 0.0
+    mean_rtf: float = 0.0
+    median_rtf: float = 0.0
+    std_rtf: float = 0.0
+    mean_audio_duration_s: float = 0.0
+    total_audio_duration_s: float = 0.0
+    audio_throughput: float = 0.0
+    request_throughput: float = 0.0
+    per_request: list = field(default_factory=list)
+
+
+def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width: int = 2) -> float:
+    return num_bytes / sample_width / sample_rate
+
+
+async def send_tts_request(
+    session: aiohttp.ClientSession,
+    api_url: str,
+    prompt: str,
+    voice: str = "vivian",
+    language: str = "English",
+    stream: bool = True,
+    pbar: tqdm | None = None,
+) -> RequestResult:
+    payload = {
+        "input": prompt,
+        "voice": voice,
+        "language": language,
+        "stream": stream,
+        "response_format": "pcm",
+    }
+
+    result = RequestResult(prompt=prompt)
+    st = time.perf_counter()
+
+    try:
+        async with session.post(api_url, json=payload) as response:
+            if response.status != 200:
+                result.error = f"HTTP {response.status}: {await response.text()}"
+                return result
+
+            first_chunk = True
+            total_bytes = 0
+
+            async for chunk in response.content.iter_any():
+                if first_chunk and len(chunk) > 0:
+                    result.ttfp = time.perf_counter() - st
+                    first_chunk = False
+                total_bytes += len(chunk)
+
+            result.e2e = time.perf_counter() - st
+            result.audio_bytes = total_bytes
+            result.audio_duration = pcm_bytes_to_duration(total_bytes)
+            if result.audio_duration > 0:
+                result.rtf = result.e2e / result.audio_duration
+            result.success = True
+
+    except Exception as e:
+        result.error = str(e)
+        result.e2e = time.perf_counter() - st
+
+    if pbar:
+        pbar.update(1)
+    return result
+
+
+async def run_benchmark(
+    host: str,
+    port: int,
+    num_prompts: int,
+    max_concurrency: int,
+    num_warmups: int = 3,
+    voice: str = "vivian",
+    language: str = "English",
+    stream: bool = True,
+) -> BenchmarkResult:
+    api_url = f"http://{host}:{port}/v1/audio/speech"
+
+    connector = aiohttp.TCPConnector(limit=max_concurrency, limit_per_host=max_concurrency, keepalive_timeout=60)
+    session = aiohttp.ClientSession(connector=connector, timeout=aiohttp.ClientTimeout(total=600))
+
+    if num_warmups > 0:
+        print(f"  Warming up with {num_warmups} requests...")
+        warmup_tasks = [
+            send_tts_request(session, api_url, PROMPTS[i % len(PROMPTS)], voice, language, stream)
+            for i in range(num_warmups)
+        ]
+        await asyncio.gather(*warmup_tasks)
+        print("  Warmup done.")
+
+    request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
+
+    print(f"  Running {num_prompts} requests with concurrency={max_concurrency}...")
+    semaphore = asyncio.Semaphore(max_concurrency)
+    pbar = tqdm(total=num_prompts, desc=f"  concurrency={max_concurrency}")
+
+    async def limited_request(prompt):
+        async with semaphore:
+            return await send_tts_request(session, api_url, prompt, voice, language, stream, pbar)
+
+    start_time = time.perf_counter()
+    tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
+    results: list[RequestResult] = await asyncio.gather(*tasks)
+    duration = time.perf_counter() - start_time
+    pbar.close()
+
+    await session.close()
+
+    successful = [r for r in results if r.success]
+    failed = [r for r in results if not r.success]
+
+    bench = BenchmarkResult(
+        concurrency=max_concurrency,
+        num_prompts=num_prompts,
+        completed=len(successful),
+        failed=len(failed),
+        duration_s=duration,
+    )
+
+    if successful:
+        ttfps = [r.ttfp * 1000 for r in successful]
+        e2es = [r.e2e * 1000 for r in successful]
+        rtfs = [r.rtf for r in successful]
+        audio_durs = [r.audio_duration for r in successful]
+
+        bench.mean_ttfp_ms = float(np.mean(ttfps))
+        bench.median_ttfp_ms = float(np.median(ttfps))
+        bench.std_ttfp_ms = float(np.std(ttfps))
+        bench.p90_ttfp_ms = float(np.percentile(ttfps, 90))
+        bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
+        bench.p99_ttfp_ms = float(np.percentile(ttfps, 99))
+
+        bench.mean_e2e_ms = float(np.mean(e2es))
+        bench.median_e2e_ms = float(np.median(e2es))
+        bench.std_e2e_ms = float(np.std(e2es))
+        bench.p90_e2e_ms = float(np.percentile(e2es, 90))
+        bench.p95_e2e_ms = float(np.percentile(e2es, 95))
+        bench.p99_e2e_ms = float(np.percentile(e2es, 99))
+
+        bench.mean_rtf = float(np.mean(rtfs))
+        bench.median_rtf = float(np.median(rtfs))
+        bench.std_rtf = float(np.std(rtfs))
+
+        bench.mean_audio_duration_s = float(np.mean(audio_durs))
+        bench.total_audio_duration_s = float(np.sum(audio_durs))
+        bench.audio_throughput = bench.total_audio_duration_s / duration
+        bench.request_throughput = len(successful) / duration
+
+        bench.per_request = [
+            {
+                "ttfp_ms": r.ttfp * 1000,
+                "e2e_ms": r.e2e * 1000,
+                "rtf": r.rtf,
+                "audio_duration_s": r.audio_duration,
+                "prompt": r.prompt,
+            }
+            for r in successful
+        ]
+
+    print(f"\n{'=' * 60}")
+    print(f"  Concurrency: {max_concurrency}  |  Completed: {bench.completed}  |  Failed: {bench.failed}")
+    print(f"  Duration: {duration:.2f}s  |  Throughput: {bench.request_throughput:.2f} req/s")
+    print(
+        f"  TTFP (ms):  mean={bench.mean_ttfp_ms:.1f}  median={bench.median_ttfp_ms:.1f}"
+        f"  p90={bench.p90_ttfp_ms:.1f}  p99={bench.p99_ttfp_ms:.1f}"
+    )
+    print(
+        f"  E2E (ms):   mean={bench.mean_e2e_ms:.1f}  median={bench.median_e2e_ms:.1f}"
+        f"  p90={bench.p90_e2e_ms:.1f}  p99={bench.p99_e2e_ms:.1f}"
+    )
+    print(f"  RTF:        mean={bench.mean_rtf:.3f}  median={bench.median_rtf:.3f}")
+    print(f"  Throughput: {bench.audio_throughput:.2f} audio-sec/wall-sec")
+    print(f"{'=' * 60}\n")
+
+    if failed:
+        for r in failed[:3]:
+            print(f"  [ERROR] {r.error[:200]}")
+
+    return bench
+
+
+async def main(args):
+    all_results = []
+
+    for concurrency in args.max_concurrency:
+        result = await run_benchmark(
+            host=args.host,
+            port=args.port,
+            num_prompts=args.num_prompts,
+            max_concurrency=concurrency,
+            num_warmups=args.num_warmups,
+            voice=args.voice,
+            language=args.language,
+            stream=args.stream,
+        )
+        result.config_name = args.config_name
+        all_results.append(asdict(result))
+
+    result_dir = Path(args.result_dir)
+    result_dir.mkdir(parents=True, exist_ok=True)
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
+
+    with open(result_file, "w") as f:
+        json.dump(all_results, f, indent=2)
+    print(f"Results saved to {result_file}")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Qwen3-TTS async_chunk benchmark client")
+    parser.add_argument("--host", type=str, default="127.0.0.1")
+    parser.add_argument("--port", type=int, default=8000)
+    parser.add_argument("--num-prompts", type=int, default=50)
+    parser.add_argument("--max-concurrency", type=int, nargs="+", default=[1, 10])
+    parser.add_argument("--num-warmups", type=int, default=3)
+    parser.add_argument("--voice", type=str, default="vivian")
+    parser.add_argument("--language", type=str, default="English")
+    parser.add_argument("--stream", action="store_true", default=True)
+    parser.add_argument("--no-stream", dest="stream", action="store_false")
+    parser.add_argument("--config-name", type=str, default="async_chunk_on")
+    parser.add_argument("--result-dir", type=str, default="results")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    asyncio.run(main(args))

benchmarks/qwen3-tts/vllm_omni/plot_async_chunk.py

@@ -0,0 +1,249 @@
+"""Plot TTFP comparison: async_chunk off vs on.
+
+Generates a bar chart with improvement arrows, matching the Qwen3-Omni
+async_chunk benchmark figure style.
+
+Usage:
+    python plot_async_chunk.py \
+        --off results/bench_async_chunk_off_*.json \
+        --on results/bench_async_chunk_on_*.json \
+        --output results/qwen3_tts_async_chunk_ttfp.png
+
+    # Also supports E2E and RTF metrics:
+    python plot_async_chunk.py \
+        --off results/bench_async_chunk_off_*.json \
+        --on results/bench_async_chunk_on_*.json \
+        --metric e2e \
+        --output results/qwen3_tts_async_chunk_e2e.png
+"""
+
+import argparse
+import json
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+METRIC_CONFIG = {
+    "ttfp": {
+        "key": "mean_ttfp_ms",
+        "ylabel": "TTFP (s)",
+        "title": "TTFP (Time to First Audio Packet) - Qwen3-TTS, by concurrency",
+        "to_seconds": True,
+    },
+    "e2e": {
+        "key": "mean_e2e_ms",
+        "ylabel": "E2E (s)",
+        "title": "E2E Latency - Qwen3-TTS, by concurrency",
+        "to_seconds": True,
+    },
+    "rtf": {
+        "key": "mean_rtf",
+        "ylabel": "RTF",
+        "title": "Real-Time Factor - Qwen3-TTS, by concurrency",
+        "to_seconds": False,
+    },
+}
+
+
+def load_results(path: str) -> list[dict]:
+    with open(path) as f:
+        return json.load(f)
+
+
+def plot_ttfp_comparison(
+    off_results: list[dict],
+    on_results: list[dict],
+    metric: str,
+    output_path: str,
+    title_override: str | None = None,
+):
+    cfg = METRIC_CONFIG[metric]
+    key = cfg["key"]
+    to_seconds = cfg["to_seconds"]
+
+    off_map = {r["concurrency"]: r for r in off_results}
+    on_map = {r["concurrency"]: r for r in on_results}
+    concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
+
+    off_vals = []
+    on_vals = []
+    for c in concurrencies:
+        v_off = off_map[c][key]
+        v_on = on_map[c][key]
+        if to_seconds:
+            v_off /= 1000.0
+            v_on /= 1000.0
+        off_vals.append(v_off)
+        on_vals.append(v_on)
+
+    fig, ax = plt.subplots(figsize=(8, 6))
+
+    x = np.arange(len(concurrencies))
+    width = 0.3
+
+    ax.bar(x - width / 2, off_vals, width, label="async_chunk off", color="#87CEEB", edgecolor="none")
+    ax.bar(x + width / 2, on_vals, width, label="async_chunk on", color="#FFF8DC", edgecolor="#DDD8B8")
+
+    # Draw improvement arrows and labels
+    for i in range(len(concurrencies)):
+        v_off = off_vals[i]
+        v_on = on_vals[i]
+        if v_on > 0:
+            improvement = v_off / v_on
+        else:
+            improvement = float("inf")
+
+        # Arrow from top of off-bar to top of on-bar
+        arrow_start_x = x[i] - width / 2
+        arrow_start_y = v_off * 0.95
+        arrow_end_x = x[i] + width / 2
+        arrow_end_y = v_on * 1.05
+
+        ax.annotate(
+            "",
+            xy=(arrow_end_x, arrow_end_y),
+            xytext=(arrow_start_x, arrow_start_y),
+            arrowprops=dict(arrowstyle="->", color="red", lw=1.5),
+        )
+
+        # Improvement label
+        label_x = (arrow_start_x + arrow_end_x) / 2
+        label_y = arrow_start_y + (v_off - v_on) * 0.15
+        ax.text(
+            label_x,
+            label_y,
+            f"{improvement:.1f}x improvement",
+            ha="center",
+            va="bottom",
+            fontsize=10,
+            color="red",
+            fontweight="bold",
+        )
+
+    title = title_override or cfg["title"]
+    ax.set_title(title, fontsize=13, fontweight="bold")
+    ax.set_ylabel(cfg["ylabel"], fontsize=12)
+    ax.set_xlabel("Max concurrency", fontsize=12)
+    ax.set_xticks(x)
+    ax.set_xticklabels([str(c) for c in concurrencies])
+    ax.set_yscale("log")
+    ax.legend(loc="upper left", fontsize=11)
+    ax.grid(axis="y", alpha=0.3, linestyle="--")
+    ax.set_axisbelow(True)
+
+    plt.tight_layout()
+    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+    plt.savefig(output_path, dpi=150, bbox_inches="tight")
+    print(f"Plot saved to {output_path}")
+    plt.close()
+
+
+def plot_all_metrics(off_results: list[dict], on_results: list[dict], output_path: str):
+    """Generate a 1x3 subplot with TTFP, E2E, and RTF comparisons."""
+    off_map = {r["concurrency"]: r for r in off_results}
+    on_map = {r["concurrency"]: r for r in on_results}
+    concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
+
+    fig, axes = plt.subplots(1, 3, figsize=(18, 6))
+    fig.suptitle("Qwen3-TTS: async_chunk on vs off", fontsize=15, fontweight="bold")
+
+    for ax, metric in zip(axes, ["ttfp", "e2e", "rtf"]):
+        cfg = METRIC_CONFIG[metric]
+        key = cfg["key"]
+        to_seconds = cfg["to_seconds"]
+
+        off_vals = []
+        on_vals = []
+        for c in concurrencies:
+            v_off = off_map[c][key]
+            v_on = on_map[c][key]
+            if to_seconds:
+                v_off /= 1000.0
+                v_on /= 1000.0
+            off_vals.append(v_off)
+            on_vals.append(v_on)
+
+        x = np.arange(len(concurrencies))
+        width = 0.3
+        ax.bar(x - width / 2, off_vals, width, label="async_chunk off", color="#87CEEB")
+        ax.bar(x + width / 2, on_vals, width, label="async_chunk on", color="#FFF8DC", edgecolor="#DDD8B8")
+
+        for i in range(len(concurrencies)):
+            if on_vals[i] > 0:
+                improvement = off_vals[i] / on_vals[i]
+                ax.annotate(
+                    "",
+                    xy=(x[i] + width / 2, on_vals[i] * 1.05),
+                    xytext=(x[i] - width / 2, off_vals[i] * 0.95),
+                    arrowprops=dict(arrowstyle="->", color="red", lw=1.5),
+                )
+                label_y = off_vals[i] * 0.85
+                ax.text(x[i], label_y, f"{improvement:.1f}x", ha="center", fontsize=10, color="red", fontweight="bold")
+
+        ax.set_title(cfg["title"].split(" - ")[0], fontsize=12, fontweight="bold")
+        ax.set_ylabel(cfg["ylabel"], fontsize=11)
+        ax.set_xlabel("Max concurrency", fontsize=11)
+        ax.set_xticks(x)
+        ax.set_xticklabels([str(c) for c in concurrencies])
+        if metric != "rtf":
+            ax.set_yscale("log")
+        ax.legend(fontsize=9)
+        ax.grid(axis="y", alpha=0.3, linestyle="--")
+        ax.set_axisbelow(True)
+
+    plt.tight_layout()
+    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+    plt.savefig(output_path, dpi=150, bbox_inches="tight")
+    print(f"Plot saved to {output_path}")
+    plt.close()
+
+
+def print_table(off_results: list[dict], on_results: list[dict]):
+    off_map = {r["concurrency"]: r for r in off_results}
+    on_map = {r["concurrency"]: r for r in on_results}
+    concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
+
+    print("\n## Benchmark Results: async_chunk off vs on\n")
+    print("| Metric | Concurrency | async_chunk off | async_chunk on | Improvement |")
+    print("| --- | --- | --- | --- | --- |")
+
+    for name, key, fmt in [
+        ("TTFP (ms)", "mean_ttfp_ms", ".1f"),
+        ("E2E (ms)", "mean_e2e_ms", ".1f"),
+        ("RTF", "mean_rtf", ".3f"),
+        ("Throughput", "audio_throughput", ".2f"),
+    ]:
+        for c in concurrencies:
+            v_off = off_map[c].get(key, 0)
+            v_on = on_map[c].get(key, 0)
+            if v_on > 0 and key != "audio_throughput":
+                ratio = f"{v_off / v_on:.1f}x"
+            elif v_off > 0 and key == "audio_throughput":
+                ratio = f"{v_on / v_off:.1f}x"
+            else:
+                ratio = "N/A"
+            print(f"| {name} | {c} | {v_off:{fmt}} | {v_on:{fmt}} | {ratio} |")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Plot async_chunk comparison for Qwen3-TTS")
+    parser.add_argument("--off", type=str, required=True, help="JSON results for async_chunk off")
+    parser.add_argument("--on", type=str, required=True, help="JSON results for async_chunk on")
+    parser.add_argument("--metric", type=str, default="ttfp", choices=["ttfp", "e2e", "rtf", "all"])
+    parser.add_argument("--output", type=str, default="results/qwen3_tts_async_chunk.png")
+    parser.add_argument("--title", type=str, default=None, help="Custom title override")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    off_results = load_results(args.off)
+    on_results = load_results(args.on)
+
+    print_table(off_results, on_results)
+
+    if args.metric == "all":
+        plot_all_metrics(off_results, on_results, args.output)
+    else:
+        plot_ttfp_comparison(off_results, on_results, args.metric, args.output, args.title)

benchmarks/qwen3-tts/vllm_omni/run_async_chunk_benchmark.sh

@@ -0,0 +1,167 @@
+#!/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)"
+
+STAGE_CONFIG_ON="vllm_omni/model_executor/stage_configs/qwen3_tts.yaml"
+STAGE_CONFIG_OFF="vllm_omni/model_executor/stage_configs/qwen3_tts_no_async_chunk.yaml"
+
+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}" \
+    --stage-configs-path "${STAGE_CONFIG_ON}" \
+    --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}" \
+    --stage-configs-path "${STAGE_CONFIG_OFF}" \
+    --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 "============================================================"

docker/Dockerfile.ci

@@ -1,14 +1,13 @@
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.16.0
+ARG VLLM_BASE_TAG=v0.17.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}
-
 COPY . .
 
 # Install system dependencies
 RUN apt-get update && \
-    apt-get install -y ffmpeg git sox libsox-fmt-all jq && \
+    apt-get install -y espeak-ng ffmpeg git sox libsox-fmt-all jq && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 

docker/Dockerfile.rocm

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.16.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.17.0
 FROM ${BASE_IMAGE} AS base
 
 # Declare a variable to know if we want to use the nightly build or the stable build.
@@ -19,7 +19,7 @@ WORKDIR ${COMMON_WORKDIR}
 
 # Step 1: Setup - Install system dependencies
 RUN apt-get update && \
-    apt-get install -y ffmpeg git sox libsox-fmt-all jq && \
+    apt-get install -y espeak-ng ffmpeg git sox libsox-fmt-all jq && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 

docker/Dockerfile.xpu

@@ -9,12 +9,12 @@ ARG PYTHON_VERSION=3.12
 ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/xpu"
 
 RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
-    echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list && \
+    echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list
-    add-apt-repository -y ppa:kobuk-team/intel-graphics
 
 RUN apt clean && apt-get update -y && \
     apt-get install -y --no-install-recommends --fix-missing \
     curl \
+    espeak-ng \
     ffmpeg \
     git \
     libsndfile1 \
@@ -31,9 +31,22 @@ RUN apt clean && apt-get update -y && \
     python3-pip
 
 RUN apt update && apt upgrade -y && \
-    apt install -y libze1 libze-dev libze-intel-gpu1 intel-opencl-icd libze-intel-gpu-raytracing intel-ocloc && \
     apt install -y intel-oneapi-compiler-dpcpp-cpp-2025.3
 
+# Install UMD with fixed version
+RUN mkdir neo && \
+    cd neo && \
+    wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-core-2_2.24.8+20344_amd64.deb && \
+    wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-opencl-2_2.24.8+20344_amd64.deb && \
+    wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-ocloc_25.48.36300.8-0_amd64.deb && \
+    wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-opencl-icd_25.48.36300.8-0_amd64.deb && \
+    wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libigdgmm12_22.8.2_amd64.deb && \
+    wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libze-intel-gpu1_25.48.36300.8-0_amd64.deb && \
+    wget https://github.com/oneapi-src/level-zero/releases/download/v1.26.0/level-zero_1.26.0+u24.04_amd64.deb && \
+    dpkg -i *.deb && \
+    cd .. && \
+    rm -rf neo
+
 ENV PATH="/root/.local/bin:$PATH"
 ENV VIRTUAL_ENV="/opt/venv"
 ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python

docs/.nav.yml

@@ -66,6 +66,7 @@ nav:
   - Model Implementation:
     - contributing/model/README.md
     - contributing/model/adding_omni_model.md
+    - contributing/model/adding_tts_model.md
     - contributing/model/adding_diffusion_model.md
   - CI: contributing/ci
   - Design Documents:

docs/configuration/README.md

@@ -2,7 +2,7 @@
 
 This section lists the most common options for running vLLM-Omni.
 
-For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.14.0/configuration/index.html)
+For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.16.0/configuration/index.html)
 
 Currently, the main options are maintained by stage configs for each model.
 
@@ -19,3 +19,4 @@ For introduction, please check [Introduction for stage config](./stage_configs.m
 - **[TeaCache Configuration](../user_guide/diffusion/teacache.md)** - Enable TeaCache adaptive caching for DiT models to achieve 1.5x-2.0x speedup with minimal quality loss
 - **[Cache-DiT Configuration](../user_guide/diffusion/cache_dit_acceleration.md)** - Enable Cache-DiT as cache acceleration backends for DiT models
 - **[Parallelism Configuration](../user_guide/diffusion/parallelism_acceleration.md)** - Enable parallelism (e.g., sequence parallelism) for for DiT models
+- **[CPU Offloading](../user_guide/diffusion/cpu_offload_diffusion.md)** - Enable CPU offloading (model-level and layerwise) for for DiT models

docs/configuration/stage_configs.md

@@ -135,6 +135,14 @@ Each stage in the `stage_args` list contains the following configuration options
 
 A unique identifier for each stage in the multi-stage pipeline. Stages are numbered sequentially starting from 0, and this ID is used to reference stages in inter-stage dependencies (e.g., `engine_input_source`).
 
+### `prompt_expand_func` (Optional)
+
+A custom Python function hook for the LLM stage (Stage 0) that expands a single incoming prompt object into multiple prompts. This is primarily used for multi-modal Classifier-Free Guidance (CFG), where it generates the necessary companion requests (like a negative text prompt) and tags them with internal roles (e.g., `cfg_text`). This ensures the upstream LLM generates the needed contextual hidden states for both the conditional and unconditional generations simultaneously.
+
+### `cfg_kv_collect_func` (Optional)
+
+A custom Python function hook for downstream diffusion stages (Stage 1+) to collect, map, and process the KV caches transferred from the companion requests fired by `prompt_expand_func`. It aggregates the hidden condition states cleanly (e.g., binding them as `cfg_text_past_key_values` and `cfg_text_kv_metadata`), allowing the diffusion runtime to perform CFG smoothly without redundantly evaluating text paths on the DiT workers.
+
 ### `runtime`
 
 Configuration for disaggregated execution of the stage, controlling how the stage is deployed and executed.

docs/contributing/README.md

@@ -71,14 +71,7 @@ For additional features and advanced configurations, refer to the:
 ### Testing
 
 vLLM-Omni uses `pytest` to test the codebase.
-
+Please refer to the [test instructions](./ci/test_guide.md) for detailed testing information.
-```bash
-# Run all tests
-pytest tests/
-
-# Run tests for a single test file with detailed output
-pytest -s -v tests/test_omni_llm.py
-```
 
 !!! warning
     Currently, not all unit tests pass when run on CPU platforms. If you don't have access to a GPU platform to run unit tests locally, rely on the continuous integration system to run the tests for now.
@@ -123,6 +116,9 @@ Only specific types of PRs will be reviewed. The PR title is prefixed appropriat
 !!! note
     If the PR spans more than one category, please include all relevant prefixes.
 
+### Local Test
+Please run the L1 and L2 test cases locally first and attach the results before contacting us to add the "ready" label. Please refer to the [test instructions](./ci/test_guide.md) for running the test cases.
+
 ### Code Quality
 
 The PR needs to meet the following code quality standards:

docs/contributing/ci/test_guide.md

@@ -0,0 +1,75 @@
+# Test Guide
+## Setting Up the Test Environment
+### Creating a Container
+vLLM-Omni provides an official Docker image for deployment. These images are built upon vLLM Docker images and are available on [Docker Hub](https://hub.docker.com/r/vllm/vllm-omni/tags). The version of vLLM-Omni indicates which vLLM release it is based on.
+For a local test environment, you can follow the steps below to create a container:
+## Installing Dependencies
+### vLLM & vLLM-Omni
+vLLM-Omni is built based on vLLM. You can follow [install guide](../../getting_started/installation/README.md) to build your local environment.
+
+### Test Case Dependencies
+When running test cases, you may need to install the following dependencies:
+
+```bash
+uv pip install ".[dev]"
+apt-get install -y ffmpeg
+```
+
+## Running Tests
+Our test scripts use the pytest framework. First, please use `git clone https://github.com/vllm-project/vllm-omni.git` to download the vllm-omni source code. Then, in the root directory of vllm-omni, you can run the following commands in your local test environment to execute the corresponding test cases.
+
+=== "L1 level"
+
+    ```bash
+    cd tests
+    pytest -s -v -m "core_model and cpu"
+    ```
+    The latest test command is available in the "Simple Unit Test" step of this [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml).
+
+=== "L2 level"
+
+    ```bash
+    cd tests
+    pytest -s -v -m "core_model and not cpu" --run-level=core_model
+    ```
+    If you only want to run a specific test case, you can use:
+    ```bash
+    pytest -s -v test_xxxx.py --run-level=core_model
+    ```
+    If you only want to run specific test cases on a particular platform, you can use:
+    ```bash
+    pytest -s -v -m "core_model and distributed_cuda and L4"  --run-level=core_model
+    ```
+    The latest test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml).
+
+=== "L3 level & L4 level"
+
+    ```bash
+    cd tests
+    pytest -s -v -m "advanced_model" --run-level=advanced_model
+    ```
+    If you only want to run L3 test case, you can use:
+    ```bash
+    pytest -s -v e2e/ --ignore-glob='*expansion.py' -m "advanced_model" --run-level=advanced_model
+    ```
+    If you only want to run a specific test case, you can use:
+    ```bash
+    pytest -s -v test_xxxx.py --run-level=advanced_model
+    ```
+    If you only want to run specific test cases on a particular platform, you can use:
+    ```bash
+    pytest -s -v -m "core_model and distributed_cuda and L4"  --run-level=core_model
+    ```
+    Note: To run performance tests, use:
+    ```bash
+    pytest -s -v perf/scripts/run_benchmark.py
+    ```
+
+    The latest L3 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-merge.yml).
+
+    The latest L4 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly.yml).
+
+You can find more information about markers in the documentation: [marker doc](./tests_markers.md)
+
+## Adding New Test Cases
+Please refer to the [L5 Layering Specification document](./CI_5levels.md).

docs/contributing/model/adding_omni_model.md

@@ -460,7 +460,7 @@ def forward(self, ...):
 These keys are then accessible in your stage transition function:
 ```python
 # In stage_input_processors/qwen3_omni.py
-thinker_embeddings = output.multimodal_output["0"]  # Access by key
+thinker_prefill_embeddings = output.multimodal_output["0"]  # Access by key
 thinker_hidden_states = output.multimodal_output["24"]
 ```
 
@@ -513,11 +513,11 @@ def thinker2talker(
     for thinker_output in thinker_outputs:
         output = thinker_output.outputs[0]
         # Extract thinker embeddings and hidden states
-        thinker_embeddings = output.multimodal_output["0"].float().clone().detach().cuda()
+        thinker_prefill_embeddings = output.multimodal_output["0"].float().clone().detach().cuda()
         thinker_hidden_states = output.multimodal_output["24"].float().clone().detach().cuda()
 
         info = {
-            "thinker_embeddings": thinker_embeddings,
+            "thinker_prefill_embeddings": thinker_prefill_embeddings,
             "thinker_hidden_states": thinker_hidden_states,
             "thinker_sequences": thinker_output.prompt_token_ids + output.token_ids,
             "thinker_input_ids": thinker_output.prompt_token_ids,

docs/contributing/model/adding_tts_model.md

@@ -0,0 +1,584 @@
+# Adding a TTS Model
+
+This guide walks through adding a new TTS model to vLLM-Omni, using **Qwen3-TTS**
+as a reference. Qwen3-TTS demonstrates the standard two-stage TTS pipeline and the
+key optimizations all TTS models in this repo should follow.
+
+## Table of Contents
+
+1. [Overview](#overview)
+2. [Directory Structure](#directory-structure)
+3. [Step-by-Step Implementation](#step-by-step-implementation)
+4. [Key Components](#key-components)
+5. [Model Registration](#model-registration)
+6. [Stage Configuration](#stage-configuration)
+7. [Stage Input Processors](#stage-input-processors)
+8. [Testing](#testing)
+9. [Summary](#summary)
+
+## Overview
+
+vLLM-Omni supports TTS models as multi-stage pipelines where each stage runs independently
+and can be placed on different devices. Qwen3-TTS has two stages:
+
+| Stage | Name | Input | Output |
+|-------|------|-------|--------|
+| 0 | Code Predictor (AR) | Text tokens | Discrete RVQ codec codes |
+| 1 | Code2Wav (Decoder) | RVQ codec codes | Audio waveform |
+
+Each stage is a separate model class configured independently via YAML. The two stages
+are connected by the `async_chunk` framework, which enables inter-stage streaming for
+low first-packet latency (see [Async Chunk Design](../../design/feature/async_chunk_design.md)).
+
+### Without async_chunk (batch mode)
+
+Stage 0 runs to completion before Stage 1 starts, resulting in long first-packet latency:
+
+```mermaid
+flowchart TB
+    subgraph stage0["Stage 0: AR Code Predictor"]
+        direction LR
+        P[Prefill] --> D1[Decode 1]
+        D1 --> D2[Decode 2]
+        D2 --> Dots1["..."]
+        Dots1 --> DN[Decode N]
+    end
+
+    subgraph stage1["Stage 1: Code2Wav"]
+        direction LR
+        DEC[Decode all codes at once]
+    end
+
+    stage0 -- "all N codes" --> stage1
+    stage1 --> FPL["First Packet Latency = Stage 0 + Stage 1"]
+
+    style stage0 fill:#dae8fc,stroke:#6c8ebf
+    style stage1 fill:#f8d7c8,stroke:#d4856a
+    style FPL fill:#e8f0fe,stroke:#3366CC,stroke-width:2px
+```
+
+### With async_chunk (streaming mode)
+
+Stage 0 sends codec codes to Stage 1 every `chunk_size=25` tokens. Stage 1 begins decoding
+immediately, reducing first-packet latency from the full AR time to just the first chunk:
+
+```mermaid
+flowchart TB
+    subgraph stage0["Stage 0: Code Predictor (AR)"]
+        direction LR
+        P[Prefill] --> D1["Decode 1-25"]
+        D1 --> D2["Decode 26-50"]
+        D2 --> Dots1["..."]
+        Dots1 --> DN["Decode N"]
+    end
+
+    subgraph stage1["Stage 1: Code2Wav"]
+        direction LR
+        C1["Chunk 1\n(25 frames)"] --> C2["Chunk 2\n(context + 25)"]
+        C2 --> Dots2["..."]
+        Dots2 --> CN["Final chunk"]
+    end
+
+    D1 -. "chunk 1 (25 codes)" .-> C1
+    D2 -. "chunk 2 (context + 25)" .-> C2
+    DN -. "final" .-> CN
+
+    stage0 --> FPL["⏱ First Packet Latency = Prefill + 25 decode steps only"]
+
+    style stage0 fill:#dae8fc,stroke:#6c8ebf
+    style stage1 fill:#e8d4f8,stroke:#8a6cad
+    style FPL fill:#e8f0fe,stroke:#3366CC,stroke-width:2px
+```
+
+Key parameters: `chunk_size=25`, `left_context_size=25` (validated defaults from Qwen3-TTS
+and Qwen3-Omni).
+
+## Directory Structure
+
+When adding a new TTS model, create the following structure:
+
+```
+vllm_omni/model_executor/models/
+  your_model_name/
+    __init__.py
+    your_model.py                    # Unified class (stage dispatch)
+    your_model_ar_stage.py           # Stage 0: AR stage
+    your_model_decoder.py            # Stage 1: audio decoder
+
+vllm_omni/model_executor/stage_input_processors/
+  your_model_name.py                 # Stage 0 -> Stage 1 transition
+
+vllm_omni/model_executor/stage_configs/
+  your_model_name.yaml               # Batch mode config
+  your_model_name_async_chunk.yaml   # Streaming mode config
+```
+
+**Qwen3-TTS reference files:**
+
+| File | Purpose |
+|------|---------|
+| `models/qwen3_tts/qwen3_tts.py` | Unified model class |
+| `models/qwen3_tts/qwen3_tts_code_predictor_vllm.py` | Stage 0 - optimized AR |
+| `models/qwen3_tts/qwen3_tts_code2wav.py` | Stage 1 - decoder |
+| `stage_configs/qwen3_tts.yaml` | Stage config (async_chunk enabled) |
+| `stage_configs/qwen3_tts_batch.yaml` | Batch mode config |
+| `stage_input_processors/qwen3_tts.py` | Stage transition processors |
+
+## Step-by-Step Implementation
+
+### Step 1: Implement Stage 0 - AR Stage
+
+Stage 0 is the autoregressive stage that generates intermediate audio representations.
+**It must use vLLM's native decoder layers with fused ops and PagedAttention** for the LLM
+backbone - this is the primary source of speedup over HuggingFace inference.
+
+#### 1.1 Use vLLM Decoder Layers Directly
+
+Build your transformer layers from the corresponding vLLM decoder layer class (e.g.
+`Qwen3DecoderLayer` for Qwen3-based backbones, or the equivalent for LLaMA, Qwen2, etc.).
+Do not wrap the HuggingFace model directly - that bypasses PagedAttention and fused kernels.
+
+```python
+# your_model_ar_stage.py
+
+from vllm.model_executor.models.qwen3 import Qwen3DecoderLayer
+
+class YourTTSARStage(nn.Module):
+
+    def __init__(self, config, vllm_config, prefix):
+        self.layers = nn.ModuleList([
+            Qwen3DecoderLayer(
+                config, vllm_config=vllm_config, prefix=f"{prefix}.layers.{i}"
+            )
+            for i in range(config.num_hidden_layers)
+        ])
+        self.lm_head = ParallelLMHead(config.codec_size, config.hidden_size)
+```
+
+See `qwen3_tts_code_predictor_vllm.py` for the full implementation.
+
+#### 1.2 Forward Pass
+
+Implement `forward()` to return an `OmniOutput` with intermediate data for Stage 1:
+
+```python
+def forward(self, input_ids, positions, intermediate_tensors=None,
+            inputs_embeds=None, **kwargs) -> OmniOutput:
+    hidden_states = self.run_layers(input_ids, positions, intermediate_tensors, inputs_embeds)
+    logits = self.lm_head(hidden_states)
+
+    return OmniOutput(
+        text_hidden_states=hidden_states,
+        multimodal_outputs={
+            "audio_codes": self.extract_codes(logits),
+        },
+    )
+```
+
+The keys in `multimodal_outputs` are what your stage input processor will read to build
+Stage 1 inputs.
+
+#### 1.3 Weight Loading with Fused QKV
+
+When using vLLM's fused `QKVParallelLinear`, pack the HF `q_proj`/`k_proj`/`v_proj` weights
+into `qkv_proj` using `stacked_params_mapping`. See the `load_weights()` method in
+`qwen3_tts_code_predictor_vllm.py` for the standard pattern - it can be reused as-is
+for any Qwen-family backbone.
+
+#### 1.4 Custom Stop Condition (if needed)
+
+Some TTS models use a learned stop head rather than an EOS token. If your model does this,
+implement it inside `sample()`:
+
+```python
+def sample(self, logits, sampling_metadata) -> SamplerOutput | None:
+    output = self.sampler(logits, sampling_metadata)
+    if self._stop_head_fired():
+        output = mark_as_finished(output)
+    return output
+```
+
+### Step 2: Implement Stage 1 - Decoder
+
+Stage 1 decodes Stage 0 output into audio. It runs outside the scheduler (no PagedAttention
+needed). Implement `chunked_decode_streaming()` to support async_chunk streaming:
+
+```python
+# your_model_decoder.py
+
+class YourTTSDecoder(nn.Module):
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        # Initialize your audio decoder (SpeechTokenizer, HiFiGAN, etc.)
+
+    def forward(self, codes: torch.Tensor, **kwargs) -> torch.Tensor:
+        return self.decoder(codes)
+
+    def chunked_decode_streaming(self, codes, chunk_size=25,
+                                  left_context_size=25) -> torch.Tensor:
+        """Decode with a sliding context window for smooth chunk boundaries."""
+        end_index = codes.shape[-1]
+        context_size = 0 if end_index <= chunk_size else left_context_size
+        wav_chunk = self(codes)
+        # Trim left context to avoid duplicate audio
+        return wav_chunk[..., context_size * self.total_upsample:]
+```
+
+### Step 3: Implement the Unified Model Class
+
+The unified class dispatches to the correct stage based on `model_stage` in the config:
+
+```python
+# your_model.py
+
+class YourTTSModelForConditionalGeneration(nn.Module, SupportsPP):
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        self.model_stage = vllm_config.model_config.model_stage
+
+        if self.model_stage == "ar_stage":
+            ar_vllm_config = vllm_config.with_hf_config(
+                vllm_config.model_config.hf_config.ar_config,
+                architectures=["YourTTSARStageForConditionalGeneration"],
+            )
+            self.ar_stage = init_vllm_registered_model(
+                vllm_config=ar_vllm_config,
+                prefix=maybe_prefix(prefix, "ar"),
+                hf_config=ar_vllm_config.model_config.hf_config,
+                architectures=["YourTTSARStageForConditionalGeneration"],
+            )
+            self.model = self.ar_stage
+
+        elif self.model_stage == "decoder":
+            self.decoder = YourTTSDecoder(vllm_config=vllm_config, prefix=prefix)
+            self.model = self.decoder
+```
+
+### Step 4: Create `__init__.py`
+
+```python
+# vllm_omni/model_executor/models/your_model_name/__init__.py
+from .your_model import YourTTSModelForConditionalGeneration
+
+__all__ = ["YourTTSModelForConditionalGeneration"]
+```
+
+## Key Components
+
+### Model Interfaces
+
+Your unified model class should implement the appropriate interfaces:
+
+- **`SupportsPP`**: Required for pipeline parallelism support (all models should implement this)
+- **`SupportsMultiModal`**: Only if your model accepts multimodal inputs (e.g. reference audio for voice cloning)
+
+### Output Format
+
+Use `OmniOutput` so the orchestrator can route intermediate data between stages:
+
+```python
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+
+return OmniOutput(
+    text_hidden_states=hidden_states,
+    multimodal_outputs={
+        "audio_codes": codec_codes,
+    },
+)
+```
+
+### Weight Loading from a Single Checkpoint
+
+If both stages load from one checkpoint, separate them by prefix in the unified class:
+
+```python
+def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+    ar_weights, decoder_weights = [], []
+    for name, tensor in weights:
+        if name.startswith("decoder."):
+            decoder_weights.append((name, tensor))
+        else:
+            ar_weights.append((name, tensor))
+
+    if self.model_stage == "ar_stage":
+        return self.ar_stage.load_weights(ar_weights)
+    elif self.model_stage == "decoder":
+        return self.decoder.load_weights(decoder_weights)
+```
+
+## Model Registration
+
+Register all stage classes in `vllm_omni/model_executor/models/registry.py`:
+
+```python
+_OMNI_MODELS = {
+    # (package_name, module_name, class_name)
+    "YourTTSModelForConditionalGeneration": (
+        "your_model_name", "your_model",
+        "YourTTSModelForConditionalGeneration",
+    ),
+    "YourTTSARStageForConditionalGeneration": (
+        "your_model_name", "your_model_ar_stage",
+        "YourTTSARStageForConditionalGeneration",
+    ),
+    "YourTTSDecoder": (
+        "your_model_name", "your_model_decoder",
+        "YourTTSDecoder",
+    ),
+}
+```
+
+The registry uses lazy loading - model classes are only imported when needed.
+
+## Stage Configuration
+
+Each stage has a `worker_type` that determines how it is scheduled:
+
+- `worker_type: ar` - autoregressive stage, uses `OmniARScheduler` with PagedAttention
+- `worker_type: generation` - non-AR stage (e.g. decoder), uses `OmniGenerationScheduler`
+
+Key configuration fields:
+
+| Field | Description |
+|-------|-------------|
+| `model_stage` | Which stage to initialize (`ar_stage`, `decoder`, etc.) |
+| `model_arch` | Architecture name, must match `registry.py` |
+| `engine_input_source` | List of upstream stage IDs that provide input (e.g. `[0]`) |
+| `engine_output_type` | Output type: `latent` for intermediate, `audio` for final |
+| `custom_process_next_stage_input_func` | Async chunk processor function path (streaming only) |
+| `final_output` | Whether this stage produces the final user-facing output |
+| `final_output_type` | Type of final output (`audio`, `text`, etc.) |
+
+### Batch mode
+
+```yaml
+# stage_configs/your_model_name.yaml
+
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    runtime:
+      devices: "0"
+      max_batch_size: 64
+    engine_args:
+      model_stage: ar_stage
+      model_arch: YourTTSModelForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      engine_output_type: latent
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 2048
+
+  - stage_id: 1
+    stage_type: llm
+    runtime:
+      devices: "0"
+    engine_args:
+      model_stage: decoder
+      model_arch: YourTTSModelForConditionalGeneration
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      engine_output_type: audio
+    engine_input_source: [0]
+    final_output: true
+    final_output_type: audio
+```
+
+### Streaming mode (async_chunk)
+
+Add `async_chunk: true` at the top level and specify `custom_process_next_stage_input_func`
+on Stage 0 to define how intermediate outputs are chunked and forwarded:
+
+```yaml
+# stage_configs/your_model_name_async_chunk.yaml
+
+async_chunk: true
+
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    runtime:
+      devices: "0"
+      max_batch_size: 64
+    engine_args:
+      model_stage: ar_stage
+      model_arch: YourTTSModelForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      engine_output_type: latent
+      custom_process_next_stage_input_func: >
+        vllm_omni.model_executor.stage_input_processors.your_model_name.ar2decoder_async_chunk
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 2048
+
+  - stage_id: 1
+    stage_type: llm
+    runtime:
+      devices: "0"
+    engine_args:
+      model_stage: decoder
+      model_arch: YourTTSModelForConditionalGeneration
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      engine_output_type: audio
+    engine_input_source: [0]
+    final_output: true
+    final_output_type: audio
+```
+
+## Stage Input Processors
+
+Stage input processors convert Stage 0 outputs into Stage 1 inputs. Create yours in
+`vllm_omni/model_executor/stage_input_processors/your_model_name.py`.
+
+See `stage_input_processors/qwen3_tts.py` for the full reference implementation.
+
+### Data structures
+
+Understanding what's available in stage outputs:
+
+- `stage_list[source_id].engine_outputs` - list of `EngineCoreOutput` objects
+- Each `EngineCoreOutput` has `outputs` - list of `RequestOutput` objects
+- Each `RequestOutput` has:
+  - `token_ids` - generated token IDs
+  - `multimodal_output` - dict with keys matching your model's `OmniOutput.multimodal_outputs`
+  - `prompt_token_ids` - original prompt token IDs
+
+### Batch mode (non-streaming)
+
+Collects all Stage 0 outputs and forwards them to Stage 1 in one shot:
+
+```python
+def ar2decoder(
+    stage_list: list[Any],
+    engine_input_source: list[int],
+    prompt: OmniTokensPrompt | TextPrompt | None = None,
+    requires_multimodal_data: bool = False,
+) -> list[OmniTokensPrompt]:
+    source_id = engine_input_source[0]
+    decoder_inputs = []
+
+    for output in stage_list[source_id].engine_outputs:
+        result = output.outputs[0]
+        codes = result.multimodal_output["audio_codes"].cpu()
+        decoder_inputs.append(
+            OmniTokensPrompt(prompt_token_ids=codes.reshape(-1).tolist())
+        )
+
+    return decoder_inputs
+```
+
+### Streaming mode (async_chunk)
+
+Buffers Stage 0 outputs and forwards a chunk to Stage 1 once `chunk_size` frames
+have accumulated. The function signature follows the `OmniChunkTransferAdapter` protocol:
+
+```python
+def ar2decoder_async_chunk(
+    transfer_manager: Any,
+    pooling_output: dict[str, Any] | None,
+    request: Any,
+    is_finished: bool = False,
+) -> dict[str, Any] | None:
+    """Forward chunks of AR output to the decoder stage."""
+    request_id = request.external_req_id
+    finished = bool(is_finished or request.is_finished())
+
+    # Extract and buffer the latest frame
+    if isinstance(pooling_output, dict):
+        frame = extract_frame(pooling_output)
+        if frame is not None:
+            transfer_manager.code_prompt_token_ids[request_id].append(
+                frame.cpu().tolist()
+            )
+    elif not finished:
+        return None
+
+    # Read chunk config from connector
+    chunk_size = 25
+    left_context_size = 25
+
+    length = len(transfer_manager.code_prompt_token_ids[request_id])
+    if length <= 0:
+        if finished:
+            return {"codes": [], "finished": torch.tensor(True, dtype=torch.bool)}
+        return None
+
+    # Wait until a full chunk is ready (or request is finished)
+    chunk_length = length % chunk_size
+    if chunk_length != 0 and not finished:
+        return None
+
+    # Build context window: left_context + chunk
+    context_length = chunk_length if chunk_length != 0 else chunk_size
+    end_index = min(length, left_context_size + context_length)
+    window = transfer_manager.code_prompt_token_ids[request_id][-end_index:]
+
+    return {
+        "codes": torch.tensor(window).transpose(0, 1).reshape(-1).tolist(),
+        "left_context_size": max(0, int(end_index - context_length)),
+        "finished": torch.tensor(finished, dtype=torch.bool),
+    }
+```
+
+Key points:
+- `transfer_manager` is the `OmniChunkTransferAdapter` that owns the chunk lifecycle
+- Each call appends one AR decode step's output; a chunk is emitted every `chunk_size` steps
+- The final (possibly partial) chunk is flushed when `is_finished` is true
+- `left_context_size` frames of overlap are included for smooth audio boundaries
+
+## Testing
+
+For general testing conventions, see [tests_style.md](../ci/tests_style.md).
+
+Recommended test cases for a new TTS model:
+
+1. **Single request** - verify waveform output shape and sample rate
+2. **Batched requests** - verify each request in the batch finishes independently
+3. **async_chunk streaming** - verify audio chunks arrive incrementally and decode correctly
+4. **Speaker conditioning** (if applicable) - verify different speaker inputs produce different outputs
+
+Reference test: `tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py`
+
+## Adding a Model Recipe
+
+After implementing and testing your model, add a model recipe to the
+[vllm-project/recipes](https://github.com/vllm-project/recipes) repository so users can
+get started quickly. See [Adding an Omni-Modality Model](./adding_omni_model.md#adding-a-model-recipe)
+for the expected format.
+
+## Summary
+
+Adding a TTS model to vLLM-Omni involves:
+
+1. **Create model directory** with AR stage, decoder stage, and unified class
+2. **AR stage** - use vLLM's native decoder layers with fused QKV; do not wrap HF directly
+3. **Decoder stage** - thin wrapper around your audio decoder; implement `chunked_decode_streaming()`
+4. **Unified class** - dispatches on `model_stage`; same structure as `Qwen3TTSModelForGeneration`
+5. **Register** all stage classes in `registry.py`
+6. **YAML configs** - provide both batch and `async_chunk` variants
+7. **Stage input processor** - buffer Stage 0 outputs and forward in chunks of 25
+8. **Tests** - cover single request, batching, and async_chunk streaming
+9. **Model recipe** - add to [vllm-project/recipes](https://github.com/vllm-project/recipes)
+
+### Qwen3-TTS Reference Files
+
+| File | Purpose |
+|------|---------|
+| `models/qwen3_tts/qwen3_tts.py` | Unified model class |
+| `models/qwen3_tts/qwen3_tts_code_predictor_vllm.py` | AR stage with vLLM fused ops |
+| `models/qwen3_tts/qwen3_tts_code2wav.py` | Decoder stage with `chunked_decode_streaming()` |
+| `stage_configs/qwen3_tts.yaml` | Stage configuration |
+| `stage_input_processors/qwen3_tts.py` | Stage transition processors |
+
+For more information, see:
+
+- [Architecture Overview](../../design/architecture_overview.md)
+- [Async Chunk Design](../../design/feature/async_chunk_design.md)
+- [Stage Configuration Guide](../../configuration/stage_configs.md)

docs/design/architecture_overview.md

@@ -92,6 +92,13 @@ The framework achieves high performance through several optimization techniques:
     * **Quantization:** Supports various quantization implementations including FP8 and AWQ.
     * **FusedOps:** Allows for custom and third-party integration.
 
+### Classifier-Free Guidance (CFG) Companion Flow
+
+vLLM-Omni natively models Classifier-Free Guidance (CFG) across disaggregated multi-stage setups via a "companion request" paradigm, eliminating redundant textual/multimodal context computation boundaries:
+1. **Prompt Expansion:** In the initial autoregressive (AR) stage, a customized `prompt_expand_func` hook intercepts incoming generation prompts and pairs them directly with negative companion prompts (e.g., a default negative prompt) on the fly, tagging the secondary prompt with a specific internal role (`cfg_text`).
+2. **Synchronized KV Cache Transfer:** The AR stage evaluates both the primary and companion sequence batches concurrently. The `OmniConnector` captures these specific structural dependencies and reliably passes the positive and negative outcome KV caches seamlessly across stage boundaries via shared memory or network protocols.
+3. **KV Cache Collection & Injection:** Upon reaching the downstream Diffusion (DiT) Engine, an assigned `cfg_kv_collect_func` automatically intercepts the mapped companion caches (`cfg_text_past_key_values`). These auxiliary dependencies are natively gathered and seamlessly bound to the primary generation sequence variables, enabling the DiT Engine to cleanly implement cross-attention CFG guidance over accurate conditioning and unconditioning structures in parallel.
+
 ### Flexibility and Usability
 
 vLLM-Omni is designed to be flexible and straightforward for users:

docs/design/feature/async_chunk_design.md

@@ -19,7 +19,7 @@ The `async_chunk` feature enables asynchronous, chunked processing of data acros
 
 For qwen3-omni:
 - **Thinker → Talker**: Per decode step (typically chunk_size=1)
-- **Talker → Code2Wav**: Accumulated to code2wav chunk_size(default=25, currently only support default, will support chunk_size soon) before sending
+- **Talker → Code2Wav**: Accumulated to `codec_chunk_frames` (default=25) before sending. Set `initial_codec_chunk_frames` to emit smaller chunks during the initial phase for reduced TTFA
 - **Code2Wav**: Streaming decode with code2wav chunk_size
 
 With `async_chunk`:

docs/design/feature/sequence_parallel.md

@@ -240,7 +240,8 @@ class TransformerWithRoPE(nn.Module):
 
 ## Approach 2: Intrusive Modification (For Complex Cases)
 
-For models with dynamic sharding logic that cannot be expressed via `_sp_plan`, manually insert shard/gather calls.
+For models with dynamic sharding logic that cannot be expressed via `_sp_plan`, manually insert shard/gather calls. Importantly, when taking this approach, be careful to ensure that you correctly manage the `_sp_shard_depth`; if the sequence parallel shard depth is 0, Ulysses will not be used.
+
 
 **When to use:**
 - Dynamic/conditional sharding logic
@@ -252,15 +253,18 @@ from vllm_omni.diffusion.distributed.sp_sharding import sp_shard, sp_gather
 
 def forward(self, hidden_states, ...):
     if self.parallel_config.sequence_parallel_size > 1:
+        # <Increment the _sp_shard_depth on the fwd context>
         hidden_states = sp_shard(hidden_states, dim=1)
 
     # ... computation ...
 
     if self.parallel_config.sequence_parallel_size > 1:
         output = sp_gather(output, dim=1)
+        # <Decrement the _sp_shard_depth on the fwd context>
 
     return output
 ```
+Note that currently, `sp_shard` / `sp_gather` do *not* automatically manage the `_sp_shard_depth`; you need to be careful to manage it yourself.
 
 ---
 

docs/design/feature/vae_parallel.md

@@ -0,0 +1,289 @@
+# VAE Patch Parallelism
+
+This document describes how to add **VAE Patch Parallelism** support to a diffusion model.
+We use **Qwen-Image** as the reference implementation.
+
+---
+
+## Table of Contents
+
+- [Overview](#overview)
+- [Step-by-Step Implementation](#step-by-step-implementation)
+- [Testing](#testing)
+- [Reference Implementations](#reference-implementations)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+### What is Vae Patch parallel?
+
+**VAE Patch Parallelism** is a decoding acceleration technique. Instead of decoding the entire latent tensor at once, the latent tensor is:
+
++ Split into multiple spatial tiles
+
++ Distributed across multiple ranks
+
++ Decoded in parallel
+
++ Merged to reconstruct the final output
+
+This approach:
+
++ Distributes computation across multiple devices
+
++ Reduces peak memory usage per device
+
++ Accelerates decoding latency
+
+### Architecture
+We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE decoding.
+
+The executor is model-agnostic and accepts three function parameters:
+
++ split – Partition the latent into tiles
+
++ exec – Decode a single tile
+
++ merge – Combine decoded tiles into the final output
+
+#### Execution Flow
+
++ Call split(z) to generate a list of TileTask and a GridSpec
+
++ Dispatch tasks across ranks using workload-based balancing
+
++ Each rank executes exec(task) on its assigned tiles
+
++ Gather decoded tile results to rank 0
+
++ Rank 0 performs merge(...)
+
++ (Optional) Broadcast final result to all ranks
+
+This design separates:
+
++ Distributed execution logic
+
++ Model-specific tiling and merging logic
+
+#### Why split / exec / merge is necessary?
+
+The latent tensor cannot be arbitrarily partitioned.
+
+During decoding:
+
++ Each output pixel may depend on neighboring pixels
+
++ The receptive field is model-dependent
+
+Therefore:
+
++ Tiles must include overlap
+
++ Merge must perform blending to avoid seams
+
+## Step-by-Step Implementation
+
+### Step 1: Implement DistributedAutoencoderKLQwenImage
+`QwenImagePipeline` use `AutoencoderKLQwenImage` for vae, so implement a distributed version:
+
+
+```
+class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMixin):
+    @classmethod
+    def from_pretrained(cls, *args: Any, **kwargs: Any):
+        model = super().from_pretrained(*args, **kwargs)
+        model.init_distributed()
+        return model
+```
+**Key points**:
++ Inherit both AutoencoderKLQwenImage and DistributedVaeMixin
++ Call init_distributed() after loading weights
+
+### Step 2: Implement split/exec/merge
+Reuse `AutoencoderKLQwenImage.tiled_decode` logic and divide it into three stages. And we need return tiles with `GridSpec` and `TileTask`:
+```
+class GridSpec:
+    split_dims: tuple[int, ...]  # Tensor dimensions being split (e.g., (2, 3) for (B, C, H, W))
+    grid_shape: tuple[int, ...]  # Tile grid layout (num_rows, num_cols)
+    tile_spec: dict = field(default_factory=dict) # Metadata required for merging
+    output_dtype: torch.dtype | None = None # Final output dtype
+```
+```
+class TileTask:
+    tile_id: int # task id
+    grid_coord: tuple[int, ...]  # Tile position in grid
+    tensor: torch.Tensor | list[torch.Tensor]  # The tile tensor
+    workload: int | float = 1 # Used for load balancing (e.g., tile area)
+```
+And tiled base split/exec/merge as follow:
+```
+def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
+    # mostly copy from AutoencoderKL
+    _, _, num_frames, height, width = z.shape
+    sample_height = height * self.spatial_compression_ratio
+    sample_width = width * self.spatial_compression_ratio
+
+    tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+    tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+    tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
+    tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+
+    blend_height = self.tile_sample_min_height - self.tile_sample_stride_height
+    blend_width = self.tile_sample_min_width - self.tile_sample_stride_width
+
+    # Split z into overlapping tiles and decode them separately.
+    # The tiles have an overlap to avoid seams between tiles.
+    tiletask_list = []
+    for i in range(0, height, tile_latent_stride_height):
+        for j in range(0, width, tile_latent_stride_width):
+            time_list = []
+            for k in range(num_frames):
+                self._conv_idx = [0]
+                tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width]
+                time_list.append(tile)
+            tiletask_list.append(
+                TileTask(
+                    len(tiletask_list),
+                    (i // tile_latent_stride_height, j // tile_latent_stride_width),
+                    time_list,
+                    workload=time_list[0].shape[3] * time_list[0].shape[4],
+                )
+            )
+    tile_spec = {
+        "sample_height": sample_height,
+        "sample_width": sample_width,
+        "blend_height": blend_height,
+        "blend_width": blend_width,
+    }
+    grid_spec = GridSpec(
+        split_dims=(3, 4),
+        grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
+        tile_spec=tile_spec,
+        output_dtype=self.dtype,
+    )
+    return tiletask_list, grid_spec
+
+def tile_exec(self, task: TileTask) -> torch.Tensor:
+    """Decode a single latent tile into RGB space."""
+    self.clear_cache()
+    time = []
+    for k in range(len(task.tensor)):
+        self._conv_idx = [0]
+        tile = self.post_quant_conv(task.tensor[k])
+        decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx)
+        time.append(decoded)
+    result = torch.cat(time, dim=2)
+    return result
+
+def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor:
+    """Merge decoded tiles into a full image."""
+    grid_h, grid_w = grid_spec.grid_shape
+    result_rows = []
+    self.clear_cache()
+
+    result_rows = []
+    for i in range(grid_h):
+        result_row = []
+        for j in range(grid_w):
+            tile = coord_tensor_map[(i, j)]
+            if i > 0:
+                tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
+            if j > 0:
+                tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
+            result_row.append(tile[:, :, :, : self.tile_sample_stride_height, : self.tile_sample_stride_width])
+        result_rows.append(torch.cat(result_row, dim=-1))
+    dec = torch.cat(result_rows, dim=3)[
+        :, :, :, : grid_spec.tile_spec["sample_height"], : grid_spec.tile_spec["sample_width"]
+    ]
+    return dec
+```
+
+### Step 3: Override tiled_decode
+We need to override tiled_decode, the main logic is:
++ check distributed is enabled
++ select split/exec/merge
++ Invoke self.distributed_decoder.execute to decode
+```
+def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
+    if not self.is_distributed_enabled():
+        return super().tiled_decode(z, return_dict=return_dict)
+
+    logger.info("Decode run with distributed executor")
+    result = self.distributed_decoder.execute(
+        z,
+        DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
+        broadcast_result=True,
+    )
+    if not return_dict:
+        return (result,)
+
+    return DecoderOutput(sample=result)
+```
+`broadcast_result` is set to True or False depending on the model; when enabled, the result will be used even on ranks other than 0.
+
+### Step 4: Modify Pipeline
+Change vae model from AutoencoderKLQwenImage to DistributedAutoencoderKLQwenImage
+```
+class YourModelPipeline(nn.Module):
+    def __init__(
+        self,
+        *,
+        od_config: OmniDiffusionConfig,
+        prefix: str = "",
+    ):
+        super().__init__()
+        ...
+-       self.vae = AutoencoderKL.from_pretrained(
+-           model, subfolder="vae", local_files_only=local_files_only).to(self.device)
++       self.vae = DistributedAutoencoderKL.from_pretrained(
++           model, subfolder="vae", local_files_only=local_files_only
++       ).to(self.device)
+```
+
+## Testing
+Verify numerical consistency between:
++ vae_patch_parallel_size = 1
+
++ vae_patch_parallel_size = N
+
+Example:
+torch.allclose(output_1, output_n, atol=1e-5)
+
+Testing requirements:
++ Fix random seed
++ Use identical tiling strategy
+
+```python
+m = Omni(
+        model=model_name,
+        vae_use_tiling=True,
+        parallel_config=DiffusionParallelConfig(
+            tensor_parallel_size=2,
+            vae_patch_parallel_size=1, # or 2
+        ),
+    )
+```
+When vae_patch_parallel_size is larger than the DiT world size, it will automatically fall back to using the DiT world size instead.
+
+## Reference Implementations
+
+Complete examples in the codebase:
+
+| Model | Path | Notes |
+|-------|------|-------|
+| **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | Distributed AutoencoderKL |
+| **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | Distributed AutoencoderKLWan |
+| **Qwen-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py` | Distributed AutoencoderKLQwenImage |
+
+---
+
+## Summary
+
+Adding Vae Patch Parallel support to diffusion model:
+
+1. **Implement Distributed Vae** - mainly copy from `diffusers` tiled_decode, and refactor into split/exec/merge
+2. **Change vae model in pipeline to Distributed Vae**
+3. **Test** - Verify with `tensor_parallel_size=N` quality

docs/features/custom_pipeline.md

@@ -115,7 +115,7 @@ python image_edit.py \
     --image cherry_blossom.jpg \
     --prompt "Let this mascot dance under the moon, surrounded by floating stars" \
     --output output_image_edit.png \
-    --num_inference_steps 10
+    --num-inference-steps 10
 ```
 
 ## Advanced Usage

docs/getting_started/installation/gpu.md

@@ -30,7 +30,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
 
 ### Pre-built wheels
 
-Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
+Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
 
 === "NVIDIA CUDA"
 

docs/getting_started/installation/gpu/cuda.inc.md

@@ -34,13 +34,13 @@ uv pip install vllm-omni
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.16.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.17.0 release version of the library
 
 ```bash
-uv pip install vllm==0.16.0 --torch-backend=auto
+uv pip install vllm==0.17.0 --torch-backend=auto
 ```
 
-The release 0.14.0 of vLLM is based on PyTorch 2.9.0 which requires CUDA 12.9 environment.
+The release 0.17.0 of vLLM is based on PyTorch 2.10.0 which requires CUDA 12.9 environment.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -56,11 +56,11 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.16.0
+git checkout v0.17.0
 ```
 Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
 ```bash
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.16.0/vllm-0.16.0-cp38-abi3-manylinux_2_31_x86_64.whl
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.17.0/vllm-0.17.0-cp38-abi3-manylinux_2_31_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -91,7 +91,7 @@ docker run --runtime nvidia --gpus 2 \
     --env "HF_TOKEN=$HF_TOKEN" \
     -p 8091:8091 \
     --ipc=host \
-    vllm/vllm-omni:v0.14.0 \
+    vllm/vllm-omni:v0.16.0 \
     --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
 ```
 

docs/getting_started/installation/gpu/rocm.inc.md

@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.16.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.16.0/rocm700
+uv pip install vllm==0.17.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.17.0/rocm700
 ```
 
 #### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm sox
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.16.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.17.0 release version of the library
 
 ```bash
-uv pip install vllm==0.16.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.16.0/rocm700
+uv pip install vllm==0.17.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.17.0/rocm700
 ```
 
-The release 0.16.0 of vLLM requires ROCm 7.0 environment.
+The release 0.17.0 of vLLM requires ROCm 7.0 environment.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.16.0
+git checkout v0.17.0
 python3 -m pip install -r requirements/rocm.txt
 python3 setup.py develop
 ```
@@ -130,7 +130,7 @@ docker run --rm \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   -p 8091:8091 \
-  vllm/vllm-omni-rocm:v0.14.0 \
+  vllm/vllm-omni-rocm:v0.16.0 \
   --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
@@ -149,7 +149,7 @@ docker run --rm -it \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   --entrypoint bash \
-  vllm/vllm-omni-rocm:v0.14.0
+  vllm/vllm-omni-rocm:v0.16.0
 ```
 
 # --8<-- [end:pre-built-images]

docs/getting_started/quickstart.md

@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
 source .venv/bin/activate
 
 # On CUDA
-uv pip install vllm==0.16.0 --torch-backend=auto
+uv pip install vllm==0.17.0 --torch-backend=auto
 
 # On ROCm
-uv pip install vllm==0.16.0 --extra-index-url https://wheels.vllm.ai/rocm/0.16.0/rocm700
+uv pip install vllm==0.17.0 --extra-index-url https://wheels.vllm.ai/rocm/0.17.0/rocm700
 
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni

docs/models/supported_models.md

@@ -30,10 +30,15 @@ th {
 |`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |
 | `WanPipeline` | Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` |
 | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` |
+| `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` |
+| `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` |
+| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` |
 | `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` |
 |`LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` |
 |`LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` |
 |`StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` |
+|`CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` |
+|`MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` |
 |`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` |
 |`FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` |
 |`OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` |

docs/serving/audio_generate_api.md

@@ -0,0 +1,338 @@
+# Audio Generate API
+
+vLLM-Omni provides an API for text-to-audio generation using diffusion-based models such as Stable Audio.
+
+Unlike the [Speech API](speech_api.md) which targets text-to-speech synthesis, the Audio Generate API is designed for general-purpose audio generation from text descriptions (sound effects, music, ambient soundscapes, etc.).
+
+Each server instance runs a single model (specified at startup via `vllm-omni serve <model> --omni`).
+
+## Quick Start
+
+### Start the Server
+
+```bash
+vllm-omni serve stabilityai/stable-audio-open-1.0 \
+    --host 0.0.0.0 \
+    --port 8000 \
+    --gpu-memory-utilization 0.9 \
+    --trust-remote-code \
+    --enforce-eager \
+    --omni
+```
+
+### Generate Audio
+
+**Using curl:**
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound of a cat purring",
+        "audio_length": 10.0
+    }' --output cat.wav
+```
+
+**Using Python:**
+
+```python
+import httpx
+
+response = httpx.post(
+    "http://localhost:8000/v1/audio/generate",
+    json={
+        "input": "The sound of a cat purring",
+        "audio_length": 10.0,
+    },
+    timeout=300.0,
+)
+
+with open("cat.wav", "wb") as f:
+    f.write(response.content)
+```
+
+## API Reference
+
+### Endpoint
+
+```
+POST /v1/audio/generate
+Content-Type: application/json
+```
+
+### Request Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `input` | string | **required** | Text prompt describing the audio to generate |
+| `model` | string | server's model | Model to use (optional, should match server if specified) |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `speed` | float | 1.0 | Playback speed (0.25 - 4.0) |
+
+#### Diffusion Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `audio_length` | float | null | Audio duration in seconds (default value is the max ~47s for `stable-audio-open-1.0`) |
+| `audio_start` | float | 0.0 | Audio start time in seconds |
+| `negative_prompt` | string | null | Text describing what to avoid in generation |
+| `guidance_scale` | float | model default | Classifier-free guidance scale (higher = more adherence to prompt) |
+| `num_inference_steps` | int | model default | Number of denoising steps (higher = better quality, slower) |
+| `seed` | int | null | Random seed for reproducible generation |
+
+### Response Format
+
+Returns binary audio data with the appropriate `Content-Type` header:
+
+| `response_format` | Content-Type |
+|--------------------|--------------|
+| `wav` | `audio/wav` |
+| `mp3` | `audio/mpeg` |
+| `flac` | `audio/flac` |
+| `pcm` | `audio/pcm` |
+| `aac` | `audio/aac` |
+| `opus` | `audio/opus` |
+
+## Examples
+
+### Basic Generation
+
+Generate audio with only a text prompt (model defaults for all other parameters):
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound of ocean waves crashing on a beach"
+    }' --output ocean.wav
+```
+
+### Custom Duration
+
+Specify an explicit audio length in seconds:
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "A dog barking",
+        "audio_length": 5.0
+    }' --output dog_5s.wav
+```
+
+### High Quality with Negative Prompt
+
+Use a negative prompt to steer generation away from undesired characteristics, and increase inference steps for higher quality:
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "A piano playing a gentle melody",
+        "audio_length": 10.0,
+        "negative_prompt": "Low quality, distorted, noisy",
+        "guidance_scale": 8.0,
+        "num_inference_steps": 150
+    }' --output piano_hq.wav
+```
+
+### Reproducible Generation
+
+Set a `seed` to get deterministic results across runs:
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Thunder and rain sounds",
+        "audio_length": 15.0,
+        "seed": 42
+    }' --output thunder.wav
+```
+
+### Full Control
+
+Combine all parameters for precise control over generation:
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Thunder and rain sounds",
+        "audio_length": 15.0,
+        "negative_prompt": "Low quality",
+        "guidance_scale": 7.0,
+        "num_inference_steps": 100,
+        "seed": 42
+    }' --output thunder_rain.wav
+```
+
+### Quick Generation (Fewer Steps)
+
+For faster generation with slightly lower quality:
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Birds chirping in a forest",
+        "audio_length": 8.0,
+        "num_inference_steps": 50
+    }' --output birds_quick.wav
+```
+
+### Python Client
+
+```python
+import httpx
+
+response = httpx.post(
+    "http://localhost:8000/v1/audio/generate",
+    json={
+        "input": "Thunder and rain",
+        "audio_length": 15.0,
+        "negative_prompt": "Low quality",
+        "guidance_scale": 7.0,
+        "num_inference_steps": 100,
+        "seed": 42,
+        "response_format": "wav",
+    },
+    timeout=300.0,
+)
+
+with open("thunder.wav", "wb") as f:
+    f.write(response.content)
+```
+
+## Parameter Tuning Guide
+
+### `guidance_scale`
+
+Controls how closely the generated audio follows the text prompt.
+
+| Range | Behaviour |
+|-------|-----------|
+| 3 - 5 | More creative / varied output |
+| 7 (default) | Balanced adherence |
+| 10+ | Strict adherence to the prompt |
+
+### `num_inference_steps`
+
+Controls the number of denoising steps in the diffusion process.
+
+| Steps | Quality | Speed | Use Case |
+|-------|---------|-------|----------|
+| 50 | Good | Fast | Quick previews |
+| 100 | Very Good | Medium | General purpose |
+| 150+ | Excellent | Slow | Final / critical audio |
+
+### `audio_length`
+
+Duration of the generated audio clip. For `stable-audio-open-1.0`, the maximum is approximately 47 seconds. If omitted, the model uses its own default length.
+
+### `negative_prompt`
+
+Describes characteristics to avoid. Common negative prompts include:
+
+- `"Low quality, distorted, noisy"`
+- `"Silence, static"`
+- `"Music"` (when generating sound effects only)
+
+## Supported Models
+
+| Model | Description |
+|-------|-------------|
+| `stabilityai/stable-audio-open-1.0` | Open-source audio generation model, up to ~47 seconds, 44.1 kHz stereo |
+
+## Error Responses
+
+### 400 Bad Request
+
+Invalid or missing parameters:
+
+```json
+{
+    "error": {
+        "message": "Audio generation model did not produce audio output.",
+        "type": "BadRequestError",
+        "param": null,
+        "code": 400
+    }
+}
+```
+
+### 404 Not Found
+
+Model mismatch:
+
+```json
+{
+    "error": {
+        "message": "The model `xxx` does not exist.",
+        "type": "NotFoundError",
+        "param": "model",
+        "code": 404
+    }
+}
+```
+
+### 422 Unprocessable Entity
+
+Pydantic validation failure (e.g. invalid `response_format`, `speed` out of range):
+
+```json
+{
+    "detail": [
+        {
+            "type": "literal_error",
+            "msg": "Input should be 'wav', 'pcm', 'flac', 'mp3', 'aac' or 'opus'",
+            ...
+        }
+    ]
+}
+```
+
+## Troubleshooting
+
+### "Audio generation model did not produce audio output"
+
+The model finished but returned no audio data. Verify the server started successfully and the model loaded without errors.
+
+### Server Not Responding
+
+```bash
+# Check if the server is healthy
+curl http://localhost:8000/health
+```
+
+### Audio Quality Issues
+
+- Increase `num_inference_steps` (e.g. 150).
+- Add a negative prompt: `"Low quality, distorted, noisy"`.
+- Increase `guidance_scale` for stronger prompt adherence.
+
+### Generation Timeout
+
+- Reduce `num_inference_steps`.
+- Reduce `audio_length`.
+- Check GPU memory with `nvidia-smi`.
+
+### Out of Memory
+
+- Lower `--gpu-memory-utilization` (e.g. 0.8).
+- Reduce `audio_length`.
+
+## Development
+
+Enable debug logging:
+
+```bash
+vllm-omni serve stabilityai/stable-audio-open-1.0 \
+    --host 0.0.0.0 \
+    --port 8000 \
+    --gpu-memory-utilization 0.9 \
+    --trust-remote-code \
+    --enforce-eager \
+    --omni \
+    --uvicorn-log-level debug
+```

docs/serving/speech_api.md

@@ -96,6 +96,7 @@ Content-Type: application/json
 | `language` | string | "Auto" | Language (see supported languages below) |
 | `instructions` | string | "" | Voice style/emotion instructions |
 | `max_new_tokens` | integer | 2048 | Maximum tokens to generate |
+| `initial_codec_chunk_frames` | integer | null | Initial chunk size for reduced TTFA (overrides stage config) |
 
 **Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian
 
@@ -103,7 +104,7 @@ Content-Type: application/json
 
 | Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
-| `ref_audio` | string | null | Reference audio (URL or base64 data URL) |
+| `ref_audio` | string | null | Reference audio (HTTP URL, base64 data URL, or `file://` URI with `--allowed-local-media-path`) |
 | `ref_text` | string | null | Transcript of reference audio |
 | `x_vector_only_mode` | bool | null | Use speaker embedding only (no ICL) |
 
@@ -124,6 +125,81 @@ Lists available voices for the loaded model.
     "voices": ["aiden", "dylan", "eric", "ono_anna", "ryan", "serena", "sohee", "uncle_fu", "vivian"]
 }
 ```
+```
+POST /v1/audio/voices
+Content-Type: multipart/form-data
+```
+
+Upload a new voice sample for voice cloning in Base task TTS requests.
+
+**Form Parameters:**
+
+| Parameter | Type | Required | Description |
+|-----------|------|----------|-------------|
+| `audio_sample` | file | Yes | Audio file (max 10MB, supported formats: wav, mp3, flac, ogg, aac, webm, mp4) |
+| `consent` | string | Yes | Consent recording ID |
+| `name` | string | Yes | Name for the new voice |
+
+**Response Example:**
+
+```json
+{
+  "success": true,
+  "voice": {
+    "name": "custom_voice_1",
+    "consent": "user_consent_id",
+    "created_at": 1738660000,
+    "mime_type": "audio/wav",
+    "file_size": 1024000
+  }
+}
+```
+
+**Usage Example:**
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/voices \
+  -F "audio_sample=@/path/to/voice_sample.wav" \
+  -F "consent=user_consent_id" \
+  -F "name=custom_voice_1"
+```
+
+
+```bash
+DELETE /v1/audio/voices/{name}
+```
+
+Delete an uploaded voice sample.
+
+**Path Parameters:**
+
+| Parameter | Type | Required | Description |
+|-----------|------|----------|-------------|
+| `name` | string | Yes | Name of the voice to delete |
+
+**Response Example:**
+
+```json
+{
+  "success": true,
+  "message": "Voice 'custom_voice_1' deleted successfully"
+}
+```
+
+**Error Response (404 Not Found):**
+
+```json
+{
+  "success": false,
+  "error": "Voice 'unknown_voice' not found"
+}
+```
+
+**Usage Example:**
+
+```bash
+curl -X DELETE http://localhost:8091/v1/audio/voices/custom_voice_1
+```
 
 ## Examples
 
@@ -184,6 +260,25 @@ curl -X POST http://localhost:8091/v1/audio/speech \
     }' --output cloned.wav
 ```
 
+upload voice
+```bash
+curl -X POST http://localhost:8091/v1/audio/voices \
+  -F "audio_sample=@/path/to/voice_sample.wav" \
+  -F "consent=user_consent_id" \
+  -F "name=custom_voice_1"
+```
+
+use upload voice
+```bash
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Hello, this is a cloned voice",
+        "task_type": "Base",
+        "voice": "custom_voice_1"
+    }' --output cloned.wav
+```
+
 ## Supported Models
 
 | Model | Task Type | Description |

docs/user_guide/diffusion/parallelism_acceleration.md

@@ -27,7 +27,7 @@ The following table shows which models are currently supported by parallelism me
 | **LongCat-Image**        | `meituan-longcat/LongCat-Image`      |     ✅      |    ✅    |      ❌       |        ✅        |         ❌          |
 | **LongCat-Image-Edit**   | `meituan-longcat/LongCat-Image-Edit` |     ✅      |    ✅    |      ❌       |        ✅        |         ❌          |
 | **Ovis-Image**           | `OvisAI/Ovis-Image`                  |     ❌      |    ❌    |      ❌       |        ❌        |         ❌          |
-| **Qwen-Image**           | `Qwen/Qwen-Image`                    |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |
+| **Qwen-Image**           | `Qwen/Qwen-Image`                    |     ✅      |    ✅    |      ✅       |        ✅        |         ✅          |
 | **Qwen-Image-Edit**      | `Qwen/Qwen-Image-Edit`               |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |
 | **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509`          |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |
 | **Qwen-Image-Layered**   | `Qwen/Qwen-Image-Layered`            |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |
@@ -51,9 +51,10 @@ The following table shows which models are currently supported by parallelism me
 
 ### VideoGen
 
-| Model | Model Identifier | Ulysses-SP | Ring-Attention | Tensor-Parallel | HSDP |
+| Model | Model Identifier | Ulysses-SP | Ring-Attention | Tensor-Parallel | HSDP | VAE-Patch-Parallel |
-|-------|------------------|:----------:|:--------------:|:---------------:|:----:|
+|-------|------------------|:----------:|:--------------:|:---------------:|:----:| :----:|
-| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ | ✅ |
+| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ | ✅ | ✅ |
+| **LTX-2** | `Lightricks/LTX-2` | ✅ | ✅ | ✅ | ❌ | ❌ |
 
 ### Tensor Parallelism
 
@@ -85,7 +86,7 @@ outputs = omni.generate(
 VAE patch parallelism distributes the VAE decode workload across multiple ranks by splitting the latent spatially. It is configured via `DiffusionParallelConfig.vae_patch_parallel_size` and can be combined with other parallelism methods (e.g., TP).
 
 !!! note "Enablement and feature gate"
-    - VAE patch parallelism is currently **enabled only for validated pipelines** (currently: `Tongyi-MAI/Z-Image-Turbo`, `stabilityai/stable-diffusion-3.5`).
+    - VAE patch parallelism is currently **enabled only for validated pipelines** (check [ImageGen](#imagegen) and [VideoGen](#videogen) for more information).
     - If `vae_patch_parallel_size > 1` is set for a validated pipeline, vLLM-Omni will automatically enable `vae_use_tiling` as a safety gate. (We use `vae_use_tiling` because it indicates the VAE supports diffusers tiling parameters like `tile_latent_min_size` and `tile_overlap_factor`.)
 
 #### Offline Inference

docs/user_guide/diffusion/quantization/fp8.md

@@ -62,6 +62,7 @@ The available `ignored_layers` names depend on the model architecture (e.g., `to
 |-------|-----------|---------------|------------------|
 | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | All layers | None |
 | Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Skip sensitive layers | `img_mlp` |
+| Flux | `black-forest-labs/FLUX.1-dev` | All layers | None |
 
 ## Combining with Other Features
 

docs/user_guide/diffusion/quantization/gguf.md

@@ -136,9 +136,9 @@ python examples/offline_inference/text_to_image/text_to_image.py \
   --height 768 \
   --width 1360 \
   --seed 42 \
-  --cfg_scale 4.0 \
+  --cfg-scale 4.0 \
-  --num_images_per_prompt 1 \
+  --num-images-per-prompt 1 \
-  --num_inference_steps 4 \
+  --num-inference-steps 4 \
   --output outputs/flux2_klein_4b.png
 ```
 
@@ -152,9 +152,9 @@ python examples/offline_inference/text_to_image/text_to_image.py \
   --height 768 \
   --width 1360 \
   --seed 42 \
-  --cfg_scale 4.0 \
+  --cfg-scale 4.0 \
-  --num_images_per_prompt 1 \
+  --num-images-per-prompt 1 \
-  --num_inference_steps 4 \
+  --num-inference-steps 4 \
   --output outputs/flux2_klein_4b_gguf.png
 ```
 

docs/user_guide/diffusion_acceleration.md

@@ -58,8 +58,8 @@ The following table shows which models are currently supported by each accelerat
 | **LongCat-Image** | `meituan-longcat/LongCat-Image` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
 | **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
 | **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
-| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
+| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
-| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
+| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
 | **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
 | **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
 | **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
@@ -67,13 +67,15 @@ The following table shows which models are currently supported by each accelerat
 | **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
 | **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
 | **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
-| **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ |
+| **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ |
+| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
 
 ### VideoGen
 
-| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel | HSDP |
+| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel | HSDP | VAE-Patch-Parallel |
-|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:----:|
+|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:----:|:----:|
-| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ |
+| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
+| **LTX-2** | `Lightricks/LTX-2` | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 
 ### Quantization
 

docs/user_guide/examples/offline_inference/qwen3_tts.md

@@ -11,7 +11,7 @@ Qwen3 TTS provides multiple task variants for speech generation:
 
 - **CustomVoice**: Generate speech with a known speaker identity (speaker ID) and optional instruction.
 - **VoiceDesign**: Generate speech from text plus a descriptive instruction that designs a new voice.
-- **Base**: Voice cloning using a reference audio + reference transcript, with optional mode selection.
+- **Base**: Voice cloning using a reference audio + reference transcript, with optional mode selection. The `ref_audio` field accepts a local file path, HTTP URL, or base64 data URL.
 
 ## Setup
 Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
@@ -98,7 +98,8 @@ Add `--streaming` to stream audio chunks progressively via `AsyncOmni` (requires
 python end2end.py --query-type CustomVoice --streaming --output-dir /tmp/out_stream
 ```
 
-Each 25-frame Code2Wav chunk is logged as it arrives. The final WAV file is written once generation
+Each Code2Wav chunk is logged as it arrives (default 25 frames; configurable via `codec_chunk_frames`
+and `initial_codec_chunk_frames` in the stage config). The final WAV file is written once generation
 completes. This demonstrates that audio data is available progressively rather than only at the end.
 
 > **Note:** Streaming uses `AsyncOmni` internally. The non-streaming path (`Omni`) is unchanged.

docs/user_guide/examples/online_serving/qwen3_tts.md

@@ -112,7 +112,7 @@ The Python client supports the following command-line arguments:
 - `--voice`: Speaker/voice name (default: `vivian`). Options: `vivian`, `ryan`, `aiden`, etc.
 - `--language`: Language. Options: `Auto`, `Chinese`, `English`, `Japanese`, `Korean`, `German`, `French`, `Russian`, `Portuguese`, `Spanish`, `Italian`
 - `--instructions`: Voice style/emotion instructions
-- `--ref-audio`: Reference audio file path or URL for voice cloning (Base task)
+- `--ref-audio`: Reference audio file path or URL for voice cloning (Base task). Local paths are automatically base64-encoded by the client before sending to the server.
 - `--ref-text`: Reference audio transcript for voice cloning (Base task)
 - `--response-format`: Audio output format (default: `wav`). Options: `wav`, `mp3`, `flac`, `pcm`, `aac`, `opus`
 - `--output` (or `-o`): Output audio file path (default: `tts_output.wav`)
@@ -220,7 +220,7 @@ Lists available voices for the loaded model:
     "task_type": "CustomVoice",
     "language": "Auto",
     "instructions": "Optional style instructions",
-    "ref_audio": "URL or base64 for voice cloning",
+    "ref_audio": "HTTP URL, base64 data URL, or file:// URI for voice cloning",
     "ref_text": "Reference audio transcript",
     "x_vector_only_mode": false,
     "max_new_tokens": 2048
@@ -261,14 +261,14 @@ Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/w
 
 | Parameter            | Type   | Required | Description                                  |
 | -------------------- | ------ | -------- | -------------------------------------------- |
-| `ref_audio`          | string | **Yes**  | Reference audio (URL or base64 data URL)     |
+| `ref_audio`          | string | **Yes**  | Reference audio (HTTP URL, base64 data URL, or `file://` URI with `--allowed-local-media-path`) |
 | `ref_text`           | string | No       | Transcript of reference audio (for ICL mode) |
 | `x_vector_only_mode` | bool   | No       | Use speaker embedding only (no ICL)          |
 
 ## Streaming
 
 Set `stream=true` with `response_format="pcm"` to receive raw PCM audio chunks as they are decoded
-(one chunk per 25-frame Code2Wav window):
+(one chunk per Code2Wav window, default 25 frames; configurable in the stage config):
 
 ```bash
 curl -X POST http://localhost:8091/v1/audio/speech \

docs/user_guide/examples/online_serving/text_to_audio.md

@@ -0,0 +1,193 @@
+# Text-To-Audio
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/stable_audio>.
+
+This example demonstrates how to deploy Stable Audio models for online text-to-audio generation using vLLM-Omni.
+
+## Supported Models
+
+| Model | Description |
+|-------|-------------|
+| `stabilityai/stable-audio-open-1.0` | Open-source audio generation, up to ~47 seconds, 44.1 kHz stereo |
+
+## Start Server
+
+### Basic Start
+
+```bash
+vllm-omni serve stabilityai/stable-audio-open-1.0 \
+    --host 0.0.0.0 \
+    --port 8000 \
+    --gpu-memory-utilization 0.9 \
+    --trust-remote-code \
+    --enforce-eager \
+    --omni
+```
+
+## API Calls
+
+### Method 1: Using curl
+
+```bash
+# Run all curl examples
+bash curl_examples.sh
+
+# Or execute directly
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound of a cat purring",
+        "audio_length": 10.0
+    }' --output cat.wav
+```
+
+### Method 2: Using Python Client
+
+```bash
+cd examples/online_serving/stable_audio
+
+# Simple generation
+python stable_audio_client.py \
+    --text "The sound of a cat purring"
+
+# With custom duration
+python stable_audio_client.py \
+    --text "A dog barking" \
+    --audio_length 5.0
+
+# With all parameters
+python stable_audio_client.py \
+    --text "Thunder and rain" \
+    --audio_length 15.0 \
+    --negative_prompt "Low quality" \
+    --guidance_scale 7.0 \
+    --num_inference_steps 100 \
+    --seed 42 \
+    --output thunder.wav
+```
+
+The Python client supports the following command-line arguments:
+
+- `--api_url`: API endpoint URL (default: `http://localhost:8000/v1/audio/generate`)
+- `--text`: Text prompt for audio generation (default: `"The sound of a cat purring"`)
+- `--audio_length`: Audio length in seconds (default: `10.0`, max ~47s for `stable-audio-open-1.0`)
+- `--audio_start`: Audio start time in seconds (default: `0.0`)
+- `--negative_prompt`: Negative prompt for classifier-free guidance (default: `"Low quality"`)
+- `--guidance_scale`: Guidance scale for diffusion (default: `7.0`)
+- `--num_inference_steps`: Number of inference steps (default: `100`)
+- `--seed`: Random seed for reproducibility (default: `None`)
+- `--response_format`: Audio output format (default: `wav`). Options: `wav`, `mp3`, `flac`, `pcm`
+- `--output`: Output file path (default: `stable_audio_output.wav`)
+
+### Method 3: Using Python httpx
+
+```python
+import httpx
+
+response = httpx.post(
+    "http://localhost:8000/v1/audio/generate",
+    json={
+        "input": "The sound of ocean waves crashing on a beach",
+        "audio_length": 10.0,
+        "negative_prompt": "Low quality, distorted",
+        "guidance_scale": 7.0,
+        "num_inference_steps": 100,
+    },
+    timeout=300.0,
+)
+
+with open("ocean.wav", "wb") as f:
+    f.write(response.content)
+```
+
+## Request Format
+
+### Simple Generation
+
+```json
+{
+    "input": "The sound of ocean waves"
+}
+```
+
+### Generation with Parameters
+
+```json
+{
+    "input": "A piano playing a gentle melody",
+    "audio_length": 10.0,
+    "negative_prompt": "Low quality, distorted, noisy",
+    "guidance_scale": 8.0,
+    "num_inference_steps": 150,
+    "seed": 42,
+    "response_format": "wav"
+}
+```
+
+## API Reference
+
+### Endpoint
+
+```
+POST /v1/audio/generate
+Content-Type: application/json
+```
+
+### Generation Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `input` | string | **required** | Text prompt describing the audio to generate |
+| `model` | string | server's model | Model to use (optional, should match server if specified) |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `speed` | float | 1.0 | Playback speed (0.25 - 4.0) |
+| `audio_length` | float | null | Audio duration in seconds (max ~47s for `stable-audio-open-1.0`) |
+| `audio_start` | float | 0.0 | Audio start time in seconds |
+| `negative_prompt` | string | null | Text describing what to avoid in generation |
+| `guidance_scale` | float | model default | Classifier-free guidance scale (higher = more adherence to prompt) |
+| `num_inference_steps` | int | model default | Number of denoising steps (higher = better quality, slower) |
+| `seed` | int | null | Random seed for reproducible generation |
+
+### Response Format
+
+Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/wav`).
+
+## Tuning Tips
+
+1. **Audio Length**: Keep under 47 seconds for `stable-audio-open-1.0`.
+2. **Quality vs Speed**:
+   - 50 steps: Fast, decent quality (quick previews)
+   - 100 steps: Good balance (general purpose)
+   - 150+ steps: High quality, slower (final / critical audio)
+3. **Guidance Scale**:
+   - Lower (3 - 5): More creative / varied output
+   - Default (7): Good balance
+   - Higher (10+): Strict adherence to the prompt
+4. **Negative Prompts**: Use to avoid unwanted characteristics such as `"Low quality"`, `"distorted"`, `"noisy"`.
+5. **Seeds**: Set a fixed seed to get deterministic, reproducible results.
+
+## File Description
+
+| File | Description |
+|------|-------------|
+| `curl_examples.sh` | Curl examples covering common use cases |
+| `stable_audio_client.py` | Python client with full CLI argument support |
+
+## Troubleshooting
+
+1. **Audio generation model did not produce audio output**: Verify the server started successfully and the model loaded without errors.
+2. **Connection refused**: Make sure the server is running on the correct port.
+3. **Generation timeout**: Reduce `num_inference_steps` or `audio_length`, and check GPU memory with `nvidia-smi`.
+4. **Out of memory**: Lower `--gpu-memory-utilization` or reduce `audio_length`.
+5. **Audio quality issues**: Increase `num_inference_steps`, add a negative prompt, or raise `guidance_scale`.
+
+## Example materials
+
+??? abstract "stable_audio_client.py"
+    ``````py
+    --8<-- "examples/online_serving/stable_audio/stable_audio_client.py"
+    ``````
+??? abstract "curl_examples.sh"
+    ``````sh
+    --8<-- "examples/online_serving/stable_audio/curl_examples.sh"
+    ``````

docs/user_guide/examples/online_serving/text_to_image.md

@@ -23,6 +23,29 @@ Or use the startup script:
 bash run_server.sh
 ```
 
+### Start with Parallelism Acceleration
+
+Enable Tensor Parallelism and VAE Patch Parallelism for faster inference:
+
+```bash
+# With Tensor Parallelism (requires >= 2 GPUs)
+vllm serve Qwen/Qwen-Image --omni --port 8091 --tensor-parallel-size 2
+
+# With Tensor Parallelism and VAE Patch Parallelism (requires >= 2 GPUs)
+vllm serve Qwen/Qwen-Image --omni --port 8091 --tensor-parallel-size 2 --vae-patch-parallel-size 2 --vae-use-tiling
+
+# With Sequence Parallelism (Ulysses-SP, requires >= 2 GPUs)
+vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2
+
+# With Ring-Attention (requires >= 2 GPUs)
+vllm serve Qwen/Qwen-Image --omni --port 8091 --ring 2
+
+# Combined: Ulysses + Ring (requires >= 4 GPUs)
+vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2 --ring 2
+```
+
+For more details on parallelism acceleration, see the [Parallelism Acceleration Guide](../../diffusion/parallelism_acceleration.md).
+
 ## API Calls
 
 ### Method 1: Using curl

examples/offline_inference/bagel/end2end.py

@@ -1,8 +1,7 @@
 import argparse
 import os
-from typing import cast
 
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
+from vllm_omni.inputs.data import OmniPromptType
 
 
 def parse_args():
@@ -49,8 +48,16 @@ def parse_args():
     parser.add_argument("--cfg-text-scale", type=float, default=4.0, help="Text CFG scale (default: 4.0)")
     parser.add_argument("--cfg-img-scale", type=float, default=1.5, help="Image CFG scale (default: 1.5)")
     parser.add_argument(
-        "--negative-prompt", type=str, default=None, help="Negative prompt (not yet supported, reserved for future)"
+        "--negative-prompt", type=str, default=None, help="Negative prompt for CFG (default: empty prompt)"
     )
+    parser.add_argument(
+        "--cfg-parallel-size",
+        type=int,
+        default=1,
+        choices=[1, 2, 3],
+        help="CFG parallel size: 1=batched (single GPU), 2=parallel with 2 branches (text CFG only), 3=parallel (3 GPUs).",
+    )
+    parser.add_argument("--seed", type=int, default=None, help="Random seed for generation.")
 
     args = parser.parse_args()
     return args
@@ -81,101 +88,80 @@ def main():
 
     from PIL import Image
 
-    if args.modality == "img2img":
+    from vllm_omni.entrypoints.omni import Omni
-        from PIL import Image
 
-        from vllm_omni.entrypoints.omni_diffusion import OmniDiffusion
+    omni_kwargs = {}
+    if args.stage_configs_path:
+        omni_kwargs["stage_configs_path"] = args.stage_configs_path
 
-        print("[Info] Running in img2img mode (Stage 1 only)")
+    omni_kwargs.update(
-        client = OmniDiffusion(model=model_name)
+        {
+            "log_stats": args.log_stats,
+            "init_sleep_seconds": args.init_sleep_seconds,
+            "batch_timeout": args.batch_timeout,
+            "init_timeout": args.init_timeout,
+            "shm_threshold_bytes": args.shm_threshold_bytes,
+            "worker_backend": args.worker_backend,
+            "ray_address": args.ray_address,
+        }
+    )
 
-        if args.image_path:
+    omni = Omni(model=model_name, **omni_kwargs)
-            if os.path.exists(args.image_path):
-                loaded_image = Image.open(args.image_path).convert("RGB")
-                prompts = [
-                    {
-                        "prompt": cast(str, p),
-                        "multi_modal_data": {"image": loaded_image},
-                    }
-                    for p in prompts
-                ]
-            else:
-                print(f"[Warning] Image path {args.image_path} does not exist.")
 
-        result = client.generate(
+    formatted_prompts = []
-            prompts,
+    for p in prompts:
-            OmniDiffusionSamplingParams(
+        if args.modality == "img2img":
-                seed=52,
+            if not args.image_path or not os.path.exists(args.image_path):
-                need_kv_receive=False,
+                raise ValueError(f"img2img requires --image-path pointing to an existing file, got: {args.image_path}")
-                num_inference_steps=args.steps,
+            loaded_image = Image.open(args.image_path).convert("RGB")
-                extra_args={
+            final_prompt_text = f"<|fim_middle|><|im_start|>{p}<|im_end|>"
-                    "cfg_text_scale": args.cfg_text_scale,
+            prompt_dict = {
-                    "cfg_img_scale": args.cfg_img_scale,
+                "prompt": final_prompt_text,
-                },
+                "multi_modal_data": {"img2img": loaded_image},
-            ),
+                "modalities": ["img2img"],
-        )
-
-        # Ensure result is a list for iteration
-        if not isinstance(result, list):
-            omni_outputs = [result]
-        else:
-            omni_outputs = result
-
-    else:
-        from vllm_omni.entrypoints.omni import Omni
-
-        omni_kwargs = {}
-        if args.stage_configs_path:
-            omni_kwargs["stage_configs_path"] = args.stage_configs_path
-
-        omni_kwargs.update(
-            {
-                "log_stats": args.log_stats,
-                "init_sleep_seconds": args.init_sleep_seconds,
-                "batch_timeout": args.batch_timeout,
-                "init_timeout": args.init_timeout,
-                "shm_threshold_bytes": args.shm_threshold_bytes,
-                "worker_backend": args.worker_backend,
-                "ray_address": args.ray_address,
             }
-        )
+            if args.negative_prompt is not None:
-
+                prompt_dict["negative_prompt"] = args.negative_prompt
-        omni = Omni(model=model_name, **omni_kwargs)
+            formatted_prompts.append(prompt_dict)
-
+        elif args.modality == "img2text":
-        formatted_prompts = []
+            if args.image_path:
-        for p in args.prompts:
+                loaded_image = Image.open(args.image_path).convert("RGB")
-            if args.modality == "img2text":
+                final_prompt_text = f"<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n"
-                if args.image_path:
+                prompt_dict = {
-                    loaded_image = Image.open(args.image_path).convert("RGB")
+                    "prompt": final_prompt_text,
-                    final_prompt_text = f"<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n"
+                    "multi_modal_data": {"image": loaded_image},
-                    prompt_dict = {
+                    "modalities": ["text"],
-                        "prompt": final_prompt_text,
-                        "multi_modal_data": {"image": loaded_image},
-                        "modalities": ["text"],
-                    }
-                    formatted_prompts.append(prompt_dict)
-            elif args.modality == "text2text":
-                final_prompt_text = f"<|im_start|>user\n{p}<|im_end|>\n<|im_start|>assistant\n"
-                prompt_dict = {"prompt": final_prompt_text, "modalities": ["text"]}
-                formatted_prompts.append(prompt_dict)
-            else:
-                # text2img
-                final_prompt_text = f"<|im_start|>{p}<|im_end|>"
-                prompt_dict = {"prompt": final_prompt_text, "modalities": ["image"]}
-                formatted_prompts.append(prompt_dict)
-
-        params_list = omni.default_sampling_params_list
-        if args.modality == "text2img":
-            params_list[0].max_tokens = 1  # type: ignore # The first stage is a SamplingParam (vllm)
-            if len(params_list) > 1:
-                diffusion_params = params_list[1]
-                diffusion_params.num_inference_steps = args.steps  # type: ignore
-                diffusion_params.extra_args = {  # type: ignore
-                    "cfg_text_scale": args.cfg_text_scale,
-                    "cfg_img_scale": args.cfg_img_scale,
                 }
+                formatted_prompts.append(prompt_dict)
+        elif args.modality == "text2text":
+            final_prompt_text = f"<|im_start|>user\n{p}<|im_end|>\n<|im_start|>assistant\n"
+            prompt_dict = {"prompt": final_prompt_text, "modalities": ["text"]}
+            formatted_prompts.append(prompt_dict)
+        else:
+            final_prompt_text = f"<|im_start|>{p}<|im_end|>"
+            prompt_dict = {"prompt": final_prompt_text, "modalities": ["image"]}
+            if args.negative_prompt is not None:
+                prompt_dict["negative_prompt"] = args.negative_prompt
+            formatted_prompts.append(prompt_dict)
 
-        omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
+    params_list = omni.default_sampling_params_list
+    if args.modality in ("text2img", "img2img"):
+        params_list[0].max_tokens = 1  # type: ignore
+        if len(params_list) > 1:
+            diffusion_params = params_list[1]
+            diffusion_params.num_inference_steps = args.steps  # type: ignore
+            diffusion_params.cfg_parallel_size = args.cfg_parallel_size  # type: ignore
+            if args.seed is not None:
+                diffusion_params.seed = args.seed  # type: ignore
+            extra = {
+                "cfg_text_scale": args.cfg_text_scale,
+                "cfg_img_scale": args.cfg_img_scale,
+            }
+            if args.negative_prompt is not None:
+                extra["negative_prompt"] = args.negative_prompt
+            diffusion_params.extra_args = extra  # type: ignore
+
+    omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
 
     for i, req_output in enumerate(omni_outputs):
         images = getattr(req_output, "images", None)

examples/offline_inference/custom_pipeline/image_to_image/image_edit.py

@@ -29,8 +29,8 @@ Accelerated with cache_dit:
     python image_edit.py \
         --image input.png \
         --prompt "Edit description" \
-        --cache_backend cache_dit \
+        --cache-backend cache_dit \
-        --cache_dit_enable_taylorseer
+        --cache-dit-enable-taylorseer
 
 Layered RGBA output:
     python image_edit.py \
@@ -66,38 +66,38 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--model", default="Qwen/Qwen-Image-Edit", help="Model name or local path.")
     parser.add_argument("--image", type=str, nargs="+", required=True, help="Input image file(s).")
     parser.add_argument("--prompt", type=str, required=True, help="Edit description prompt.")
-    parser.add_argument("--negative_prompt", type=str, default=None)
+    parser.add_argument("--negative-prompt", type=str, default=None)
     parser.add_argument("--seed", type=int, default=0)
-    parser.add_argument("--cfg_scale", type=float, default=4.0)
+    parser.add_argument("--cfg-scale", type=float, default=4.0)
-    parser.add_argument("--guidance_scale", type=float, default=1.0)
+    parser.add_argument("--guidance-scale", type=float, default=1.0)
     parser.add_argument("--output", type=str, default="output.png")
-    parser.add_argument("--num_outputs_per_prompt", type=int, default=1)
+    parser.add_argument("--num-outputs-per-prompt", type=int, default=1)
-    parser.add_argument("--num_inference_steps", type=int, default=50)
+    parser.add_argument("--num-inference-steps", type=int, default=50)
-    parser.add_argument("--cache_backend", type=str, default=None, choices=["cache_dit", "tea_cache"])
+    parser.add_argument("--cache-backend", type=str, default=None, choices=["cache_dit", "tea_cache"])
-    parser.add_argument("--ulysses_degree", type=int, default=1)
+    parser.add_argument("--ulysses-degree", type=int, default=1)
-    parser.add_argument("--ring_degree", type=int, default=1)
+    parser.add_argument("--ring-degree", type=int, default=1)
-    parser.add_argument("--tensor_parallel_size", type=int, default=1)
+    parser.add_argument("--tensor-parallel-size", type=int, default=1)
     parser.add_argument("--layers", type=int, default=4)
     parser.add_argument("--resolution", type=int, default=640)
     parser.add_argument("--color-format", type=str, default="RGB")
 
     # Acceleration + Optimization Options
-    parser.add_argument("--cache_dit_fn_compute_blocks", type=int, default=1)
+    parser.add_argument("--cache-dit-fn-compute-blocks", type=int, default=1)
-    parser.add_argument("--cache_dit_bn_compute_blocks", type=int, default=0)
+    parser.add_argument("--cache-dit-bn-compute-blocks", type=int, default=0)
-    parser.add_argument("--cache_dit_max_warmup_steps", type=int, default=4)
+    parser.add_argument("--cache-dit-max-warmup-steps", type=int, default=4)
-    parser.add_argument("--cache_dit_residual_diff_threshold", type=float, default=0.24)
+    parser.add_argument("--cache-dit-residual-diff-threshold", type=float, default=0.24)
-    parser.add_argument("--cache_dit_max_continuous_cached_steps", type=int, default=3)
+    parser.add_argument("--cache-dit-max-continuous-cached-steps", type=int, default=3)
-    parser.add_argument("--cache_dit_enable_taylorseer", action="store_true", default=False)
+    parser.add_argument("--cache-dit-enable-taylorseer", action="store_true", default=False)
-    parser.add_argument("--cache_dit_taylorseer_order", type=int, default=1)
+    parser.add_argument("--cache-dit-taylorseer-order", type=int, default=1)
     parser.add_argument(
-        "--cache_dit_scm_steps_mask_policy", type=str, default=None, choices=[None, "slow", "medium", "fast", "ultra"]
+        "--cache-dit-scm-steps-mask-policy", type=str, default=None, choices=[None, "slow", "medium", "fast", "ultra"]
     )
-    parser.add_argument("--cache_dit_scm_steps_policy", type=str, default="dynamic", choices=["dynamic", "static"])
+    parser.add_argument("--cache-dit-scm-steps-policy", type=str, default="dynamic", choices=["dynamic", "static"])
-    parser.add_argument("--tea_cache_rel_l1_thresh", type=float, default=0.2)
+    parser.add_argument("--tea-cache-rel-l1-thresh", type=float, default=0.2)
-    parser.add_argument("--cfg_parallel_size", type=int, default=1, choices=[1, 2])
+    parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2])
-    parser.add_argument("--enforce_eager", action="store_true")
+    parser.add_argument("--enforce-eager", action="store_true")
-    parser.add_argument("--vae_use_slicing", action="store_true")
+    parser.add_argument("--vae-use-slicing", action="store_true")
-    parser.add_argument("--vae_use_tiling", action="store_true")
+    parser.add_argument("--vae-use-tiling", action="store_true")
     parser.add_argument("--enable-cpu-offload", action="store_true")
 
     return parser.parse_args()

examples/offline_inference/custom_pipeline/image_to_image/run.sh

@@ -3,4 +3,4 @@ python image_edit.py \
         --image cherry_blossom.jpg \
         --prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'" \
         --output output_image_edit.png \
-        --num_inference_steps 10
+        --num-inference-steps 10

examples/offline_inference/glm_image/README.md

@@ -40,8 +40,7 @@ GLM-Image uses a 2-stage pipeline:
 
 ```bash
 python end2end.py \
-    --model-path /path/to/glm-image \
+    --config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
-    --config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
     --prompt "A beautiful sunset over the ocean with sailing boats" \
     --height 1024 \
     --width 1024 \
@@ -52,8 +51,7 @@ python end2end.py \
 
 ```bash
 python end2end.py \
-    --model-path /path/to/glm-image \
+    --config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
-    --config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
     --prompt "Transform this scene into a winter wonderland" \
     --image input.png \
     --output output_i2i.png
@@ -64,7 +62,7 @@ python end2end.py \
 ```bash
 python end2end.py \
     --model-path /path/to/glm-image \
-    --config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
+    --config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
     --prompt "A photorealistic cat sitting on a window sill" \
     --height 1024 \
     --width 1024 \

examples/offline_inference/glm_image/end2end.py

@@ -178,9 +178,6 @@ def main(args: argparse.Namespace) -> None:
     print("=" * 60)
 
     # Validate arguments
-    if not args.model_path:
-        raise ValueError("--model-path is required")
-
     if not args.prompt:
         raise ValueError("--prompt is required")
 
@@ -365,7 +362,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--model-path",
         type=str,
-        required=True,
+        default="zai-org/GLM-Image",
         help="Path to GLM-Image model directory or HuggingFace model ID",
     )
     parser.add_argument(

examples/offline_inference/helios/README.md

@@ -0,0 +1,193 @@
+# Helios Video Generation
+
+Helios is a text-to-video (T2V), image-to-video (I2V), and video-to-video (V2V) diffusion model. This example demonstrates end-to-end video generation using vLLM-Omni with three model variants:
+
+| Variant | Description | Key Features |
+|---------|-------------|--------------|
+| **Helios-Base** | Base model, Stage 1 only | Single-stage denoising, `guidance_scale=5.0` |
+| **Helios-Mid** | Mid model, Stage 2 pyramid | Multi-stage pyramid denoising, CFG-Zero* support |
+| **Helios-Distilled** | Distilled model, Stage 2+3 | Few-step inference with DMD, `guidance_scale=1.0` |
+
+## Setup
+
+Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
+
+## Run Examples
+
+Get into the example folder:
+```bash
+cd examples/offline_inference/helios
+```
+
+### Text-to-Video (T2V)
+
+**Helios-Base** (Stage 1 only):
+```bash
+python end2end.py \
+    --model BestWishYsh/Helios-Base \
+    --sample-type t2v \
+    --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
+    --guidance-scale 5.0 \
+    --output helios_t2v_base.mp4
+```
+
+**Helios-Mid** (Stage 2 + CFG-Zero*):
+```bash
+python end2end.py \
+    --model BestWishYsh/Helios-Mid \
+    --sample-type t2v \
+    --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
+    --guidance-scale 5.0 \
+    --is-enable-stage2 \
+    --pyramid-num-inference-steps-list 20 20 20 \
+    --use-cfg-zero-star --use-zero-init --zero-steps 1 \
+    --output helios_t2v_mid.mp4
+```
+
+**Helios-Distilled** (Stage 2 pyramid + DMD):
+```bash
+python end2end.py \
+    --model BestWishYsh/Helios-Distilled \
+    --sample-type t2v \
+    --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
+    --num-frames 240 \
+    --guidance-scale 1.0 \
+    --is-enable-stage2 \
+    --pyramid-num-inference-steps-list 2 2 2 \
+    --is-amplify-first-chunk \
+    --output helios_t2v_distilled.mp4
+```
+
+### Image-to-Video (I2V)
+
+**Helios-Base**:
+```bash
+python end2end.py \
+    --model BestWishYsh/Helios-Base \
+    --sample-type i2v \
+    --image-path /path/to/image.jpg \
+    --prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
+    --guidance-scale 5.0 \
+    --output helios_i2v_base.mp4
+```
+
+**Helios-Mid** (Stage 2 + CFG-Zero*):
+```bash
+python end2end.py \
+    --model BestWishYsh/Helios-Mid \
+    --sample-type i2v \
+    --image-path /path/to/image.jpg \
+    --prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
+    --guidance-scale 5.0 \
+    --is-enable-stage2 \
+    --pyramid-num-inference-steps-list 20 20 20 \
+    --use-cfg-zero-star --use-zero-init --zero-steps 1 \
+    --output helios_i2v_mid.mp4
+```
+
+**Helios-Distilled**:
+```bash
+python end2end.py \
+    --model BestWishYsh/Helios-Distilled \
+    --sample-type i2v \
+    --image-path /path/to/image.jpg \
+    --prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
+    --num-frames 240 \
+    --guidance-scale 1.0 \
+    --is-enable-stage2 \
+    --pyramid-num-inference-steps-list 2 2 2 \
+    --is-amplify-first-chunk \
+    --output helios_i2v_distilled.mp4
+```
+
+### Video-to-Video (V2V)
+
+**Helios-Base**:
+```bash
+python end2end.py \
+    --model BestWishYsh/Helios-Base \
+    --sample-type v2v \
+    --video-path /path/to/video.mp4 \
+    --prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
+    --guidance-scale 5.0 \
+    --output helios_v2v_base.mp4
+```
+
+**Helios-Mid** (Stage 2 + CFG-Zero*):
+```bash
+python end2end.py \
+    --model BestWishYsh/Helios-Mid \
+    --sample-type v2v \
+    --video-path /path/to/video.mp4 \
+    --prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
+    --guidance-scale 5.0 \
+    --is-enable-stage2 \
+    --pyramid-num-inference-steps-list 20 20 20 \
+    --use-cfg-zero-star --use-zero-init --zero-steps 1 \
+    --output helios_v2v_mid.mp4
+```
+
+**Helios-Distilled**:
+```bash
+python end2end.py \
+    --model BestWishYsh/Helios-Distilled \
+    --sample-type v2v \
+    --video-path /path/to/video.mp4 \
+    --prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
+    --num-frames 240 \
+    --guidance-scale 1.0 \
+    --is-enable-stage2 \
+    --pyramid-num-inference-steps-list 2 2 2 \
+    --is-amplify-first-chunk \
+    --output helios_v2v_distilled.mp4
+```
+
+## Common Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `--model` | `BestWishYsh/Helios-Base` | Model ID or local path |
+| `--sample-type` | `t2v` | Generation mode: `t2v`, `i2v`, or `v2v` |
+| `--prompt` | — | Text prompt describing the video |
+| `--negative-prompt` | *(see source)* | Negative prompt for CFG (includes anti-static terms) |
+| `--image-path` | — | Input image (required for `i2v`) |
+| `--video-path` | — | Input video (required for `v2v`) |
+| `--height` | `384` | Video height in pixels |
+| `--width` | `640` | Video width in pixels |
+| `--num-frames` | `99` | Number of output frames |
+| `--num-inference-steps` | `50` | Denoising steps (Stage 1 only) |
+| `--guidance-scale` | `5.0` | CFG scale (`1.0` for Distilled) |
+| `--seed` | `42` | Random seed |
+| `--fps` | `16` | Output video frame rate |
+| `--output` | `helios_output.mp4` | Output file path |
+
+### Stage 2 / Pyramid Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `--is-enable-stage2` | off | Enable pyramid multi-stage denoising |
+| `--pyramid-num-stages` | `3` | Number of pyramid stages |
+| `--pyramid-num-inference-steps-list` | `10 10 10` | Steps per pyramid stage |
+| `--is-amplify-first-chunk` | off | DMD amplification (Distilled) |
+
+### CFG-Zero* Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `--use-cfg-zero-star` | off | Enable CFG-Zero* guidance (Mid) |
+| `--use-zero-init` | off | Zero init for first steps |
+| `--zero-steps` | `1` | Number of zero-init steps |
+
+### Memory & Parallelism
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `--vae-use-slicing` | off | Enable VAE slicing |
+| `--vae-use-tiling` | off | Enable VAE tiling |
+| `--enforce-eager` | off | Disable torch.compile |
+| `--enable-cpu-offload` | off | CPU offloading |
+| `--enable-layerwise-offload` | off | Layerwise offloading |
+| `--ulysses-degree` | `1` | Ulysses SP degree |
+| `--ring-degree` | `1` | Ring SP degree |
+| `--cfg-parallel-size` | `1` | CFG parallel size (1 or 2) |
+| `--tensor-parallel-size` | `1` | Tensor parallelism size |

examples/offline_inference/helios/end2end.py

@@ -0,0 +1,343 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Helios video generation example supporting T2V, I2V, and V2V.
+
+Usage (T2V, Helios-Base, Stage 1 only):
+    python end2end.py \
+        --model /path/to/Helios-Base --sample-type t2v \
+        --prompt "A serene lakeside sunrise with mist over the water." \
+        --height 384 --width 640 --num-frames 99 \
+        --num-inference-steps 50 --guidance-scale 5.0
+
+Usage (I2V, Helios-Base):
+    python end2end.py \
+        --model /path/to/Helios-Base --sample-type i2v \
+        --image-path /path/to/image.jpg \
+        --prompt "Description of desired animation." \
+        --guidance-scale 5.0
+
+Usage (V2V, Helios-Base):
+    python end2end.py \
+        --model /path/to/Helios-Base --sample-type v2v \
+        --video-path /path/to/video.mp4 \
+        --prompt "Description of desired transformation." \
+        --guidance-scale 5.0
+
+Usage (Helios-Mid, Stage 2 + CFG-Zero*):
+    python end2end.py \
+        --model /path/to/Helios-Mid --sample-type t2v \
+        --prompt "A serene lakeside sunrise with mist over the water." \
+        --guidance-scale 5.0 \
+        --is-enable-stage2 \
+        --pyramid-num-inference-steps-list 20 20 20 \
+        --use-cfg-zero-star --use-zero-init --zero-steps 1
+
+Usage (Helios-Distilled, Stage 2 pyramid + DMD):
+    python end2end.py \
+        --model /path/to/Helios-Distilled --sample-type t2v \
+        --prompt "A serene lakeside sunrise with mist over the water." \
+        --guidance-scale 1.0 \
+        --is-enable-stage2 \
+        --pyramid-num-inference-steps-list 2 2 2 \
+        --is-amplify-first-chunk
+"""
+
+import argparse
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.platforms import current_omni_platform
+
+
+def load_image_as_tensor(image_path: str, height: int, width: int) -> torch.Tensor:
+    """Load an image and return a (1, C, H, W) tensor normalized to [-1, 1]."""
+    from PIL import Image
+    from torchvision import transforms
+
+    image = Image.open(image_path).convert("RGB").resize((width, height))
+    transform = transforms.Compose(
+        [
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+        ]
+    )
+    return transform(image).unsqueeze(0)
+
+
+def load_video_as_tensor(video_path: str, height: int, width: int) -> torch.Tensor:
+    """Load a video and return a (1, C, T, H, W) tensor normalized to [-1, 1]."""
+    from torchvision import transforms
+    from torchvision.io import read_video
+
+    video_frames, _, _ = read_video(video_path, output_format="TCHW")
+    video_frames = video_frames.float() / 255.0
+
+    transform = transforms.Compose(
+        [
+            transforms.Resize((height, width), antialias=True),
+            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+        ]
+    )
+    video_frames = torch.stack([transform(f) for f in video_frames])
+    return video_frames.permute(1, 0, 2, 3).unsqueeze(0)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Generate a video with Helios (T2V / I2V / V2V).")
+    parser.add_argument(
+        "--model",
+        default="BestWishYsh/Helios-Base",
+        help="Helios model ID or local path (e.g. Helios-Base, Helios-Mid, Helios-Distilled).",
+    )
+    parser.add_argument(
+        "--sample-type",
+        choices=["t2v", "i2v", "v2v"],
+        default="t2v",
+        help="Generation mode: t2v (text-to-video), i2v (image-to-video), v2v (video-to-video).",
+    )
+    parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.")
+    parser.add_argument(
+        "--negative-prompt",
+        default=(
+            "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, "
+            "images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, "
+            "incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, "
+            "misshapen limbs, fused fingers, still picture, messy background, three legs, many people "
+            "in the background, walking backwards"
+        ),
+        help="Negative prompt.",
+    )
+    parser.add_argument("--seed", type=int, default=42, help="Random seed.")
+    parser.add_argument("--guidance-scale", type=float, default=5.0, help="CFG scale.")
+    parser.add_argument("--height", type=int, default=384, help="Video height.")
+    parser.add_argument("--width", type=int, default=640, help="Video width.")
+    parser.add_argument("--num-frames", type=int, default=99, help="Number of video frames.")
+    parser.add_argument("--num-inference-steps", type=int, default=50, help="Sampling steps (Stage 1 only).")
+    parser.add_argument("--output", type=str, default="helios_output.mp4", help="Output video path.")
+    parser.add_argument("--fps", type=int, default=16, help="Frames per second for the output video.")
+
+    # I2V / V2V
+    parser.add_argument("--image-path", type=str, default=None, help="Input image path for I2V mode.")
+    parser.add_argument("--video-path", type=str, default=None, help="Input video path for V2V mode.")
+
+    # Stage 2 (pyramid multi-stage denoising)
+    parser.add_argument(
+        "--is-enable-stage2",
+        action="store_true",
+        help="Enable pyramid multi-stage denoising (Stage 2). Required for Helios-Distilled.",
+    )
+    parser.add_argument(
+        "--pyramid-num-stages",
+        type=int,
+        default=3,
+        help="Number of pyramid stages for Stage 2.",
+    )
+    parser.add_argument(
+        "--pyramid-num-inference-steps-list",
+        type=int,
+        nargs="+",
+        default=[10, 10, 10],
+        help="Inference steps per pyramid stage.",
+    )
+
+    # DMD
+    parser.add_argument(
+        "--is-amplify-first-chunk",
+        action="store_true",
+        help="Enable DMD amplification for the first chunk (Helios-Distilled).",
+    )
+
+    # CFG Zero Star
+    parser.add_argument(
+        "--use-cfg-zero-star",
+        action="store_true",
+        help="Enable CFG Zero Star guidance (recommended for Helios-Mid).",
+    )
+    parser.add_argument(
+        "--use-zero-init",
+        action="store_true",
+        help="Use zero initialization for the first denoising steps with CFG-Zero*.",
+    )
+    parser.add_argument(
+        "--zero-steps",
+        type=int,
+        default=1,
+        help="Number of initial denoising steps using zero prediction (default: 1).",
+    )
+
+    # Memory & parallelism
+    parser.add_argument("--vae-use-slicing", action="store_true", help="Enable VAE slicing.")
+    parser.add_argument("--vae-use-tiling", action="store_true", help="Enable VAE tiling.")
+    parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.")
+    parser.add_argument("--enable-cpu-offload", action="store_true", help="Enable CPU offloading.")
+    parser.add_argument("--enable-layerwise-offload", action="store_true", help="Enable layerwise offloading.")
+    parser.add_argument("--ulysses-degree", type=int, default=1, help="Ulysses SP degree.")
+    parser.add_argument("--ring-degree", type=int, default=1, help="Ring SP degree.")
+    parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2], help="CFG parallel size.")
+    parser.add_argument("--tensor-parallel-size", type=int, default=1, help="Tensor parallelism size.")
+
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
+
+    parallel_config = DiffusionParallelConfig(
+        ulysses_degree=args.ulysses_degree,
+        ring_degree=args.ring_degree,
+        cfg_parallel_size=args.cfg_parallel_size,
+        tensor_parallel_size=args.tensor_parallel_size,
+    )
+
+    omni = Omni(
+        model=args.model,
+        enable_layerwise_offload=args.enable_layerwise_offload,
+        vae_use_slicing=args.vae_use_slicing,
+        vae_use_tiling=args.vae_use_tiling,
+        enable_cpu_offload=args.enable_cpu_offload,
+        parallel_config=parallel_config,
+        enforce_eager=args.enforce_eager,
+    )
+
+    # Validate I2V / V2V arguments
+    if args.sample_type == "i2v" and not args.image_path:
+        raise ValueError("--image-path is required for I2V mode.")
+    if args.sample_type == "v2v" and not args.video_path:
+        raise ValueError("--video-path is required for V2V mode.")
+
+    # Build extra_args for Helios-specific parameters
+    extra_args = {}
+    if args.is_enable_stage2:
+        extra_args["is_enable_stage2"] = True
+        extra_args["pyramid_num_stages"] = args.pyramid_num_stages
+        extra_args["pyramid_num_inference_steps_list"] = args.pyramid_num_inference_steps_list
+    if args.is_amplify_first_chunk:
+        extra_args["is_amplify_first_chunk"] = True
+    if args.use_cfg_zero_star:
+        extra_args["use_cfg_zero_star"] = True
+    if args.use_zero_init:
+        extra_args["use_zero_init"] = True
+        extra_args["zero_steps"] = args.zero_steps
+
+    if args.sample_type == "i2v":
+        image_tensor = load_image_as_tensor(args.image_path, args.height, args.width)
+        extra_args["image"] = image_tensor
+    elif args.sample_type == "v2v":
+        video_tensor = load_video_as_tensor(args.video_path, args.height, args.width)
+        extra_args["video"] = video_tensor
+
+    # Print generation configuration
+    print(f"\n{'=' * 60}")
+    print("Helios Generation Configuration:")
+    print(f"  Model: {args.model}")
+    print(f"  Sample type: {args.sample_type.upper()}")
+    if args.sample_type == "i2v":
+        print(f"  Image: {args.image_path}")
+    elif args.sample_type == "v2v":
+        print(f"  Video: {args.video_path}")
+    print(f"  Prompt: {args.prompt[:80]}{'...' if len(args.prompt) > 80 else ''}")
+    print(f"  Video size: {args.width}x{args.height}, {args.num_frames} frames")
+    print(f"  Inference steps: {args.num_inference_steps}")
+    print(f"  Guidance scale: {args.guidance_scale}")
+    if args.is_enable_stage2:
+        print(f"  Stage 2: enabled (stages={args.pyramid_num_stages}, steps={args.pyramid_num_inference_steps_list})")
+        if args.is_amplify_first_chunk:
+            print("  DMD amplify first chunk: enabled")
+        if args.use_cfg_zero_star:
+            print(f"  CFG Zero Star: enabled (zero_init={args.use_zero_init}, zero_steps={args.zero_steps})")
+    else:
+        if args.use_cfg_zero_star:
+            print(f"  CFG Zero Star: enabled (zero_init={args.use_zero_init}, zero_steps={args.zero_steps})")
+        print("  Stage 2: disabled (Stage 1 only)")
+    print(f"{'=' * 60}\n")
+
+    generation_start = time.perf_counter()
+    frames = omni.generate(
+        {
+            "prompt": args.prompt,
+            "negative_prompt": args.negative_prompt,
+        },
+        OmniDiffusionSamplingParams(
+            height=args.height,
+            width=args.width,
+            generator=generator,
+            guidance_scale=args.guidance_scale,
+            num_inference_steps=args.num_inference_steps,
+            num_frames=args.num_frames,
+            extra_args=extra_args,
+        ),
+    )
+    generation_end = time.perf_counter()
+    generation_time = generation_end - generation_start
+
+    print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
+
+    # Extract video frames from OmniRequestOutput
+    if isinstance(frames, list) and len(frames) > 0:
+        first_item = frames[0]
+
+        if hasattr(first_item, "final_output_type"):
+            if first_item.final_output_type != "image":
+                raise ValueError(
+                    f"Unexpected output type '{first_item.final_output_type}', expected 'image' for video generation."
+                )
+
+            if hasattr(first_item, "is_pipeline_output") and first_item.is_pipeline_output:
+                if isinstance(first_item.request_output, list) and len(first_item.request_output) > 0:
+                    inner_output = first_item.request_output[0]
+                    if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"):
+                        frames = inner_output.images[0] if inner_output.images else None
+                        if frames is None:
+                            raise ValueError("No video frames found in output.")
+            elif hasattr(first_item, "images") and first_item.images:
+                frames = first_item.images
+            else:
+                raise ValueError("No video frames found in OmniRequestOutput.")
+
+    # Unwrap batch list from postprocess_video: [numpy(T,H,W,C)] -> numpy(T,H,W,C)
+    if isinstance(frames, list) and len(frames) > 0 and isinstance(frames[0], np.ndarray):
+        frames = frames[0]
+
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    try:
+        from diffusers.utils import export_to_video
+    except ImportError:
+        raise ImportError("diffusers is required for export_to_video.")
+
+    if isinstance(frames, torch.Tensor):
+        video_tensor = frames.detach().cpu()
+        if video_tensor.dim() == 5:
+            if video_tensor.shape[1] in (3, 4):
+                video_tensor = video_tensor[0].permute(1, 2, 3, 0)
+            else:
+                video_tensor = video_tensor[0]
+        elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4):
+            video_tensor = video_tensor.permute(1, 2, 3, 0)
+        if video_tensor.is_floating_point():
+            video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5
+        video_array = video_tensor.float().numpy()
+    elif isinstance(frames, np.ndarray):
+        video_array = frames
+    else:
+        video_array = frames
+
+    if isinstance(video_array, np.ndarray) and video_array.ndim == 5:
+        video_array = video_array[0]
+
+    export_to_video(video_array, str(output_path), fps=args.fps)
+    print(f"Saved generated video to {output_path}")
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/hunyuan_image3/README.md

@@ -0,0 +1,25 @@
+# HunyuanImage-3.0 Image-to-Text Inference
+
+This example demonstrates how to run HunyuanImage-3.0 Image-to-Text with the vLLM-Omni.
+
+## Local CLI Usage
+
+Download the example image:
+
+```bash
+wget https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg
+```
+
+Run example:
+
+```bash
+python image_to_text.py \
+  --image cherry_blossom.jpg \
+  --prompt "<|startoftext|>You are an assistant that understands images and outputs text.<img>Describe the content of the picture."
+```
+
+Key arguments:
+
+- `--model`: Model used. Default is: tencent/HunyuanImage-3.0-Instruct (Optional).
+- `--image`: Path to input image (required).
+- `--prompt`: Text description used to guide image understanding (required).

examples/offline_inference/hunyuan_image3/image_to_text.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import argparse
+import os
+
+from PIL import Image
+
+from vllm_omni.entrypoints.omni import Omni
+
+"""
+The tencent/HunyuanImage-3.0-Instruct base model is built on the Hunyuan v1 architecture, specifically the tencent/Hunyuan-A13B-Instruct model. It utilizes two tokenizer delimiter templates:
+
+1) Pretrained template (default for gen_text mode), which concatenates system, image
+   tokens, and user question WITHOUT role delimiters:
+"<|startoftext|>{system_prompt}{image_tokens}{user_question}"
+
+   Example (before image token expansion):
+"<|startoftext|>You are an assistant that understands images and outputs text.<img>Describe the content of the picture."
+
+2) Instruct template, which uses explicit role prefixes and separators.
+"""
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Generate text from image using HunyuanImage-3.0-Instruct.")
+    parser.add_argument(
+        "--model",
+        default="tencent/HunyuanImage-3.0-Instruct",
+        help="Model name or local path.",
+    )
+    parser.add_argument(
+        "--image",
+        type=str,
+        required=True,
+        help="Path to input image file (PNG, JPG, etc.).",
+    )
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        required=True,
+        help="Pretrain template prompt: <|startoftext|>{system}<img>{question}",
+    )
+    return parser.parse_args()
+
+
+def load_image(image_path: str) -> Image.Image:
+    """Load an image from file path."""
+    if not os.path.exists(image_path):
+        raise FileNotFoundError(f"Image file not found: {image_path}")
+    return Image.open(image_path).convert("RGB")
+
+
+def main(args: argparse.Namespace) -> None:
+    omni = Omni(model=args.model)
+
+    prompt_dict = {
+        "prompt": args.prompt,
+        "modalities": ["text"],
+    }
+
+    # Add image input if provided
+    if args.image:
+        if not os.path.exists(args.image):
+            raise FileNotFoundError(f"Input image not found: {args.image}")
+
+        input_image = load_image(args.image)
+        prompt_dict["multi_modal_data"] = {"image": input_image}
+
+    prompts = [prompt_dict]
+    omni_outputs = omni.generate(prompts=prompts)
+
+    prompt_text = omni_outputs[0].request_output[0].prompt
+    generated_text = omni_outputs[0].request_output[0].outputs[0].text
+    print(f"Prompt: {prompt_text}")
+    print(f"Text: {generated_text}")
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)

examples/offline_inference/image_to_video/README.md

@@ -65,6 +65,7 @@ Key arguments:
 - `--vae-use-slicing`: Enable VAE slicing for memory optimization.
 - `--vae-use-tiling`: Enable VAE tiling for memory optimization.
 - `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--tensor-parallel-size`: tensor parallel size (effective for models that support TP, e.g. LTX2).
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
 - `--use-hsdp`: Enable Hybrid Sharded Data Parallel to shard model weights across GPUs.
 - `--hsdp-shard-size`: Number of GPUs to shard model weights across within each replica group. -1 (default) auto-calculates as world_size / replicate_size.

examples/offline_inference/image_to_video/image_to_video.py

@@ -2,11 +2,12 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 """
-Image-to-Video generation example using Wan2.2 I2V/TI2V models.
+Image-to-Video generation example using Wan2.2 I2V/TI2V models or LTX2.
 
 Supports:
 - Wan2.2-I2V-A14B-Diffusers: MoE model with CLIP image encoder
 - Wan2.2-TI2V-5B-Diffusers: Unified T2V+I2V model (dense 5B)
+- LTX2 image-to-video pipeline
 
 Usage:
     # I2V-A14B (MoE)
@@ -16,10 +17,18 @@ Usage:
     # TI2V-5B (unified)
     python image_to_video.py --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
         --image input.jpg --prompt "A cat playing with yarn"
+
+    # LTX2 image-to-video
+    python image_to_video.py --model /path/to/LTX-2 \
+        --model-class-name LTX2ImageToVideoPipeline \
+        --image input.jpg --prompt "A cinematic dolly shot of a boat" \
+        --num-frames 121 --num-inference-steps 40 --guidance-scale 4.0 \
+        --frame-rate 24 --fps 24 --output ltx2_i2v.mp4
 """
 
 import argparse
 import os
+import time
 from pathlib import Path
 
 import numpy as np
@@ -34,12 +43,17 @@ from vllm_omni.platforms import current_omni_platform
 
 
 def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Generate a video from an image with Wan2.2 I2V/TI2V.")
+    parser = argparse.ArgumentParser(description="Generate a video from an image with Wan2.2 or LTX2.")
     parser.add_argument(
         "--model",
         default="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
         help="Diffusers Wan2.2 I2V model ID or local path.",
     )
+    parser.add_argument(
+        "--model-class-name",
+        default=None,
+        help="Override model class name (e.g., LTX2ImageToVideoPipeline).",
+    )
     parser.add_argument("--image", required=True, help="Path to input image.")
     parser.add_argument("--prompt", default="", help="Text prompt describing the desired motion.")
     parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
@@ -55,11 +69,17 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--num-frames", type=int, default=81, help="Number of frames.")
     parser.add_argument("--num-inference-steps", type=int, default=50, help="Sampling steps.")
     parser.add_argument("--boundary-ratio", type=float, default=0.875, help="Boundary split ratio for MoE models.")
+    parser.add_argument(
+        "--frame-rate",
+        type=float,
+        default=None,
+        help="Optional generation frame rate (used by models like LTX2). Defaults to --fps.",
+    )
     parser.add_argument(
         "--flow-shift", type=float, default=5.0, help="Scheduler flow_shift (5.0 for 720p, 12.0 for 480p)."
     )
     parser.add_argument("--output", type=str, default="i2v_output.mp4", help="Path to save the video (mp4).")
-    parser.add_argument("--fps", type=int, default=16, help="Frames per second for the output video.")
+    parser.add_argument("--fps", type=int, default=None, help="Frames per second for the output video.")
     parser.add_argument(
         "--vae-use-slicing",
         action="store_true",
@@ -105,11 +125,34 @@ def parse_args() -> argparse.Namespace:
         choices=[1, 2],
         help="Number of GPUs used for classifier free guidance parallel size.",
     )
+    parser.add_argument(
+        "--vae-patch-parallel-size",
+        type=int,
+        default=1,
+        help="Number of GPUs used for VAE patch/tile parallelism (decode).",
+    )
     parser.add_argument(
         "--enforce-eager",
         action="store_true",
         help="Disable torch.compile and force eager execution.",
     )
+    parser.add_argument(
+        "--audio-sample-rate",
+        type=int,
+        default=24000,
+        help="Sample rate for audio output when saved (default: 24000 for LTX2).",
+    )
+    parser.add_argument(
+        "--cache-backend",
+        type=str,
+        default=None,
+        choices=["cache_dit", "tea_cache"],
+        help=(
+            "Cache backend to use for acceleration. "
+            "Options: 'cache_dit' (DBCache + SCM + TaylorSeer), 'tea_cache' (Timestep Embedding Aware Cache). "
+            "Default: None (no cache acceleration)."
+        ),
+    )
     parser.add_argument(
         "--use-hsdp",
         action="store_true",
@@ -136,10 +179,13 @@ def parse_args() -> argparse.Namespace:
     return parser.parse_args()
 
 
-def calculate_dimensions(image: PIL.Image.Image, max_area: int = 480 * 832) -> tuple[int, int]:
+def calculate_dimensions(
+    image: PIL.Image.Image,
+    max_area: int = 480 * 832,
+    mod_value: int = 16,
+) -> tuple[int, int]:
     """Calculate output dimensions maintaining aspect ratio."""
     aspect_ratio = image.height / image.width
-    mod_value = 16  # Must be divisible by 16
 
     height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value
     width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value
@@ -150,22 +196,66 @@ def calculate_dimensions(image: PIL.Image.Image, max_area: int = 480 * 832) -> t
 def main():
     args = parse_args()
     generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
+    model_name = str(args.model).lower() if args.model is not None else ""
+    model_class_name = args.model_class_name
+    is_ltx2 = "ltx2" in model_name or (model_class_name and "ltx2" in model_class_name.lower())
+    if model_class_name is None and is_ltx2:
+        model_class_name = "LTX2ImageToVideoPipeline"
 
     # Load input image
     image = PIL.Image.open(args.image).convert("RGB")
 
+    fps = args.fps if args.fps is not None else (24 if is_ltx2 else 16)
+    frame_rate = args.frame_rate if args.frame_rate is not None else float(fps)
+    guidance_scale = args.guidance_scale if args.guidance_scale is not None else (4.0 if is_ltx2 else 5.0)
+    num_frames = args.num_frames if args.num_frames is not None else (121 if is_ltx2 else 81)
+    num_inference_steps = args.num_inference_steps if args.num_inference_steps is not None else (40 if is_ltx2 else 50)
+
     # Calculate dimensions if not provided
     height = args.height
     width = args.width
     if height is None or width is None:
-        # Default to 480P area for I2V
+        # Default to 480P area for Wan2.2 I2V, 512x768 area for LTX2
-        calc_height, calc_width = calculate_dimensions(image, max_area=480 * 832)
+        max_area = 512 * 768 if is_ltx2 else 480 * 832
+        mod_value = 32 if is_ltx2 else 16
+        calc_height, calc_width = calculate_dimensions(image, max_area=max_area, mod_value=mod_value)
         height = height or calc_height
         width = width or calc_width
 
     # Resize image to target dimensions
     image = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
 
+    # Configure cache based on backend type
+    cache_config = None
+    if args.cache_backend == "cache_dit":
+        if is_ltx2:
+            cache_config = {
+                "Fn_compute_blocks": 2,
+                "Bn_compute_blocks": 0,
+                "max_warmup_steps": 8,
+                "residual_diff_threshold": 0.12,
+                "max_continuous_cached_steps": 1,
+                "max_cached_steps": 20,
+                "enable_taylorseer": False,
+                "scm_steps_mask_policy": None,
+            }
+        else:
+            cache_config = {
+                "Fn_compute_blocks": 1,
+                "Bn_compute_blocks": 0,
+                "max_warmup_steps": 4,
+                "residual_diff_threshold": 0.24,
+                "max_continuous_cached_steps": 3,
+                "enable_taylorseer": False,
+                "taylorseer_order": 1,
+                "scm_steps_mask_policy": None,
+                "scm_steps_policy": "dynamic",
+            }
+    elif args.cache_backend == "tea_cache":
+        cache_config = {
+            "rel_l1_thresh": 0.2,
+        }
+
     # Check if profiling is requested via environment variable
     profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
     parallel_config = DiffusionParallelConfig(
@@ -173,6 +263,7 @@ def main():
         ring_degree=args.ring_degree,
         cfg_parallel_size=args.cfg_parallel_size,
         tensor_parallel_size=args.tensor_parallel_size,
+        vae_patch_parallel_size=args.vae_patch_parallel_size,
         use_hsdp=args.use_hsdp,
         hsdp_shard_size=args.hsdp_shard_size,
         hsdp_replicate_size=args.hsdp_replicate_size,
@@ -187,6 +278,9 @@ def main():
         enable_cpu_offload=args.enable_cpu_offload,
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
+        model_class_name=model_class_name,
+        cache_backend=args.cache_backend,
+        cache_config=cache_config,
     )
 
     if profiler_enabled:
@@ -199,10 +293,14 @@ def main():
     print(f"  Model: {args.model}")
     print(f"  Inference steps: {args.num_inference_steps}")
     print(f"  Frames: {args.num_frames}")
-    print(f"  Parallel configuration: cfg_parallel_size={args.cfg_parallel_size}")
+    print(
+        f"  Parallel configuration: cfg_parallel_size={args.cfg_parallel_size},"
+        f" tensor_parallel_size={args.tensor_parallel_size}, vae_patch_parallel_size={args.vae_patch_parallel_size}"
+    )
     print(f"  Video size: {args.width}x{args.height}")
     print(f"{'=' * 60}\n")
 
+    generation_start = time.perf_counter()
     # omni.generate() returns Generator[OmniRequestOutput, None, None]
     frames = omni.generate(
         {
@@ -214,38 +312,69 @@ def main():
             height=height,
             width=width,
             generator=generator,
-            guidance_scale=args.guidance_scale,
+            guidance_scale=guidance_scale,
             guidance_scale_2=args.guidance_scale_high,
-            num_inference_steps=args.num_inference_steps,
+            num_inference_steps=num_inference_steps,
-            num_frames=args.num_frames,
+            num_frames=num_frames,
+            frame_rate=frame_rate,
         ),
     )
+    generation_end = time.perf_counter()
+    generation_time = generation_end - generation_start
 
-    # Extract video frames from OmniRequestOutput
+    # Print profiling results
-    if isinstance(frames, list) and len(frames) > 0:
+    print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
-        first_item = frames[0]
 
-        # Check if it's an OmniRequestOutput
+    audio = None
-        if hasattr(first_item, "final_output_type"):
+    if isinstance(frames, list):
-            if first_item.final_output_type != "image":
+        frames = frames[0] if frames else None
-                raise ValueError(
-                    f"Unexpected output type '{first_item.final_output_type}', expected 'image' for video generation."
-                )
 
-            # Pipeline mode: extract from nested request_output
+    if isinstance(frames, OmniRequestOutput):
-            if hasattr(first_item, "is_pipeline_output") and first_item.is_pipeline_output:
+        if frames.final_output_type != "image":
-                if isinstance(first_item.request_output, list) and len(first_item.request_output) > 0:
+            raise ValueError(
-                    inner_output = first_item.request_output[0]
+                f"Unexpected output type '{frames.final_output_type}', expected 'image' for video generation."
-                    if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"):
+            )
-                        frames = inner_output.images[0] if inner_output.images else None
+        if frames.multimodal_output and "audio" in frames.multimodal_output:
-                        if frames is None:
+            audio = frames.multimodal_output["audio"]
-                            raise ValueError("No video frames found in output.")
+        if frames.is_pipeline_output and frames.request_output is not None:
-            # Diffusion mode: use direct images field
+            inner_output = frames.request_output
-            elif hasattr(first_item, "images") and first_item.images:
+            if isinstance(inner_output, list):
-                frames = first_item.images
+                inner_output = inner_output[0] if inner_output else None
+            if isinstance(inner_output, OmniRequestOutput):
+                if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
+                    audio = inner_output.multimodal_output["audio"]
+                frames = inner_output
+        if isinstance(frames, OmniRequestOutput):
+            if frames.images:
+                if len(frames.images) == 1 and isinstance(frames.images[0], tuple) and len(frames.images[0]) == 2:
+                    frames, audio = frames.images[0]
+                elif len(frames.images) == 1 and isinstance(frames.images[0], dict):
+                    audio = frames.images[0].get("audio")
+                    frames = frames.images[0].get("frames") or frames.images[0].get("video")
+                else:
+                    frames = frames.images
             else:
                 raise ValueError("No video frames found in OmniRequestOutput.")
 
+    if isinstance(frames, list) and frames:
+        first_item = frames[0]
+        if isinstance(first_item, tuple) and len(first_item) == 2:
+            frames, audio = first_item
+        elif isinstance(first_item, dict):
+            audio = first_item.get("audio")
+            frames = first_item.get("frames") or first_item.get("video")
+        elif isinstance(first_item, list):
+            frames = first_item
+
+    if isinstance(frames, tuple) and len(frames) == 2:
+        frames, audio = frames
+    elif isinstance(frames, dict):
+        audio = frames.get("audio")
+        frames = frames.get("frames") or frames.get("video")
+
+    if frames is None:
+        raise ValueError("No video frames found in output.")
+
     output_path = Path(args.output)
     output_path.parent.mkdir(parents=True, exist_ok=True)
 
@@ -254,32 +383,128 @@ def main():
     except ImportError:
         raise ImportError("diffusers is required for export_to_video.")
 
-    # frames may be np.ndarray (preferred) or torch.Tensor
+    def _normalize_frame(frame):
+        if isinstance(frame, torch.Tensor):
+            frame_tensor = frame.detach().cpu()
+            if frame_tensor.dim() == 4 and frame_tensor.shape[0] == 1:
+                frame_tensor = frame_tensor[0]
+            if frame_tensor.dim() == 3 and frame_tensor.shape[0] in (3, 4):
+                frame_tensor = frame_tensor.permute(1, 2, 0)
+            if frame_tensor.is_floating_point():
+                frame_tensor = frame_tensor.clamp(-1, 1) * 0.5 + 0.5
+            return frame_tensor.float().numpy()
+        if isinstance(frame, np.ndarray):
+            frame_array = frame
+            if frame_array.ndim == 4 and frame_array.shape[0] == 1:
+                frame_array = frame_array[0]
+            if np.issubdtype(frame_array.dtype, np.integer):
+                frame_array = frame_array.astype(np.float32) / 255.0
+            return frame_array
+        try:
+            from PIL import Image
+        except ImportError:
+            Image = None
+        if Image is not None and isinstance(frame, Image.Image):
+            return np.asarray(frame).astype(np.float32) / 255.0
+        return frame
+
+    def _ensure_frame_list(video_array):
+        if isinstance(video_array, list):
+            if len(video_array) == 0:
+                return video_array
+            first_item = video_array[0]
+            if isinstance(first_item, np.ndarray):
+                if first_item.ndim == 5:
+                    return list(first_item[0])
+                if first_item.ndim == 4:
+                    if len(video_array) == 1:
+                        return list(first_item)
+                    return list(first_item)
+                if first_item.ndim == 3:
+                    return video_array
+            return video_array
+        if isinstance(video_array, np.ndarray):
+            if video_array.ndim == 5:
+                return list(video_array[0])
+            if video_array.ndim == 4:
+                return list(video_array)
+            if video_array.ndim == 3:
+                return [video_array]
+        return video_array
+
+    # frames may be np.ndarray, torch.Tensor, or list of tensors/arrays/images
     # export_to_video expects a list of frames with values in [0, 1]
     if isinstance(frames, torch.Tensor):
         video_tensor = frames.detach().cpu()
         if video_tensor.dim() == 5:
-            # [B, C, F, H, W] or [B, F, H, W, C]
             if video_tensor.shape[1] in (3, 4):
                 video_tensor = video_tensor[0].permute(1, 2, 3, 0)
             else:
                 video_tensor = video_tensor[0]
         elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4):
             video_tensor = video_tensor.permute(1, 2, 3, 0)
-        # If float, assume [-1,1] and normalize to [0,1]
         if video_tensor.is_floating_point():
             video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5
         video_array = video_tensor.float().numpy()
+    elif isinstance(frames, np.ndarray):
+        video_array = frames
+        if video_array.ndim == 5:
+            video_array = video_array[0]
+        if np.issubdtype(video_array.dtype, np.integer):
+            video_array = video_array.astype(np.float32) / 255.0
+    elif isinstance(frames, list):
+        if len(frames) == 0:
+            raise ValueError("No video frames found in output.")
+        video_array = [_normalize_frame(frame) for frame in frames]
     else:
         video_array = frames
-        if hasattr(video_array, "shape") and video_array.ndim == 5:
-            video_array = video_array[0]
 
-    # Convert 4D array (frames, H, W, C) to list of frames for export_to_video
+    video_array = _ensure_frame_list(video_array)
-    if isinstance(video_array, np.ndarray) and video_array.ndim == 4:
-        video_array = list(video_array)
 
-    export_to_video(video_array, str(output_path), fps=args.fps)
+    use_ltx2_export = is_ltx2
+    encode_video = None
+    if use_ltx2_export:
+        try:
+            from diffusers.pipelines.ltx2.export_utils import encode_video
+        except ImportError:
+            encode_video = None
+
+    if use_ltx2_export and encode_video is not None:
+        if isinstance(video_array, list):
+            frames_np = np.stack(video_array, axis=0)
+        elif isinstance(video_array, np.ndarray):
+            frames_np = video_array
+        else:
+            frames_np = np.asarray(video_array)
+
+        if frames_np.ndim == 4 and frames_np.shape[-1] == 4:
+            frames_np = frames_np[..., :3]
+
+        frames_np = np.clip(frames_np, 0.0, 1.0)
+        frames_u8 = (frames_np * 255).round().clip(0, 255).astype("uint8")
+        video_tensor = torch.from_numpy(frames_u8)
+
+        audio_out = None
+        if audio is not None:
+            if isinstance(audio, list):
+                audio = audio[0] if audio else None
+            if isinstance(audio, np.ndarray):
+                audio = torch.from_numpy(audio)
+            if isinstance(audio, torch.Tensor):
+                audio_out = audio
+                if audio_out.dim() > 1:
+                    audio_out = audio_out[0]
+                audio_out = audio_out.float().cpu()
+
+        encode_video(
+            video_tensor,
+            fps=fps,
+            audio=audio_out,
+            audio_sample_rate=args.audio_sample_rate if audio_out is not None else None,
+            output_path=str(output_path),
+        )
+    else:
+        export_to_video(video_array, str(output_path), fps=fps)
     print(f"Saved generated video to {output_path}")
 
     if profiler_enabled:

examples/offline_inference/mammothmodal2_preview/README.md

@@ -0,0 +1,32 @@
+# MammothModa2-Preview
+
+## Run examples (MammothModa2-Preview)
+
+Download model
+```bash
+hf bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-Preview
+```
+
+### Text-to-Image (T2I)
+
+```bash
+python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py \
+  --model ./MammothModa2-Preview \
+  --stage-config ./vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml \
+  --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
+  --height 1024 \
+  --width 1024 \
+  --num-inference-steps 50 \
+  --text-guidance-scale 4.0 \
+  --out output.png
+```
+
+### Image Summary
+
+```bash
+python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
+  --model ./MammothModa2-Preview \
+  --stage-config ./vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml \
+  --question "Summarize this image." \
+  --image ./image.png
+```

examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py

@@ -0,0 +1,115 @@
+"""
+Offline inference example: MammothModa2 image summarization (single AR stage).
+
+Example:
+  uv run python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
+    --model path/to/MammothModa2-Preview \
+    --stage-config vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml \
+    --image /path/to/input.jpg \
+    --question "Please summarize the content of this image."
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+
+from PIL import Image
+from vllm import SamplingParams
+from vllm.multimodal.image import convert_image_mode
+
+from vllm_omni import Omni
+
+DEFAULT_SYSTEM = "You are a helpful assistant."
+DEFAULT_QUESTION = "Please summarize the content of this image."
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="MammothModa2 image summarization (offline, AR only).")
+    parser.add_argument("--model", type=str, required=True, help="Path to model directory or model id.")
+    parser.add_argument(
+        "--stage-config", type=str, required=True, help="Path to stage config yaml (single-stage AR->text)."
+    )
+    parser.add_argument("--image", type=str, required=True, help="Path to input image.")
+    parser.add_argument("--question", type=str, default=DEFAULT_QUESTION, help="Question/instruction for the model.")
+    parser.add_argument("--system", type=str, default=DEFAULT_SYSTEM, help="System prompt.")
+    parser.add_argument(
+        "--max-tokens",
+        type=int,
+        default=512,
+        help="Max new tokens to generate.",
+    )
+    parser.add_argument("--temperature", type=float, default=0.2)
+    parser.add_argument("--top-p", type=float, default=0.9)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--trust-remote-code", action="store_true")
+    return parser.parse_args()
+
+
+def build_prompt(system: str, question: str) -> str:
+    return (
+        f"<|im_start|>system\n{system}<|im_end|>\n"
+        "<|im_start|>user\n"
+        "<|vision_start|><|image_pad|><|vision_end|>"
+        f"{question}<|im_end|>\n"
+        "<|im_start|>assistant\n"
+    )
+
+
+def main() -> None:
+    args = parse_args()
+
+    if not os.path.exists(args.image):
+        raise FileNotFoundError(f"Image file not found: {args.image}")
+
+    pil_image = Image.open(args.image)
+    image_data = convert_image_mode(pil_image, "RGB")
+    prompt = build_prompt(args.system, args.question)
+
+    omni = Omni(
+        model=args.model,
+        stage_configs_path=args.stage_config,
+        trust_remote_code=args.trust_remote_code,
+    )
+    try:
+        sp = SamplingParams(
+            temperature=float(args.temperature),
+            top_p=float(args.top_p),
+            top_k=-1,
+            max_tokens=int(args.max_tokens),
+            seed=int(args.seed),
+            detokenize=True,
+        )
+        # NOTE: omni.generate() returns a Generator[OmniRequestOutput, None, None].
+        # Consume it inside the try block so the worker isn't closed early.
+        outputs = list(
+            omni.generate(
+                [
+                    {
+                        "prompt": prompt,
+                        "multi_modal_data": {"image": image_data},
+                        "additional_information": {"omni_task": ["chat"]},
+                    }
+                ],
+                [sp],
+            )
+        )
+    finally:
+        omni.close()
+
+    lines: list[str] = []
+    for stage_outputs in outputs:
+        req_outputs = getattr(stage_outputs, "request_output", stage_outputs)
+        req_outputs = req_outputs if isinstance(req_outputs, list) else [req_outputs]
+        for ro in req_outputs:
+            text = ro.outputs[0].text if getattr(ro, "outputs", None) else str(ro)
+            lines.append(f"request_id: {getattr(ro, 'request_id', 'unknown')}\n")
+            lines.append("answer:\n")
+            lines.append(text.strip() + "\n")
+            lines.append("\n")
+
+    print("\n".join(lines))
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py

@@ -0,0 +1,251 @@
+"""
+Offline inference example for MammothModa2 Text-to-Image (T2I) generation.
+This script uses the vllm_omni.Omni pipeline with a multi-stage configuration.
+
+Workflow:
+1. Stage 0 (AR): Generates visual tokens and their corresponding hidden states.
+2. Stage 1 (DiT): Consumes the hidden states as conditions to perform diffusion
+   and VAE decoding to produce the final image.
+
+Example Usage:
+    uv run python examples/offline_inference/run_mammothmoda2_t2i.py \
+        --model path/to/MammothModa2-Preview \
+        --stage-config vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml \
+        --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
+        --out output.png
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import os
+from pathlib import Path
+from typing import NamedTuple
+
+import torch
+from PIL import Image
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni import Omni
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+logger = logging.getLogger(__name__)
+
+# ---------------------------------------------------------------------------
+# Constants
+# ---------------------------------------------------------------------------
+_PATCH_SIZE = 16  # AR image grid patch size (pixels per token)
+
+
+class T2IGenConfig(NamedTuple):
+    eol_token_id: int
+    visual_token_start_id: int
+    visual_token_end_id: int
+    top_k: int  # AR sampling top-k (covers the full visual generation vocabulary)
+    # Qwen2.5-VL special vision tokens: <|image_pad|>, <|video_pad|>, <|vision_start|>, <|vision_end|>
+    visual_ids: list[int]
+
+
+def load_t2i_generation_config(model_dir: str) -> T2IGenConfig:
+    """Load T2I token IDs from t2i_generation_config.json and config.json.
+
+    Supports both local directory paths and HuggingFace Hub model IDs.
+    """
+    model_path = Path(model_dir)
+
+    def _read_json(filename: str) -> dict:
+        local = model_path / filename
+        if local.exists():
+            with local.open(encoding="utf-8") as f:
+                return json.load(f)
+        # Fall back to HuggingFace Hub when model_dir is a repo ID.
+        try:
+            from huggingface_hub import hf_hub_download
+        except ImportError as exc:
+            raise ImportError(
+                "huggingface_hub is required to load configs from HF Hub. Install it with: pip install huggingface_hub"
+            ) from exc
+        cached = hf_hub_download(repo_id=model_dir, filename=filename)
+        with open(cached, encoding="utf-8") as f:
+            return json.load(f)
+
+    gen_cfg = _read_json("t2i_generation_config.json")
+    llm_cfg = _read_json("config.json").get("llm_config", {})
+
+    return T2IGenConfig(
+        eol_token_id=int(gen_cfg["eol_token_id"]),
+        visual_token_start_id=int(gen_cfg["visual_token_start_id"]),
+        visual_token_end_id=int(gen_cfg["visual_token_end_id"]),
+        top_k=int(gen_cfg["top_k"]),
+        visual_ids=[
+            int(llm_cfg["image_token_id"]),
+            int(llm_cfg["video_token_id"]),
+            int(llm_cfg["vision_start_token_id"]),
+            int(llm_cfg["vision_end_token_id"]),
+        ],
+    )
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description="Run MammothModa2 T2I (AR -> DiT) with vLLM-Omni.")
+    p.add_argument("--model", type=str, required=True, help="Path to the model directory.")
+    p.add_argument("--stage-config", type=str, required=True, help="Path to the multi-stage YAML configuration.")
+    p.add_argument(
+        "--prompt",
+        type=str,
+        action="append",
+        default=None,
+        help=(
+            "Text prompt for image generation. Can be provided multiple times "
+            "to generate multiple images with shared height/width/CFG settings."
+        ),
+    )
+    p.add_argument("--height", type=int, default=1024, help="Output image height (must be a multiple of 16).")
+    p.add_argument("--width", type=int, default=1024, help="Output image width (must be a multiple of 16).")
+    p.add_argument("--num-inference-steps", type=int, default=50, help="Number of diffusion steps for the DiT stage.")
+    p.add_argument(
+        "--text-guidance-scale", type=float, default=9.0, help="Classifier-Free Guidance (CFG) scale for DiT."
+    )
+    p.add_argument(
+        "--cfg-range",
+        type=float,
+        nargs=2,
+        default=(0.0, 1.0),
+        help="Relative step range [start, end] where CFG is active.",
+    )
+    p.add_argument("--out", type=str, default="output.png", help="Path to save the generated image.")
+    p.add_argument("--trust-remote-code", action="store_true", help="Trust remote code when loading the model.")
+    args = p.parse_args()
+    if not args.prompt:
+        args.prompt = ["A stylish woman with sunglasses riding a motorcycle in NYC."]
+    return args
+
+
+def tensor_to_pil(image: torch.Tensor) -> Image.Image:
+    """Convert a normalized torch tensor [-1, 1] to a PIL Image."""
+    if image.ndim == 4:
+        image = image[0]
+    image = image.detach().to("cpu")
+    image = (image / 2 + 0.5).clamp(0, 1)
+    image = (image * 255).to(torch.uint8)
+    image = image.permute(1, 2, 0).contiguous().numpy()
+    return Image.fromarray(image)
+
+
+def _format_prompt(user_prompt: str, ar_width: int, ar_height: int) -> str:
+    """Build the AR-stage prompt string including the image grid header."""
+    return (
+        "<|im_start|>system\nYou are a helpful image generator.<|im_end|>\n"
+        f"<|im_start|>user\n{user_prompt}<|im_end|>\n"
+        "<|im_start|>assistant\n"
+        f"<|image start|>{ar_width}*{ar_height}<|image token|>"
+    )
+
+
+def _collect_images(outputs: list) -> list[torch.Tensor]:
+    """Extract all image tensors produced by the final (DiT) stage."""
+    images: list[torch.Tensor] = []
+    for out in outputs:
+        ro_list = getattr(out, "request_output", out)
+        if not isinstance(ro_list, list):
+            ro_list = [ro_list]
+        for ro_item in ro_list:
+            for completion in getattr(ro_item, "outputs", None) or []:
+                mm = getattr(completion, "multimodal_output", None)
+                if not isinstance(mm, dict) or "image" not in mm:
+                    raise RuntimeError(f"Missing image in multimodal output: {mm}")
+                payload = mm["image"]
+                for tensor in payload if isinstance(payload, list) else [payload]:
+                    if not isinstance(tensor, torch.Tensor):
+                        raise TypeError(f"Expected image tensor, got {type(tensor)}")
+                    images.append(tensor)
+    return images
+
+
+def _save_images(images: list[torch.Tensor], out_path: str) -> list[str]:
+    """Save image tensors to disk.
+
+    Single image: written to *out_path* exactly.
+    Multiple images: suffixed as ``<base>_0<ext>``, ``<base>_1<ext>``, …
+    """
+    if not images:
+        raise RuntimeError("No images to save.")
+    base, ext = os.path.splitext(out_path)
+    ext = ext or ".png"
+    paths = []
+    for i, tensor in enumerate(images):
+        path = out_path if len(images) == 1 else f"{base}_{i}{ext}"
+        tensor_to_pil(tensor).save(path)
+        paths.append(path)
+    return paths
+
+
+def main() -> None:
+    args = parse_args()
+    os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
+
+    if args.height <= 0 or args.width <= 0:
+        raise ValueError(f"Height and width must be positive, got {args.height}x{args.width}")
+    if args.height % _PATCH_SIZE != 0 or args.width % _PATCH_SIZE != 0:
+        raise ValueError(f"Height and width must be multiples of {_PATCH_SIZE}, got {args.height}x{args.width}")
+
+    ar_height = args.height // _PATCH_SIZE
+    ar_width = args.width // _PATCH_SIZE
+    gen_cfg = load_t2i_generation_config(args.model)
+    expected_grid_tokens = ar_height * (ar_width + 1)
+
+    logger.info("Initializing Omni pipeline...")
+    omni = Omni(model=args.model, stage_configs_path=args.stage_config, trust_remote_code=args.trust_remote_code)
+    try:
+        ar_sampling = SamplingParams(
+            temperature=1.0,
+            top_p=1.0,
+            top_k=gen_cfg.top_k,
+            max_tokens=max(1, expected_grid_tokens + 1),  # +1 for hidden state of eoi
+            detokenize=False,
+        )
+        dit_sampling = SamplingParams(
+            temperature=0.0,
+            top_p=1.0,
+            top_k=-1,
+            max_tokens=1,
+            detokenize=False,
+        )
+
+        additional_information = {
+            "omni_task": ["t2i"],
+            "ar_width": [ar_width],
+            "ar_height": [ar_height],
+            "eol_token_id": [gen_cfg.eol_token_id],
+            "visual_token_start_id": [gen_cfg.visual_token_start_id],
+            "visual_token_end_id": [gen_cfg.visual_token_end_id],
+            "image_height": [args.height],
+            "image_width": [args.width],
+            "num_inference_steps": [args.num_inference_steps],
+            "text_guidance_scale": [args.text_guidance_scale],
+            "cfg_range": [args.cfg_range[0], args.cfg_range[1]],
+            "visual_ids": gen_cfg.visual_ids,
+        }
+        inputs = [
+            {
+                "prompt": _format_prompt(p, ar_width, ar_height),
+                "additional_information": dict(additional_information),
+            }
+            for p in args.prompt
+        ]
+
+        logger.info("Starting generation...")
+        # omni.generate() returns a Generator; consume it to run the full pipeline.
+        outputs = list(omni.generate(inputs, [ar_sampling, dit_sampling]))
+
+        logger.info("Post-processing and saving image(s)...")
+        for path in _save_images(_collect_images(outputs), args.out):
+            logger.info(f"Saved: {path}")
+    finally:
+        omni.close()
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/qwen3_omni/README.md

@@ -64,6 +64,51 @@ If media file paths are not provided, the script will use default assets. Suppor
 - `multi_audios`: Multiple audio inputs
 - `mixed_modalities`: Combination of video, image, and audio inputs
 
+### Async-chunk (offline)
+
+For true stage-level concurrency -- where downstream stages (Talker, Code2Wav)
+start **before** the upstream stage (Thinker) finishes -- use the async_chunk
+example. This requires:
+
+1. A stage config YAML with ``async_chunk: true`` (e.g.
+   ``qwen3_omni_moe_async_chunk.yaml``).
+2. Hardware that matches the config (e.g. 2x H100 for the default 3-stage
+   config).
+
+The async_chunk example uses ``AsyncOmni`` instead of the synchronous ``Omni``
+class, which enables the async orchestrator to receive stage-0 intermediate
+outputs and trigger downstream stages early. Chunk data flows directly between
+stage workers via the in-worker ``OmniChunkTransferAdapter`` / connector,
+**not** through the orchestrator.
+
+#### Single prompt
+```bash
+cd examples/offline_inference/qwen3_omni
+bash run_single_prompt_async_chunk.sh
+```
+
+#### Multiple prompts with concurrency control
+```bash
+bash run_multiple_prompts_async_chunk.sh --max-in-flight 4
+```
+
+#### Text-only output (skip audio generation)
+```bash
+python end2end_async_chunk.py --query-type text --modalities text
+```
+
+#### Custom stage config
+```bash
+python end2end_async_chunk.py \
+    --query-type use_audio \
+    --stage-configs-path /path/to/your_async_chunk.yaml
+```
+
+> **Note**: The synchronous ``end2end.py`` (using ``Omni``) is still the
+> recommended entry point for non-async-chunk workflows. Only use the
+> async_chunk example when you need the stage-level concurrency semantics
+> described in PR #962 / #1151.
+
 ### FAQ
 
 If you encounter error about backend of librosa, try to install ffmpeg with command below.

examples/offline_inference/qwen3_omni/end2end_async_chunk.py

@@ -0,0 +1,614 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline inference with async_chunk enabled via AsyncOmni.
+
+This script uses AsyncOmni (the async orchestrator) to run offline inference
+with async_chunk semantics: downstream stages (Talker, Code2Wav) start
+*before* upstream stages finish, consuming chunks as they arrive via
+the in-worker OmniChunkTransferAdapter / connector.
+
+Compared to the synchronous ``end2end.py`` (which uses ``Omni``), this
+entry point achieves true stage-level concurrency -- stage-1/2 are
+actively processing while stage-0 is still generating.
+
+Usage
+-----
+    python end2end_async_chunk.py --query-type use_audio \
+        --stage-configs-path <path-to-async-chunk-yaml>
+
+See ``--help`` for all options.
+"""
+
+import asyncio
+import logging
+import os
+import time
+import uuid
+from typing import NamedTuple
+
+import numpy as np
+import soundfile as sf
+import torch
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+import librosa
+from PIL import Image
+from vllm import SamplingParams
+from vllm.assets.audio import AudioAsset
+from vllm.assets.image import ImageAsset
+from vllm.assets.video import VideoAsset, video_to_ndarrays
+from vllm.multimodal.image import convert_image_mode
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+logger = logging.getLogger(__name__)
+
+# ---------------------------------------------------------------------------
+# Query builders (reuse the patterns from end2end.py)
+# ---------------------------------------------------------------------------
+
+default_system = (
+    "You are Qwen, a virtual human developed by the Qwen Team, Alibaba "
+    "Group, capable of perceiving auditory and visual inputs, as well as "
+    "generating text and speech."
+)
+
+
+class QueryResult(NamedTuple):
+    inputs: dict
+    limit_mm_per_prompt: dict[str, int]
+
+
+def get_text_query(question: str = None) -> QueryResult:
+    if question is None:
+        question = "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
+    prompt = (
+        f"<|im_start|>system\n{default_system}<|im_end|>\n"
+        "<|im_start|>user\n"
+        f"{question}<|im_end|>\n"
+        "<|im_start|>assistant\n"
+    )
+    return QueryResult(inputs={"prompt": prompt}, limit_mm_per_prompt={})
+
+
+def get_audio_query(
+    question: str = None,
+    audio_path: str | None = None,
+    sampling_rate: int = 16000,
+) -> QueryResult:
+    if question is None:
+        question = "What is the content of this audio?"
+    prompt = (
+        f"<|im_start|>system\n{default_system}<|im_end|>\n"
+        "<|im_start|>user\n<|audio_start|><|audio_pad|><|audio_end|>"
+        f"{question}<|im_end|>\n"
+        "<|im_start|>assistant\n"
+    )
+    if audio_path:
+        if not os.path.exists(audio_path):
+            raise FileNotFoundError(f"Audio file not found: {audio_path}")
+        audio_signal, sr = librosa.load(audio_path, sr=sampling_rate)
+        audio_data = (audio_signal.astype(np.float32), sr)
+    else:
+        audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
+    return QueryResult(
+        inputs={
+            "prompt": prompt,
+            "multi_modal_data": {"audio": audio_data},
+        },
+        limit_mm_per_prompt={"audio": 1},
+    )
+
+
+def get_image_query(question: str = None, image_path: str | None = None) -> QueryResult:
+    if question is None:
+        question = "What is the content of this image?"
+    prompt = (
+        f"<|im_start|>system\n{default_system}<|im_end|>\n"
+        "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>"
+        f"{question}<|im_end|>\n"
+        "<|im_start|>assistant\n"
+    )
+    if image_path:
+        if not os.path.exists(image_path):
+            raise FileNotFoundError(f"Image file not found: {image_path}")
+        pil_image = Image.open(image_path)
+        image_data = convert_image_mode(pil_image, "RGB")
+    else:
+        image_data = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB")
+    return QueryResult(
+        inputs={
+            "prompt": prompt,
+            "multi_modal_data": {"image": image_data},
+        },
+        limit_mm_per_prompt={"image": 1},
+    )
+
+
+def get_video_query(
+    question: str = None,
+    video_path: str | None = None,
+    num_frames: int = 16,
+) -> QueryResult:
+    if question is None:
+        question = "Why is this video funny?"
+    prompt = (
+        f"<|im_start|>system\n{default_system}<|im_end|>\n"
+        "<|im_start|>user\n<|vision_start|><|video_pad|><|vision_end|>"
+        f"{question}<|im_end|>\n"
+        "<|im_start|>assistant\n"
+    )
+    if video_path:
+        if not os.path.exists(video_path):
+            raise FileNotFoundError(f"Video file not found: {video_path}")
+        video_frames = video_to_ndarrays(video_path, num_frames=num_frames)
+    else:
+        video_frames = VideoAsset(name="baby_reading", num_frames=num_frames).np_ndarrays
+    return QueryResult(
+        inputs={
+            "prompt": prompt,
+            "multi_modal_data": {"video": video_frames},
+        },
+        limit_mm_per_prompt={"video": 1},
+    )
+
+
+query_map = {
+    "text": get_text_query,
+    "use_audio": get_audio_query,
+    "use_image": get_image_query,
+    "use_video": get_video_query,
+}
+
+# ---------------------------------------------------------------------------
+# Core async routine
+# ---------------------------------------------------------------------------
+
+
+def clone_prompt_for_request(template: dict) -> dict:
+    """Shallow-clone prompt dict so concurrent requests own independent containers."""
+    cloned = dict(template)
+    for key in ("multi_modal_data", "mm_processor_kwargs", "additional_information"):
+        value = template.get(key)
+        if isinstance(value, dict):
+            cloned[key] = dict(value)
+        elif isinstance(value, list):
+            cloned[key] = list(value)
+    return cloned
+
+
+def _default_async_chunk_stage_configs_path() -> str | None:
+    """Best-effort default stage config for running Qwen3-Omni with async_chunk.
+
+    When this example is executed from within the repository, we resolve the
+    default YAML path relative to this file. When installed elsewhere, the
+    file may not exist and callers should pass --stage-configs-path explicitly.
+    """
+    repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))
+    candidate = os.path.join(
+        repo_root,
+        "vllm_omni",
+        "model_executor",
+        "stage_configs",
+        "qwen3_omni_moe_async_chunk.yaml",
+    )
+    return candidate if os.path.exists(candidate) else None
+
+
+async def run_single_request(
+    async_omni: AsyncOmni,
+    prompt: dict,
+    request_id: str,
+    sampling_params_list: list[SamplingParams] | None,
+    output_dir: str,
+    output_modalities: list[str] | None = None,
+    stream_audio_to_disk: bool = False,
+) -> dict:
+    """Run one request through AsyncOmni and collect outputs.
+
+    Returns a dict with timing information and saved file paths.
+    """
+    t_start = time.perf_counter()
+    text_parts: list[str] = []
+    audio_chunks: list[torch.Tensor] = []
+    audio_sr: int | None = None
+    first_audio_ts: float | None = None
+    audio_list_consumed: int = 0
+    audio_last_tensor: torch.Tensor | None = None
+    stage_0_first_output_ts: float | None = None
+
+    samplerate = 24000
+    wav_file = os.path.join(output_dir, f"output_{request_id}.wav")
+    sf_writer: sf.SoundFile | None = None
+    audio_samples_written: int = 0
+
+    try:
+        async for omni_output in async_omni.generate(
+            prompt=prompt,
+            request_id=request_id,
+            sampling_params_list=sampling_params_list,
+            output_modalities=output_modalities,
+        ):
+            if not isinstance(omni_output.request_output, list):
+                outputs_list = [omni_output.request_output]
+            else:
+                outputs_list = omni_output.request_output
+
+            for output in outputs_list:
+                if omni_output.final_output_type == "text":
+                    if stage_0_first_output_ts is None:
+                        stage_0_first_output_ts = time.perf_counter()
+                    text_output = output.outputs[0].text
+                    if output.finished:
+                        text_parts.append(text_output)
+                elif omni_output.final_output_type == "audio":
+                    mm_out = output.outputs[0].multimodal_output
+                    if mm_out and "audio" in mm_out:
+                        if first_audio_ts is None:
+                            first_audio_ts = time.perf_counter()
+                        if audio_sr is None and "sr" in mm_out:
+                            sr_val = mm_out["sr"]
+                            audio_sr = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
+                            samplerate = audio_sr
+                        audio_data = mm_out["audio"]
+                        if isinstance(audio_data, list):
+                            new_chunks = audio_data[audio_list_consumed:]
+                            audio_list_consumed = len(audio_data)
+                        elif isinstance(audio_data, torch.Tensor):
+                            new_chunks = [audio_data]
+                            audio_last_tensor = audio_data
+                        else:
+                            new_chunks = []
+
+                        if stream_audio_to_disk and new_chunks:
+                            if sf_writer is None:
+                                sf_writer = sf.SoundFile(
+                                    wav_file,
+                                    mode="w",
+                                    samplerate=samplerate,
+                                    channels=1,
+                                    subtype="FLOAT",
+                                )
+                            for chunk in new_chunks:
+                                chunk_np = chunk.float().detach().cpu().numpy().flatten()
+                                sf_writer.write(chunk_np)
+                                audio_samples_written += len(chunk_np)
+                        else:
+                            audio_chunks.extend(new_chunks)
+    finally:
+        if sf_writer is not None:
+            sf_writer.close()
+
+    t_end = time.perf_counter()
+    result = {
+        "request_id": request_id,
+        "e2e_latency_s": t_end - t_start,
+        "saved_files": [],
+    }
+
+    # Save text output
+    if text_parts:
+        text_file = os.path.join(output_dir, f"{request_id}.txt")
+        with open(text_file, "w", encoding="utf-8") as f:
+            f.write("\n".join(text_parts))
+        result["saved_files"].append(text_file)
+        print(
+            f"[Request {request_id}] Text saved to {text_file} "
+            f"(stage-0 first output at {stage_0_first_output_ts - t_start:.3f}s)"
+        )
+
+    # Save audio output
+    if stream_audio_to_disk and audio_samples_written > 0:
+        result["saved_files"].append(wav_file)
+        result["audio_duration_s"] = audio_samples_written / samplerate
+        result["num_audio_chunks"] = audio_list_consumed
+        ttfa = (first_audio_ts - t_start) if first_audio_ts else None
+        result["time_to_first_audio_s"] = ttfa
+        ttfa_str = f"{ttfa:.3f}s" if ttfa is not None else "N/A"
+        print(
+            f"[Request {request_id}] Audio streamed to {wav_file} "
+            f"(duration={result['audio_duration_s']:.2f}s, "
+            f"TTFA={ttfa_str}, "
+            f"e2e={result['e2e_latency_s']:.3f}s)"
+        )
+    elif audio_chunks or audio_last_tensor is not None:
+        if audio_chunks:
+            if len(audio_chunks) > 1:
+                audio_tensor = torch.cat(audio_chunks, dim=-1)
+            else:
+                audio_tensor = audio_chunks[0]
+        else:
+            audio_tensor = audio_last_tensor
+        audio_numpy = audio_tensor.float().detach().cpu().numpy()
+        if audio_numpy.ndim > 1:
+            audio_numpy = audio_numpy.flatten()
+        sf.write(wav_file, audio_numpy, samplerate=samplerate, format="WAV")
+        result["saved_files"].append(wav_file)
+        result["audio_duration_s"] = len(audio_numpy) / samplerate
+        result["num_audio_chunks"] = len(audio_chunks)
+        ttfa = (first_audio_ts - t_start) if first_audio_ts else None
+        result["time_to_first_audio_s"] = ttfa
+        ttfa_str = f"{ttfa:.3f}s" if ttfa is not None else "N/A"
+        print(
+            f"[Request {request_id}] Audio saved to {wav_file} "
+            f"({len(audio_chunks)} chunks, "
+            f"duration={result['audio_duration_s']:.2f}s, "
+            f"TTFA={ttfa_str}, "
+            f"e2e={result['e2e_latency_s']:.3f}s)"
+        )
+
+    return result
+
+
+async def run_all(args):
+    """Main async entry: build prompts, create AsyncOmni, run requests."""
+    # Build query
+    query_func = query_map[args.query_type]
+    if args.query_type == "use_video":
+        query_result = query_func(
+            video_path=getattr(args, "video_path", None),
+            num_frames=getattr(args, "num_frames", 16),
+        )
+    elif args.query_type == "use_image":
+        query_result = query_func(image_path=getattr(args, "image_path", None))
+    elif args.query_type == "use_audio":
+        query_result = query_func(
+            audio_path=getattr(args, "audio_path", None),
+            sampling_rate=getattr(args, "sampling_rate", 16000),
+        )
+    else:
+        query_result = query_func()
+
+    # Build prompt list
+    if args.txt_prompts is not None:
+        assert args.query_type == "text", "txt-prompts is only supported for text query type"
+        with open(args.txt_prompts, encoding="utf-8") as f:
+            lines = [ln.strip() for ln in f if ln.strip()]
+        prompts = [get_text_query(ln).inputs for ln in lines]
+        print(f"[Info] Loaded {len(prompts)} prompts from {args.txt_prompts}")
+    else:
+        prompts = [clone_prompt_for_request(query_result.inputs) for _ in range(args.num_prompts)]
+
+    # Inject output modalities if specified
+    output_modalities = None
+    if args.modalities is not None:
+        output_modalities = args.modalities.split(",")
+        for prompt in prompts:
+            prompt["modalities"] = output_modalities
+
+    # Create AsyncOmni
+    print(f"[Info] Creating AsyncOmni with stage_configs_path={args.stage_configs_path}")
+    async_omni = None
+    try:
+        async_omni = AsyncOmni(
+            model=args.model,
+            stage_configs_path=args.stage_configs_path,
+            log_stats=args.log_stats,
+            stage_init_timeout=args.stage_init_timeout,
+        )
+
+        # Use default sampling params from stage config (they are pre-configured
+        # in the YAML for each stage).
+        sampling_params_list = None
+
+        output_dir = args.output_dir
+        os.makedirs(output_dir, exist_ok=True)
+
+        # Run requests with concurrency control
+        semaphore = asyncio.Semaphore(args.max_in_flight)
+        request_timeout = getattr(args, "request_timeout_s", None)
+        stream_audio = getattr(args, "stream_audio_to_disk", False)
+
+        async def _run_one(idx: int, prompt: dict):
+            async with semaphore:
+                request_id = f"req_{idx}_{uuid.uuid4().hex[:8]}"
+                coro = run_single_request(
+                    async_omni=async_omni,
+                    prompt=prompt,
+                    request_id=request_id,
+                    sampling_params_list=sampling_params_list,
+                    output_dir=output_dir,
+                    output_modalities=output_modalities,
+                    stream_audio_to_disk=stream_audio,
+                )
+                if request_timeout and request_timeout > 0:
+                    return await asyncio.wait_for(coro, timeout=request_timeout)
+                return await coro
+
+        wall_start = time.perf_counter()
+        tasks = [_run_one(i, p) for i, p in enumerate(prompts)]
+        all_results = await asyncio.gather(*tasks, return_exceptions=True)
+        wall_end = time.perf_counter()
+
+        # Print summary
+        print("\n" + "=" * 60)
+        print("Summary")
+        print("=" * 60)
+        success_count = 0
+        total_audio_dur = 0.0
+        for r in all_results:
+            if isinstance(r, Exception):
+                print(f"  [ERROR] {type(r).__name__}: {r}")
+            else:
+                success_count += 1
+                total_audio_dur += r.get("audio_duration_s", 0.0)
+                print(f"  [{r['request_id']}] e2e={r['e2e_latency_s']:.3f}s  files={r['saved_files']}")
+        wall_time = wall_end - wall_start
+        print(f"\nTotal: {success_count}/{len(prompts)} succeeded")
+        print(f"Wall time: {wall_time:.3f}s")
+        if total_audio_dur > 0:
+            print(f"Total audio duration: {total_audio_dur:.2f}s")
+            print(f"Real-time factor: {total_audio_dur / wall_time:.2f}x")
+        print("=" * 60)
+    finally:
+        if async_omni is not None:
+            async_omni.shutdown()
+
+
+# ---------------------------------------------------------------------------
+# CLI
+# ---------------------------------------------------------------------------
+
+
+def parse_args():
+    parser = FlexibleArgumentParser(
+        description=(
+            "Offline inference with async_chunk enabled via AsyncOmni. "
+            "Downstream stages start before upstream stages finish, "
+            "achieving true stage-level concurrency."
+        )
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="Qwen/Qwen3-Omni-30B-A3B-Instruct",
+        help="Model name or path.",
+    )
+    parser.add_argument(
+        "--query-type",
+        "-q",
+        type=str,
+        default="use_audio",
+        choices=query_map.keys(),
+        help="Query type.",
+    )
+    parser.add_argument(
+        "--stage-configs-path",
+        type=str,
+        default=_default_async_chunk_stage_configs_path(),
+        help=(
+            "Path to an async_chunk stage config YAML. "
+            "If not set, uses the model's default config "
+            "(make sure it has async_chunk: true)."
+        ),
+    )
+    parser.add_argument(
+        "--log-stats",
+        action="store_true",
+        default=False,
+        help="Enable writing detailed statistics.",
+    )
+    parser.add_argument(
+        "--stage-init-timeout",
+        type=int,
+        default=300,
+        help="Timeout for initializing a single stage (seconds).",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default="output_audio_async_chunk",
+        help="Directory to save output files.",
+    )
+    parser.add_argument(
+        "--num-prompts",
+        type=int,
+        default=1,
+        help="Number of prompts to generate (duplicated from query).",
+    )
+    parser.add_argument(
+        "--txt-prompts",
+        type=str,
+        default=None,
+        help="Path to a .txt file with one prompt per line.",
+    )
+    parser.add_argument(
+        "--max-in-flight",
+        type=int,
+        default=1,
+        help="Maximum concurrent requests (default: 1).",
+    )
+    parser.add_argument(
+        "--request-timeout-s",
+        type=float,
+        default=None,
+        help=(
+            "Per-request timeout in seconds. When set, a request that "
+            "exceeds this duration is cancelled and reported as an error. "
+            "Default: None (no timeout)."
+        ),
+    )
+    parser.add_argument(
+        "--batch-timeout-s",
+        type=float,
+        default=None,
+        help=(
+            "Global timeout for the entire batch in seconds. When set, "
+            "the whole run_all() is cancelled if it exceeds this duration. "
+            "Default: None (no global timeout)."
+        ),
+    )
+    parser.add_argument(
+        "--stream-audio-to-disk",
+        action="store_true",
+        default=False,
+        help=(
+            "Write audio chunks to WAV incrementally instead of "
+            "accumulating in memory. Useful for very long audio or "
+            "high --max-in-flight to reduce memory footprint."
+        ),
+    )
+    parser.add_argument(
+        "--modalities",
+        type=str,
+        default=None,
+        help="Comma-separated output modalities filter (e.g. 'text', 'audio', 'text,audio').",
+    )
+    parser.add_argument(
+        "--audio-path",
+        "-a",
+        type=str,
+        default=None,
+        help="Path to local audio file.",
+    )
+    parser.add_argument(
+        "--image-path",
+        "-i",
+        type=str,
+        default=None,
+        help="Path to local image file.",
+    )
+    parser.add_argument(
+        "--video-path",
+        "-v",
+        type=str,
+        default=None,
+        help="Path to local video file.",
+    )
+    parser.add_argument(
+        "--num-frames",
+        type=int,
+        default=16,
+        help="Number of frames to extract from video.",
+    )
+    parser.add_argument(
+        "--sampling-rate",
+        type=int,
+        default=16000,
+        help="Sampling rate for audio loading.",
+    )
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+
+    async def _main():
+        batch_timeout = getattr(args, "batch_timeout_s", None)
+        if batch_timeout and batch_timeout > 0:
+            await asyncio.wait_for(run_all(args), timeout=batch_timeout)
+        else:
+            await run_all(args)
+
+    try:
+        asyncio.run(_main())
+    except asyncio.TimeoutError:
+        print(
+            f"\n[TIMEOUT] Batch exceeded --batch-timeout-s="
+            f"{args.batch_timeout_s}s. AsyncOmni shutdown was handled "
+            f"by finally block."
+        )
+    except KeyboardInterrupt:
+        print("\nInterrupted by user. AsyncOmni shutdown was handled by finally block.")

examples/offline_inference/qwen3_omni/run_multiple_prompts_async_chunk.sh

@@ -0,0 +1,23 @@
+#!/bin/bash
+# Run multiple Qwen3-Omni requests with async_chunk enabled.
+#
+# Uses AsyncOmni with --max-in-flight to control request-level
+# concurrency (each request still gets true stage-level concurrency
+# via async_chunk).
+#
+# Usage:
+#   bash run_multiple_prompts_async_chunk.sh
+#   bash run_multiple_prompts_async_chunk.sh --max-in-flight 4
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
+
+python "${SCRIPT_DIR}/end2end_async_chunk.py" \
+    --query-type text \
+    --txt-prompts "${SCRIPT_DIR}/text_prompts_10.txt" \
+    --stage-configs-path "${REPO_ROOT}/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml" \
+    --output-dir output_audio_async_chunk \
+    --max-in-flight 2 \
+    "$@"

examples/offline_inference/qwen3_omni/run_single_prompt_async_chunk.sh

@@ -0,0 +1,26 @@
+#!/bin/bash
+# Run a single Qwen3-Omni request with async_chunk enabled.
+#
+# This uses AsyncOmni (async orchestrator) so that downstream stages
+# (Talker, Code2Wav) start *before* stage-0 (Thinker) finishes,
+# achieving true stage-level concurrency via chunk-level streaming.
+#
+# Prerequisites:
+#   - An async_chunk stage config YAML (e.g. qwen3_omni_moe_async_chunk.yaml)
+#   - Hardware matching the config (e.g. 2x H100 for the default 3-stage config)
+#
+# Usage:
+#   bash run_single_prompt_async_chunk.sh
+#   bash run_single_prompt_async_chunk.sh --query-type text --modalities text
+#   bash run_single_prompt_async_chunk.sh --stage-configs-path /path/to/custom.yaml
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
+
+python "${SCRIPT_DIR}/end2end_async_chunk.py" \
+    --query-type use_audio \
+    --stage-configs-path "${REPO_ROOT}/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml" \
+    --output-dir output_audio_async_chunk \
+    "$@"

examples/offline_inference/qwen3_tts/README.md

@@ -95,7 +95,8 @@ Add `--streaming` to stream audio chunks progressively via `AsyncOmni` (requires
 python end2end.py --query-type CustomVoice --streaming --output-dir /tmp/out_stream
 ```
 
-Each 25-frame Code2Wav chunk is logged as it arrives. The final WAV file is written once generation
+Each Code2Wav chunk is logged as it arrives (default 25 frames; configurable via `codec_chunk_frames`
+and `initial_codec_chunk_frames` in the stage config). The final WAV file is written once generation
 completes. This demonstrates that audio data is available progressively rather than only at the end.
 
 > **Note:** Streaming uses `AsyncOmni` internally. The non-streaming path (`Omni`) is unchanged.

examples/offline_inference/qwen3_tts/end2end.py

@@ -7,6 +7,7 @@ tasks, then runs Omni generation and saves output wav files.
 import asyncio
 import logging
 import os
+import time
 from typing import Any, NamedTuple
 
 import soundfile as sf
@@ -78,10 +79,16 @@ def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult:
     task_type = "CustomVoice"
     model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
     if use_batch_sample:
-        texts = ["其实我真的有发现，我是一个特别善于观察别人情绪的人。", "She said she would be here by noon."]
+        texts = [
-        instructs = ["", "Very happy."]
+            "其实我真的有发现，我是一个特别善于观察别人情绪的人。",
-        languages = ["Chinese", "English"]
+            "She said she would be here by noon.",
-        speakers = ["Vivian", "Ryan"]
+            "I like you very much.",
+            "Really, you do?",
+            "Yes, absolutely.",
+        ]
+        instructs = ["", "Very happy.", "Very happy.", "Very happy.", "Very happy."]
+        languages = ["Chinese", "English", "English", "English", "English"]
+        speakers = ["Vivian", "Ryan", "Ryan", "Ryan", "Ryan"]
         inputs = []
         for text, instruct, language, speaker in zip(texts, instructs, languages, speakers):
             additional_information = {
@@ -337,13 +344,27 @@ async def main_streaming(args):
 
     for i, prompt in enumerate(inputs):
         request_id = str(i)
+        t_start = time.perf_counter()
+        t_prev = t_start
+        chunk_idx = 0
         async for stage_output in omni.generate(prompt, request_id=request_id):
             mm = stage_output.request_output.outputs[0].multimodal_output
             if not stage_output.finished:
+                t_now = time.perf_counter()
                 audio = mm.get("audio")
                 n = len(audio) if isinstance(audio, list) else (0 if audio is None else 1)
-                logger.info(f"Request {request_id}: received chunk {n}")
+                dt_ms = (t_now - t_prev) * 1000
+                ttfa_ms = (t_now - t_start) * 1000
+                if chunk_idx == 0:
+                    logger.info(f"Request {request_id}: chunk {chunk_idx} samples={n} TTFA={ttfa_ms:.1f}ms")
+                else:
+                    logger.info(f"Request {request_id}: chunk {chunk_idx} samples={n} inter_chunk={dt_ms:.1f}ms")
+                t_prev = t_now
+                chunk_idx += 1
             else:
+                t_end = time.perf_counter()
+                total_ms = (t_end - t_start) * 1000
+                logger.info(f"Request {request_id}: done total={total_ms:.1f}ms chunks={chunk_idx}")
                 _save_wav(output_dir, request_id, mm)
 
 

examples/offline_inference/text_to_video/text_to_video.md

@@ -1,6 +1,7 @@
 # Text-To-Video
 
-The `Wan-AI/Wan2.2-T2V-A14B-Diffusers` pipeline generates short videos from text prompts.
+The `Wan-AI/Wan2.2-T2V-A14B-Diffusers` pipeline generates short videos from text prompts. This script can also be used
+for `Lightricks/LTX-2` to generate video+audio.
 
 ## Local CLI Usage
 
@@ -19,6 +20,22 @@ python text_to_video.py \
   --output t2v_out.mp4
 ```
 
+LTX2 example:
+
+```bash
+python text_to_video.py \
+  --model "Lightricks/LTX-2" \
+  --prompt "A cinematic close-up of ocean waves at golden hour." \
+  --negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \
+  --height 512 \
+  --width 768 \
+  --num-frames 121 \
+  --num-inference-steps 40 \
+  --guidance-scale 4.0 \
+  --frame-rate 24 \
+  --output ltx2_out.mp4
+```
+
 Key arguments:
 
 - `--prompt`: text description (string).
@@ -32,6 +49,9 @@ Key arguments:
 - `--vae-use-slicing`: enable VAE slicing for memory optimization.
 - `--vae-use-tiling`: enable VAE tiling for memory optimization.
 - `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--tensor-parallel-size`: tensor parallel size (effective for models that support TP, e.g. LTX2).
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+- `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2).
+- `--audio-sample-rate`: audio sample rate for embedded audio (when the pipeline returns audio).
 
 > ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.

examples/offline_inference/text_to_video/text_to_video.py

@@ -31,6 +31,12 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--height", type=int, default=720, help="Video height.")
     parser.add_argument("--width", type=int, default=1280, help="Video width.")
     parser.add_argument("--num-frames", type=int, default=81, help="Number of frames (Wan default is 81).")
+    parser.add_argument(
+        "--frame-rate",
+        type=float,
+        default=None,
+        help="Optional generation frame rate (used by models like LTX2). Defaults to --fps.",
+    )
     parser.add_argument("--num-inference-steps", type=int, default=40, help="Sampling steps.")
     parser.add_argument(
         "--boundary-ratio",
@@ -109,12 +115,25 @@ def parse_args() -> argparse.Namespace:
         default=1,
         help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
     )
+    parser.add_argument(
+        "--audio-sample-rate",
+        type=int,
+        default=24000,
+        help="Sample rate for audio output when saved (default: 24000 for LTX2).",
+    )
+    parser.add_argument(
+        "--vae-patch-parallel-size",
+        type=int,
+        default=1,
+        help="Number of GPUs used for VAE patch/tile parallelism (decode).",
+    )
     return parser.parse_args()
 
 
 def main():
     args = parse_args()
     generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
+    frame_rate = args.frame_rate if args.frame_rate is not None else float(args.fps)
 
     # Wan2.2 cache-dit tuning (from cache-dit examples and cache_alignment).
     cache_config = None
@@ -134,13 +153,13 @@ def main():
             "scm_steps_mask_policy": None,  # SCM mask policy: None (disabled), "slow", "medium", "fast", "ultra"
             "scm_steps_policy": "dynamic",  # SCM steps policy: "dynamic" or "static"
         }
-    # Configure parallel settings (only SP is supported for Wan)
+    # Configure parallel settings
-    # Note: cfg_parallel and tensor_parallel are not implemented for Wan models
     parallel_config = DiffusionParallelConfig(
         ulysses_degree=args.ulysses_degree,
         ring_degree=args.ring_degree,
         cfg_parallel_size=args.cfg_parallel_size,
         tensor_parallel_size=args.tensor_parallel_size,
+        vae_patch_parallel_size=args.vae_patch_parallel_size,
     )
 
     # Check if profiling is requested via environment variable
@@ -153,12 +172,12 @@ def main():
         vae_use_tiling=args.vae_use_tiling,
         boundary_ratio=args.boundary_ratio,
         flow_shift=args.flow_shift,
-        cache_backend=args.cache_backend,
-        cache_config=cache_config,
         enable_cache_dit_summary=args.enable_cache_dit_summary,
         enable_cpu_offload=args.enable_cpu_offload,
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
+        cache_backend=args.cache_backend,
+        cache_config=cache_config,
     )
 
     if profiler_enabled:
@@ -172,7 +191,9 @@ def main():
     print(f"  Inference steps: {args.num_inference_steps}")
     print(f"  Frames: {args.num_frames}")
     print(
-        f"  Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size}"
+        f"  Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree},"
+        f" cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size},"
+        f" vae_patch_parallel_size={args.vae_patch_parallel_size}"
     )
     print(f"  Video size: {args.width}x{args.height}")
     print(f"{'=' * 60}\n")
@@ -191,6 +212,7 @@ def main():
             guidance_scale_2=args.guidance_scale_high,
             num_inference_steps=args.num_inference_steps,
             num_frames=args.num_frames,
+            frame_rate=frame_rate,
         ),
     )
     generation_end = time.perf_counter()
@@ -199,31 +221,56 @@ def main():
     # Print profiling results
     print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
 
-    # Extract video frames from OmniRequestOutput
+    audio = None
-    if isinstance(frames, list) and len(frames) > 0:
+    if isinstance(frames, list):
-        first_item = frames[0]
+        frames = frames[0] if frames else None
 
-        # Check if it's an OmniRequestOutput
+    if isinstance(frames, OmniRequestOutput):
-        if hasattr(first_item, "final_output_type"):
+        if frames.final_output_type != "image":
-            if first_item.final_output_type != "image":
+            raise ValueError(
-                raise ValueError(
+                f"Unexpected output type '{frames.final_output_type}', expected 'image' for video generation."
-                    f"Unexpected output type '{first_item.final_output_type}', expected 'image' for video generation."
+            )
-                )
+        if frames.multimodal_output and "audio" in frames.multimodal_output:
-
+            audio = frames.multimodal_output["audio"]
-            # Pipeline mode: extract from nested request_output
+        if frames.is_pipeline_output and frames.request_output is not None:
-            if hasattr(first_item, "is_pipeline_output") and first_item.is_pipeline_output:
+            inner_output = frames.request_output
-                if isinstance(first_item.request_output, list) and len(first_item.request_output) > 0:
+            if isinstance(inner_output, list):
-                    inner_output = first_item.request_output[0]
+                inner_output = inner_output[0] if inner_output else None
-                    if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"):
+            if isinstance(inner_output, OmniRequestOutput):
-                        frames = inner_output.images[0] if inner_output.images else None
+                if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
-                        if frames is None:
+                    audio = inner_output.multimodal_output["audio"]
-                            raise ValueError("No video frames found in output.")
+                frames = inner_output
-            # Diffusion mode: use direct images field
+        if isinstance(frames, OmniRequestOutput):
-            elif hasattr(first_item, "images") and first_item.images:
+            if frames.images:
-                frames = first_item.images
+                if len(frames.images) == 1 and isinstance(frames.images[0], tuple) and len(frames.images[0]) == 2:
+                    frames, audio = frames.images[0]
+                elif len(frames.images) == 1 and isinstance(frames.images[0], dict):
+                    audio = frames.images[0].get("audio")
+                    frames = frames.images[0].get("frames") or frames.images[0].get("video")
+                else:
+                    frames = frames.images
             else:
                 raise ValueError("No video frames found in OmniRequestOutput.")
 
+    if isinstance(frames, list) and frames:
+        first_item = frames[0]
+        if isinstance(first_item, tuple) and len(first_item) == 2:
+            frames, audio = first_item
+        elif isinstance(first_item, dict):
+            audio = first_item.get("audio")
+            frames = first_item.get("frames") or first_item.get("video")
+        elif isinstance(first_item, list):
+            frames = first_item
+
+    if isinstance(frames, tuple) and len(frames) == 2:
+        frames, audio = frames
+    elif isinstance(frames, dict):
+        audio = frames.get("audio")
+        frames = frames.get("frames") or frames.get("video")
+
+    if frames is None:
+        raise ValueError("No video frames found in output.")
+
     output_path = Path(args.output)
     output_path.parent.mkdir(parents=True, exist_ok=True)
     try:
@@ -231,32 +278,127 @@ def main():
     except ImportError:
         raise ImportError("diffusers is required for export_to_video.")
 
-    # frames may be np.ndarray (preferred) or torch.Tensor
+    def _normalize_frame(frame):
+        if isinstance(frame, torch.Tensor):
+            frame_tensor = frame.detach().cpu()
+            if frame_tensor.dim() == 4 and frame_tensor.shape[0] == 1:
+                frame_tensor = frame_tensor[0]
+            if frame_tensor.dim() == 3 and frame_tensor.shape[0] in (3, 4):
+                frame_tensor = frame_tensor.permute(1, 2, 0)
+            if frame_tensor.is_floating_point():
+                frame_tensor = frame_tensor.clamp(-1, 1) * 0.5 + 0.5
+            return frame_tensor.float().numpy()
+        if isinstance(frame, np.ndarray):
+            frame_array = frame
+            if frame_array.ndim == 4 and frame_array.shape[0] == 1:
+                frame_array = frame_array[0]
+            if np.issubdtype(frame_array.dtype, np.integer):
+                frame_array = frame_array.astype(np.float32) / 255.0
+            return frame_array
+        try:
+            from PIL import Image
+        except ImportError:
+            Image = None
+        if Image is not None and isinstance(frame, Image.Image):
+            return np.asarray(frame).astype(np.float32) / 255.0
+        return frame
+
+    def _ensure_frame_list(video_array):
+        if isinstance(video_array, list):
+            if len(video_array) == 0:
+                return video_array
+            first_item = video_array[0]
+            if isinstance(first_item, np.ndarray):
+                if first_item.ndim == 5:
+                    return list(first_item[0])
+                if first_item.ndim == 4:
+                    if len(video_array) == 1:
+                        return list(first_item)
+                    return list(first_item)
+                if first_item.ndim == 3:
+                    return video_array
+            return video_array
+        if isinstance(video_array, np.ndarray):
+            if video_array.ndim == 5:
+                return list(video_array[0])
+            if video_array.ndim == 4:
+                return list(video_array)
+            if video_array.ndim == 3:
+                return [video_array]
+        return video_array
+
+    # frames may be np.ndarray, torch.Tensor, or list of tensors/arrays/images
     # export_to_video expects a list of frames with values in [0, 1]
     if isinstance(frames, torch.Tensor):
         video_tensor = frames.detach().cpu()
         if video_tensor.dim() == 5:
-            # [B, C, F, H, W] or [B, F, H, W, C]
             if video_tensor.shape[1] in (3, 4):
                 video_tensor = video_tensor[0].permute(1, 2, 3, 0)
             else:
                 video_tensor = video_tensor[0]
         elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4):
             video_tensor = video_tensor.permute(1, 2, 3, 0)
-        # If float, assume [-1,1] and normalize to [0,1]
         if video_tensor.is_floating_point():
             video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5
         video_array = video_tensor.float().numpy()
+    elif isinstance(frames, np.ndarray):
+        video_array = frames
+        if video_array.ndim == 5:
+            video_array = video_array[0]
+        if np.issubdtype(video_array.dtype, np.integer):
+            video_array = video_array.astype(np.float32) / 255.0
+    elif isinstance(frames, list):
+        if len(frames) == 0:
+            raise ValueError("No video frames found in output.")
+        video_array = [_normalize_frame(frame) for frame in frames]
     else:
         video_array = frames
-        if hasattr(video_array, "shape") and video_array.ndim == 5:
-            video_array = video_array[0]
 
-    # Convert 4D array (frames, H, W, C) to list of frames for export_to_video
+    video_array = _ensure_frame_list(video_array)
-    if isinstance(video_array, np.ndarray) and video_array.ndim == 4:
-        video_array = list(video_array)
 
-    export_to_video(video_array, str(output_path), fps=args.fps)
+    use_ltx2_export = False
+    if args.model and "ltx" in str(args.model).lower():
+        use_ltx2_export = True
+    if audio is not None:
+        use_ltx2_export = True
+
+    if use_ltx2_export:
+        try:
+            from diffusers.pipelines.ltx2.export_utils import encode_video
+        except ImportError:
+            raise ImportError("diffusers is required for LTX2 encode_video.")
+
+        if isinstance(video_array, list):
+            frames_np = np.stack(video_array, axis=0)
+        elif isinstance(video_array, np.ndarray):
+            frames_np = video_array
+        else:
+            frames_np = np.asarray(video_array)
+
+        frames_u8 = (frames_np * 255).round().clip(0, 255).astype("uint8")
+        video_tensor = torch.from_numpy(frames_u8)
+
+        audio_out = None
+        if audio is not None:
+            if isinstance(audio, list):
+                audio = audio[0] if audio else None
+            if isinstance(audio, np.ndarray):
+                audio = torch.from_numpy(audio)
+            if isinstance(audio, torch.Tensor):
+                audio_out = audio
+                if audio_out.dim() > 1:
+                    audio_out = audio_out[0]
+                audio_out = audio_out.float().cpu()
+
+        encode_video(
+            video_tensor,
+            fps=args.fps,
+            audio=audio_out,
+            audio_sample_rate=args.audio_sample_rate if audio_out is not None else None,
+            output_path=str(output_path),
+        )
+    else:
+        export_to_video(video_array, str(output_path), fps=args.fps)
     print(f"Saved generated video to {output_path}")
 
     if profiler_enabled:

examples/online_serving/chart-helm/.helmignore

@@ -0,0 +1,6 @@
+*.png
+.git/
+ct.yaml
+lintconf.yaml
+values.schema.json
+/workflows

examples/online_serving/chart-helm/Chart.yaml

@@ -0,0 +1,15 @@
+apiVersion: v2
+name: chart-vllm-omni
+description: A Helm chart for deploying vLLM-Omni on Kubernetes for omni-modality model serving
+
+# Application chart that can be packaged and deployed
+type: application
+
+# Chart version - increment on each change to the chart
+version: 0.1.0
+
+# vllm-omni application version — keep in sync with image.tag in values.yaml
+appVersion: "0.16.0"
+
+maintainers:
+  - name: vllm-omni-team

examples/online_serving/chart-helm/README.md

@@ -0,0 +1,144 @@
+# vLLM-Omni Helm Chart
+
+Helm chart for deploying [vLLM-Omni](https://github.com/vllm-project/vllm-omni) on Kubernetes. vLLM-Omni extends vLLM with omni-modality model serving, supporting text-to-image, multimodal chat, text-to-speech, and more.
+
+## Prerequisites
+
+- Kubernetes 1.24+
+- Helm 3.x
+- NVIDIA GPU nodes with [NVIDIA Device Plugin](https://github.com/NVIDIA/k8s-device-plugin)
+
+## Quick Start
+
+```bash
+helm install my-release ./chart-helm \
+  --set model=Tongyi-MAI/Z-Image-Turbo
+```
+
+## Configuration
+
+### Model Selection
+
+Set the `model` value to any supported HuggingFace model ID:
+
+| Model | Type | GPUs | Notes |
+|-------|------|------|-------|
+| `Tongyi-MAI/Z-Image-Turbo` | text-to-image | 1 | Small, fast (default) |
+| `stabilityai/stable-diffusion-3.5-medium` | text-to-image | 1 | ~6GB VRAM |
+| `Qwen/Qwen-Image` | text-to-image | 1 | Large, ~40GB+ VRAM |
+| `Qwen/Qwen2.5-Omni-7B` | multimodal | 2 | Text + audio + image + video |
+| `Qwen/Qwen3-Omni-7B-Chat` | multimodal | 2 | Latest omni model |
+| `Qwen/Qwen3-TTS` | text-to-speech | 1 | TTS |
+
+### HuggingFace Token
+
+For gated models that require authentication:
+
+```bash
+helm install my-release ./chart-helm \
+  --set model=Qwen/Qwen2.5-Omni-7B \
+  --set hfToken=hf_xxxxx \
+  --set resources.requests."nvidia\.com/gpu"=2 \
+  --set resources.limits."nvidia\.com/gpu"=2
+```
+
+### Omni-Specific Flags
+
+Enable VAE memory optimizations for diffusion models:
+
+```bash
+helm install my-release ./chart-helm \
+  --set model=Qwen/Qwen-Image \
+  --set omniArgs.vaeUseSlicing=true \
+  --set omniArgs.vaeUseTiling=true
+```
+
+Enable CPU offloading:
+
+```bash
+helm install my-release ./chart-helm \
+  --set model=Qwen/Qwen-Image \
+  --set omniArgs.enableCpuOffload=true
+```
+
+Pass additional raw CLI flags:
+
+```bash
+helm install my-release ./chart-helm \
+  --set model=Qwen/Qwen-Image \
+  --set omniArgs.extraArgs[0]="--enable-layerwise-offload"
+```
+
+### Model Cache
+
+By default, a PersistentVolumeClaim is created for the HuggingFace model cache to avoid re-downloading models on pod restarts:
+
+```yaml
+modelCache:
+  enabled: true
+  storageSize: "50Gi"
+  storageClassName: ""
+```
+
+To use an ephemeral volume instead:
+
+```bash
+helm install my-release ./chart-helm \
+  --set modelCache.enabled=false
+```
+
+### Custom Command Override
+
+To fully override the container command:
+
+```bash
+helm install my-release ./chart-helm \
+  --set image.command[0]=vllm \
+  --set image.command[1]=serve \
+  --set image.command[2]=my-model \
+  --set image.command[3]=--omni \
+  --set image.command[4]=--host \
+  --set image.command[5]=0.0.0.0
+```
+
+## API Endpoints
+
+Once deployed, vLLM-Omni exposes the following OpenAI-compatible endpoints:
+
+| Endpoint | Method | Description |
+|----------|--------|-------------|
+| `/health` | GET | Health check |
+| `/v1/models` | GET | List available models |
+| `/v1/chat/completions` | POST | Chat completions (text/multimodal) |
+| `/v1/images/generations` | POST | Image generation |
+| `/v1/images/edits` | POST | Image editing |
+| `/v1/audio/speech` | POST | Text-to-speech |
+
+## Files
+
+| File | Description |
+|------|-------------|
+| `Chart.yaml` | Chart metadata (name, version, maintainers) |
+| `values.yaml` | Default configuration values |
+| `values.schema.json` | JSON schema for validating values |
+| `templates/_helpers.tpl` | Helper templates for common configurations |
+| `templates/deployment.yaml` | Kubernetes Deployment |
+| `templates/service.yaml` | Kubernetes Service (ClusterIP) |
+| `templates/secrets.yaml` | Secrets (generic + HuggingFace token) |
+| `templates/pvc.yaml` | PersistentVolumeClaim for model cache |
+| `templates/configmap.yaml` | Optional ConfigMap |
+| `templates/hpa.yaml` | HorizontalPodAutoscaler |
+| `templates/poddisruptionbudget.yaml` | PodDisruptionBudget |
+| `templates/custom-objects.yaml` | Custom Kubernetes objects |
+
+## Running Tests
+
+This chart includes unit tests using [helm-unittest](https://github.com/helm-unittest/helm-unittest). Install the plugin and run tests:
+
+```bash
+# Install plugin
+helm plugin install https://github.com/helm-unittest/helm-unittest
+
+# Run tests
+helm unittest .
+```

examples/online_serving/chart-helm/ct.yaml

@@ -0,0 +1,3 @@
+chart-dirs:
+  - charts
+validate-maintainers: false

examples/online_serving/chart-helm/lintconf.yaml

@@ -0,0 +1,42 @@
+---
+rules:
+  braces:
+    min-spaces-inside: 0
+    max-spaces-inside: 0
+    min-spaces-inside-empty: -1
+    max-spaces-inside-empty: -1
+  brackets:
+    min-spaces-inside: 0
+    max-spaces-inside: 0
+    min-spaces-inside-empty: -1
+    max-spaces-inside-empty: -1
+  colons:
+    max-spaces-before: 0
+    max-spaces-after: 1
+  commas:
+    max-spaces-before: 0
+    min-spaces-after: 1
+    max-spaces-after: 1
+  comments:
+    require-starting-space: true
+    min-spaces-from-content: 2
+  document-end: disable
+  document-start: disable
+  empty-lines:
+    max: 2
+    max-start: 0
+    max-end: 0
+  hyphens:
+    max-spaces-after: 1
+  indentation:
+    spaces: consistent
+    indent-sequences: whatever
+    check-multi-line-strings: false
+  key-duplicates: enable
+  line-length: disable
+  new-line-at-end-of-file: disable
+  new-lines:
+    type: unix
+  trailing-spaces: enable
+  truthy:
+    level: warning

examples/online_serving/chart-helm/templates/_helpers.tpl

@@ -0,0 +1,203 @@
+{{/*
+Define the vllm-omni serve command from model + omniArgs.
+If image.command is set, uses that as a full override.
+*/}}
+{{- define "chart.omni-command" -}}
+{{- if .Values.image.command }}
+{{-   toYaml .Values.image.command }}
+{{- else }}
+- "vllm"
+- "serve"
+- {{ .Values.model | quote }}
+- "--omni"
+- "--host"
+- "0.0.0.0"
+- "--port"
+- {{ include "chart.container-port" . | quote }}
+{{- if .Values.omniArgs.vaeUseSlicing }}
+- "--vae-use-slicing"
+{{- end }}
+{{- if .Values.omniArgs.vaeUseTiling }}
+- "--vae-use-tiling"
+{{- end }}
+{{- if .Values.omniArgs.enableCpuOffload }}
+- "--enable-cpu-offload"
+{{- end }}
+{{- if .Values.omniArgs.numGpus }}
+- "--num-gpus"
+- {{ .Values.omniArgs.numGpus | quote }}
+{{- end }}
+{{- if .Values.omniArgs.stageConfigsPath }}
+- "--stage-configs-path"
+- {{ .Values.omniArgs.stageConfigsPath | quote }}
+{{- end }}
+{{- if and .Values.omniArgs.cacheBackend (ne .Values.omniArgs.cacheBackend "none") }}
+- "--cache-backend"
+- {{ .Values.omniArgs.cacheBackend | quote }}
+{{- end }}
+{{- if .Values.omniArgs.defaultSamplingParams }}
+- "--default-sampling-params"
+- {{ .Values.omniArgs.defaultSamplingParams | quote }}
+{{- end }}
+{{- if and .Values.omniArgs.workerBackend (ne .Values.omniArgs.workerBackend "multi_process") }}
+- "--worker-backend"
+- {{ .Values.omniArgs.workerBackend | quote }}
+{{- end }}
+{{- range .Values.omniArgs.extraArgs }}
+- {{ . | quote }}
+{{- end }}
+{{- end }}
+{{- end }}
+
+{{/*
+Define HuggingFace environment variables
+*/}}
+{{- define "chart.hf-env" -}}
+- name: HF_HOME
+  value: "/cache/huggingface"
+- name: HOME
+  value: "/cache"
+{{- if .Values.hfToken }}
+- name: HF_TOKEN
+  valueFrom:
+    secretKeyRef:
+      name: {{ .Release.Name }}-hf-token
+      key: token
+{{- end }}
+{{- end }}
+
+{{/*
+Define ports for the pods
+*/}}
+{{- define "chart.container-port" -}}
+{{-  default "8000" .Values.containerPort }}
+{{- end }}
+
+{{/*
+Define service name
+*/}}
+{{- define "chart.service-name" -}}
+{{-  if .Values.serviceName -}}
+{{ .Values.serviceName | lower | trim }}
+{{-  else -}}
+{{ printf "%s-service" .Release.Name }}
+{{-  end -}}
+{{- end }}
+
+{{/*
+Define service port
+*/}}
+{{- define "chart.service-port" -}}
+{{-  if .Values.servicePort }}
+{{-    .Values.servicePort }}
+{{-  else }}
+{{-    include "chart.container-port" . }}
+{{-  end }}
+{{- end }}
+
+{{/*
+Define service port name
+*/}}
+{{- define "chart.service-port-name" -}}
+"service-port"
+{{- end }}
+
+{{/*
+Define container port name
+*/}}
+{{- define "chart.container-port-name" -}}
+"container-port"
+{{- end }}
+
+{{/*
+Define deployment strategy
+*/}}
+{{- define "chart.strategy" -}}
+strategy:
+{{-   if not .Values.deploymentStrategy }}
+  type: Recreate
+{{-   else }}
+{{      toYaml .Values.deploymentStrategy | indent 2 }}
+{{-   end }}
+{{- end }}
+
+{{/*
+Define additional ports
+*/}}
+{{- define "chart.extraPorts" }}
+{{-   with .Values.extraPorts }}
+{{      toYaml . }}
+{{-   end }}
+{{- end }}
+
+{{/*
+Define chart external ConfigMaps and Secrets
+*/}}
+{{- define "chart.externalConfigs" -}}
+{{-   with .Values.externalConfigs -}}
+{{      toYaml . }}
+{{-   end }}
+{{- end }}
+
+{{/*
+Define startup, liveness and readiness probes
+*/}}
+{{- define "chart.probes" -}}
+{{-   if .Values.startupProbe  }}
+startupProbe:
+{{-     with .Values.startupProbe }}
+{{-       toYaml . | nindent 2 }}
+{{-     end }}
+{{-   end }}
+{{-   if .Values.readinessProbe  }}
+readinessProbe:
+{{-     with .Values.readinessProbe }}
+{{-       toYaml . | nindent 2 }}
+{{-     end }}
+{{-   end }}
+{{-   if .Values.livenessProbe  }}
+livenessProbe:
+{{-     with .Values.livenessProbe }}
+{{-       toYaml . | nindent 2 }}
+{{-     end }}
+{{-   end }}
+{{- end }}
+
+{{/*
+Define resources
+*/}}
+{{- define "chart.resources" -}}
+requests:
+  memory: {{ required "Value 'resources.requests.memory' must be defined !" .Values.resources.requests.memory | quote }}
+  cpu: {{ required "Value 'resources.requests.cpu' must be defined !" .Values.resources.requests.cpu | quote }}
+  {{- if and (gt (int (index .Values.resources.requests "nvidia.com/gpu")) 0) (gt (int (index .Values.resources.limits "nvidia.com/gpu")) 0) }}
+  nvidia.com/gpu: {{ required "Value 'resources.requests.nvidia.com/gpu' must be defined !" (index .Values.resources.requests "nvidia.com/gpu") }}
+  {{- end }}
+limits:
+  memory: {{ required "Value 'resources.limits.memory' must be defined !" .Values.resources.limits.memory | quote }}
+  cpu: {{ required "Value 'resources.limits.cpu' must be defined !" .Values.resources.limits.cpu | quote }}
+  {{- if and (gt (int (index .Values.resources.requests "nvidia.com/gpu")) 0) (gt (int (index .Values.resources.limits "nvidia.com/gpu")) 0) }}
+  nvidia.com/gpu: {{ required "Value 'resources.limits.nvidia.com/gpu' must be defined !" (index .Values.resources.limits "nvidia.com/gpu") }}
+  {{- end }}
+{{- end }}
+
+{{/*
+Define user for the main container
+*/}}
+{{- define "chart.user" }}
+{{-   if .Values.image.runAsUser  }}
+runAsUser:
+{{-     with .Values.image.runAsUser }}
+{{-       toYaml . | nindent 2 }}
+{{-     end }}
+{{-   end }}
+{{- end }}
+
+{{/*
+Define chart labels
+*/}}
+{{- define "chart.labels" -}}
+{{-   with .Values.labels -}}
+{{      toYaml . }}
+{{-   end }}
+{{- end }}

examples/online_serving/chart-helm/templates/configmap.yaml

@@ -0,0 +1,11 @@
+{{- if .Values.configs -}}
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: "{{ .Release.Name }}-configs"
+  namespace: {{ .Release.Namespace }}
+data:
+  {{- with .Values.configs }}
+  {{- toYaml . | nindent 2 }}
+  {{- end }}
+{{- end -}}

examples/online_serving/chart-helm/templates/custom-objects.yaml

@@ -0,0 +1,6 @@
+{{- if .Values.customObjects }}
+{{- range .Values.customObjects }}
+{{- tpl (. | toYaml) $ }}
+---
+{{- end }}
+{{- end }}

examples/online_serving/chart-helm/templates/deployment.yaml

@@ -0,0 +1,106 @@
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: "{{ .Release.Name }}-deployment-vllm-omni"
+  namespace: {{ .Release.Namespace }}
+  labels:
+  {{- include "chart.labels" . | nindent 4 }}
+spec:
+  {{- if not .Values.autoscaling.enabled }}
+  replicas: {{ .Values.replicaCount }}
+  {{- end }}
+  {{- include "chart.strategy" . | nindent 2 }}
+  selector:
+    matchLabels:
+    {{- include "chart.labels" . | nindent 6 }}
+  progressDeadlineSeconds: 1800
+  template:
+    metadata:
+      labels:
+      {{- include "chart.labels" . | nindent 8 }}
+    spec:
+      containers:
+        - name: "vllm-omni"
+          image: "{{ required "Required value 'image.repository' must be defined !" .Values.image.repository }}:{{ required "Required value 'image.tag' must be defined !" .Values.image.tag }}"
+          command:
+            {{- include "chart.omni-command" . | nindent 12 }}
+          securityContext:
+            {{- if .Values.image.securityContext }}
+              {{- with .Values.image.securityContext }}
+              {{- toYaml . | nindent 12 }}
+              {{- end }}
+            {{- else }}
+            runAsNonRoot: false
+              {{- include "chart.user" . | indent 12 }}
+            {{- end }}
+          imagePullPolicy: IfNotPresent
+          env:
+            {{- include "chart.hf-env" . | nindent 12 }}
+            {{- if .Values.image.env }}
+            {{- toYaml .Values.image.env | nindent 12 }}
+            {{- end }}
+          {{- if or .Values.externalConfigs .Values.configs .Values.secrets }}
+          envFrom:
+            {{- if .Values.configs }}
+            - configMapRef:
+                name: "{{ .Release.Name }}-configs"
+            {{- end }}
+            {{- if .Values.secrets }}
+            - secretRef:
+                name: "{{ .Release.Name }}-secrets"
+            {{- end }}
+            {{- include "chart.externalConfigs" . | nindent 12 }}
+          {{- end }}
+          ports:
+            - name: {{ include "chart.container-port-name" . }}
+              containerPort: {{ include "chart.container-port" . }}
+            {{- include "chart.extraPorts" . | nindent 12 }}
+          {{- include "chart.probes" . | indent 10 }}
+          resources: {{- include "chart.resources" . | nindent 12 }}
+          volumeMounts:
+          - name: shm
+            mountPath: /dev/shm
+          - name: model-cache
+            mountPath: /cache
+
+        {{- with .Values.extraContainers }}
+        {{ toYaml . | nindent 8 }}
+        {{- end }}
+
+      {{- if .Values.extraInit.initContainers }}
+      initContainers:
+        {{- toYaml .Values.extraInit.initContainers | nindent 8 }}
+      {{- end }}
+
+      volumes:
+        - name: shm
+          emptyDir:
+            medium: Memory
+            sizeLimit: {{ .Values.shmSize }}
+        - name: model-cache
+          {{- if .Values.modelCache.enabled }}
+          persistentVolumeClaim:
+            claimName: {{ .Release.Name }}-model-cache
+          {{- else }}
+          emptyDir: {}
+          {{- end }}
+
+      {{- with .Values.nodeSelector }}
+      nodeSelector:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
+      {{- with .Values.tolerations }}
+      tolerations:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
+      {{- if .Values.gpuModels }}
+      affinity:
+        nodeAffinity:
+          requiredDuringSchedulingIgnoredDuringExecution:
+            nodeSelectorTerms:
+              - matchExpressions:
+                - key: nvidia.com/gpu.product
+                  operator: In
+                  values:
+                    {{- toYaml .Values.gpuModels | nindent 20 }}
+      {{- end }}

examples/online_serving/chart-helm/templates/hpa.yaml

@@ -0,0 +1,31 @@
+{{- if .Values.autoscaling.enabled }}
+apiVersion: autoscaling/v2
+kind: HorizontalPodAutoscaler
+metadata:
+  name: "{{ .Release.Name }}-hpa"
+  namespace: {{ .Release.Namespace }}
+spec:
+  scaleTargetRef:
+    apiVersion: apps/v1
+    kind: Deployment
+    name: "{{ .Release.Name }}-deployment-vllm-omni"
+  minReplicas: {{ .Values.autoscaling.minReplicas }}
+  maxReplicas: {{ .Values.autoscaling.maxReplicas }}
+  metrics:
+    {{- if .Values.autoscaling.targetCPUUtilizationPercentage }}
+    - type: Resource
+      resource:
+        name: cpu
+        target:
+          type: Utilization
+          averageUtilization: {{ .Values.autoscaling.targetCPUUtilizationPercentage }}
+    {{- end }}
+    {{- if .Values.autoscaling.targetMemoryUtilizationPercentage }}
+    - type: Resource
+      resource:
+        name: memory
+        target:
+          type: Utilization
+          averageUtilization: {{ .Values.autoscaling.targetMemoryUtilizationPercentage }}
+    {{- end }}
+{{- end }}

examples/online_serving/chart-helm/templates/ingress.yaml

@@ -0,0 +1,32 @@
+{{- if .Values.ingress.enabled }}
+apiVersion: networking.k8s.io/v1
+kind: Ingress
+metadata:
+  name: "{{ .Release.Name }}-ingress"
+  namespace: {{ .Release.Namespace }}
+  labels:
+  {{- include "chart.labels" . | nindent 4 }}
+  {{- with .Values.ingress.annotations }}
+  annotations:
+    {{- toYaml . | nindent 4 }}
+  {{- end }}
+spec:
+  {{- if .Values.ingress.ingressClassName }}
+  ingressClassName: {{ .Values.ingress.ingressClassName }}
+  {{- end }}
+  {{- if .Values.ingress.tls }}
+  tls:
+    {{- toYaml .Values.ingress.tls | nindent 4 }}
+  {{- end }}
+  rules:
+    - host: {{ .Values.ingress.host }}
+      http:
+        paths:
+          - path: /
+            pathType: Prefix
+            backend:
+              service:
+                name: "{{ include "chart.service-name" . }}"
+                port:
+                  number: {{ include "chart.service-port" . }}
+{{- end }}

examples/online_serving/chart-helm/templates/poddisruptionbudget.yaml

@@ -0,0 +1,12 @@
+{{- if gt (int .Values.replicaCount) 1 }}
+apiVersion: policy/v1
+kind: PodDisruptionBudget
+metadata:
+  name: "{{ .Release.Name }}-pdb"
+  namespace: {{ .Release.Namespace }}
+spec:
+  maxUnavailable: {{ default 1 .Values.maxUnavailablePodDisruptionBudget }}
+  selector:
+    matchLabels:
+    {{- include "chart.labels" . | nindent 6 }}
+{{- end }}

examples/online_serving/chart-helm/templates/pvc.yaml

@@ -0,0 +1,16 @@
+{{- if .Values.modelCache.enabled }}
+apiVersion: v1
+kind: PersistentVolumeClaim
+metadata:
+  name: "{{ .Release.Name }}-model-cache"
+  namespace: {{ .Release.Namespace }}
+spec:
+  accessModes:
+    {{- toYaml .Values.modelCache.accessModes | nindent 4 }}
+  {{- if .Values.modelCache.storageClassName }}
+  storageClassName: {{ .Values.modelCache.storageClassName }}
+  {{- end }}
+  resources:
+    requests:
+      storage: {{ .Values.modelCache.storageSize }}
+{{- end }}

examples/online_serving/chart-helm/templates/secrets.yaml

@@ -0,0 +1,23 @@
+{{- if .Values.secrets }}
+apiVersion: v1
+kind: Secret
+metadata:
+  name: "{{ .Release.Name }}-secrets"
+  namespace: {{ .Release.Namespace }}
+type: Opaque
+data:
+  {{- range $key, $val := .Values.secrets }}
+  {{ $key }}: {{ $val | b64enc | quote }}
+  {{- end }}
+---
+{{- end }}
+{{- if .Values.hfToken }}
+apiVersion: v1
+kind: Secret
+metadata:
+  name: "{{ .Release.Name }}-hf-token"
+  namespace: {{ .Release.Namespace }}
+type: Opaque
+data:
+  token: {{ .Values.hfToken | b64enc | quote }}
+{{- end }}

examples/online_serving/chart-helm/templates/service.yaml

@@ -0,0 +1,14 @@
+apiVersion: v1
+kind: Service
+metadata:
+  name: "{{ include "chart.service-name" . }}"
+  namespace: {{ .Release.Namespace }}
+spec:
+  type: ClusterIP
+  ports:
+    - name: {{ include "chart.service-port-name" . }}
+      port: {{ include "chart.service-port" . }}
+      targetPort: {{ include "chart.container-port-name" . }}
+      protocol: TCP
+  selector:
+  {{- include "chart.labels" . | nindent 4 }}

examples/online_serving/chart-helm/tests/deployment_test.yaml

@@ -0,0 +1,181 @@
+suite: test deployment
+templates:
+  - deployment.yaml
+tests:
+  - it: should create deployment with default omni command
+    asserts:
+      - hasDocuments:
+          count: 1
+      - isKind:
+          of: Deployment
+      - equal:
+          path: spec.template.spec.containers[0].name
+          value: vllm-omni
+      - equal:
+          path: spec.template.spec.containers[0].command[0]
+          value: vllm
+      - equal:
+          path: spec.template.spec.containers[0].command[1]
+          value: serve
+      - equal:
+          path: spec.template.spec.containers[0].command[2]
+          value: Tongyi-MAI/Z-Image-Turbo
+      - equal:
+          path: spec.template.spec.containers[0].command[3]
+          value: "--omni"
+      - equal:
+          path: spec.template.spec.containers[0].command[4]
+          value: "--host"
+      - equal:
+          path: spec.template.spec.containers[0].command[5]
+          value: "0.0.0.0"
+      - equal:
+          path: spec.template.spec.containers[0].command[6]
+          value: "--port"
+      - equal:
+          path: spec.template.spec.containers[0].command[7]
+          value: "8000"
+
+  - it: should use custom model when set
+    set:
+      model: "Qwen/Qwen2.5-Omni-7B"
+    asserts:
+      - equal:
+          path: spec.template.spec.containers[0].command[2]
+          value: Qwen/Qwen2.5-Omni-7B
+
+  - it: should include omniArgs flags in command
+    set:
+      omniArgs:
+        vaeUseSlicing: true
+        vaeUseTiling: true
+        enableCpuOffload: true
+        numGpus: 2
+        cacheBackend: "tea_cache"
+        workerBackend: "ray"
+        extraArgs:
+          - "--enable-layerwise-offload"
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].command
+          content: "--vae-use-slicing"
+      - contains:
+          path: spec.template.spec.containers[0].command
+          content: "--vae-use-tiling"
+      - contains:
+          path: spec.template.spec.containers[0].command
+          content: "--enable-cpu-offload"
+      - contains:
+          path: spec.template.spec.containers[0].command
+          content: "--num-gpus"
+      - contains:
+          path: spec.template.spec.containers[0].command
+          content: "--cache-backend"
+      - contains:
+          path: spec.template.spec.containers[0].command
+          content: "--worker-backend"
+      - contains:
+          path: spec.template.spec.containers[0].command
+          content: "--enable-layerwise-offload"
+
+  - it: should use full command override when image.command is set
+    set:
+      image:
+        command:
+          - "custom-binary"
+          - "--flag"
+    asserts:
+      - equal:
+          path: spec.template.spec.containers[0].command[0]
+          value: custom-binary
+      - equal:
+          path: spec.template.spec.containers[0].command[1]
+          value: "--flag"
+
+  - it: should mount shm and model-cache volumes
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].volumeMounts
+          content:
+            name: shm
+            mountPath: /dev/shm
+      - contains:
+          path: spec.template.spec.containers[0].volumeMounts
+          content:
+            name: model-cache
+            mountPath: /cache
+      - contains:
+          path: spec.template.spec.volumes
+          content:
+            name: shm
+            emptyDir:
+              medium: Memory
+              sizeLimit: 8Gi
+
+  - it: should use PVC for model-cache when modelCache is enabled
+    set:
+      modelCache:
+        enabled: true
+        storageSize: "50Gi"
+        accessModes:
+          - ReadWriteOnce
+    asserts:
+      - contains:
+          path: spec.template.spec.volumes
+          content:
+            name: model-cache
+            persistentVolumeClaim:
+              claimName: RELEASE-NAME-model-cache
+
+  - it: should use emptyDir for model-cache when modelCache is disabled
+    set:
+      modelCache:
+        enabled: false
+        storageSize: "50Gi"
+        accessModes:
+          - ReadWriteOnce
+    asserts:
+      - contains:
+          path: spec.template.spec.volumes
+          content:
+            name: model-cache
+            emptyDir: {}
+
+  - it: should set HF_HOME environment variable
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: HF_HOME
+            value: /cache/huggingface
+
+  - it: should set HF_TOKEN env when hfToken is provided
+    set:
+      hfToken: "hf_test_token_123"
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: HF_TOKEN
+            valueFrom:
+              secretKeyRef:
+                name: RELEASE-NAME-hf-token
+                key: token
+
+  - it: should create custom init containers when specified
+    set:
+      extraInit:
+        initContainers:
+          - name: my-init
+            image: busybox:latest
+            command: ["echo", "init"]
+    asserts:
+      - lengthEqual:
+          path: spec.template.spec.initContainers
+          count: 1
+      - equal:
+          path: spec.template.spec.initContainers[0].name
+          value: my-init
+      - equal:
+          path: spec.template.spec.initContainers[0].image
+          value: busybox:latest

examples/online_serving/chart-helm/tests/ingress_test.yaml

@@ -0,0 +1,86 @@
+suite: test ingress
+templates:
+  - ingress.yaml
+tests:
+  - it: should not create Ingress when disabled
+    set:
+      ingress:
+        enabled: false
+    asserts:
+      - hasDocuments:
+          count: 0
+
+  - it: should create Ingress when enabled
+    set:
+      ingress:
+        enabled: true
+        host: "vllm-omni.example.com"
+    asserts:
+      - hasDocuments:
+          count: 1
+      - isKind:
+          of: Ingress
+      - equal:
+          path: spec.rules[0].host
+          value: vllm-omni.example.com
+      - equal:
+          path: spec.rules[0].http.paths[0].path
+          value: /
+      - equal:
+          path: spec.rules[0].http.paths[0].pathType
+          value: Prefix
+      - equal:
+          path: spec.rules[0].http.paths[0].backend.service.name
+          value: RELEASE-NAME-service
+      - equal:
+          path: spec.rules[0].http.paths[0].backend.service.port.number
+          value: 80
+
+  - it: should set ingressClassName when specified
+    set:
+      ingress:
+        enabled: true
+        host: "vllm-omni.example.com"
+        ingressClassName: "nginx"
+    asserts:
+      - equal:
+          path: spec.ingressClassName
+          value: nginx
+
+  - it: should not set ingressClassName when empty
+    set:
+      ingress:
+        enabled: true
+        host: "vllm-omni.example.com"
+        ingressClassName: ""
+    asserts:
+      - notExists:
+          path: spec.ingressClassName
+
+  - it: should include annotations when specified
+    set:
+      ingress:
+        enabled: true
+        host: "vllm-omni.example.com"
+        annotations:
+          nginx.ingress.kubernetes.io/proxy-body-size: "100m"
+    asserts:
+      - isNotEmpty:
+          path: metadata.annotations
+
+  - it: should include TLS configuration when specified
+    set:
+      ingress:
+        enabled: true
+        host: "vllm-omni.example.com"
+        tls:
+          - secretName: vllm-omni-tls
+            hosts:
+              - vllm-omni.example.com
+    asserts:
+      - equal:
+          path: spec.tls[0].secretName
+          value: vllm-omni-tls
+      - equal:
+          path: spec.tls[0].hosts[0]
+          value: vllm-omni.example.com

examples/online_serving/chart-helm/tests/pvc_test.yaml

@@ -0,0 +1,73 @@
+suite: test pvc
+templates:
+  - pvc.yaml
+tests:
+  - it: should create PVC when modelCache is enabled
+    set:
+      modelCache:
+        enabled: true
+        storageSize: "50Gi"
+        storageClassName: ""
+        accessModes:
+          - ReadWriteOnce
+    asserts:
+      - hasDocuments:
+          count: 1
+      - isKind:
+          of: PersistentVolumeClaim
+      - equal:
+          path: spec.accessModes[0]
+          value: ReadWriteOnce
+      - equal:
+          path: spec.resources.requests.storage
+          value: 50Gi
+
+  - it: should not create PVC when modelCache is disabled
+    set:
+      modelCache:
+        enabled: false
+        storageSize: "50Gi"
+        accessModes:
+          - ReadWriteOnce
+    asserts:
+      - hasDocuments:
+          count: 0
+
+  - it: should use custom storage size
+    set:
+      modelCache:
+        enabled: true
+        storageSize: "100Gi"
+        storageClassName: ""
+        accessModes:
+          - ReadWriteOnce
+    asserts:
+      - equal:
+          path: spec.resources.requests.storage
+          value: 100Gi
+
+  - it: should set storageClassName when specified
+    set:
+      modelCache:
+        enabled: true
+        storageSize: "50Gi"
+        storageClassName: "fast-ssd"
+        accessModes:
+          - ReadWriteOnce
+    asserts:
+      - equal:
+          path: spec.storageClassName
+          value: fast-ssd
+
+  - it: should support ReadWriteMany access mode
+    set:
+      modelCache:
+        enabled: true
+        storageSize: "50Gi"
+        storageClassName: ""
+        accessModes:
+          - ReadWriteMany
+    asserts:
+      - equal:
+          path: spec.accessModes[0]
+          value: ReadWriteMany

examples/online_serving/chart-helm/tests/secrets_test.yaml

@@ -0,0 +1,48 @@
+suite: test secrets
+templates:
+  - secrets.yaml
+tests:
+  - it: should create HF token secret when hfToken is provided
+    set:
+      hfToken: "hf_test_token_123"
+    asserts:
+      - hasDocuments:
+          count: 1
+      - isKind:
+          of: Secret
+      - equal:
+          path: metadata.name
+          value: RELEASE-NAME-hf-token
+      - exists:
+          path: data.token
+
+  - it: should not create any secrets when hfToken is empty and no secrets defined
+    set:
+      hfToken: ""
+      secrets: {}
+    asserts:
+      - hasDocuments:
+          count: 0
+
+  - it: should create both generic and HF token secrets
+    set:
+      hfToken: "hf_test_token_123"
+      secrets:
+        mykey: myvalue
+    asserts:
+      - hasDocuments:
+          count: 2
+
+  - it: should create only generic secrets when no hfToken
+    set:
+      hfToken: ""
+      secrets:
+        mykey: myvalue
+    asserts:
+      - hasDocuments:
+          count: 1
+      - isKind:
+          of: Secret
+      - equal:
+          path: metadata.name
+          value: RELEASE-NAME-secrets

examples/online_serving/chart-helm/values.yaml

@@ -0,0 +1,197 @@
+# -- Default values for chart vllm-omni
+# -- Declare variables to be passed into your templates.
+
+# -- HuggingFace model ID to serve
+model: "Tongyi-MAI/Z-Image-Turbo"
+
+# -- Image configuration
+image:
+  # -- Image repository
+  repository: "vllm/vllm-omni"
+  # -- Image tag
+  tag: "v0.16.0"
+  # -- Override the container command entirely. If empty, the chart constructs
+  # -- it automatically from model + omniArgs.
+  command: []
+  # -- Optional environment variables for the container
+  env: []
+  # -- Security context override
+  securityContext: {}
+
+# -- Container port
+containerPort: 8000
+# -- Service name (auto-generated from release name if empty)
+serviceName:
+# -- Service port
+servicePort: 80
+# -- Additional ports configuration
+extraPorts: []
+
+# -- Number of replicas
+replicaCount: 1
+
+# -- Deployment strategy configuration
+deploymentStrategy: {}
+
+# -- Resource configuration
+resources:
+  requests:
+    # -- Number of CPUs
+    cpu: 4
+    # -- CPU memory
+    memory: 24Gi
+    # -- Number of GPUs
+    nvidia.com/gpu: 1
+  limits:
+    # -- Number of CPUs
+    cpu: 4
+    # -- CPU memory
+    memory: 24Gi
+    # -- Number of GPUs
+    nvidia.com/gpu: 1
+
+# -- GPU model types for node affinity scheduling (optional)
+# -- Leave empty to schedule on any GPU node. Set to restrict to specific GPU types.
+# -- Example: ["NVIDIA-A100-SXM4-40GB"]
+gpuModels: []
+
+# -- vllm-omni specific CLI arguments
+omniArgs:
+  # -- Additional raw CLI flags appended to the serve command (list of strings)
+  extraArgs: []
+  # -- Enable VAE slicing for memory optimization
+  vaeUseSlicing: false
+  # -- Enable VAE tiling for memory optimization
+  vaeUseTiling: false
+  # -- Enable CPU offloading for diffusion models
+  enableCpuOffload: false
+  # -- Number of GPUs for diffusion inference (empty = auto)
+  numGpus:
+  # -- Path to stage configs file (empty = auto-detected from model)
+  stageConfigsPath:
+  # -- Cache backend: none, tea_cache, or cache_dit
+  cacheBackend: "none"
+  # -- Default sampling params as a JSON string
+  defaultSamplingParams:
+  # -- Worker backend: multi_process or ray
+  workerBackend: "multi_process"
+
+# -- HuggingFace token for gated models (optional)
+hfToken: ""
+
+# -- Shared memory size for PyTorch multiprocessing
+shmSize: "8Gi"
+
+# -- Model cache configuration (HuggingFace downloads)
+modelCache:
+  # -- Use a PersistentVolumeClaim for model cache (recommended for production)
+  enabled: true
+  # -- Storage size for the PVC
+  storageSize: "50Gi"
+  # -- Storage class name (empty = cluster default)
+  storageClassName: ""
+  # -- Access modes
+  accessModes:
+    - ReadWriteOnce
+
+# -- Autoscaling configuration
+autoscaling:
+  # -- Enable autoscaling
+  enabled: false
+  # -- Minimum replicas
+  minReplicas: 1
+  # -- Maximum replicas
+  maxReplicas: 10
+  # -- Target CPU utilization for autoscaling
+  targetCPUUtilizationPercentage: 80
+  # targetMemoryUtilizationPercentage: 80
+
+# -- ConfigMap data (key-value pairs injected as environment variables)
+configs: {}
+
+# -- Secrets data (key-value pairs, base64-encoded automatically)
+secrets: {}
+
+# -- External ConfigMaps/Secrets references
+externalConfigs: []
+
+# -- Custom Kubernetes objects
+customObjects: []
+
+# -- PodDisruptionBudget max unavailable
+maxUnavailablePodDisruptionBudget: ""
+
+# -- Additional init containers
+extraInit:
+  initContainers: []
+
+# -- Additional sidecar containers
+extraContainers: []
+
+# -- Startup probe configuration
+# -- Protects slow-starting containers. Liveness and readiness probes
+# -- do not start until the startup probe succeeds.
+# -- failureThreshold * periodSeconds = max startup time (40 * 30 = 1200s = 20 min)
+startupProbe:
+  # -- HTTP check configuration
+  httpGet:
+    # -- Path to check
+    path: /health
+    # -- Port to check (uses named port to stay in sync with containerPort)
+    port: "container-port"
+  # -- Failures before killing the container
+  failureThreshold: 40
+  # -- How often to perform the probe
+  periodSeconds: 30
+
+# -- Readiness probe configuration
+readinessProbe:
+  # -- How often to perform the probe
+  periodSeconds: 10
+  # -- Failures before marking unready
+  failureThreshold: 3
+  # -- HTTP check configuration
+  httpGet:
+    # -- Path to check
+    path: /health
+    # -- Port to check (uses named port to stay in sync with containerPort)
+    port: "container-port"
+
+# -- Liveness probe configuration
+livenessProbe:
+  # -- Failures before restarting
+  failureThreshold: 3
+  # -- How often to perform the probe
+  periodSeconds: 15
+  # -- HTTP check configuration
+  httpGet:
+    # -- Path to check
+    path: /health
+    # -- Port to check (uses named port to stay in sync with containerPort)
+    port: "container-port"
+
+# -- Ingress configuration
+ingress:
+  # -- Enable Ingress
+  enabled: false
+  # -- Ingress class name (e.g., nginx, traefik)
+  ingressClassName: ""
+  # -- Hostname for the Ingress rule
+  host: "vllm-omni.example.com"
+  # -- Additional annotations
+  annotations: {}
+  # -- TLS configuration
+  tls: []
+  #  - secretName: vllm-omni-tls
+  #    hosts:
+  #      - vllm-omni.example.com
+
+# -- Node selector for pod scheduling
+nodeSelector: {}
+
+# -- Tolerations for pod scheduling
+tolerations: []
+
+# -- Labels applied to all resources
+labels:
+  app: "vllm-omni"

examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh

@@ -160,10 +160,7 @@ esac
 echo "Running query type: $QUERY_TYPE"
 echo ""
 
-
+request_body=$(cat <<EOF
-output=$(curl -sS -X POST http://localhost:8091/v1/chat/completions \
-    -H "Content-Type: application/json" \
-    -d @- <<EOF
 {
   "model": "Qwen/Qwen2.5-Omni-7B",
   "sampling_params_list": $sampling_params_list,
@@ -186,7 +183,12 @@ output=$(curl -sS -X POST http://localhost:8091/v1/chat/completions \
   ]
 }
 EOF
-  )
+)
+
+output=$(curl -sS --retry 3 --retry-delay 3 --retry-connrefused \
+    -X POST http://localhost:8091/v1/chat/completions \
+    -H "Content-Type: application/json" \
+    -d "$request_body")
 
 # Here it only shows the text content of the first choice. Audio content has many binaries, so it's not displayed here.
 echo "Output of request: $(echo "$output" | jq '.choices[0].message.content')"

examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh

@@ -138,10 +138,7 @@ esac
 echo "Running query type: $QUERY_TYPE"
 echo ""
 
-
+request_body=$(cat <<EOF
-output=$(curl -sS -X POST http://localhost:8091/v1/chat/completions \
-    -H "Content-Type: application/json" \
-    -d @- <<EOF
 {
   "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
   "sampling_params_list": $sampling_params_list,
@@ -164,7 +161,12 @@ output=$(curl -sS -X POST http://localhost:8091/v1/chat/completions \
   ]
 }
 EOF
-  )
+)
+
+output=$(curl -sS --retry 3 --retry-delay 3 --retry-connrefused \
+    -X POST http://localhost:8091/v1/chat/completions \
+    -H "Content-Type: application/json" \
+    -d "$request_body")
 
 # Here it only shows the text content of the first choice. Audio content has many binaries, so it's not displayed here.
 echo "Output of request: $(echo "$output" | jq '.choices[0].message.content')"

examples/online_serving/qwen3_tts/README.md

@@ -184,6 +184,59 @@ sudo apt install ffmpeg
 
 ## API Reference
 
+### Voices Endpoint
+
+#### GET /v1/audio/voices
+
+List all available voices/speakers from the loaded model, including both built-in model voices and uploaded custom voices.
+
+**Response Example:**
+```json
+{
+  "voices": ["vivian", "ryan", "custom_voice_1"],
+  "uploaded_voices": [
+    {
+      "name": "custom_voice_1",
+      "consent": "user_consent_id",
+      "created_at": 1738660000,
+      "file_size": 1024000,
+      "mime_type": "audio/wav"
+    }
+  ]
+}
+```
+
+#### POST /v1/audio/voices
+
+Upload a new voice sample for voice cloning in Base task TTS requests.
+
+**Form Parameters:**
+- `audio_sample` (required): Audio file (max 10MB, supported formats: wav, mp3, flac, ogg, aac, webm, mp4)
+- `consent` (required): Consent recording ID
+- `name` (required): Name for the new voice
+
+**Response Example:**
+```json
+{
+  "success": true,
+  "voice": {
+    "name": "custom_voice_1",
+    "consent": "user_consent_id",
+    "created_at": 1738660000,
+    "mime_type": "audio/wav",
+    "file_size": 1024000
+  }
+}
+```
+
+**Usage Example:**
+```bash
+curl -X POST http://localhost:8000/v1/audio/voices \
+  -F "audio_sample=@/path/to/voice_sample.wav" \
+  -F "consent=user_consent_id" \
+  -F "name=custom_voice_1"
+```
+
 ### Endpoint
 
 ```
@@ -193,20 +246,6 @@ Content-Type: application/json
 
 This endpoint follows the [OpenAI Audio Speech API](https://platform.openai.com/docs/api-reference/audio/createSpeech) format with additional Qwen3-TTS parameters.
 
-### Voices Endpoint
-
-```
-GET /v1/audio/voices
-```
-
-Lists available voices for the loaded model:
-
-```json
-{
-    "voices": ["aiden", "dylan", "eric", "one_anna", "ryan", "serena", "sohee", "uncle_fu", "vivian"]
-}
-```
-
 ### Request Body
 
 ```json
@@ -250,6 +289,7 @@ Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/w
 | `language`       | string | "Auto"        | Language (see supported languages below)                                     |
 | `instructions`   | string | ""            | Voice style/emotion instructions                                             |
 | `max_new_tokens` | int    | 2048          | Maximum tokens to generate                                                   |
+| `initial_codec_chunk_frames` | int | null | Initial chunk size for reduced TTFA (overrides stage config)                 |
 | `stream`         | bool   | false         | Stream raw PCM chunks as they are decoded (requires `response_format="pcm"`) |
 
 **Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian
@@ -265,7 +305,7 @@ Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/w
 ## Streaming
 
 Set `stream=true` with `response_format="pcm"` to receive raw PCM audio chunks as they are decoded
-(one chunk per 25-frame Code2Wav window):
+(one chunk per Code2Wav window, default 25 frames; configurable in the stage config):
 
 ```bash
 curl -X POST http://localhost:8091/v1/audio/speech \

examples/online_serving/stable_audio/README.md

@@ -0,0 +1,234 @@
+# Stable Audio Online Serving
+
+Generate audio from text prompts using Stable Audio models via an OpenAI-compatible API endpoint.
+
+## Features
+
+- **OpenAI-compatible API**: Use `/v1/audio/generate` endpoint
+- **Flexible control**: Adjust audio length, guidance scale, inference steps
+- **Quality control**: Use negative prompts to avoid unwanted characteristics
+- **Reproducible**: Set random seed for deterministic generation
+
+## Quick Start
+
+### 1. Start the Server
+
+```bash
+vllm-omni serve stabilityai/stable-audio-open-1.0 \
+    --host 0.0.0.0 \
+    --port 8000 \
+    --gpu-memory-utilization 0.9 \
+    --trust-remote-code \
+    --enforce-eager \
+    --omni
+```
+
+### 2. Generate Audio
+
+#### Using curl
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound of a cat purring",
+        "audio_length": 10.0
+    }' --output cat.wav
+```
+
+#### Using Python Client
+
+```bash
+python stable_audio_client.py \
+    --text "The sound of a cat purring" \
+    --audio_length 10.0 \
+    --output cat.wav
+```
+
+#### Using Bash Script
+
+```bash
+bash curl_examples.sh
+```
+
+## API Reference
+
+### Endpoint
+
+```
+POST /v1/audio/generate
+```
+
+### Request Body
+
+```json
+{
+    "input": "Text description of the audio",
+    "audio_length": 10.0,
+    "audio_start": 0.0,
+    "negative_prompt": "Low quality",
+    "guidance_scale": 7.0,
+    "num_inference_steps": 100,
+    "seed": 42,
+    "response_format": "wav"
+}
+```
+
+### Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `input` | string | **required** | Text prompt describing the audio to generate |
+| `audio_length` | float | ~47s | Audio duration in seconds (max ~47s for stable-audio-open-1.0) |
+| `audio_start` | float | 0.0 | Audio start time in seconds |
+| `negative_prompt` | string | null | Text describing what to avoid in generation |
+| `guidance_scale` | float | 7.0 | Classifier-free guidance scale (higher = more adherence to prompt) |
+| `num_inference_steps` | int | 50 | Number of denoising steps (higher = better quality, slower) |
+| `seed` | int | null | Random seed for reproducibility |
+| `response_format` | string | "wav" | Output format: wav, mp3, flac, pcm |
+
+### Response
+
+Returns audio data in the requested format (default: WAV).
+
+## Usage Examples
+
+### Basic Generation
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound of ocean waves"
+    }' --output ocean.wav
+```
+
+### Custom Duration
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "A dog barking",
+        "audio_length": 5.0
+    }' --output dog_5s.wav
+```
+
+### High Quality with Negative Prompt
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "A piano playing a gentle melody",
+        "audio_length": 10.0,
+        "negative_prompt": "Low quality, distorted, noisy",
+        "guidance_scale": 8.0,
+        "num_inference_steps": 150
+    }' --output piano_hq.wav
+```
+
+### Reproducible Generation
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Thunder and rain sounds",
+        "audio_length": 15.0,
+        "seed": 42
+    }' --output thunder.wav
+```
+
+### Quick Generation (Fewer Steps)
+
+For faster generation with slightly lower quality:
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Birds chirping in a forest",
+        "audio_length": 8.0,
+        "num_inference_steps": 50
+    }' --output birds_quick.wav
+```
+
+## Python Client Examples
+
+### Simple Generation
+
+```bash
+python stable_audio_client.py \
+    --text "The sound of a cat purring"
+```
+
+### Custom Parameters
+
+```bash
+python stable_audio_client.py \
+    --text "Thunder and rain" \
+    --audio_length 15.0 \
+    --negative_prompt "Low quality" \
+    --guidance_scale 7.0 \
+    --num_inference_steps 100 \
+    --seed 42 \
+    --output thunder.wav
+```
+
+### Different Output Format
+
+```bash
+python stable_audio_client.py \
+    --text "Guitar playing" \
+    --response_format mp3 \
+    --output guitar.mp3
+```
+
+## Tips
+
+1. **Audio Length**: Keep under 47 seconds for `stable-audio-open-1.0`
+2. **Quality vs Speed**:
+   - 50 steps: Fast, decent quality
+   - 100 steps: Good balance (default)
+   - 150+ steps: High quality, slower
+3. **Guidance Scale**:
+   - Lower (3-5): More creative/varied
+   - Default (7): Good balance
+   - Higher (10+): More literal to prompt
+4. **Negative Prompts**: Use to avoid "Low quality", "distorted", "noisy", etc.
+5. **Seeds**: Use same seed for reproducible results
+
+## Performance
+
+| Inference Steps | Quality | Speed | Use Case |
+|----------------|---------|-------|----------|
+| 50 | Good | Fast | Quick previews |
+| 100 (default) | Very Good | Medium | Production |
+| 150+ | Excellent | Slow | Final/critical audio |
+
+## Troubleshooting
+
+### Server not responding
+- Check if server is running: `curl http://localhost:8000/health`
+- Check server logs for errors
+
+### Audio quality issues
+- Increase `num_inference_steps` (e.g., 150)
+- Add negative prompts: `"Low quality, distorted, noisy"`
+- Increase `guidance_scale` for more prompt adherence
+
+### Generation timeout
+- Reduce `num_inference_steps`
+- Reduce `audio_length`
+- Check GPU memory with `nvidia-smi`
+
+### Wrong audio length
+- Ensure `audio_length` is within model limits (~47s max)
+- Adjust `audio_start` if trimming is needed
+
+## See Also
+
+- [Offline Inference Example](../../offline_inference/text_to_audio/README.md)
+- [Stable Audio Model Card](https://huggingface.co/stabilityai/stable-audio-open-1.0)
+- [vLLM-Omni Documentation](https://github.com/vllm-project/vllm-omni)

examples/online_serving/stable_audio/curl_examples.sh

@@ -0,0 +1,54 @@
+#!/bin/bash
+# Examples for using Stable Audio with curl via /v1/audio/generate endpoint
+
+# Example 1: Simple request with default parameters
+echo "Example 1: Simple request with default parameters"
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound audience clapping and cheering in a stadium"
+    }' --output stadium.wav
+
+# Example 2: Request with custom audio_length
+echo "Example 2: Custom audio length (5 seconds)"
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound of a dog barking",
+        "audio_length": 5.0
+    }' --output dog_5s.wav
+
+# Example 3: Request with negative prompt for quality control
+echo "Example 3: With negative prompt"
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "A piano playing a gentle melody",
+        "audio_length": 10.0,
+        "negative_prompt": "Low quality, distorted, noisy"
+    }' --output piano.wav
+
+# Example 4: Full control with all parameters
+echo "Example 4: Full control (custom length, guidance, steps, seed)"
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Thunder and rain sounds",
+        "audio_length": 15.0,
+        "negative_prompt": "Low quality",
+        "guidance_scale": 7.0,
+        "num_inference_steps": 100,
+        "seed": 42
+    }' --output thunder_rain.wav
+
+# Example 5: Quick generation with fewer steps (faster but lower quality)
+echo "Example 5: Quick generation (fewer steps)"
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Ocean waves crashing on a beach",
+        "audio_length": 8.0,
+        "num_inference_steps": 50
+    }' --output ocean.wav
+
+echo "All examples completed!"

examples/online_serving/stable_audio/stable_audio_client.py

@@ -0,0 +1,170 @@
+#!/usr/bin/env python3
+"""
+OpenAI-compatible client for Stable Audio via /v1/audio/generate endpoint.
+
+This script demonstrates how to use the OpenAI-compatible speech API
+to generate audio from text using Stable Audio models.
+
+Examples:
+    # Simple generation
+    python stable_audio_client.py --text "The sound of a cat purring"
+
+    # With custom duration
+    python stable_audio_client.py --text "A dog barking" --audio_length 5.0
+
+    # With all parameters
+    python stable_audio_client.py --text "Thunder and rain" \
+        --audio_length 15.0 \
+        --negative_prompt "Low quality" \
+        --guidance_scale 7.0 \
+        --num_inference_steps 100 \
+        --seed 42 \
+        --output thunder.wav
+"""
+
+import argparse
+import sys
+
+import requests
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Generate audio with Stable Audio via OpenAI-compatible API")
+    parser.add_argument(
+        "--api_url",
+        default="http://localhost:8000/v1/audio/generate",
+        help="API endpoint URL",
+    )
+    parser.add_argument(
+        "--text",
+        default="The sound of a cat purring",
+        help="Text prompt for audio generation",
+    )
+    parser.add_argument(
+        "--audio_length",
+        type=float,
+        default=10.0,
+        help="Audio length in seconds (max ~47s for stable-audio-open-1.0)",
+    )
+    parser.add_argument(
+        "--audio_start",
+        type=float,
+        default=0.0,
+        help="Audio start time in seconds",
+    )
+    parser.add_argument(
+        "--negative_prompt",
+        default="Low quality",
+        help="Negative prompt for classifier-free guidance",
+    )
+    parser.add_argument(
+        "--guidance_scale",
+        type=float,
+        default=7.0,
+        help="Guidance scale for diffusion (higher = more adherence to prompt)",
+    )
+    parser.add_argument(
+        "--num_inference_steps",
+        type=int,
+        default=100,
+        help="Number of inference steps (higher = better quality, slower)",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=None,
+        help="Random seed for reproducibility",
+    )
+    parser.add_argument(
+        "--output",
+        default="stable_audio_output.wav",
+        help="Output file path",
+    )
+    parser.add_argument(
+        "--response_format",
+        default="wav",
+        choices=["wav", "mp3", "flac", "pcm"],
+        help="Audio output format",
+    )
+    return parser.parse_args()
+
+
+def generate_audio(args):
+    """Generate audio using the API."""
+
+    # Build request payload
+    payload = {
+        "input": args.text,
+        "audio_length": args.audio_length,
+        "audio_start": args.audio_start,
+        "response_format": args.response_format,
+    }
+
+    # Add optional parameters
+    if args.negative_prompt:
+        payload["negative_prompt"] = args.negative_prompt
+    if args.guidance_scale:
+        payload["guidance_scale"] = args.guidance_scale
+    if args.num_inference_steps:
+        payload["num_inference_steps"] = args.num_inference_steps
+    if args.seed is not None:
+        payload["seed"] = args.seed
+
+    print(f"\n{'=' * 60}")
+    print("Stable Audio - Text-to-Audio Generation")
+    print(f"{'=' * 60}")
+    print(f"API URL: {args.api_url}")
+    print(f"Prompt: {args.text}")
+    print(f"Audio length: {args.audio_length}s")
+    print(f"Negative prompt: {args.negative_prompt}")
+    print(f"Guidance scale: {args.guidance_scale}")
+    print(f"Inference steps: {args.num_inference_steps}")
+    if args.seed is not None:
+        print(f"Seed: {args.seed}")
+    print(f"Output: {args.output}")
+    print(f"{'=' * 60}\n")
+
+    try:
+        # Make the API request
+        print("Generating audio...")
+        response = requests.post(
+            args.api_url,
+            json=payload,
+            headers={"Content-Type": "application/json"},
+            timeout=300,  # 5 minute timeout for long generations
+        )
+
+        # Check for errors
+        if response.status_code != 200:
+            print(f"Error: API returned status code {response.status_code}")
+            print(f"Response: {response.text}")
+            return False
+
+        # Save the audio
+        with open(args.output, "wb") as f:
+            f.write(response.content)
+
+        print(f"✓ Audio saved to {args.output}")
+        print(f"  File size: {len(response.content) / 1024:.1f} KB")
+        return True
+
+    except requests.exceptions.Timeout:
+        print("Error: Request timed out. Try reducing inference steps or audio length.")
+        return False
+    except requests.exceptions.ConnectionError:
+        print(f"Error: Could not connect to {args.api_url}")
+        print("Make sure the server is running.")
+        return False
+    except Exception as e:
+        print(f"Error: {e}")
+        return False
+
+
+def main():
+    args = parse_args()
+    success = generate_audio(args)
+    sys.exit(0 if success else 1)
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -49,7 +49,9 @@ dev = [
     "imageio[ffmpeg]>=0.6.0",
     "opencv-python>=4.12.0.88",
     "mooncake-transfer-engine==0.3.8.post1",
-    "av" # for ComfyUI tests
+    "av", # for ComfyUI tests
+    "openpyxl>=3.0.0", # for nightly CI
+    "pyttsx3>=2.99"
 ]
 
 docs = [
@@ -202,3 +204,5 @@ extend-ignore-identifiers-re = [
     ".*[Oo]no_[Aa]nna.*",
     ".*cann.*",
 ]
+[tool.typos.default.extend-words]
+ue = "ue"

requirements/common.txt

@@ -12,6 +12,8 @@ torchsde>=0.2.6
 openai-whisper>=20250625
 imageio[ffmpeg]>=2.37.2
 sox>=1.5.0
+x-transformers>=2.12.2
+einops>=0.8.1
 prettytable>=3.8.0
 aenum==3.1.16
 pyzmq>=25.0.0

tests/conftest.py

@@ -12,6 +12,7 @@ if "VLLM_TARGET_DEVICE" not in os.environ:
 
 import concurrent.futures
 import gc
+import multiprocessing
 import socket
 import subprocess
 import sys
@@ -217,102 +218,193 @@ def generate_synthetic_audio(
     sample_rate: int = 48000,  # Default use 48000Hz.
     save_to_file: bool = False,
 ) -> dict[str, Any]:
-    """ "Generate synthetic audio with rain."""
+    """
+    Generate TTS speech with pyttsx3 and return base64 string.
+    """
+    import tempfile
+
+    import pyttsx3
     import soundfile as sf
 
-    # Initialize audio data array
+    def _pick_voice(engine: pyttsx3.Engine) -> str | None:
-    num_samples = int(sample_rate * duration)
+        voices = engine.getProperty("voices")
+        if not voices:
+            return None
+
+        preferred_tokens = (
+            "natural",
+            "jenny",
+            "sonia",
+            "susan",
+            "zira",
+            "aria",
+            "hazel",
+            "samantha",
+            "ava",
+            "allison",
+            "female",
+            "woman",
+            "english-us",
+            "en-us",
+            "english",
+        )
+        discouraged_tokens = (
+            "espeak",
+            "robot",
+            "mbrola",
+            "microsoft david",
+            "male",
+            "man",
+        )
+
+        best_voice = voices[0]
+        best_score = float("-inf")
+        for voice in voices:
+            voice_text = f"{getattr(voice, 'id', '')} {getattr(voice, 'name', '')}".lower()
+            voice_languages = " ".join(
+                lang.decode(errors="ignore") if isinstance(lang, bytes) else str(lang)
+                for lang in getattr(voice, "languages", [])
+            ).lower()
+            combined_text = f"{voice_text} {voice_languages}"
+            score = 0
+            for idx, token in enumerate(preferred_tokens):
+                if token in combined_text:
+                    score += 20 - idx
+            for token in discouraged_tokens:
+                if token in combined_text:
+                    score -= 10
+            if "english" in combined_text or "en_" in combined_text or "en-" in combined_text:
+                score += 4
+            if "en-us" in combined_text or "english-us" in combined_text:
+                score += 4
+            if score > best_score:
+                best_score = score
+                best_voice = voice
+
+        return best_voice.id
+
+    def _resample_audio(audio: np.ndarray, src_sr: int, dst_sr: int) -> np.ndarray:
+        if src_sr == dst_sr or len(audio) == 0:
+            return audio.astype(np.float32)
+
+        src_len = audio.shape[0]
+        dst_len = max(1, int(round(src_len * float(dst_sr) / float(src_sr))))
+        src_idx = np.arange(src_len, dtype=np.float32)
+        dst_idx = np.linspace(0, src_len - 1, dst_len, dtype=np.float32)
+
+        resampled_channels: list[np.ndarray] = []
+        for ch in range(audio.shape[1]):
+            resampled_channels.append(np.interp(dst_idx, src_idx, audio[:, ch]).astype(np.float32))
+        return np.stack(resampled_channels, axis=1)
+
+    def _match_channels(audio: np.ndarray, target_channels: int) -> np.ndarray:
+        current_channels = audio.shape[1]
+        if current_channels == target_channels:
+            return audio.astype(np.float32)
+        if target_channels == 1:
+            return np.mean(audio, axis=1, keepdims=True, dtype=np.float32)
+        if current_channels == 1:
+            return np.repeat(audio, target_channels, axis=1).astype(np.float32)
+
+        collapsed = np.mean(audio, axis=1, keepdims=True, dtype=np.float32)
+        return np.repeat(collapsed, target_channels, axis=1).astype(np.float32)
+
+    def _trim_silence(audio: np.ndarray, threshold: float = 0.01) -> np.ndarray:
+        if len(audio) == 0:
+            return audio
+        energy = np.max(np.abs(audio), axis=1)
+        voiced = np.where(energy > threshold)[0]
+        if len(voiced) == 0:
+            return audio
+        start = max(0, int(voiced[0]) - int(sample_rate * 0.02))
+        end = min(len(audio), int(voiced[-1]) + int(sample_rate * 0.04) + 1)
+        return audio[start:end]
+
+    def _enhance_speech(audio: np.ndarray) -> np.ndarray:
+        if len(audio) == 0:
+            return audio.astype(np.float32)
+        enhanced = audio.astype(np.float32).copy()
+        enhanced -= np.mean(enhanced, axis=0, keepdims=True, dtype=np.float32)
+        if len(enhanced) > 1:
+            preemphasis = enhanced.copy()
+            preemphasis[1:] = enhanced[1:] - 0.94 * enhanced[:-1]
+            enhanced = 0.7 * enhanced + 0.3 * preemphasis
+        # Mild dynamic-range compression for ASR/TTS robustness.
+        enhanced = np.sign(enhanced) * np.sqrt(np.abs(enhanced))
+        # Light fade to avoid clicks after trimming/repeating.
+        fade = min(len(enhanced) // 4, max(1, int(sample_rate * 0.01)))
+        if fade > 1:
+            ramp_in = np.linspace(0.0, 1.0, fade, dtype=np.float32)
+            ramp_out = np.linspace(1.0, 0.0, fade, dtype=np.float32)
+            enhanced[:fade] *= ramp_in[:, None]
+            enhanced[-fade:] *= ramp_out[:, None]
+        peak = float(np.max(np.abs(enhanced)))
+        if peak > 1e-8:
+            enhanced = enhanced / peak * 0.95
+        return enhanced.astype(np.float32)
+
+    phrase_text = "test"
+    num_samples = int(sample_rate * max(1, duration))
     audio_data = np.zeros((num_samples, num_channels), dtype=np.float32)
 
-    # Configure parameters based on rain intensity
+    engine = pyttsx3.init()
-    drop_density = 10  # Number of raindrops per second
+    engine.setProperty("rate", 112)
-    drop_volume = 0.15  # Volume of individual raindrops
+    engine.setProperty("volume", 1.0)
-    background_volume = 0.02  # Volume of background rain noise
+    selected_voice = _pick_voice(engine)
+    if selected_voice is not None:
+        engine.setProperty("voice", selected_voice)
 
-    # Pink noise sounds more natural than white noise for rain
+    temp_wav = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
-    white_noise = np.random.randn(num_samples)
+    temp_wav.close()
-    pink_noise = np.convolve(white_noise, np.ones(8) / 8, mode="same")
-    pink_noise = pink_noise / np.max(np.abs(pink_noise)) if np.max(np.abs(pink_noise)) > 0 else pink_noise
-    bg_noise = pink_noise * background_volume
 
-    # Add background noise to all channels
+    try:
-    for ch in range(num_channels):
+        engine.save_to_file(phrase_text, temp_wav.name)
-        audio_data[:, ch] += bg_noise
+        engine.runAndWait()
+        engine.stop()
 
-    # Total number of raindrops = density × duration × channels for stereo effect
+        ready = False
-    total_drops = int(drop_density * duration * num_channels)
+        for _ in range(50):
+            if os.path.exists(temp_wav.name) and os.path.getsize(temp_wav.name) > 44:
+                ready = True
+                break
+            time.sleep(0.1)
 
-    for _ in range(total_drops):
+        if not ready:
-        # Random timing for raindrop
+            raise RuntimeError("pyttsx3 did not produce a WAV file in time.")
-        drop_time = random.uniform(0, duration)
 
-        # Random duration of raindrop sound (0.01-0.05 seconds)
+        tts_audio, tts_sr = sf.read(temp_wav.name, dtype="float32", always_2d=True)
-        drop_duration = random.uniform(0.01, 0.05)
+    finally:
+        if os.path.exists(temp_wav.name):
+            os.unlink(temp_wav.name)
 
-        # Random frequency gives variation in raindrop pitch
+    if len(tts_audio) == 0:
-        drop_freq = random.uniform(500, 5000)  # Hz
+        raise RuntimeError("pyttsx3 produced an empty WAV file.")
 
-        # Random channel selection for stereo positioning
+    tts_audio = _resample_audio(tts_audio, tts_sr, sample_rate)
-        channel = random.randint(0, num_channels - 1)
+    tts_audio = _match_channels(tts_audio, num_channels)
+    tts_audio = _trim_silence(tts_audio, threshold=0.012)
+    tts_audio = _enhance_speech(tts_audio)
 
-        # Calculate sample positions for this raindrop
+    lead_silence = min(int(sample_rate * 0.02), num_samples // 8)
-        start_sample = int(drop_time * sample_rate)
+    pause_samples = int(sample_rate * 0.18)
-        drop_samples = int(drop_duration * sample_rate)
+    start = lead_silence
-        end_sample = min(start_sample + drop_samples, num_samples)
+    phrase_len = tts_audio.shape[0]
 
-        if start_sample < end_sample:
+    while start < num_samples:
-            # Generate the raindrop sound
+        take = min(phrase_len, num_samples - start)
-            num_drop_samples = end_sample - start_sample
+        audio_data[start : start + take] = tts_audio[:take]
-            t = np.arange(num_drop_samples) / sample_rate
+        start += phrase_len + pause_samples
 
-            # Basic sine wave for raindrop sound
+    max_amp = float(np.max(np.abs(audio_data)))
-            drop_sound = drop_volume * np.sin(2 * math.pi * drop_freq * t)
-
-            # Apply envelope for natural attack and decay
-            envelope = np.ones(num_drop_samples)
-            attack_samples = int(num_drop_samples * 0.1)  # 10% of samples for attack
-            decay_samples = num_drop_samples - attack_samples
-
-            if attack_samples > 0:
-                # Linear attack: volume increases from 0 to 1
-                envelope[:attack_samples] = np.linspace(0, 1, attack_samples)
-
-            if decay_samples > 0:
-                # Exponential decay for natural sound fade
-                decay = np.exp(-8 * t[attack_samples:] / drop_duration)
-                envelope[attack_samples:] = decay
-
-            # Apply envelope to raindrop sound
-            drop_sound *= envelope
-
-            # Add raindrop sound to selected channel
-            audio_data[start_sample:end_sample, channel] += drop_sound
-
-    # Step 3: Add simple reverb effect for realism
-    # Reverb simulates sound reflections in environment
-    if duration > 2:
-        # Single delay reverb (100ms delay)
-        delay_samples = int(0.1 * sample_rate)
-        if delay_samples < num_samples:
-            for ch in range(num_channels):
-                delayed = np.zeros(num_samples)
-                delayed[delay_samples:] = audio_data[:-delay_samples, ch] * 0.3
-                audio_data[:, ch] += delayed
-
-    # Step 4: Normalize audio to prevent clipping
-    # Find maximum amplitude and scale to 80% of maximum volume
-    max_amp = np.max(np.abs(audio_data))
     if max_amp > 0:
-        audio_data = audio_data / max_amp * 0.8
+        audio_data = audio_data / max_amp * 0.95
 
-    audio_array = audio_data.copy()
+    audio_bytes: bytes | None = None
-    result = {
+    output_path: str | None = None
-        "np_array": audio_array,
+    result: dict[str, Any] = {
+        "np_array": audio_data.copy(),
     }
 
-    # Handle file saving
-    audio_bytes = None
-
     if save_to_file:
         timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
         output_path = f"audio_{num_channels}ch_{timestamp}.wav"
@@ -617,24 +709,38 @@ def convert_audio_to_text(audio_data):
     return text
 
 
-def convert_audio_file_to_text(output_path):
+def _whisper_transcribe_in_current_process(output_path: str) -> str:
     import whisper
 
-    model = whisper.load_model("small")
+    device = "cuda" if torch.cuda.is_available() else "cpu"
-    text = model.transcribe(
+    model = whisper.load_model("small", device=device)
-        output_path,
+    try:
-        temperature=0.0,
+        text = model.transcribe(
-        word_timestamps=True,
+            output_path,
-        condition_on_previous_text=False,
+            temperature=0.0,
-    )["text"]
+            word_timestamps=True,
-    del model
+            condition_on_previous_text=False,
-    if torch.cuda.is_available():
+        )["text"]
+    finally:
+        # Sync GPU so in-flight ops finish before we free the model; otherwise
+        # freed memory may not show up until those ops complete.
+        if torch is not None and torch.cuda.is_available():
+            torch.cuda.synchronize()
+        del model
         gc.collect()
-        torch.cuda.empty_cache()
+        if torch is not None and torch.cuda.is_available():
-    if text:
+            torch.cuda.empty_cache()
-        return text
+
-    else:
+    return text or ""
-        return ""
+
+
+def convert_audio_file_to_text(output_path: str) -> str:
+    """Convert an audio file to text in an isolated subprocess."""
+    # Import locally to avoid impacting test module import time.
+    ctx = multiprocessing.get_context("spawn")
+    with concurrent.futures.ProcessPoolExecutor(max_workers=1, mp_context=ctx) as executor:
+        future = executor.submit(_whisper_transcribe_in_current_process, output_path)
+        return future.result()
 
 
 def merge_base64_and_convert_to_text(base64_list):
@@ -1105,7 +1211,7 @@ def assert_omni_response(response: OmniResponse, request_config: dict[str, Any],
     assert response.success, "The request failed."
     e2e_latency = response.e2e_latency
     if e2e_latency is not None:
-        print(f"the avg e2e is: {e2e_latency}")
+        print(f"the avg e2e latency is: {e2e_latency}")
 
     modalities = request_config.get("modalities", ["text", "audio"])
 

tests/diffusion/attention/test_attention_sp.py

@@ -414,7 +414,8 @@ def ulysses_attention_on_test_model(
             # Generate and save full input data with fixed seed
             # Reinitialize RNG to ensure reproducibility
             torch.manual_seed(42)
-            current_omni_platform.seed_everything(42)
+            if torch.cuda.is_available():
+                torch.cuda.manual_seed_all(42)
             full_hidden_states = torch.randn(
                 (batch_size, seq_len, hidden_size),
                 dtype=dtype,

tests/diffusion/distributed/test_distributed_vae_executor.py

@@ -0,0 +1,165 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+import torch.distributed as dist
+
+from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+    DistributedVaeExecutor,
+    DistributedVaeMixin,
+    GridSpec,
+    TileTask,
+)
+
+
+class E2EOperator:
+    """tiles with (2, 3) -- (H,W)"""
+
+    def split(self, z):
+        rows_num = 2
+        cols_num = 3
+
+        h_size = z.shape[0]
+        w_size = z.shape[1]
+
+        tasks = []
+        for i in range(rows_num):
+            for j in range(cols_num):
+                tasks.append(
+                    TileTask(
+                        tile_id=len(tasks),
+                        grid_coord=(i, j),
+                        tensor=z[
+                            ((i * h_size) // rows_num) : (((i + 1) * h_size) // rows_num),
+                            ((j * w_size) // cols_num) : (((j + 1) * w_size) // cols_num),
+                        ],
+                    )
+                )
+
+        grid_spec = GridSpec(
+            split_dims=(0, 1),
+            grid_shape=(2, 3),
+        )
+        return tasks, grid_spec
+
+    def exec(self, task: TileTask):
+        return torch.full_like(task.tensor, fill_value=task.tile_id)
+
+    def merge(self, coord_tensor_map, grid_spec):
+        tiles = []
+        for r in range(grid_spec.grid_shape[0]):
+            row_tiles = []
+            for c in range(grid_spec.grid_shape[1]):
+                coord = (r, c)
+                row_tiles.append(coord_tensor_map[coord])
+            tiles.append(torch.cat(row_tiles, dim=1))
+        return torch.cat(tiles, dim=0)
+
+
+class DummyMixin(DistributedVaeMixin):
+    def __init__(self):
+        self.use_tiling = True
+        self.distributed_decoder = MagicMock()
+        self.distributed_decoder.parallel_size = 2
+        self.distributed_decoder.group = None
+
+
+@pytest.fixture(autouse=True)
+def mock_dist():
+    with (
+        patch.object(dist, "get_world_size", return_value=2),
+        patch.object(dist, "get_rank", return_value=0),
+        patch.object(dist, "is_initialized", return_value=True),
+        patch.object(dist, "gather", return_value=None),
+        patch.object(dist, "broadcast", return_value=None),
+    ):
+        yield
+
+
+@pytest.fixture(autouse=True)
+def mock_dit_group():
+    with patch(
+        "vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor.get_dit_group",
+        new=MagicMock(return_value=None),
+    ):
+        yield
+
+
+@pytest.fixture(autouse=True)
+def mock_dist_vae_executor():
+    with (
+        patch.object(DistributedVaeExecutor, "gather_tensors", side_effect=lambda x: [x]),
+        patch.object(DistributedVaeExecutor, "broadcast_tensor", side_effect=lambda x: x),
+    ):
+        yield
+
+
+# ============================
+# Unitest
+# ============================
+
+
+def test_balance_tasks():
+    executor = DistributedVaeExecutor()
+    workloads = [2, 5, 13, 8, 2, 4]  # each is 17
+    tasks = [TileTask(0, (i,), torch.tensor([i]), workload=load) for i, load in enumerate(workloads)]
+    assigned = executor._balance_tasks(tasks, num_rank=2)
+    assert len(assigned) == 2
+    total_work = [sum(t.workload for t in group) for group in assigned]
+    assert total_work[0] == total_work[1]
+
+
+def test_compute_global_padding_shape():
+    executor = DistributedVaeExecutor()
+    executor.rank = 0
+
+    local_results = [(0, torch.zeros((2, 3))), (1, torch.zeros((4, 2)))]
+    shape = executor._compute_global_padding_shape(local_results, 2, "cpu")
+
+    assert shape == [2, 4, 3]
+
+
+def test_pack_and_unpack():
+    executor = DistributedVaeExecutor()
+    executor.world_size = 1
+
+    grid_spec = GridSpec(split_dims=(0, 1), grid_shape=(2, 2))
+
+    # ======================
+    # pack
+    # ======================
+    local_results = [(0, torch.tensor([[1, 2], [3, 4]]))]
+
+    global_shape = [3, 3, 3]  # (tiles, H, W)
+
+    tile_tensor, meta_tensor = executor._pack_local_tiles(
+        local_results, global_shape, grid_spec, device="cpu", dtype=torch.int64
+    )
+
+    # check pack
+    assert tile_tensor.shape == torch.Size(global_shape)
+    assert meta_tensor.shape == (global_shape[0], len(grid_spec.split_dims) + 1)
+    assert meta_tensor[0, 0] == 0
+    assert meta_tensor[0, 1] == 2
+    assert meta_tensor[0, 2] == 2
+
+    # ======================
+    # unpack
+    # ======================
+    meta_gather = [meta_tensor]
+    tile_gather = [tile_tensor]
+
+    tid_coord_map = {0: (0, 0)}
+
+    coord_tensor_map = executor._unpack_tiles(meta_gather, tile_gather, grid_spec, tid_coord_map)
+
+    # check unpack
+    assert torch.equal(coord_tensor_map[(0, 0)], torch.tensor([[1, 2], [3, 4]]))
+
+
+def test_is_distributed_enabled():
+    mixin = DummyMixin()
+    assert mixin.is_distributed_enabled() is True
+
+    mixin.use_tiling = False
+    assert mixin.is_distributed_enabled() is False

tests/diffusion/distributed/test_hsdp.py

@@ -156,13 +156,22 @@ class TestDiffusionParallelConfigHSDP:
 
     def test_hsdp_cannot_use_with_tp(self):
         """Test that HSDP and Tensor Parallelism cannot be used together."""
-        with pytest.raises(ValueError, match="cannot be used with Tensor Parallelism"):
+        with pytest.raises(ValueError, match="cannot be used with TP or DP"):
             DiffusionParallelConfig(
                 tensor_parallel_size=2,
                 use_hsdp=True,
                 hsdp_shard_size=4,
             )
 
+    def test_hsdp_cannot_use_with_dp(self):
+        """Test that HSDP and Data Parallelism cannot be used together."""
+        with pytest.raises(ValueError, match="cannot be used with TP or DP"):
+            DiffusionParallelConfig(
+                data_parallel_size=2,
+                use_hsdp=True,
+                hsdp_shard_size=4,
+            )
+
     def test_from_dict_with_hsdp(self):
         """Test creating config from dict with HSDP settings."""
         config = DiffusionParallelConfig.from_dict(
@@ -177,6 +186,158 @@ class TestDiffusionParallelConfigHSDP:
         assert config.hsdp_shard_size == 2  # auto: 4 // 2
 
 
+class TestStandaloneHSDPDetection:
+    """Tests for standalone HSDP detection and dit_parallel_size calculation.
+
+    These tests verify the logic used in initialize_model_parallel() to detect
+    standalone HSDP mode and calculate effective parallel sizes.
+
+    Standalone HSDP is when all non-HSDP parallelism dimensions are 1.
+    """
+
+    @staticmethod
+    def compute_standalone_hsdp_params(
+        data_parallel_size: int = 1,
+        cfg_parallel_size: int = 1,
+        sequence_parallel_size: int = 1,
+        pipeline_parallel_size: int = 1,
+        tensor_parallel_size: int = 1,
+        fully_shard_degree: int = 1,
+        hsdp_replicate_size: int = 1,
+    ) -> dict:
+        """Compute standalone HSDP detection parameters.
+
+        This mirrors the logic in initialize_model_parallel().
+        """
+        dit_parallel_size = (
+            data_parallel_size
+            * cfg_parallel_size
+            * sequence_parallel_size
+            * pipeline_parallel_size
+            * tensor_parallel_size
+        )
+
+        # Check for standalone HSDP: all non-HSDP parallelism dimensions are 1
+        is_standalone_hsdp = dit_parallel_size == 1 and fully_shard_degree > 1
+
+        # For standalone HSDP: use (fully_shard_degree * hsdp_replicate_size)
+        if is_standalone_hsdp:
+            effective_dit_parallel_size = fully_shard_degree * hsdp_replicate_size
+        else:
+            effective_dit_parallel_size = dit_parallel_size
+
+        effective_dp_size = (fully_shard_degree * hsdp_replicate_size) if is_standalone_hsdp else data_parallel_size
+
+        return {
+            "original_dit_parallel_size": dit_parallel_size,
+            "is_standalone_hsdp": is_standalone_hsdp,
+            "effective_dit_parallel_size": effective_dit_parallel_size,
+            "effective_dp_size": effective_dp_size,
+        }
+
+    def test_standalone_hsdp_basic(self):
+        """Test basic standalone HSDP detection (shard_size=4, replicate=1)."""
+        result = self.compute_standalone_hsdp_params(
+            fully_shard_degree=4,
+            hsdp_replicate_size=1,
+        )
+        assert result["original_dit_parallel_size"] == 1
+        assert result["is_standalone_hsdp"] is True
+        assert result["effective_dit_parallel_size"] == 4
+        assert result["effective_dp_size"] == 4
+
+    def test_standalone_hsdp_with_replicate(self):
+        """Test standalone HSDP with replication (shard_size=4, replicate=2)."""
+        result = self.compute_standalone_hsdp_params(
+            fully_shard_degree=4,
+            hsdp_replicate_size=2,
+        )
+        assert result["original_dit_parallel_size"] == 1
+        assert result["is_standalone_hsdp"] is True
+        assert result["effective_dit_parallel_size"] == 8  # 4 * 2
+        assert result["effective_dp_size"] == 8
+
+    def test_combined_hsdp_sp_not_standalone(self):
+        """Test HSDP combined with SP is NOT detected as standalone.
+
+        This is a regression test for the bug where the condition
+        `dit_parallel_size == fully_shard_degree` incorrectly matched
+        combined modes like SP=4 + HSDP=4.
+        """
+        result = self.compute_standalone_hsdp_params(
+            sequence_parallel_size=4,
+            fully_shard_degree=4,
+            hsdp_replicate_size=1,
+        )
+        assert result["original_dit_parallel_size"] == 4
+        assert result["is_standalone_hsdp"] is False
+        # Should NOT override dp_size for combined mode
+        assert result["effective_dp_size"] == 1  # original data_parallel_size
+
+    def test_combined_hsdp_cfg_not_standalone(self):
+        """Test HSDP combined with CFG is NOT detected as standalone."""
+        result = self.compute_standalone_hsdp_params(
+            cfg_parallel_size=2,
+            fully_shard_degree=4,
+            hsdp_replicate_size=1,
+        )
+        assert result["original_dit_parallel_size"] == 2
+        assert result["is_standalone_hsdp"] is False
+        assert result["effective_dp_size"] == 1
+
+    def test_combined_hsdp_pp_not_standalone(self):
+        """Test HSDP combined with PP is NOT detected as standalone."""
+        result = self.compute_standalone_hsdp_params(
+            pipeline_parallel_size=2,
+            fully_shard_degree=4,
+            hsdp_replicate_size=1,
+        )
+        assert result["original_dit_parallel_size"] == 2
+        assert result["is_standalone_hsdp"] is False
+        assert result["effective_dp_size"] == 1
+
+    def test_no_hsdp_not_standalone(self):
+        """Test that no HSDP (fully_shard_degree=1) is NOT standalone."""
+        result = self.compute_standalone_hsdp_params(
+            fully_shard_degree=1,
+        )
+        assert result["original_dit_parallel_size"] == 1
+        assert result["is_standalone_hsdp"] is False
+        assert result["effective_dp_size"] == 1
+
+    def test_combined_multiple_parallelism_not_standalone(self):
+        """Test HSDP combined with multiple parallelism is NOT standalone."""
+        result = self.compute_standalone_hsdp_params(
+            sequence_parallel_size=2,
+            cfg_parallel_size=2,
+            fully_shard_degree=4,
+            hsdp_replicate_size=1,
+        )
+        assert result["original_dit_parallel_size"] == 4  # 2 * 2
+        assert result["is_standalone_hsdp"] is False
+        assert result["effective_dp_size"] == 1
+
+    def test_standalone_hsdp_large_shard(self):
+        """Test standalone HSDP with large shard size."""
+        result = self.compute_standalone_hsdp_params(
+            fully_shard_degree=8,
+            hsdp_replicate_size=1,
+        )
+        assert result["is_standalone_hsdp"] is True
+        assert result["effective_dit_parallel_size"] == 8
+        assert result["effective_dp_size"] == 8
+
+    def test_standalone_hsdp_large_replicate(self):
+        """Test standalone HSDP with large replicate size."""
+        result = self.compute_standalone_hsdp_params(
+            fully_shard_degree=4,
+            hsdp_replicate_size=4,
+        )
+        assert result["is_standalone_hsdp"] is True
+        assert result["effective_dit_parallel_size"] == 16  # 4 * 4
+        assert result["effective_dp_size"] == 16
+
+
 class TestHSDPShardConditions:
     """Tests for _hsdp_shard_conditions matching logic."""
 

tests/diffusion/lora/test_lora_manager.py

@@ -182,7 +182,11 @@ def test_lora_manager_replaces_packed_layer_when_targeting_sublayers(monkeypatch
     monkeypatch.setattr(manager_mod, "replace_submodule", _fake_replace_submodule)
 
     pipeline = torch.nn.Module()
-    pipeline.packed_modules_mapping = {"to_qkv": ["to_q", "to_k", "to_v"]}
+    pipeline.stacked_params_mapping = [
+        (".to_qkv.", ".to_q.", "q"),
+        (".to_qkv.", ".to_k.", "k"),
+        (".to_qkv.", ".to_v.", "v"),
+    ]
     pipeline.transformer = torch.nn.Module()
     pipeline.transformer.to_qkv = _FakeLinearBase()
 
@@ -194,7 +198,7 @@ def test_lora_manager_replaces_packed_layer_when_targeting_sublayers(monkeypatch
     )
 
     # Treat the dummy layer as a packed 3-slice projection so the manager uses
-    # `packed_modules_mapping` to decide replacement based on target_modules.
+    # `stacked_params_mapping` to decide replacement based on target_modules.
     monkeypatch.setattr(manager, "_get_packed_modules_list", lambda _module: ["q", "k", "v"])
 
     peft_helper = type("_PH", (), {"r": 1, "target_modules": ["to_q"]})()
@@ -254,7 +258,11 @@ def test_lora_manager_activates_fused_lora_on_packed_layer():
 
 def test_lora_manager_activates_packed_lora_from_sublayers():
     pipeline = torch.nn.Module()
-    pipeline.packed_modules_mapping = {"to_qkv": ["to_q", "to_k", "to_v"]}
+    pipeline.stacked_params_mapping = [
+        (".to_qkv", ".to_q", "q"),
+        (".to_qkv", ".to_k", "k"),
+        (".to_qkv", ".to_v", "v"),
+    ]
     manager = DiffusionLoRAManager(
         pipeline=pipeline,
         device=torch.device("cpu"),

tests/diffusion/models/bagel/test_combine_cfg.py

@@ -0,0 +1,316 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for Bagel._combine_cfg logic."""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.bagel.bagel_transformer import Bagel
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestCombineCfg:
+    """Tests for the _combine_cfg static method."""
+
+    def _make_tensors(self, shape=(10, 64), seed=42):
+        """Create deterministic test tensors."""
+        gen = torch.Generator().manual_seed(seed)
+        v_t = torch.randn(shape, generator=gen)
+        cfg_text_v_t = torch.randn(shape, generator=gen)
+        cfg_img_v_t = torch.randn(shape, generator=gen)
+        return v_t, cfg_text_v_t, cfg_img_v_t
+
+    def test_text_channel_renorm_preserves_direction(self):
+        """text_channel renorm should change direction but constrain magnitude."""
+        v_t, cfg_text_v_t, _ = self._make_tensors()
+
+        result = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            None,
+            cfg_text_scale=4.0,
+            cfg_img_scale=1.0,
+            cfg_renorm_type="text_channel",
+            cfg_renorm_min=0.0,
+        )
+
+        # Result norm per token should be <= original v_t norm (clamp max=1.0)
+        result_norm = torch.norm(result, dim=-1)
+        v_t_norm = torch.norm(v_t, dim=-1)
+        assert torch.all(result_norm <= v_t_norm + 1e-6), "text_channel renorm should not increase per-token norm"
+
+    def test_scale_1_returns_v_t(self):
+        """cfg_text_scale=1.0 means no CFG: result should equal v_t."""
+        v_t, cfg_text_v_t, _ = self._make_tensors()
+
+        result = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            None,
+            cfg_text_scale=1.0,
+            cfg_img_scale=1.0,
+            cfg_renorm_type="text_channel",
+            cfg_renorm_min=0.0,
+        )
+
+        # scale=1 → v_t_text_ = cfg_text + 1*(v_t - cfg_text) = v_t
+        # renorm scale = norm(v_t)/norm(v_t) = 1.0, so result = v_t
+        assert torch.allclose(result, v_t, atol=1e-6)
+
+    def test_img_cfg_applied_when_scale_gt_1(self):
+        """When cfg_img_scale > 1.0, result should differ from text-only CFG."""
+        v_t, cfg_text_v_t, cfg_img_v_t = self._make_tensors()
+
+        text_only = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            None,
+            cfg_text_scale=4.0,
+            cfg_img_scale=1.0,
+            cfg_renorm_type="text_channel",
+            cfg_renorm_min=0.0,
+        )
+
+        with_img = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            cfg_img_v_t,
+            cfg_text_scale=4.0,
+            cfg_img_scale=1.5,
+            cfg_renorm_type="text_channel",
+            cfg_renorm_min=0.0,
+        )
+
+        assert not torch.allclose(text_only, with_img, atol=1e-6), (
+            "Image CFG should produce different result from text-only CFG"
+        )
+
+    def test_img_cfg_none_ignored(self):
+        """cfg_img_v_t=None should be equivalent to cfg_img_scale <= 1.0."""
+        v_t, cfg_text_v_t, cfg_img_v_t = self._make_tensors()
+
+        result_none = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            None,
+            cfg_text_scale=4.0,
+            cfg_img_scale=1.5,
+            cfg_renorm_type="text_channel",
+            cfg_renorm_min=0.0,
+        )
+
+        result_low_scale = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            cfg_img_v_t,
+            cfg_text_scale=4.0,
+            cfg_img_scale=0.5,
+            cfg_renorm_type="text_channel",
+            cfg_renorm_min=0.0,
+        )
+
+        assert torch.allclose(result_none, result_low_scale, atol=1e-6), (
+            "cfg_img_v_t=None and cfg_img_scale<=1.0 should give same result"
+        )
+
+    def test_global_renorm(self):
+        """global renorm should produce valid output without error."""
+        v_t, cfg_text_v_t, cfg_img_v_t = self._make_tensors()
+
+        result = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            cfg_img_v_t,
+            cfg_text_scale=4.0,
+            cfg_img_scale=1.5,
+            cfg_renorm_type="global",
+            cfg_renorm_min=0.0,
+        )
+
+        assert result.shape == v_t.shape
+        assert not torch.any(torch.isnan(result))
+
+    def test_channel_renorm(self):
+        """channel renorm should produce valid output without error."""
+        v_t, cfg_text_v_t, cfg_img_v_t = self._make_tensors()
+
+        result = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            cfg_img_v_t,
+            cfg_text_scale=4.0,
+            cfg_img_scale=1.5,
+            cfg_renorm_type="channel",
+            cfg_renorm_min=0.0,
+        )
+
+        assert result.shape == v_t.shape
+        assert not torch.any(torch.isnan(result))
+
+    def test_invalid_renorm_type_raises(self):
+        """Unknown renorm type should raise NotImplementedError."""
+        v_t, cfg_text_v_t, _ = self._make_tensors()
+
+        with pytest.raises(NotImplementedError):
+            Bagel._combine_cfg(
+                v_t,
+                cfg_text_v_t,
+                None,
+                cfg_text_scale=4.0,
+                cfg_img_scale=1.0,
+                cfg_renorm_type="unknown",
+                cfg_renorm_min=0.0,
+            )
+
+    def test_renorm_min_clamps_scale(self):
+        """cfg_renorm_min should prevent scale from going too low."""
+        v_t = torch.ones(10, 64)
+        # Make cfg_text_v_t very different so CFG amplifies heavily
+        cfg_text_v_t = torch.zeros(10, 64)
+
+        result_no_min = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            None,
+            cfg_text_scale=100.0,
+            cfg_img_scale=1.0,
+            cfg_renorm_type="text_channel",
+            cfg_renorm_min=0.0,
+        )
+
+        result_with_min = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            None,
+            cfg_text_scale=100.0,
+            cfg_img_scale=1.0,
+            cfg_renorm_type="text_channel",
+            cfg_renorm_min=0.5,
+        )
+
+        # With higher renorm_min, result magnitude should be larger
+        # (scale is clamped to at least 0.5 instead of going near 0)
+        norm_no_min = torch.norm(result_no_min)
+        norm_with_min = torch.norm(result_with_min)
+        assert norm_with_min >= norm_no_min - 1e-6, "Higher cfg_renorm_min should preserve more magnitude"
+
+    def test_global_renorm_with_img_cfg(self):
+        """global renorm + img CFG should produce valid, different output."""
+        v_t, cfg_text_v_t, cfg_img_v_t = self._make_tensors()
+
+        text_only = Bagel._combine_cfg(
+            v_t.clone(),
+            cfg_text_v_t.clone(),
+            None,
+            cfg_text_scale=4.0,
+            cfg_img_scale=1.0,
+            cfg_renorm_type="global",
+            cfg_renorm_min=0.0,
+        )
+
+        with_img = Bagel._combine_cfg(
+            v_t.clone(),
+            cfg_text_v_t.clone(),
+            cfg_img_v_t.clone(),
+            cfg_text_scale=4.0,
+            cfg_img_scale=1.5,
+            cfg_renorm_type="global",
+            cfg_renorm_min=0.0,
+        )
+
+        assert not torch.allclose(text_only, with_img, atol=1e-6), (
+            "global renorm + img CFG should differ from text-only"
+        )
+        assert not torch.any(torch.isnan(with_img))
+
+    def test_channel_renorm_with_img_cfg(self):
+        """channel renorm + img CFG should produce valid, different output."""
+        v_t, cfg_text_v_t, cfg_img_v_t = self._make_tensors()
+
+        text_only = Bagel._combine_cfg(
+            v_t.clone(),
+            cfg_text_v_t.clone(),
+            None,
+            cfg_text_scale=4.0,
+            cfg_img_scale=1.0,
+            cfg_renorm_type="channel",
+            cfg_renorm_min=0.0,
+        )
+
+        with_img = Bagel._combine_cfg(
+            v_t.clone(),
+            cfg_text_v_t.clone(),
+            cfg_img_v_t.clone(),
+            cfg_text_scale=4.0,
+            cfg_img_scale=1.5,
+            cfg_renorm_type="channel",
+            cfg_renorm_min=0.0,
+        )
+
+        assert not torch.allclose(text_only, with_img, atol=1e-6), (
+            "channel renorm + img CFG should differ from text-only"
+        )
+        assert not torch.any(torch.isnan(with_img))
+
+    def test_global_channel_renorm_constrains_norm(self):
+        """global and channel renorm should not increase overall norm."""
+        v_t, cfg_text_v_t, cfg_img_v_t = self._make_tensors()
+
+        for renorm_type in ["global", "channel"]:
+            result = Bagel._combine_cfg(
+                v_t.clone(),
+                cfg_text_v_t.clone(),
+                cfg_img_v_t.clone(),
+                cfg_text_scale=4.0,
+                cfg_img_scale=1.5,
+                cfg_renorm_type=renorm_type,
+                cfg_renorm_min=0.0,
+            )
+            # Global norm of result should be <= global norm of v_t (clamp max=1.0)
+            assert torch.norm(result) <= torch.norm(v_t) + 1e-5, f"{renorm_type} renorm should not increase global norm"
+
+    def test_text_channel_img_cfg_no_second_renorm(self):
+        """text_channel mode: img CFG is applied AFTER renorm, without a second renorm.
+        So the result norm can exceed v_t norm when img_scale > 1."""
+        v_t, cfg_text_v_t, cfg_img_v_t = self._make_tensors()
+
+        result = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            cfg_img_v_t,
+            cfg_text_scale=4.0,
+            cfg_img_scale=2.0,
+            cfg_renorm_type="text_channel",
+            cfg_renorm_min=0.0,
+        )
+
+        # text_channel renorms after text CFG, then applies img CFG without renorm
+        # So result norm CAN exceed v_t norm — this is expected behavior
+        assert result.shape == v_t.shape
+        assert not torch.any(torch.isnan(result))
+
+    def test_all_renorm_types_consistent_direction(self):
+        """All renorm types should guide in the same general direction."""
+        v_t, cfg_text_v_t, _ = self._make_tensors()
+
+        results = {}
+        for renorm_type in ["text_channel", "global", "channel"]:
+            results[renorm_type] = Bagel._combine_cfg(
+                v_t.clone(),
+                cfg_text_v_t.clone(),
+                None,
+                cfg_text_scale=4.0,
+                cfg_img_scale=1.0,
+                cfg_renorm_type=renorm_type,
+                cfg_renorm_min=0.0,
+            )
+
+        # All results should have positive cosine similarity with each other
+        for a_name, a in results.items():
+            for b_name, b in results.items():
+                cos_sim = torch.nn.functional.cosine_similarity(a.flatten().unsqueeze(0), b.flatten().unsqueeze(0))
+                assert cos_sim > 0.5, (
+                    f"{a_name} and {b_name} should point in similar direction, "
+                    f"but cosine similarity = {cos_sim.item():.4f}"
+                )

tests/diffusion/test_diffusion_model_runner.py

@@ -56,7 +56,10 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar
         enable_cache_dit_summary=enable_cache_dit_summary,
         parallel_config=SimpleNamespace(use_hsdp=False),
     )
-    runner.kv_transfer_manager = SimpleNamespace(receive_kv_cache=lambda req, target_device: None)
+    runner.kv_transfer_manager = SimpleNamespace(
+        receive_kv_cache=lambda req, target_device=None: None,
+        receive_multi_kv_cache=lambda req, cfg_kv_collect_func=None, target_device=None: None,
+    )
     return runner
 
 
@@ -105,8 +108,8 @@ def test_execute_model_emits_cache_summary_with_active_cache_dit_backend(monkeyp
 
 def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
     class _DummyLoader:
-        def __init__(self, load_config):
+        def __init__(self, load_config, od_config=None):
-            del load_config
+            del load_config, od_config
 
         def load_model(self, **kwargs):
             del kwargs

tests/diffusion/test_worker_wrapper_base.py

@@ -96,6 +96,7 @@ class TestWorkerWrapperBaseInitialization:
             local_rank=0,
             rank=0,
             od_config=mock_od_config,
+            skip_load_model=False,
         )
 
 

tests/distributed/omni_connectors/test_basic_connectors.py

@@ -66,14 +66,13 @@ def shm_connector():
 
 
 def test_put_get_inline(shm_connector):
-    """Test inline transfer for small data."""
+    """Test transfer for small data (currently always uses SHM)."""
     data = {"small": "data"}
-    # Ensure data is smaller than threshold (100 bytes)
 
     success, size, metadata = shm_connector.put("stage_0", "stage_1", "req_1", data)
     assert success is True
-    assert "inline_bytes" in metadata
+    assert "shm" in metadata
-    assert "shm" not in metadata
+    assert "size" in metadata
 
     # Retrieve
     retrieved_data, ret_size = shm_connector.get("stage_0", "stage_1", "req_1", metadata)
@@ -87,7 +86,7 @@ def test_put_get_shm(mocker: MockerFixture, shm_connector, monkeypatch: pytest.M
     data = {"large": "x" * 200}
 
     # Mock SHM return values
-    mock_handle = {"name": "test_shm", "size": 200}
+    mock_handle = {"name": "req_2", "size": 200}
     mock_write = mocker.MagicMock(return_value=mock_handle)
     monkeypatch.setattr("vllm_omni.distributed.omni_connectors.connectors.shm_connector.shm_write_bytes", mock_write)
 

tests/distributed/omni_connectors/test_kv_flow.py

@@ -6,6 +6,7 @@ from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
     OmniKVCacheConfig,
     OmniKVTransferManager,
 )
+from vllm_omni.distributed.omni_connectors.utils.kv_utils import normalize_layer_kv
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.cache]
@@ -192,8 +193,7 @@ def test_normalize_layer_kv_rejects_invalid_inputs(kv_config, common_constants,
     else:
         layer_kv = (torch.randn(2, block_size, num_heads, head_dim), "not-a-tensor")
 
-    manager = OmniKVTransferManager(kv_config)
+    normalized = normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=0)
-    normalized = manager._normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=0)
     assert normalized is None