[skip ci]Update slides links in README&Docs (#1828)

Signed-off-by: hsliu_ustc <hsliu_ustc@noreply.gitcode.com>
Co-authored-by: hsliu_ustc <hsliu_ustc@noreply.gitcode.com>

.buildkite/test-nightly.yml

@@ -54,7 +54,6 @@ steps:
       - docker#v5.2.0:
           image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
           always-pull: true
-          shm-size: "8gb"
           propagate-environment: true
           shm-size: "8gb"
           environment:
@@ -62,6 +61,42 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
+  - label: ":full_moon: Diffusion Model Test with H100"
+    timeout_in_minutes: 60
+    depends_on: upload-nightly-pipeline
+    # if: build.env("NIGHTLY") == "1"
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and diffusion and H100" --run-level "advanced_model"
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 2
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
   - label: ":full_moon: Qwen3-TTS Non-Async-Chunk E2E Test"
     timeout_in_minutes: 30
     depends_on: upload-nightly-pipeline

.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

@@ -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://drive.google.com/file/d/12k_DOo2xoEROvd4xQxYNuEi6C2GECoRx/view?usp=drive_link"><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://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true"><b>Slides</b></a> |
 </p>
 
 

apps/ComfyUI-vLLM-Omni/README.md

@@ -47,6 +47,7 @@ If no, check your shell running the ComfyUI process. There may be some error mes
 This extension offers the following nodes based on the output modalities (at **ComfyUI sidebar -> Node Library**):
 
 - **Generate Image** for text-to-image and image-to-image tasks
+- **Generate Video** for text-to-video and image-to-video tasks
 - **Multimodality Understanding** for multimodality-to-text and multimodality-to-audio tasks
 - **TTS** and **TTS Voice Clone** for TTS tasks
 
@@ -59,7 +60,7 @@ To build a simple workflow yourself,
 
 - Drag a generation node onto the canvas.
 - Depending on your need, grab built-in multimedia file loader nodes, such as **image->Load Image**, **image->video->Load Video**, **audio->Load Audio**
-- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**. For text output, you can install [ComfyUI-Custom-Scripts plugin](https://github.com/pythongosssss/ComfyUI-Custom-Scripts/) and grab its **utils->Show Text 🐍** node.
+- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**, **utils->Preview as Text**.
 - If you want to tune sampling parameters, grab corresponding nodes from **vLLM-Omni-> Sampling Params**.
     - For multi-stage models, you can connect multiple **AR Sampling Params** and **Diffusion Sampling Params** nodes to a **Multi-Stage Sampling Params List** node, and connect this node to the generation node.
     - For some multi-stage models like BAGEL, [only one stage's sampling parameters are exposed and tunable via vLLM-Omni's online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/). Thus, these models are treated as single-stage ones. Please check the vLLM-Omni documentation on how to correctly set each model's sampling parameters.
@@ -104,6 +105,20 @@ You can configure per-stage sampling parameters for multi-stage models.
 > [!TIP]
 > The node automatically choose text-to-image or image-to-image API endpoints depending on whether you connect an image input or not.
 
+### Text-to-video and image-to-video generation (e.g., Wan)
+
+(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Video Generation**)
+
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-video-generation.jpg">
+    <img alt="vLLM-Omni video generation" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-video-generation.jpg" width=55%>
+  </picture>
+</p>
+
+> [!TIP]
+> The node automatically choose text-to-video or image-to-video API endpoints depending on whether you connect an image input or not.
+
 ### TTS (e.g., Qwen TTS series)
 
 (Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni TTS**)

apps/ComfyUI-vLLM-Omni/__init__.py

@@ -14,17 +14,21 @@ from .comfyui_vllm_omni.nodes import (
     VLLMOmniARSampling,
     VLLMOmniDiffusionSampling,
     VLLMOmniGenerateImage,
+    VLLMOmniGenerateVideo,
     VLLMOmniQwenTTSParams,
+    VLLMOmniRemoteLoRA,
     VLLMOmniSamplingParamsList,
     VLLMOmniTTS,
     VLLMOmniUnderstanding,
     VLLMOmniVoiceClone,
+    VLLMOmniWanParams,
 )
 
 # A dictionary that contains all nodes you want to export with their names
 NODE_CLASS_MAPPINGS = {
     # === Generation ===
     "VLLMOmniGenerateImage": VLLMOmniGenerateImage,
+    "VLLMOmniGenerateVideo": VLLMOmniGenerateVideo,
     "VLLMOmniUnderstanding": VLLMOmniUnderstanding,
     "VLLMOmniTTS": VLLMOmniTTS,
     "VLLMOmniVoiceClone": VLLMOmniVoiceClone,
@@ -32,13 +36,16 @@ NODE_CLASS_MAPPINGS = {
     "VLLMOmniARSampling": VLLMOmniARSampling,
     "VLLMOmniDiffusionSampling": VLLMOmniDiffusionSampling,
     "VLLMOmniSamplingParamsList": VLLMOmniSamplingParamsList,
+    "VLLMOmniRemoteLoRA": VLLMOmniRemoteLoRA,
     "VLLMOmniQwenTTSParams": VLLMOmniQwenTTSParams,
+    "VLLMOmniWanParams": VLLMOmniWanParams,
 }
 
 # A dictionary that contains the friendly/humanly readable titles for the nodes
 NODE_DISPLAY_NAME_MAPPINGS = {
     # === Generation ===
     "VLLMOmniGenerateImage": "Generate Image",
+    "VLLMOmniGenerateVideo": "Generate Video",
     "VLLMOmniUnderstanding": "Multimodality Understanding",
     "VLLMOmniTTS": "TTS (Text to Speech)",
     "VLLMOmniVoiceClone": "TTS Voice Cloning",
@@ -46,7 +53,9 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "VLLMOmniARSampling": "AR Sampling Params",
     "VLLMOmniDiffusionSampling": "Diffusion Sampling Params",
     "VLLMOmniSamplingParamsList": "Multi-Stage Sampling Params List",
+    "VLLMOmniRemoteLoRA": "LoRA",
     "VLLMOmniQwenTTSParams": "Qwen TTS Params",
+    "VLLMOmniWanParams": "Wan Video Params",
 }
 
 WEB_DIRECTORY = "./web"

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py

@@ -11,6 +11,7 @@ from .utils.types import (
     AutoregressionSamplingParams,
     DiffusionSamplingParams,
     QwenTTSModelSpecificParams,
+    WanModelSpecificParams,
 )
 from .utils.validators import (
     add_sampling_parameters_to_stage,
@@ -56,6 +57,7 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
                 # "video": ("VIDEO",),
                 # "audio": ("AUDIO",),
                 "sampling_params": ("SAMPLING_PARAMS",),
+                "lora": ("REMOTE_LORA",),
             },
         }
 
@@ -76,9 +78,11 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
         audio: AudioInput | None = None,  # Hidden & unused
         video: VideoInput | None = None,  # Hidden & unused
         sampling_params: dict | list[dict] | None = None,
+        lora: dict | None = None,
         **kwargs,
     ):
-        logger.info("Uncaught kwargs: %s", kwargs)
+        if kwargs:
+            logger.info("Uncaught kwargs: %s", kwargs)
         logger.debug("Got sampling params: %s", sampling_params)
         validate_model_and_sampling_params_types(model, sampling_params)
         if image is None and mask is not None:
@@ -105,6 +109,7 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
                     height=height,
                     negative_prompt=negative_prompt,
                     sampling_params=sampling_params,
+                    lora=lora,
                 )
                 return (output,)
             elif image is not None and audio is None and video is None:
@@ -119,6 +124,7 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
                     negative_prompt=negative_prompt,
                     mask=mask,
                     sampling_params=sampling_params,
+                    lora=lora,
                 )
                 return (output,)
 
@@ -136,11 +142,91 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
             audio=audio,
             video=video,
             sampling_params=sampling_params,
+            lora=lora,
         )
 
         return (output,)
 
 
+class VLLMOmniGenerateVideo(_VLLMOmniGenerateBase):
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "url": ("STRING", {"default": "http://localhost:8000/v1"}),
+                "model": ("STRING", {"default": "Wan-AI/Wan2.2-T2V-A14B-Diffusers"}),
+                "prompt": ("STRING", {"multiline": True}),
+                "negative_prompt": ("STRING", {"multiline": True, "default": ""}),
+                "width": ("INT", {"default": 832, "min": 1}),
+                "height": ("INT", {"default": 480, "min": 1}),
+                "fps": ("INT", {"default": 16, "min": 1}),
+                "num_frames": ("INT", {"default": 41, "min": 1}),
+            },
+            "optional": {
+                "image": ("IMAGE",),
+                "sampling_params": ("SAMPLING_PARAMS",),
+                "lora": ("REMOTE_LORA",),
+                "model_params": ("VIDEO_PARAMS",),
+            },
+        }
+
+    RETURN_TYPES = ("VIDEO",)
+    RETURN_NAMES = ("video",)
+    FUNCTION = "generate"
+
+    async def generate(
+        self,
+        url: str,
+        model: str,
+        prompt: str,
+        width: int,
+        height: int,
+        fps: int,
+        num_frames: int,
+        negative_prompt: str | None = None,
+        image: torch.Tensor | None = None,
+        sampling_params: dict | list[dict] | None = None,
+        model_params: dict | None = None,
+        lora: dict | None = None,
+        **kwargs,
+    ):
+        if kwargs:
+            logger.info("Uncaught kwargs: %s", kwargs)
+        logger.debug("Got sampling params: %s", sampling_params)
+        logger.debug("Got model params: %s", model_params)
+        validate_model_and_sampling_params_types(model, sampling_params)
+
+        # Currently, all video generation models are single-stage diffusion models
+        if isinstance(sampling_params, list):
+            if len(sampling_params) != 1:
+                raise ValueError(
+                    "Video generation expects a single sampling params group. "
+                    "Please provide one Diffusion sampling node."
+                )
+            sampling_params = sampling_params[0]
+
+        if sampling_params is not None:
+            sampling_params.pop("type", None)  # internal fields
+        if model_params is not None:
+            model_params.pop("type", None)  # internal fields
+
+        client = VLLMOmniClient(url)
+        output = await client.generate_video(
+            model=model,
+            prompt=prompt,
+            image=image,  # image present => i2v, absent => t2v
+            width=width,
+            height=height,
+            num_frames=num_frames,
+            fps=fps,
+            negative_prompt=negative_prompt,
+            sampling_params=sampling_params,
+            lora=lora,
+            model_params=model_params,
+        )
+        return (output,)
+
+
 class VLLMOmniUnderstanding(_VLLMOmniGenerateBase):
     @classmethod
     def INPUT_TYPES(cls):
@@ -188,7 +274,8 @@ class VLLMOmniUnderstanding(_VLLMOmniGenerateBase):
         use_audio_in_video: bool = True,
         **kwargs,
     ) -> tuple[str, AudioInput]:
-        logger.info("Uncaught kwargs: %s", kwargs)
+        if kwargs:
+            logger.info("Uncaught kwargs: %s", kwargs)
         logger.debug("Got sampling params: %s", sampling_params)
         validate_model_and_sampling_params_types(model, sampling_params)
 
@@ -547,6 +634,52 @@ class VLLMOmniSamplingParamsList:
         return (params,)
 
 
+class VLLMOmniRemoteLoRA:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "local_path": ("STRING", {"default": ""}),
+                "name": ("STRING", {"default": ""}),
+                "scale": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.1},
+                ),
+                "int_id": (
+                    "INT",
+                    {
+                        "default": 0,
+                        "min": 0,
+                        "step": 1,
+                        "tooltip": "0 means it is not set and the server can derive it.",
+                    },
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("REMOTE_LORA",)
+    RETURN_NAMES = ("lora",)
+    FUNCTION = "get_lora"
+    CATEGORY = "vLLM-Omni"
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, local_path, name) -> str | Literal[True]:
+        if not local_path.strip() or not name.strip():
+            return "Both local_path and name must be provided."
+        return True
+
+    def get_lora(self, local_path: str, name: str, scale: float, int_id: int):
+        local_path = local_path.strip()
+        name = name.strip()
+        lora = {
+            "local_path": local_path or None,
+            "name": name or None,
+            "scale": float(scale),
+            "int_id": int(int_id) if int_id > 0 else None,
+        }
+        return (lora,)
+
+
 class VLLMOmniQwenTTSParams:
     @classmethod
     def INPUT_TYPES(cls):
@@ -572,3 +705,32 @@ class VLLMOmniQwenTTSParams:
 
     def get_params(self, **kwargs):
         return (QwenTTSModelSpecificParams(kwargs),)
+
+
+class VLLMOmniWanParams:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "guidance_scale_2": (
+                    "FLOAT",
+                    {"default": 4.0, "min": 0.0, "max": 20.0, "step": 0.1},
+                ),
+                "boundary_ratio": (
+                    "FLOAT",
+                    {"default": 0.875, "min": 0.0, "max": 1.0, "step": 0.001},
+                ),
+                "flow_shift": (
+                    "FLOAT",
+                    {"default": 5.0, "min": 0.0, "max": 100.0, "step": 0.1},
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("VIDEO_PARAMS",)
+    RETURN_NAMES = ("Wan video params",)
+    FUNCTION = "get_params"
+    CATEGORY = "vLLM-Omni/Video Params"
+
+    def get_params(self, **kwargs):
+        return (WanModelSpecificParams(kwargs),)

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/api_client.py

@@ -6,6 +6,7 @@ The image generation part is derived from dougbtv/comfyui-vllm-omni by Doug (@do
 Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed under the MIT License.
 """
 
+import json
 from typing import Any
 
 import aiohttp
@@ -17,6 +18,7 @@ from .format import (
     audio_to_base64,
     base64_to_audio,
     base64_to_image_tensor,
+    base64_to_video,
     bytes_to_audio,
     image_tensor_to_base64,
     image_tensor_to_png_bytes,
@@ -30,7 +32,7 @@ logger = get_logger(__name__)
 
 
 class VLLMOmniClient:
-    def __init__(self, base_url: str, timeout: float = 300.0):
+    def __init__(self, base_url: str, timeout: float | None = None):
         self.base_url = base_url
         self.timeout = aiohttp.ClientTimeout(total=timeout)
 
@@ -43,12 +45,13 @@ class VLLMOmniClient:
         height: int,
         negative_prompt: str | None = None,
         sampling_params: dict | None = None,
+        lora: dict | None = None,
     ) -> torch.Tensor:
         """Run text-to-image generation via DALLE API"""
         await self._check_model_exist(model)
 
         size = f"{width}x{height}"
-        payload = {
+        payload: dict[str, Any] = {
             "model": model,
             "prompt": prompt,
             "size": size,
@@ -58,6 +61,8 @@ class VLLMOmniClient:
             payload["negative_prompt"] = negative_prompt
         if sampling_params is not None:
             payload.update(sampling_params)
+        if lora is not None:
+            payload["lora"] = lora
         logger.debug("img gen payload: %s", payload)
 
         url = self.base_url + "/images/generations"
@@ -110,6 +115,7 @@ class VLLMOmniClient:
         negative_prompt: str | None = None,
         mask: torch.Tensor | None = None,
         sampling_params: dict | None = None,
+        lora: dict | None = None,
     ) -> torch.Tensor:
         """Run image editing via DALLE API"""
         await self._check_model_exist(model)
@@ -131,6 +137,8 @@ class VLLMOmniClient:
         if sampling_params is not None:
             for k, v in sampling_params.items():
                 form.add_field(k, str(v))
+        if lora is not None:
+            form.add_field("lora", json.dumps(lora, ensure_ascii=False))
         if mask is not None:
             mask_filename = "mask.png"
             form.add_field(
@@ -183,6 +191,7 @@ class VLLMOmniClient:
         audio: AudioInput | None = None,
         video: VideoInput | None = None,
         sampling_params: dict | list[dict] | None = None,
+        lora: dict | None = None,
     ) -> torch.Tensor:
         payload = VLLMOmniClient._prepare_chat_completion_messages(
             model=model,
@@ -193,6 +202,8 @@ class VLLMOmniClient:
             video=video,
             sampling_params=sampling_params,
             modalities=["image"],
+            # === below are additional `extra_body` fields, handled by **kwargs ===
+            lora=lora,
         )
         choices = await self._generate_base_chat_completion(model, payload)
 
@@ -206,6 +217,79 @@ class VLLMOmniClient:
 
         return torch.stack(image_tensors, dim=0)
 
+    async def generate_video(
+        self,
+        *,
+        model: str,
+        prompt: str,
+        width: int,
+        height: int,
+        num_frames: int,
+        fps: int,
+        negative_prompt: str | None = None,
+        image: torch.Tensor | None = None,
+        sampling_params: dict | None = None,
+        model_params: dict | None = None,
+        lora: dict | None = None,
+        **extra_body,
+    ) -> VideoInput:
+        form = aiohttp.FormData()
+        form.add_field("model", model)
+        form.add_field("prompt", prompt)
+        form.add_field("width", str(width))
+        form.add_field("height", str(height))
+        form.add_field("num_frames", str(num_frames))
+        form.add_field("fps", str(fps))
+        if negative_prompt:
+            form.add_field("negative_prompt", negative_prompt)
+        if sampling_params is not None:
+            for k, v in sampling_params.items():
+                form.add_field(k, str(v))
+        if model_params is not None:
+            for k, v in model_params.items():
+                form.add_field(k, str(v))
+        if lora is not None:
+            form.add_field("lora", json.dumps(lora, ensure_ascii=False))
+        if extra_body:
+            form.add_field("extra_body", json.dumps(extra_body, ensure_ascii=False))
+
+        if image is not None:
+            image_filename = "image.png"  # Required for multipart form
+            form.add_field(
+                "input_reference",
+                image_tensor_to_png_bytes(image, image_filename),
+                filename=image_filename,
+                content_type="image/png",
+            )
+
+        url = self.base_url + "/videos"
+        async with aiohttp.ClientSession(timeout=self.timeout) as session:
+            try:
+                async with session.post(url, data=form) as response:
+                    if not response.ok:
+                        error_text = await response.text()
+                        raise (ValueError if response.status < 500 else RuntimeError)(
+                            f"vLLM-Omni API returned status {response.status}: {error_text}"
+                        )
+
+                    try:
+                        data = await response.json()
+                    except aiohttp.ContentTypeError as e:
+                        raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}")
+            except aiohttp.ClientError as e:
+                raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
+
+        if "data" not in data:
+            raise RuntimeError("API response missing 'data' field - expected OpenAI DALL-E format")
+        if not data["data"]:
+            raise RuntimeError("API returned empty data array")
+        try:
+            base64_str = data["data"][0]["b64_json"]
+        except (KeyError, IndexError):
+            raise RuntimeError("API response missing 'b64_json' field in first data item")
+
+        return base64_to_video(base64_str)
+
     async def generate_understanding_chat_completion(
         self,
         *,

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/format.py

@@ -6,12 +6,17 @@ Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed und
 
 import base64
 import mimetypes
+from fractions import Fraction
 from io import BytesIO
 
 import av
 import numpy as np
 import torch
+from av.audio.frame import AudioFrame
+from av.audio.resampler import AudioResampler
+from av.video.frame import VideoFrame
 from comfy_api.input import AudioInput, VideoInput
+from comfy_api.latest import InputImpl, Types
 from comfy_extras import nodes_audio
 from PIL import Image
 
@@ -138,6 +143,71 @@ def video_to_base64(video: VideoInput, filename: str = "video.mp4") -> str:
     return f"data:{mime_type};base64,{base64_str}"
 
 
+def base64_to_video(base64_str: str) -> VideoInput:
+    if base64_str.startswith("data:video"):
+        _, base64_str = base64_str.split(",", 1)
+
+    try:
+        video_bytes = base64.b64decode(base64_str)
+    except Exception as e:
+        raise ValueError(f"Invalid base64 string: {e}")
+
+    video_buffer = BytesIO(video_bytes)
+
+    try:
+        with av.open(video_buffer, mode="r") as container:
+            video_stream = next((s for s in container.streams if s.type == "video"), None)
+            if video_stream is None:
+                raise ValueError("No video stream found in decoded payload.")
+
+            frames = []
+            for frame in container.decode(video_stream):
+                if not isinstance(frame, VideoFrame):
+                    continue
+                image = frame.to_ndarray(format="rgb24")
+                frames.append(torch.from_numpy(image).float() / 255.0)
+
+            if len(frames) == 0:
+                raise ValueError("No video frames found in decoded payload.")
+
+            images = torch.stack(frames, dim=0)
+            frame_rate = Fraction(video_stream.average_rate) if video_stream.average_rate else Fraction(1)
+
+            audio: AudioInput | None = None
+            if len(container.streams.audio):
+                audio_stream = container.streams.audio[-1]
+                audio_frames = []
+                resampler = AudioResampler(format="fltp")
+                for frame in container.decode(audio_stream):
+                    if not isinstance(frame, AudioFrame):
+                        continue
+                    resampled = resampler.resample(frame)
+                    if not isinstance(resampled, list):
+                        resampled = [resampled]
+                    for audio_frame in resampled:
+                        if audio_frame is not None:
+                            audio_frames.append(audio_frame.to_ndarray())
+
+                if len(audio_frames) > 0:
+                    audio_data = np.concatenate(audio_frames, axis=1)
+                    sample_rate = int(audio_stream.sample_rate) if audio_stream.sample_rate else 1
+                    audio = {
+                        "waveform": torch.from_numpy(audio_data).unsqueeze(0),
+                        "sample_rate": sample_rate,
+                    }
+
+            components = Types.VideoComponents(
+                images=images,
+                frame_rate=frame_rate,
+                audio=audio,
+                metadata=container.metadata if container.metadata else None,
+            )
+    except Exception as e:
+        raise RuntimeError(f"Failed to decode video: {e}")
+
+    return InputImpl.VideoFromComponents(components)
+
+
 def audio_to_bytes(audio: AudioInput, filename: str = "audio.mp3", quality: str = "128k") -> BytesIO:
     waveform = audio["waveform"][0]  # Shape: (C, T)
     sample_rate = audio["sample_rate"]

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py

@@ -23,6 +23,10 @@ class QwenTTSModelSpecificParams(dict):
     pass
 
 
+class WanModelSpecificParams(dict):
+    pass
+
+
 class ModelMode(Enum):
     IMAGE_GENERATION = auto()
     VIDEO_GENERATION = auto()

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/validators.py

@@ -16,7 +16,7 @@ def validate_model_and_sampling_params_types(
     # Skip if no spec or no sampling params
     pipeline_spec, _ = lookup_model_spec(model_name)
     if pipeline_spec is None:
-        logger.info("skipping sampling params check because spec is not found")
+        logger.info(f"skipping sampling params check because spec for {model_name} is not found")
         return
     if sampling_param_list is None:
         return

apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Video Generation.json

@@ -0,0 +1,513 @@
+{
+    "id": "3aea1d08-7de5-4997-ba71-b2f66ad72a21",
+    "revision": 0,
+    "last_node_id": 14,
+    "last_link_id": 28,
+    "nodes": [
+        {
+            "id": 10,
+            "type": "MarkdownNote",
+            "pos": [
+                227.99306819575278,
+                -231.24143306069843
+            ],
+            "size": [
+                240.3326478922474,
+                136.3505820791881
+            ],
+            "flags": {},
+            "order": 0,
+            "mode": 0,
+            "inputs": [],
+            "outputs": [],
+            "title": "Note: Task and Input",
+            "properties": {},
+            "widgets_values": [
+                "vLLM-Omni nodes are categorized based on the output modality. The \"Generate Video\" node supports both text-to-video generation and image-to-video generation. The node will create corresponding payload depending on whether an input image is present or not."
+            ],
+            "color": "#432",
+            "bgcolor": "#000"
+        },
+        {
+            "id": 4,
+            "type": "LoadImage",
+            "pos": [
+                496.31859627609606,
+                -229.71277089860084
+            ],
+            "size": [
+                270,
+                314
+            ],
+            "flags": {},
+            "order": 1,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "type": "COMBO",
+                    "widget": {
+                        "name": "image"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "choose file to upload",
+                    "name": "upload",
+                    "type": "IMAGEUPLOAD",
+                    "widget": {
+                        "name": "upload"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "IMAGE",
+                    "name": "IMAGE",
+                    "type": "IMAGE",
+                    "links": [
+                        25
+                    ]
+                },
+                {
+                    "localized_name": "MASK",
+                    "name": "MASK",
+                    "type": "MASK",
+                    "links": null
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "LoadImage"
+            },
+            "widgets_values": [
+                "cute-cat.jpg",
+                "image"
+            ]
+        },
+        {
+            "id": 8,
+            "type": "VLLMOmniDiffusionSampling",
+            "pos": [
+                465.5140218220927,
+                148.80072646828987
+            ],
+            "size": [
+                284.205078125,
+                226
+            ],
+            "flags": {},
+            "order": 2,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "n",
+                    "name": "n",
+                    "type": "INT",
+                    "widget": {
+                        "name": "n"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "num_inference_steps",
+                    "name": "num_inference_steps",
+                    "type": "INT",
+                    "widget": {
+                        "name": "num_inference_steps"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "guidance_scale",
+                    "name": "guidance_scale",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "guidance_scale"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "true_cfg_scale",
+                    "name": "true_cfg_scale",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "true_cfg_scale"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "vae_use_slicing",
+                    "name": "vae_use_slicing",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "vae_use_slicing"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "vae_use_tiling",
+                    "name": "vae_use_tiling",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "vae_use_tiling"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "seed",
+                    "name": "seed",
+                    "type": "INT",
+                    "widget": {
+                        "name": "seed"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "diffusion sampling params",
+                    "name": "diffusion sampling params",
+                    "type": "SAMPLING_PARAMS",
+                    "links": [
+                        27
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniDiffusionSampling"
+            },
+            "widgets_values": [
+                1,
+                50,
+                1,
+                1,
+                false,
+                false,
+                1992,
+                "randomize"
+            ]
+        },
+        {
+            "id": 14,
+            "type": "VLLMOmniWanParams",
+            "pos": [
+                518.6005521845747,
+                428.4455663719664
+            ],
+            "size": [
+                270,
+                106
+            ],
+            "flags": {},
+            "order": 3,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "guidance_scale_2",
+                    "name": "guidance_scale_2",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "guidance_scale_2"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "boundary_ratio",
+                    "name": "boundary_ratio",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "boundary_ratio"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "flow_shift",
+                    "name": "flow_shift",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "flow_shift"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "Wan video params",
+                    "name": "Wan video params",
+                    "type": "VIDEO_PARAMS",
+                    "links": [
+                        28
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniWanParams"
+            },
+            "widgets_values": [
+                4,
+                0.875,
+                5
+            ]
+        },
+        {
+            "id": 12,
+            "type": "VLLMOmniGenerateVideo",
+            "pos": [
+                827.2566336087867,
+                -73.77449831827545
+            ],
+            "size": [
+                400,
+                346
+            ],
+            "flags": {},
+            "order": 4,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "shape": 7,
+                    "type": "IMAGE",
+                    "link": 25
+                },
+                {
+                    "localized_name": "sampling_params",
+                    "name": "sampling_params",
+                    "shape": 7,
+                    "type": "SAMPLING_PARAMS",
+                    "link": 27
+                },
+                {
+                    "localized_name": "lora",
+                    "name": "lora",
+                    "shape": 7,
+                    "type": "REMOTE_LORA",
+                    "link": null
+                },
+                {
+                    "localized_name": "model_params",
+                    "name": "model_params",
+                    "shape": 7,
+                    "type": "VIDEO_PARAMS",
+                    "link": 28
+                },
+                {
+                    "localized_name": "url",
+                    "name": "url",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "url"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "model",
+                    "name": "model",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "model"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "prompt",
+                    "name": "prompt",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "prompt"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "negative_prompt",
+                    "name": "negative_prompt",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "negative_prompt"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "width",
+                    "name": "width",
+                    "type": "INT",
+                    "widget": {
+                        "name": "width"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "height",
+                    "name": "height",
+                    "type": "INT",
+                    "widget": {
+                        "name": "height"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "fps",
+                    "name": "fps",
+                    "type": "INT",
+                    "widget": {
+                        "name": "fps"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "num_frames",
+                    "name": "num_frames",
+                    "type": "INT",
+                    "widget": {
+                        "name": "num_frames"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "video",
+                    "name": "video",
+                    "type": "VIDEO",
+                    "links": [
+                        26
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniGenerateVideo"
+            },
+            "widgets_values": [
+                "http://localhost:8000/v1",
+                "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
+                "Make the cat in the attached video yawning. Eye closed and mouth wide open. Keep the original realistic style.",
+                "Cartoonish",
+                900,
+                1200,
+                12,
+                41
+            ]
+        },
+        {
+            "id": 13,
+            "type": "SaveVideo",
+            "pos": [
+                1257.5441369501657,
+                -44.12956394419863
+            ],
+            "size": [
+                374,
+                372.34896826553796
+            ],
+            "flags": {},
+            "order": 5,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "video",
+                    "name": "video",
+                    "type": "VIDEO",
+                    "link": 26
+                },
+                {
+                    "localized_name": "filename_prefix",
+                    "name": "filename_prefix",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "filename_prefix"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "format",
+                    "name": "format",
+                    "type": "COMBO",
+                    "widget": {
+                        "name": "format"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "codec",
+                    "name": "codec",
+                    "type": "COMBO",
+                    "widget": {
+                        "name": "codec"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [],
+            "properties": {},
+            "widgets_values": [
+                "video/ComfyUI",
+                "auto",
+                "auto"
+            ]
+        }
+    ],
+    "links": [
+        [
+            25,
+            4,
+            0,
+            12,
+            0,
+            "IMAGE"
+        ],
+        [
+            26,
+            12,
+            0,
+            13,
+            0,
+            "VIDEO"
+        ],
+        [
+            27,
+            8,
+            0,
+            12,
+            1,
+            "SAMPLING_PARAMS"
+        ],
+        [
+            28,
+            14,
+            0,
+            12,
+            3,
+            "VIDEO_PARAMS"
+        ]
+    ],
+    "groups": [
+        {
+            "id": 1,
+            "title": "Input",
+            "bounding": [
+                213.1706010087147,
+                -313.20095750667554,
+                560.8246471437027,
+                407.4246532472182
+            ],
+            "color": "#3f789e",
+            "font_size": 24,
+            "flags": {}
+        }
+    ],
+    "config": {},
+    "extra": {
+        "workflowRendererVersion": "LG",
+        "ds": {
+            "scale": 1.1469075819486894,
+            "offset": [
+                101.7632950847089,
+                492.29121889347806
+            ]
+        }
+    },
+    "version": 0.4
+}

benchmarks/qwen3-tts/README.md

@@ -0,0 +1,100 @@
+# Qwen3-TTS Benchmark
+
+Benchmarks for Qwen3-TTS text-to-speech models, comparing vLLM-Omni streaming serving against HuggingFace Transformers offline inference.
+
+## Prerequisites
+
+```bash
+pip install matplotlib aiohttp soundfile numpy tqdm
+pip install qwen_tts  # for HF baseline
+```
+
+## Quick Start
+
+Run the full benchmark (vllm-omni + HF baseline) with a single command:
+
+```bash
+cd benchmarks/qwen3-tts
+bash run_benchmark.sh
+```
+
+Results (JSON + PNG plots) are saved to `results/`.
+
+### Common options
+
+```bash
+# Only vllm-omni (skip HF baseline)
+bash run_benchmark.sh --async-only
+
+# Only HF baseline
+bash run_benchmark.sh --hf-only
+
+# Use a different model (e.g. 1.7B)
+MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
+
+# Use batch_size=4 config for higher throughput
+STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs4.yaml bash run_benchmark.sh --async-only
+
+# Custom GPU, prompt count, concurrency levels
+GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
+```
+
+## Manual Steps
+
+### 1) Start the vLLM-Omni server
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python -m vllm_omni.entrypoints.cli.main serve \
+    "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
+    --omni --host 127.0.0.1 --port 8000 \
+    --stage-configs-path benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs1.yaml \
+    --trust-remote-code
+```
+
+### 2) Run online serving benchmark
+
+```bash
+python vllm_omni/bench_tts_serve.py \
+    --port 8000 \
+    --num-prompts 50 \
+    --max-concurrency 1 4 10 \
+    --config-name "async_chunk" \
+    --result-dir results/
+```
+
+### 3) Run HuggingFace baseline
+
+```bash
+python transformers/bench_tts_hf.py \
+    --model "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
+    --num-prompts 50 \
+    --gpu-device 0 \
+    --result-dir results/
+```
+
+### 4) Generate comparison plots
+
+```bash
+python plot_results.py \
+    --results results/bench_async_chunk_*.json results/bench_hf_transformers_*.json \
+    --labels "vllm-omni" "hf_transformers" \
+    --output results/comparison.png
+```
+
+## Stage Configs
+
+| Config | Batch Size | Description |
+|--------|:----------:|-------------|
+| `vllm_omni/configs/qwen3_tts_bs1.yaml` | 1 | Single-request processing (lower latency) |
+| `vllm_omni/configs/qwen3_tts_bs4.yaml` | 4 | Concurrent request processing (higher throughput) |
+
+Both configs use a 2-stage pipeline (Talker -> Code2Wav) with `async_chunk` streaming enabled. The `SharedMemoryConnector` streams codec frames (25-frame chunks with 25-frame context overlap) between stages.
+
+The model is specified via the CLI `--model` flag (or `MODEL` env var in `run_benchmark.sh`), so the same configs work for both the 0.6B and 1.7B model variants.
+
+## Metrics
+
+- **TTFP (Time to First Audio Packet)**: Time from request to first audio chunk (streaming latency)
+- **E2E (End-to-End Latency)**: Total time from request to complete audio response
+- **RTF (Real-Time Factor)**: E2E latency / audio duration. RTF < 1.0 means faster-than-real-time synthesis
+- **Throughput**: Total audio seconds generated per wall-clock second

benchmarks/qwen3-tts/plot_results.py

@@ -0,0 +1,254 @@
+"""Plot Qwen3-TTS benchmark results.
+
+Generates comparison bar charts similar to the async_chunk design doc:
+- TTFP (Time-to-First-Packet) across concurrency levels
+- E2E latency across concurrency levels
+- RTF (Real-Time Factor) across concurrency levels
+
+Usage:
+    # Compare two configs (async_chunk vs no_async_chunk):
+    python plot_results.py \
+        --results results/bench_async_chunk_*.json results/bench_no_async_chunk_*.json \
+        --labels "async_chunk" "no_async_chunk" \
+        --output results/qwen3_tts_benchmark.png
+
+    # Single config:
+    python plot_results.py \
+        --results results/bench_async_chunk_*.json \
+        --labels "async_chunk" \
+        --output results/qwen3_tts_benchmark.png
+"""
+
+import argparse
+import json
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+def load_results(result_files: list[str]) -> list[list[dict]]:
+    """Load benchmark results from JSON files."""
+    all_results = []
+    for f in result_files:
+        with open(f) as fh:
+            data = json.load(fh)
+        all_results.append(data)
+    return all_results
+
+
+def plot_comparison(
+    all_results: list[list[dict]],
+    labels: list[str],
+    output_path: str,
+    title_prefix: str = "Qwen3-TTS",
+):
+    """Generate comparison bar charts."""
+    n_configs = len(all_results)
+
+    # Collect concurrency levels present in ALL configs (skip missing data)
+    all_concurrencies = [set(r["concurrency"] for r in results) for results in all_results]
+    concurrencies = sorted(set.union(*all_concurrencies))
+
+    # Build data arrays, using None for missing concurrency levels
+    ttfp_data = {label: [] for label in labels}
+    e2e_data = {label: [] for label in labels}
+    rtf_data = {label: [] for label in labels}
+    throughput_data = {label: [] for label in labels}
+
+    for results, label in zip(all_results, labels):
+        conc_map = {r["concurrency"]: r for r in results}
+        for c in concurrencies:
+            r = conc_map.get(c)
+            ttfp_data[label].append(r["mean_ttfp_ms"] if r else None)
+            e2e_data[label].append(r["mean_e2e_ms"] if r else None)
+            rtf_data[label].append(r["mean_rtf"] if r else None)
+            throughput_data[label].append(r["audio_throughput"] if r else None)
+
+    fig, axes = plt.subplots(2, 2, figsize=(14, 10))
+    fig.suptitle(f"{title_prefix} Performance Benchmark", fontsize=16, fontweight="bold")
+
+    x = np.arange(len(concurrencies))
+    width = 0.35 if n_configs == 2 else 0.5
+    if n_configs > 1:
+        offsets = np.linspace(-width / 2 * (n_configs - 1), width / 2 * (n_configs - 1), n_configs)
+    else:
+        offsets = [0]
+
+    colors = ["#2196F3", "#FF5722", "#4CAF50", "#FFC107"]
+
+    def plot_metric(ax, data_dict, ylabel, title, fmt=".1f"):
+        bars = []
+        for i, (label, values) in enumerate(data_dict.items()):
+            # Replace None with 0 for plotting, but track which are missing
+            plot_values = [v if v is not None else 0 for v in values]
+            color = colors[i % len(colors)]
+            bar = ax.bar(x + offsets[i], plot_values, width, label=label, color=color, alpha=0.85)
+            bars.append(bar)
+            # Add value labels on bars (skip None/missing data)
+            max_val = max((v for v in values if v is not None), default=1)
+            for rect, val in zip(bar, values):
+                if val is not None and val > 0:
+                    ax.text(
+                        rect.get_x() + rect.get_width() / 2,
+                        rect.get_height() + max_val * 0.02,
+                        f"{val:{fmt}}",
+                        ha="center",
+                        va="bottom",
+                        fontsize=9,
+                        fontweight="bold",
+                    )
+        ax.set_xlabel("Concurrency", fontsize=12)
+        ax.set_ylabel(ylabel, fontsize=12)
+        ax.set_title(title, fontsize=13, fontweight="bold")
+        ax.set_xticks(x)
+        ax.set_xticklabels([str(c) for c in concurrencies])
+        ax.legend(fontsize=10)
+        ax.grid(axis="y", alpha=0.3)
+        ax.set_axisbelow(True)
+
+    plot_metric(axes[0, 0], ttfp_data, "TTFP (ms)", "Time to First Audio Packet (TTFP)")
+    plot_metric(axes[0, 1], e2e_data, "E2E Latency (ms)", "End-to-End Latency (E2E)")
+    plot_metric(axes[1, 0], rtf_data, "RTF", "Real-Time Factor (RTF)", fmt=".3f")
+    plot_metric(axes[1, 1], throughput_data, "Audio-sec / Wall-sec", "Audio Throughput", fmt=".2f")
+
+    plt.tight_layout()
+    plt.savefig(output_path, dpi=150, bbox_inches="tight")
+    print(f"Plot saved to {output_path}")
+    plt.close()
+
+
+def plot_single_summary(results: list[dict], label: str, output_path: str):
+    """Generate a single-config summary with percentile breakdown."""
+    concurrencies = [r["concurrency"] for r in results]
+
+    fig, axes = plt.subplots(1, 3, figsize=(16, 5))
+    fig.suptitle(f"Qwen3-TTS Benchmark - {label}", fontsize=15, fontweight="bold")
+
+    # TTFP breakdown
+    ax = axes[0]
+    means = [r["mean_ttfp_ms"] for r in results]
+    medians = [r["median_ttfp_ms"] for r in results]
+    p90s = [r["p90_ttfp_ms"] for r in results]
+    p99s = [r["p99_ttfp_ms"] for r in results]
+    x = np.arange(len(concurrencies))
+    w = 0.2
+    ax.bar(x - 1.5 * w, means, w, label="mean", color="#2196F3")
+    ax.bar(x - 0.5 * w, medians, w, label="median", color="#4CAF50")
+    ax.bar(x + 0.5 * w, p90s, w, label="p90", color="#FF9800")
+    ax.bar(x + 1.5 * w, p99s, w, label="p99", color="#F44336")
+    ax.set_xticks(x)
+    ax.set_xticklabels([str(c) for c in concurrencies])
+    ax.set_xlabel("Concurrency")
+    ax.set_ylabel("TTFP (ms)")
+    ax.set_title("Time to First Audio Packet")
+    ax.legend(fontsize=9)
+    ax.grid(axis="y", alpha=0.3)
+
+    # E2E breakdown
+    ax = axes[1]
+    means = [r["mean_e2e_ms"] for r in results]
+    medians = [r["median_e2e_ms"] for r in results]
+    p90s = [r["p90_e2e_ms"] for r in results]
+    p99s = [r["p99_e2e_ms"] for r in results]
+    ax.bar(x - 1.5 * w, means, w, label="mean", color="#2196F3")
+    ax.bar(x - 0.5 * w, medians, w, label="median", color="#4CAF50")
+    ax.bar(x + 0.5 * w, p90s, w, label="p90", color="#FF9800")
+    ax.bar(x + 1.5 * w, p99s, w, label="p99", color="#F44336")
+    ax.set_xticks(x)
+    ax.set_xticklabels([str(c) for c in concurrencies])
+    ax.set_xlabel("Concurrency")
+    ax.set_ylabel("E2E Latency (ms)")
+    ax.set_title("End-to-End Latency")
+    ax.legend(fontsize=9)
+    ax.grid(axis="y", alpha=0.3)
+
+    # RTF
+    ax = axes[2]
+    means = [r["mean_rtf"] for r in results]
+    medians = [r["median_rtf"] for r in results]
+    ax.bar(x - 0.15, means, 0.3, label="mean", color="#2196F3")
+    ax.bar(x + 0.15, medians, 0.3, label="median", color="#4CAF50")
+    ax.set_xticks(x)
+    ax.set_xticklabels([str(c) for c in concurrencies])
+    ax.set_xlabel("Concurrency")
+    ax.set_ylabel("RTF")
+    ax.set_title("Real-Time Factor")
+    ax.legend(fontsize=9)
+    ax.grid(axis="y", alpha=0.3)
+
+    plt.tight_layout()
+    plt.savefig(output_path, dpi=150, bbox_inches="tight")
+    print(f"Plot saved to {output_path}")
+    plt.close()
+
+
+def print_comparison_table(all_results: list[list[dict]], labels: list[str]):
+    """Print a markdown-formatted comparison table."""
+    concurrencies = sorted(set(r["concurrency"] for r in all_results[0]))
+
+    print("\n## Benchmark Results\n")
+    header = "| Metric | Concurrency |"
+    sep = "| --- | --- |"
+    for label in labels:
+        header += f" {label} |"
+        sep += " --- |"
+    print(header)
+    print(sep)
+
+    for metric, key, fmt in [
+        ("TTFP (ms)", "mean_ttfp_ms", ".1f"),
+        ("E2E (ms)", "mean_e2e_ms", ".1f"),
+        ("RTF", "mean_rtf", ".3f"),
+        ("Throughput (audio-s/s)", "audio_throughput", ".2f"),
+    ]:
+        for c in concurrencies:
+            row = f"| {metric} | {c} |"
+            for results in all_results:
+                conc_map = {r["concurrency"]: r for r in results}
+                val = conc_map.get(c, {}).get(key, 0)
+                row += f" {val:{fmt}} |"
+            print(row)
+
+    # Improvement calculation (only if 2 configs)
+    if len(all_results) == 2:
+        print(f"\n## Improvement ({labels[0]} vs {labels[1]})\n")
+        print("| Metric | Concurrency | Improvement |")
+        print("| --- | --- | --- |")
+        for metric, key in [("TTFP", "mean_ttfp_ms"), ("E2E", "mean_e2e_ms"), ("RTF", "mean_rtf")]:
+            for c in concurrencies:
+                m0 = {r["concurrency"]: r for r in all_results[0]}
+                m1 = {r["concurrency"]: r for r in all_results[1]}
+                v0 = m0.get(c, {}).get(key, 0)
+                v1 = m1.get(c, {}).get(key, 0)
+                if v1 > 0:
+                    pct = (v1 - v0) / v1 * 100
+                    print(f"| {metric} | {c} | {pct:+.1f}% |")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Plot Qwen3-TTS benchmark results")
+    parser.add_argument(
+        "--results", type=str, nargs="+", required=True, help="Path(s) to result JSON files (one per config)"
+    )
+    parser.add_argument(
+        "--labels", type=str, nargs="+", required=True, help="Labels for each config (must match --results count)"
+    )
+    parser.add_argument("--output", type=str, default="results/qwen3_tts_benchmark.png", help="Output image path")
+    parser.add_argument("--title", type=str, default="Qwen3-TTS", help="Title prefix for the plot")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    assert len(args.results) == len(args.labels), "--results and --labels must have the same count"
+
+    all_results = load_results(args.results)
+    print_comparison_table(all_results, args.labels)
+
+    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
+
+    if len(all_results) == 1:
+        plot_single_summary(all_results[0], args.labels[0], args.output)
+    else:
+        plot_comparison(all_results, args.labels, args.output, title_prefix=args.title)

benchmarks/qwen3-tts/results/.gitignore

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

benchmarks/qwen3-tts/run_benchmark.sh

@@ -0,0 +1,272 @@
+#!/bin/bash
+# Qwen3-TTS Benchmark Runner
+#
+# Compares vllm-omni streaming serving vs HuggingFace transformers offline inference.
+# Produces JSON results and comparison plots.
+#
+# Usage:
+#   # Full comparison (vllm-omni + HF):
+#   bash run_benchmark.sh
+#
+#   # Only vllm-omni async_chunk config:
+#   bash run_benchmark.sh --async-only
+#
+#   # Only HuggingFace baseline:
+#   bash run_benchmark.sh --hf-only
+#
+#   # vllm-omni only (skip HF):
+#   bash run_benchmark.sh --skip-hf
+#
+#   # Custom settings:
+#   GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
+#
+#   # Use 1.7B model:
+#   MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
+#
+#   # Use batch_size=4 config:
+#   STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs4.yaml bash run_benchmark.sh --async-only
+#
+# Environment variables:
+#   GPU_DEVICE       - GPU index to use (default: 0)
+#   NUM_PROMPTS      - Number of prompts per concurrency level (default: 50)
+#   CONCURRENCY      - Space-separated concurrency levels (default: "1 4 10")
+#   MODEL            - Model name (default: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice)
+#   PORT             - Server port (default: 8000)
+#   GPU_MEM_TALKER   - gpu_memory_utilization for talker stage (default: 0.3)
+#   GPU_MEM_CODE2WAV - gpu_memory_utilization for code2wav stage (default: 0.2)
+#   STAGE_CONFIG     - Path to stage config YAML (default: configs/qwen3_tts_bs1.yaml)
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
+
+# Defaults
+GPU_DEVICE="${GPU_DEVICE:-0}"
+NUM_PROMPTS="${NUM_PROMPTS:-50}"
+CONCURRENCY="${CONCURRENCY:-1 4 10}"
+MODEL="${MODEL:-Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice}"
+PORT="${PORT:-8000}"
+GPU_MEM_TALKER="${GPU_MEM_TALKER:-0.3}"
+GPU_MEM_CODE2WAV="${GPU_MEM_CODE2WAV:-0.2}"
+NUM_WARMUPS="${NUM_WARMUPS:-3}"
+STAGE_CONFIG="${STAGE_CONFIG:-vllm_omni/configs/qwen3_tts_bs1.yaml}"
+RESULT_DIR="${SCRIPT_DIR}/results"
+TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
+
+# Parse args
+RUN_ASYNC=true
+RUN_HF=true
+for arg in "$@"; do
+    case "$arg" in
+        --async-only) RUN_HF=false ;;
+        --hf-only) RUN_ASYNC=false ;;
+        --skip-hf) RUN_HF=false ;;
+    esac
+done
+
+mkdir -p "${RESULT_DIR}"
+
+echo "============================================================"
+echo " Qwen3-TTS Benchmark"
+echo "============================================================"
+echo " GPU:          ${GPU_DEVICE}"
+echo " Model:        ${MODEL}"
+echo " Prompts:      ${NUM_PROMPTS}"
+echo " Concurrency:  ${CONCURRENCY}"
+echo " Port:         ${PORT}"
+echo " Stage config: ${STAGE_CONFIG}"
+echo " Results:      ${RESULT_DIR}"
+echo "============================================================"
+
+# Prepare stage config with correct GPU device and memory settings
+prepare_config() {
+    local config_template="$1"
+    local config_name="$2"
+    local output_path="${RESULT_DIR}/${config_name}_stage_config.yaml"
+
+    # Use sed to patch GPU device and memory utilization
+    sed \
+        -e "s/devices: \"0\"/devices: \"${GPU_DEVICE}\"/g" \
+        -e "s/gpu_memory_utilization: 0.3/gpu_memory_utilization: ${GPU_MEM_TALKER}/g" \
+        -e "s/gpu_memory_utilization: 0.2/gpu_memory_utilization: ${GPU_MEM_CODE2WAV}/g" \
+        "${config_template}" > "${output_path}"
+
+    echo "${output_path}"
+}
+
+# Start server and wait for it to be ready
+start_server() {
+    local stage_config="$1"
+    local config_name="$2"
+    local log_file="${RESULT_DIR}/server_${config_name}_${TIMESTAMP}.log"
+
+    echo ""
+    echo "Starting server with config: ${config_name}"
+    echo "  Stage config: ${stage_config}"
+    echo "  Log file: ${log_file}"
+
+    VLLM_WORKER_MULTIPROC_METHOD=spawn \
+    CUDA_VISIBLE_DEVICES="${GPU_DEVICE}" \
+    python -m vllm_omni.entrypoints.cli.main serve "${MODEL}" \
+        --omni \
+        --host 127.0.0.1 \
+        --port "${PORT}" \
+        --stage-configs-path "${stage_config}" \
+        --stage-init-timeout 120 \
+        --trust-remote-code \
+        --disable-log-stats \
+        > "${log_file}" 2>&1 &
+
+    SERVER_PID=$!
+    echo "  Server PID: ${SERVER_PID}"
+
+    # Wait for server to be ready
+    echo "  Waiting for server to be ready..."
+    local max_wait=300
+    local waited=0
+    while [ ${waited} -lt ${max_wait} ]; do
+        if curl -sf "http://127.0.0.1:${PORT}/v1/models" > /dev/null 2>&1; then
+            echo "  Server is ready! (waited ${waited}s)"
+            return 0
+        fi
+        # Check if process is still alive
+        if ! kill -0 ${SERVER_PID} 2>/dev/null; then
+            echo "  ERROR: Server process died. Check log: ${log_file}"
+            tail -20 "${log_file}"
+            return 1
+        fi
+        sleep 2
+        waited=$((waited + 2))
+    done
+
+    echo "  ERROR: Server did not start within ${max_wait}s. Check log: ${log_file}"
+    kill ${SERVER_PID} 2>/dev/null || true
+    return 1
+}
+
+# Stop the server
+stop_server() {
+    if [ -n "${SERVER_PID:-}" ]; then
+        echo "  Stopping server (PID: ${SERVER_PID})..."
+        kill ${SERVER_PID} 2>/dev/null || true
+        wait ${SERVER_PID} 2>/dev/null || true
+        # Kill any remaining child processes on the port
+        local pids
+        pids=$(lsof -ti:${PORT} 2>/dev/null || true)
+        if [ -n "${pids}" ]; then
+            echo "  Cleaning up remaining processes on port ${PORT}..."
+            echo "${pids}" | xargs kill -9 2>/dev/null || true
+        fi
+        echo "  Server stopped."
+        SERVER_PID=""
+    fi
+}
+
+# Cleanup on exit
+trap 'stop_server' EXIT
+
+# Run benchmark for a given config
+run_bench() {
+    local config_name="$1"
+    local config_template="$2"
+
+    echo ""
+    echo "============================================================"
+    echo " Benchmarking: ${config_name}"
+    echo "============================================================"
+
+    local stage_config
+    stage_config=$(prepare_config "${config_template}" "${config_name}")
+
+    start_server "${stage_config}" "${config_name}"
+
+    # Convert concurrency string to args
+    local conc_args=""
+    for c in ${CONCURRENCY}; do
+        conc_args="${conc_args} ${c}"
+    done
+
+    cd "${PROJECT_ROOT}"
+    python "${SCRIPT_DIR}/vllm_omni/bench_tts_serve.py" \
+        --host 127.0.0.1 \
+        --port "${PORT}" \
+        --num-prompts "${NUM_PROMPTS}" \
+        --max-concurrency ${conc_args} \
+        --num-warmups "${NUM_WARMUPS}" \
+        --config-name "${config_name}" \
+        --result-dir "${RESULT_DIR}"
+
+    stop_server
+
+    # Allow GPU memory to settle
+    sleep 5
+}
+
+# Run vllm-omni benchmark
+if [ "${RUN_ASYNC}" = true ]; then
+    run_bench "async_chunk" "${SCRIPT_DIR}/${STAGE_CONFIG}"
+fi
+
+# Run HuggingFace baseline benchmark
+if [ "${RUN_HF}" = true ]; then
+    echo ""
+    echo "============================================================"
+    echo " Benchmarking: HuggingFace transformers (offline)"
+    echo "============================================================"
+
+    cd "${PROJECT_ROOT}"
+    python "${SCRIPT_DIR}/transformers/bench_tts_hf.py" \
+        --model "${MODEL}" \
+        --num-prompts "${NUM_PROMPTS}" \
+        --num-warmups "${NUM_WARMUPS}" \
+        --gpu-device "${GPU_DEVICE}" \
+        --config-name "hf_transformers" \
+        --result-dir "${RESULT_DIR}"
+
+    # Allow GPU memory to settle
+    sleep 5
+fi
+
+# Plot results
+echo ""
+echo "============================================================"
+echo " Generating plots..."
+echo "============================================================"
+
+RESULT_FILES=""
+LABELS=""
+
+if [ "${RUN_ASYNC}" = true ]; then
+    ASYNC_FILE=$(ls -t "${RESULT_DIR}"/bench_async_chunk_*.json 2>/dev/null | head -1)
+    if [ -n "${ASYNC_FILE}" ]; then
+        RESULT_FILES="${ASYNC_FILE}"
+        LABELS="async_chunk"
+    fi
+fi
+
+if [ "${RUN_HF}" = true ]; then
+    HF_FILE=$(ls -t "${RESULT_DIR}"/bench_hf_transformers_*.json 2>/dev/null | head -1)
+    if [ -n "${HF_FILE}" ]; then
+        if [ -n "${RESULT_FILES}" ]; then
+            RESULT_FILES="${RESULT_FILES} ${HF_FILE}"
+            LABELS="${LABELS} hf_transformers"
+        else
+            RESULT_FILES="${HF_FILE}"
+            LABELS="hf_transformers"
+        fi
+    fi
+fi
+
+if [ -n "${RESULT_FILES}" ]; then
+    python "${SCRIPT_DIR}/plot_results.py" \
+        --results ${RESULT_FILES} \
+        --labels ${LABELS} \
+        --output "${RESULT_DIR}/qwen3_tts_benchmark_${TIMESTAMP}.png"
+fi
+
+echo ""
+echo "============================================================"
+echo " Benchmark complete!"
+echo " Results: ${RESULT_DIR}"
+echo "============================================================"

benchmarks/qwen3-tts/transformers/bench_tts_hf.py

@@ -0,0 +1,281 @@
+"""Benchmark Qwen3-TTS using HuggingFace transformers (qwen_tts library).
+
+Measures E2E latency, RTF, and audio duration for offline (non-serving) inference.
+Results are saved in the same JSON format as bench_tts_serve.py for unified plotting.
+
+Usage:
+    python bench_tts_hf.py \
+        --model Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice \
+        --num-prompts 50 \
+        --num-warmups 3 \
+        --gpu-device 0 \
+        --result-dir results/
+"""
+
+import argparse
+import json
+import time
+from dataclasses import asdict, dataclass, field
+from datetime import datetime
+from pathlib import Path
+
+import numpy as np
+import soundfile as sf
+import torch
+
+PROMPTS = [
+    "Hello, welcome to the voice synthesis benchmark test.",
+    "She said she would be here by noon, but nobody showed up.",
+    "The quick brown fox jumps over the lazy dog near the riverbank.",
+    "I can't believe how beautiful the sunset looks from up here on the mountain.",
+    "Please remember to bring your identification documents to the appointment tomorrow morning.",
+    "Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
+    "The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
+    "After the meeting, we should discuss the quarterly results and plan for the next phase.",
+    "Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
+    "The train leaves at half past seven, so we need to arrive at the station before then.",
+    "Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
+    "It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
+]
+
+
+@dataclass
+class BenchmarkResult:
+    config_name: str = ""
+    concurrency: int = 1  # always 1 for offline
+    num_prompts: int = 0
+    completed: int = 0
+    failed: int = 0
+    duration_s: float = 0.0
+    # TTFP stats - not applicable for HF offline, set to E2E for compatibility
+    mean_ttfp_ms: float = 0.0
+    median_ttfp_ms: float = 0.0
+    std_ttfp_ms: float = 0.0
+    p90_ttfp_ms: float = 0.0
+    p95_ttfp_ms: float = 0.0
+    p99_ttfp_ms: float = 0.0
+    # E2E stats (ms)
+    mean_e2e_ms: float = 0.0
+    median_e2e_ms: float = 0.0
+    std_e2e_ms: float = 0.0
+    p90_e2e_ms: float = 0.0
+    p95_e2e_ms: float = 0.0
+    p99_e2e_ms: float = 0.0
+    # RTF stats
+    mean_rtf: float = 0.0
+    median_rtf: float = 0.0
+    std_rtf: float = 0.0
+    p99_rtf: float = 0.0
+    # Audio stats
+    mean_audio_duration_s: float = 0.0
+    total_audio_duration_s: float = 0.0
+    audio_throughput: float = 0.0
+    request_throughput: float = 0.0
+    # Per-request details
+    per_request: list = field(default_factory=list)
+
+
+def run_benchmark(args):
+    from qwen_tts import Qwen3TTSModel
+
+    device = f"cuda:{args.gpu_device}"
+    print(f"Loading model: {args.model} on {device}")
+    model = Qwen3TTSModel.from_pretrained(
+        args.model,
+        device_map=device,
+        dtype=torch.bfloat16,
+    )
+    print("Model loaded.")
+
+    # Build prompt list
+    prompts = [PROMPTS[i % len(PROMPTS)] for i in range(args.num_prompts)]
+
+    # Warmup
+    if args.num_warmups > 0:
+        print(f"Warming up with {args.num_warmups} requests...")
+        for i in range(args.num_warmups):
+            p = PROMPTS[i % len(PROMPTS)]
+            wavs, sr = model.generate_custom_voice(
+                text=p,
+                language=args.language,
+                speaker=args.voice,
+            )
+        # Sync GPU
+        torch.cuda.synchronize(device)
+        print("Warmup done.")
+
+    # Benchmark
+    print(f"Running {args.num_prompts} requests sequentially...")
+    e2e_times = []
+    rtfs = []
+    audio_durations = []
+    per_request = []
+    failed = 0
+
+    audio_dir = None
+    if args.save_audio:
+        audio_dir = Path(args.result_dir) / "audio_hf"
+        audio_dir.mkdir(parents=True, exist_ok=True)
+
+    total_start = time.perf_counter()
+
+    for i, prompt in enumerate(prompts):
+        try:
+            torch.cuda.synchronize(device)
+            st = time.perf_counter()
+
+            wavs, sr = model.generate_custom_voice(
+                text=prompt,
+                language=args.language,
+                speaker=args.voice,
+            )
+
+            torch.cuda.synchronize(device)
+            elapsed = time.perf_counter() - st
+
+            # Compute audio duration
+            audio_samples = wavs[0]
+            if isinstance(audio_samples, torch.Tensor):
+                audio_samples = audio_samples.cpu().numpy()
+            audio_dur = len(audio_samples) / sr
+
+            rtf = elapsed / audio_dur if audio_dur > 0 else 0.0
+
+            e2e_times.append(elapsed)
+            rtfs.append(rtf)
+            audio_durations.append(audio_dur)
+            per_request.append(
+                {
+                    "e2e_ms": elapsed * 1000,
+                    "ttfp_ms": elapsed * 1000,  # no streaming, TTFP = E2E
+                    "rtf": rtf,
+                    "audio_duration_s": audio_dur,
+                    "prompt": prompt,
+                }
+            )
+
+            if audio_dir:
+                sf.write(str(audio_dir / f"output_{i:04d}.wav"), audio_samples, sr)
+
+            if (i + 1) % 10 == 0 or i == 0:
+                print(
+                    f"  [{i + 1}/{args.num_prompts}] e2e={elapsed * 1000:.0f}ms  rtf={rtf:.3f}  audio={audio_dur:.2f}s"
+                )
+
+        except Exception as e:
+            print(f"  [{i + 1}/{args.num_prompts}] FAILED: {e}")
+            failed += 1
+
+    total_duration = time.perf_counter() - total_start
+    completed = len(e2e_times)
+
+    # Compute stats
+    result = BenchmarkResult(
+        config_name=args.config_name,
+        concurrency=1,
+        num_prompts=args.num_prompts,
+        completed=completed,
+        failed=failed,
+        duration_s=total_duration,
+    )
+
+    if e2e_times:
+        e2e_ms = [t * 1000 for t in e2e_times]
+
+        result.mean_e2e_ms = float(np.mean(e2e_ms))
+        result.median_e2e_ms = float(np.median(e2e_ms))
+        result.std_e2e_ms = float(np.std(e2e_ms))
+        result.p90_e2e_ms = float(np.percentile(e2e_ms, 90))
+        result.p95_e2e_ms = float(np.percentile(e2e_ms, 95))
+        result.p99_e2e_ms = float(np.percentile(e2e_ms, 99))
+
+        # For HF offline, TTFP = E2E (no streaming)
+        result.mean_ttfp_ms = result.mean_e2e_ms
+        result.median_ttfp_ms = result.median_e2e_ms
+        result.std_ttfp_ms = result.std_e2e_ms
+        result.p90_ttfp_ms = result.p90_e2e_ms
+        result.p95_ttfp_ms = result.p95_e2e_ms
+        result.p99_ttfp_ms = result.p99_e2e_ms
+
+        result.mean_rtf = float(np.mean(rtfs))
+        result.median_rtf = float(np.median(rtfs))
+        result.std_rtf = float(np.std(rtfs))
+        result.p99_rtf = float(np.percentile(rtfs, 99))
+
+        result.mean_audio_duration_s = float(np.mean(audio_durations))
+        result.total_audio_duration_s = float(np.sum(audio_durations))
+        result.audio_throughput = result.total_audio_duration_s / total_duration
+        result.request_throughput = completed / total_duration
+        result.per_request = per_request
+
+    # Print summary in standardized performance template
+    W = 50
+    print("")
+    print(f"{'=' * W}")
+    print(f"{'Serving Benchmark Result':^{W}}")
+    print(f"{'=' * W}")
+    print(f"{'Successful requests:':<40}{completed:<10}")
+    print(f"{'Failed requests:':<40}{failed:<10}")
+    print(f"{'Maximum request concurrency:':<40}{1:<10}")
+    print(f"{'Benchmark duration (s):':<40}{total_duration:<10.2f}")
+    print(f"{'Request throughput (req/s):':<40}{result.request_throughput:<10.2f}")
+    print(f"{'-' * W}")
+    print(f"{'End-to-end Latency':^{W}}")
+    print(f"{'-' * W}")
+    print(f"{'Mean E2EL (ms):':<40}{result.mean_e2e_ms:<10.2f}")
+    print(f"{'Median E2EL (ms):':<40}{result.median_e2e_ms:<10.2f}")
+    print(f"{'P99 E2EL (ms):':<40}{result.p99_e2e_ms:<10.2f}")
+    print(f"{'=' * W}")
+    print(f"{'Audio Result':^{W}}")
+    print(f"{'=' * W}")
+    print(f"{'Total audio duration generated (s):':<40}{result.total_audio_duration_s:<10.2f}")
+    print(f"{'Audio throughput (audio duration/s):':<40}{result.audio_throughput:<10.2f}")
+    print(f"{'-' * W}")
+    print(f"{'Time to First Packet':^{W}}")
+    print(f"{'-' * W}")
+    print(f"{'Mean AUDIO_TTFP (ms):':<40}{result.mean_ttfp_ms:<10.2f}")
+    print(f"{'Median AUDIO_TTFP (ms):':<40}{result.median_ttfp_ms:<10.2f}")
+    print(f"{'P99 AUDIO_TTFP (ms):':<40}{result.p99_ttfp_ms:<10.2f}")
+    print(f"{'-' * W}")
+    print(f"{'Real Time Factor':^{W}}")
+    print(f"{'-' * W}")
+    print(f"{'Mean AUDIO_RTF:':<40}{result.mean_rtf:<10.3f}")
+    print(f"{'Median AUDIO_RTF:':<40}{result.median_rtf:<10.3f}")
+    print(f"{'P99 AUDIO_RTF:':<40}{result.p99_rtf:<10.3f}")
+    print(f"{'=' * W}")
+    print("")
+
+    # Save results (as a list with single concurrency=1 entry, matching serve format)
+    result_dir = Path(args.result_dir)
+    result_dir.mkdir(parents=True, exist_ok=True)
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
+
+    with open(result_file, "w") as f:
+        json.dump([asdict(result)], f, indent=2)
+    print(f"Results saved to {result_file}")
+
+    return result
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Qwen3-TTS HuggingFace Benchmark")
+    parser.add_argument(
+        "--model", type=str, default="Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice", help="HuggingFace model name or path"
+    )
+    parser.add_argument("--num-prompts", type=int, default=50)
+    parser.add_argument("--num-warmups", type=int, default=3)
+    parser.add_argument("--gpu-device", type=int, default=0)
+    parser.add_argument("--voice", type=str, default="Vivian")
+    parser.add_argument("--language", type=str, default="English")
+    parser.add_argument(
+        "--config-name", type=str, default="hf_transformers", help="Label for this config (used in filenames)"
+    )
+    parser.add_argument("--result-dir", type=str, default="results")
+    parser.add_argument("--save-audio", action="store_true", help="Save generated audio files")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    run_benchmark(args)

benchmarks/qwen3-tts/vllm_omni/bench_tts_serve.py

@@ -0,0 +1,348 @@
+"""Benchmark client for Qwen3-TTS via /v1/audio/speech endpoint.
+
+Measures TTFP (Time-to-First-Packet), E2E latency, and RTF (Real-Time Factor)
+across configurable concurrency levels. Saves results as JSON for plotting.
+
+Usage:
+    python bench_tts_serve.py \
+        --host 127.0.0.1 --port 8000 \
+        --num-prompts 50 \
+        --max-concurrency 1 4 10 \
+        --result-dir results/
+"""
+
+import argparse
+import asyncio
+import json
+import time
+from dataclasses import asdict, dataclass, field
+from datetime import datetime
+from pathlib import Path
+
+import aiohttp
+import numpy as np
+from tqdm.asyncio import tqdm
+
+PROMPTS = [
+    "Hello, welcome to the voice synthesis benchmark test.",
+    "She said she would be here by noon, but nobody showed up.",
+    "The quick brown fox jumps over the lazy dog near the riverbank.",
+    "I can't believe how beautiful the sunset looks from up here on the mountain.",
+    "Please remember to bring your identification documents to the appointment tomorrow morning.",
+    "Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
+    "The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
+    "After the meeting, we should discuss the quarterly results and plan for the next phase.",
+    "Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
+    "The train leaves at half past seven, so we need to arrive at the station before then.",
+    "Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
+    "It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
+]
+
+
+@dataclass
+class RequestResult:
+    success: bool = False
+    ttfp: float = 0.0  # Time to first audio packet (seconds)
+    e2e: float = 0.0  # End-to-end latency (seconds)
+    audio_bytes: int = 0  # Total audio bytes received
+    audio_duration: float = 0.0  # Audio duration in seconds (estimated from PCM)
+    rtf: float = 0.0  # Real-time factor = e2e / audio_duration
+    prompt: str = ""
+    error: str = ""
+
+
+@dataclass
+class BenchmarkResult:
+    config_name: str = ""
+    concurrency: int = 0
+    num_prompts: int = 0
+    completed: int = 0
+    failed: int = 0
+    duration_s: float = 0.0
+    # TTFP stats (ms)
+    mean_ttfp_ms: float = 0.0
+    median_ttfp_ms: float = 0.0
+    std_ttfp_ms: float = 0.0
+    p90_ttfp_ms: float = 0.0
+    p95_ttfp_ms: float = 0.0
+    p99_ttfp_ms: float = 0.0
+    # E2E stats (ms)
+    mean_e2e_ms: float = 0.0
+    median_e2e_ms: float = 0.0
+    std_e2e_ms: float = 0.0
+    p90_e2e_ms: float = 0.0
+    p95_e2e_ms: float = 0.0
+    p99_e2e_ms: float = 0.0
+    # RTF stats
+    mean_rtf: float = 0.0
+    median_rtf: float = 0.0
+    std_rtf: float = 0.0
+    p99_rtf: float = 0.0
+    # Audio stats
+    mean_audio_duration_s: float = 0.0
+    total_audio_duration_s: float = 0.0
+    audio_throughput: float = 0.0  # audio_duration / wall_time
+    request_throughput: float = 0.0  # requests / second
+    # Per-request details
+    per_request: list = field(default_factory=list)
+
+
+def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width: int = 2) -> float:
+    """Convert raw PCM byte count to duration in seconds."""
+    num_samples = num_bytes / sample_width
+    return num_samples / sample_rate
+
+
+async def send_tts_request(
+    session: aiohttp.ClientSession,
+    api_url: str,
+    prompt: str,
+    voice: str = "vivian",
+    language: str = "English",
+    pbar: tqdm | None = None,
+) -> RequestResult:
+    """Send a streaming TTS request and measure latency metrics."""
+    payload = {
+        "input": prompt,
+        "voice": voice,
+        "language": language,
+        "stream": True,
+        "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()}"
+                result.success = False
+                return result
+
+            first_chunk = True
+            total_bytes = 0
+
+            async for chunk in response.content.iter_any():
+                if first_chunk and len(chunk) > 0:
+                    result.ttfp = time.perf_counter() - st
+                    first_chunk = False
+                total_bytes += len(chunk)
+
+            result.e2e = time.perf_counter() - st
+            result.audio_bytes = total_bytes
+            result.audio_duration = pcm_bytes_to_duration(total_bytes)
+
+            if result.audio_duration > 0:
+                result.rtf = result.e2e / result.audio_duration
+            result.success = True
+
+    except Exception as e:
+        result.error = str(e)
+        result.success = False
+        result.e2e = time.perf_counter() - st
+
+    if pbar:
+        pbar.update(1)
+    return result
+
+
+async def run_benchmark(
+    host: str,
+    port: int,
+    num_prompts: int,
+    max_concurrency: int,
+    num_warmups: int = 3,
+    voice: str = "vivian",
+    language: str = "English",
+) -> BenchmarkResult:
+    """Run benchmark at a given concurrency level."""
+    api_url = f"http://{host}:{port}/v1/audio/speech"
+
+    connector = aiohttp.TCPConnector(
+        limit=max_concurrency,
+        limit_per_host=max_concurrency,
+        keepalive_timeout=60,
+    )
+    session = aiohttp.ClientSession(
+        connector=connector,
+        timeout=aiohttp.ClientTimeout(total=600),
+    )
+
+    # Warmup
+    if num_warmups > 0:
+        print(f"  Warming up with {num_warmups} requests...")
+        warmup_tasks = []
+        for i in range(num_warmups):
+            prompt = PROMPTS[i % len(PROMPTS)]
+            warmup_tasks.append(send_tts_request(session, api_url, prompt, voice, language))
+        await asyncio.gather(*warmup_tasks)
+        print("  Warmup done.")
+
+    # Build request list
+    request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
+
+    # Run benchmark
+    print(f"  Running {num_prompts} requests with concurrency={max_concurrency}...")
+    semaphore = asyncio.Semaphore(max_concurrency)
+    pbar = tqdm(total=num_prompts, desc=f"  concurrency={max_concurrency}")
+
+    async def limited_request(prompt):
+        async with semaphore:
+            return await send_tts_request(session, api_url, prompt, voice, language, pbar)
+
+    start_time = time.perf_counter()
+    tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
+    results: list[RequestResult] = await asyncio.gather(*tasks)
+    duration = time.perf_counter() - start_time
+    pbar.close()
+
+    await session.close()
+
+    # Compute stats
+    successful = [r for r in results if r.success]
+    failed = [r for r in results if not r.success]
+
+    bench = BenchmarkResult(
+        concurrency=max_concurrency,
+        num_prompts=num_prompts,
+        completed=len(successful),
+        failed=len(failed),
+        duration_s=duration,
+    )
+
+    if successful:
+        ttfps = [r.ttfp * 1000 for r in successful]  # convert to ms
+        e2es = [r.e2e * 1000 for r in successful]
+        rtfs = [r.rtf for r in successful]
+        audio_durs = [r.audio_duration for r in successful]
+
+        bench.mean_ttfp_ms = float(np.mean(ttfps))
+        bench.median_ttfp_ms = float(np.median(ttfps))
+        bench.std_ttfp_ms = float(np.std(ttfps))
+        bench.p90_ttfp_ms = float(np.percentile(ttfps, 90))
+        bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
+        bench.p99_ttfp_ms = float(np.percentile(ttfps, 99))
+
+        bench.mean_e2e_ms = float(np.mean(e2es))
+        bench.median_e2e_ms = float(np.median(e2es))
+        bench.std_e2e_ms = float(np.std(e2es))
+        bench.p90_e2e_ms = float(np.percentile(e2es, 90))
+        bench.p95_e2e_ms = float(np.percentile(e2es, 95))
+        bench.p99_e2e_ms = float(np.percentile(e2es, 99))
+
+        bench.mean_rtf = float(np.mean(rtfs))
+        bench.median_rtf = float(np.median(rtfs))
+        bench.std_rtf = float(np.std(rtfs))
+        bench.p99_rtf = float(np.percentile(rtfs, 99))
+
+        bench.mean_audio_duration_s = float(np.mean(audio_durs))
+        bench.total_audio_duration_s = float(np.sum(audio_durs))
+        bench.audio_throughput = bench.total_audio_duration_s / duration
+        bench.request_throughput = len(successful) / duration
+
+        bench.per_request = [
+            {
+                "ttfp_ms": r.ttfp * 1000,
+                "e2e_ms": r.e2e * 1000,
+                "rtf": r.rtf,
+                "audio_duration_s": r.audio_duration,
+                "prompt": r.prompt,
+            }
+            for r in successful
+        ]
+
+    # Print summary in standardized performance template
+    W = 50
+    print("")
+    print(f"{'=' * W}")
+    print(f"{'Serving Benchmark Result':^{W}}")
+    print(f"{'=' * W}")
+    print(f"{'Successful requests:':<40}{bench.completed:<10}")
+    print(f"{'Failed requests:':<40}{bench.failed:<10}")
+    print(f"{'Maximum request concurrency:':<40}{max_concurrency:<10}")
+    print(f"{'Benchmark duration (s):':<40}{duration:<10.2f}")
+    print(f"{'Request throughput (req/s):':<40}{bench.request_throughput:<10.2f}")
+    print(f"{'-' * W}")
+    print(f"{'End-to-end Latency':^{W}}")
+    print(f"{'-' * W}")
+    print(f"{'Mean E2EL (ms):':<40}{bench.mean_e2e_ms:<10.2f}")
+    print(f"{'Median E2EL (ms):':<40}{bench.median_e2e_ms:<10.2f}")
+    print(f"{'P99 E2EL (ms):':<40}{bench.p99_e2e_ms:<10.2f}")
+    print(f"{'=' * W}")
+    print(f"{'Audio Result':^{W}}")
+    print(f"{'=' * W}")
+    print(f"{'Total audio duration generated (s):':<40}{bench.total_audio_duration_s:<10.2f}")
+    print(f"{'Audio throughput (audio duration/s):':<40}{bench.audio_throughput:<10.2f}")
+    print(f"{'-' * W}")
+    print(f"{'Time to First Packet':^{W}}")
+    print(f"{'-' * W}")
+    print(f"{'Mean AUDIO_TTFP (ms):':<40}{bench.mean_ttfp_ms:<10.2f}")
+    print(f"{'Median AUDIO_TTFP (ms):':<40}{bench.median_ttfp_ms:<10.2f}")
+    print(f"{'P99 AUDIO_TTFP (ms):':<40}{bench.p99_ttfp_ms:<10.2f}")
+    print(f"{'-' * W}")
+    print(f"{'Real Time Factor':^{W}}")
+    print(f"{'-' * W}")
+    print(f"{'Mean AUDIO_RTF:':<40}{bench.mean_rtf:<10.3f}")
+    print(f"{'Median AUDIO_RTF:':<40}{bench.median_rtf:<10.3f}")
+    print(f"{'P99 AUDIO_RTF:':<40}{bench.p99_rtf:<10.3f}")
+    print(f"{'=' * W}")
+    print("")
+
+    if failed:
+        for r in failed[:3]:
+            print(f"  [ERROR] {r.error[:200]}")
+
+    return bench
+
+
+async def main(args):
+    all_results = []
+
+    for concurrency in args.max_concurrency:
+        result = await run_benchmark(
+            host=args.host,
+            port=args.port,
+            num_prompts=args.num_prompts,
+            max_concurrency=concurrency,
+            num_warmups=args.num_warmups,
+            voice=args.voice,
+            language=args.language,
+        )
+        result.config_name = args.config_name
+        all_results.append(asdict(result))
+
+    # Save results
+    result_dir = Path(args.result_dir)
+    result_dir.mkdir(parents=True, exist_ok=True)
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
+
+    with open(result_file, "w") as f:
+        json.dump(all_results, f, indent=2)
+    print(f"Results saved to {result_file}")
+
+    return all_results
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Qwen3-TTS Benchmark Client")
+    parser.add_argument("--host", type=str, default="127.0.0.1")
+    parser.add_argument("--port", type=int, default=8000)
+    parser.add_argument("--num-prompts", type=int, default=50, help="Number of prompts per concurrency level")
+    parser.add_argument(  # noqa: E501
+        "--max-concurrency", type=int, nargs="+", default=[1, 4, 10], help="Concurrency levels to test"
+    )
+    parser.add_argument("--num-warmups", type=int, default=3)
+    parser.add_argument("--voice", type=str, default="vivian")
+    parser.add_argument("--language", type=str, default="English")
+    parser.add_argument(
+        "--config-name", type=str, default="async_chunk", help="Label for this config (used in filenames)"
+    )
+    parser.add_argument("--result-dir", type=str, default="results")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    asyncio.run(main(args))

benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs1.yaml

@@ -0,0 +1,96 @@
+# Qwen3-TTS batch_size=1 config (streaming with async_chunk)
+# 2-stage pipeline: Talker -> Code2Wav
+async_chunk: true
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    runtime:
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: qwen3_tts
+      model_arch: Qwen3TTSTalkerForConditionalGeneration
+      hf_overrides:
+        architectures: [Qwen3TTSTalkerForConditionalGeneration]
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      enforce_eager: false
+      trust_remote_code: true
+      async_scheduling: false
+      enable_prefix_caching: false
+      engine_output_type: latent
+      gpu_memory_utilization: 0.3
+      distributed_executor_backend: "mp"
+      max_num_batched_tokens: 512
+      max_model_len: 4096
+      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 4096
+      seed: 42
+      detokenize: false
+      repetition_penalty: 1.05
+      stop_token_ids: [2150]
+
+  - stage_id: 1
+    stage_type: llm
+    runtime:
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: code2wav
+      model_arch: Qwen3TTSCode2Wav
+      hf_overrides:
+        architectures: [Qwen3TTSCode2Wav]
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      enforce_eager: true
+      trust_remote_code: true
+      async_scheduling: false
+      enable_prefix_caching: false
+      engine_output_type: audio
+      gpu_memory_utilization: 0.2
+      distributed_executor_backend: "mp"
+      max_num_batched_tokens: 8192
+      max_model_len: 32768
+    engine_input_source: [0]
+    final_output: true
+    final_output_type: audio
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    tts_args:
+      max_instructions_length: 500
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      detokenize: true
+      repetition_penalty: 1.0
+
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1
+
+  connectors:
+    connector_of_shared_memory:
+      name: SharedMemoryConnector
+      extra:
+        shm_threshold_bytes: 65536
+        codec_streaming: true
+        connector_get_sleep_s: 0.01
+        connector_get_max_wait_first_chunk: 3000
+        connector_get_max_wait: 300
+        codec_chunk_frames: 25
+        codec_left_context_frames: 25
+
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1

benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs4.yaml

@@ -0,0 +1,97 @@
+# Qwen3-TTS batch_size=4 config (streaming with async_chunk)
+# Enables concurrent request processing with max_inflight=4
+# 2-stage pipeline: Talker -> Code2Wav
+async_chunk: true
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    runtime:
+      devices: "0"
+      max_batch_size: 4
+    engine_args:
+      model_stage: qwen3_tts
+      model_arch: Qwen3TTSTalkerForConditionalGeneration
+      hf_overrides:
+        architectures: [Qwen3TTSTalkerForConditionalGeneration]
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      enforce_eager: false
+      trust_remote_code: true
+      async_scheduling: false
+      enable_prefix_caching: false
+      engine_output_type: latent
+      gpu_memory_utilization: 0.3
+      distributed_executor_backend: "mp"
+      max_num_batched_tokens: 512
+      max_model_len: 4096
+      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 4096
+      seed: 42
+      detokenize: false
+      repetition_penalty: 1.05
+      stop_token_ids: [2150]
+
+  - stage_id: 1
+    stage_type: llm
+    runtime:
+      devices: "0"
+      max_batch_size: 4
+    engine_args:
+      model_stage: code2wav
+      model_arch: Qwen3TTSCode2Wav
+      hf_overrides:
+        architectures: [Qwen3TTSCode2Wav]
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      enforce_eager: true
+      trust_remote_code: true
+      async_scheduling: false
+      enable_prefix_caching: false
+      engine_output_type: audio
+      gpu_memory_utilization: 0.2
+      distributed_executor_backend: "mp"
+      max_num_batched_tokens: 8192
+      max_model_len: 32768
+    engine_input_source: [0]
+    final_output: true
+    final_output_type: audio
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    tts_args:
+      max_instructions_length: 500
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      detokenize: true
+      repetition_penalty: 1.0
+
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 4
+
+  connectors:
+    connector_of_shared_memory:
+      name: SharedMemoryConnector
+      extra:
+        shm_threshold_bytes: 65536
+        codec_streaming: true
+        connector_get_sleep_s: 0.01
+        connector_get_max_wait_first_chunk: 3000
+        connector_get_max_wait: 300
+        codec_chunk_frames: 25
+        codec_left_context_frames: 25
+
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1

docker/Dockerfile.ci

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

@@ -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

@@ -14,6 +14,7 @@ RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRO
 RUN apt clean && apt-get update -y && \
     apt-get install -y --no-install-recommends --fix-missing \
     curl \
+    espeak-ng \
     ffmpeg \
     git \
     libsndfile1 \

docs/README.md

@@ -60,5 +60,5 @@ vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, in
 
 For more information, checkout the following:
 
-- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true)
+- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true)
 - [vllm-omni announcement blogpost](https://blog.vllm.ai/2025/11/30/vllm-omni.html)

docs/contributing/ci/CI_5levels.md

@@ -67,8 +67,8 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
       <td>Online & Offline (multiple deployment scenarios):<br>real model, normal inference function, normal accuracy</td>
       <td>&lt;30min</td>
       <td>
-        /tests/e2e/online_serving/test_{model_name}_expansion.py<br>
-        /tests/e2e/offline_inference/test_{model_name}_expansion.py
+        /tests/e2e/online_serving/test_{model_name}.py<br>
+        /tests/e2e/offline_inference/test_{model_name}.py
       </td>
       <td>
         <a href="#chapter-2-l3-level-testing-core-integration-performance-and-accuracy-verification">Chapter 2</a><br>
@@ -508,7 +508,9 @@ L3 level testing executes after code is merged into the main branch. Its core pu
     **2.4.5 Request Execution**
 
     ```python
-    openai_client.send_request(request_config, request_num=1)
+    openai_client.send_omni_request(request_config, request_num=1)  # for omni-understanding models
+    # or
+    openai_client.send_diffusion_request(request_config, request_num=1)  # for diffusion models
     ```
 
     **Explanation**:
@@ -517,7 +519,7 @@ L3 level testing executes after code is merged into the main branch. Its core pu
 
     **Single Request**: The comment clearly states this is a single-request completion test. For concurrent testing, it can be extended to multiple requests using request_num = n.
 
-    **Implicit Validation**: The send_request method internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model performs deep validation.
+    **Implicit Validation**: The `send_omni_request` and `send_diffusion_request` methods internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model performs deep validation.
 
 -   ***Run Command***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py -m advanced_model --run-level=advanced_model`
 

docs/contributing/ci/tests_markers.md

@@ -32,7 +32,7 @@ Defined in `pyproject.toml`:
 | `slow`             | Slow tests (may skip in quick CI)                         |
 | `benchmark`        | Benchmark tests                                           |
 
-\* Means those markers are auto-added, and they will be added by the `@hardware_test` decorator.
+\* Means those markers are auto-added by `@hardware_test` (parametrization decorator) or `hardware_marks` (only returning the list of marks for flexibility).
 
 ### Example usage for markers
 
@@ -49,6 +49,7 @@ from tests.utils import hardware_test
 def test_video_to_audio()
     ...
 ```
+
 ### Decorator: `@hardware_test`
 
 This decorator is intended to make hardware-aware, cross-platform test authoring easier and more robust for CI/CD environments. The `hardware_test` decorator in `vllm-omni/tests/utils.py` performs the following actions:
@@ -98,6 +99,24 @@ This decorator is intended to make hardware-aware, cross-platform test authoring
     - ROCm only: `pytest -m "distributed_rocm and MI325"`
     - NPU only: `pytest -m "distributed_npu"`
 
+### Function: `hardware_marks`
+
+`hardware_marks` returns a list of pytest mark objects with the same signature as `@hardware_test`. Use it when you need more flexibility, such as attaching hardware marks to individual `pytest.param` entries rather than an entire test function.
+
+```python
+from tests.utils import hardware_marks
+
+MULTI_CARD_MARKS = hardware_marks(
+    res={"cuda": "H100", "rocm": "MI325", "npu": "A2"}, num_cards=2
+)
+
+@pytest.mark.parametrize("omni_server", [
+    pytest.param(OmniServerParams(...), id="case_001", marks=MULTI_CARD_MARKS),
+], indirect=True)
+def test_feature(omni_server):
+    ...
+```
+
 ## Add Support for a New Platform
 
 If you want to add support for a new platform (e.g., "tpu" for a new accelerator), follow these steps:
@@ -132,12 +151,13 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator
            # Optionally: add skipif_tpu when implemented
            return [test_platform, test_resource, test_distributed]
    ```
-3. **Update `hardware_test` to recognize your new platform**:
-    In the relevant place (see the `hardware_test` implementation), add:
+3. **Update `hardware_marks` to recognize your new platform**:
+    In the relevant place (see the `hardware_marks` implementation), add:
     ```python
     if platform == "tpu":
         marks = tpu_marks(res=resource, num_cards=cards)
     ```
+    (`hardware_test` calls `hardware_marks` internally, so both will pick up the change.)
 4. **(Recommended) Add a test using your new markers**:
    ```python
    @hardware_test(
@@ -151,7 +171,7 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator
 **Summary**:  
 - Add pytest markers for your new platform/resources  
 - Implement a marker function (`xxx_marks`)  
-- Plug into `hardware_test`  
-- You're done: tests decorated with `@hardware_test` using your platform now automatically get the correct markers, distribution, and isolation!
+- Plug into `hardware_marks`  
+- You're done: tests using `@hardware_test` or `hardware_marks` with your platform now automatically get the correct markers, distribution, and isolation!
 
 See code in `vllm-omni/tests/utils.py` for existing examples (`cuda_marks`, `rocm_marks`, `npu_marks`).

docs/contributing/ci/tests_style.md

@@ -158,6 +158,7 @@ vllm_omni/                                    tests/
                                                    ├── test_diffusion_layerwise_offload.py
                                                    ├── test_diffusion_lora.py
                                                    ├── test_sequence_parallel.py
+                                                   ├── test_qwen_image_edit_expansion.py
                                                    └── stage_configs/
                                                        ├── qwen2_5_omni_ci.yaml
                                                        ├── qwen3_omni_ci.yaml

docs/contributing/model/adding_diffusion_model.md

@@ -651,8 +651,27 @@ For a fair comparison, keep the same **prompt**, **seed**, **resolution**, **num
 
 #### 5.2 Functionality Check in CI
 
-To ensure project maintainability and sustainable development, we encourage contributors to submit test code (unit tests, system tests, or end-to-end tests) alongside their code changes.
-For comprehensive testing guidelines, please refer to the [Test File Structure and Style Guide](../ci/tests_style.md).
+To ensure project maintainability and sustainable development, please submit test code (unit tests, system tests, or end-to-end tests) alongside their code changes.
+
+For comprehensive testing guidelines and the definition of test levels (L1-L5), please refer to the [Test File Structure and Style Guide](../ci/tests_style.md).
+The following tests are required to add:
+
+- L4 test of the model's full *functionality* (i.e., all the *diffusion features* that are supported by this model), including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload_diffusion/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/quantization/overview/).
+  - Test cases: Currently all the features are available in online serving mode. Hence, only need to add `tests/e2e/online_serving/test_{model}_expansion.py`. The following test cases shall cover all features:
+    - 1 GPU: TeaCache & GGUF (or fallback to FP8, or disable it) & Layer-wise CPU offloading (or fallback to Module-wise)
+    - 2 GPUs: Cache-DiT & FP8 (or fallback to GGUF, or disable it) & Ulysses = 2
+    - 2 GPUs: Cache-DiT & GGUF (or fallback to FP8, or disable it) & Ring = 2
+    - 2 GPUs: TeaCache & FP8 (or fallback to GGUF, or disable it) & CFG Parallel = 2
+    - 2 GPUs: Cache-DiT & FP8 (or fallback to GGUF, or disable it) & Tensor Parallel = 2 & VAE Patch Parallel = 2
+    - 2 GPUs: Cache-DiT & GGUF (or fallback to FP8, or disable it) & HSDP = 2 & VAE Patch Parallel = 2
+  - Validation: test that the multimodal output files of your model have the correct shapes.
+  - Test marks: always add `advanced_model` and `diffusion`. Add `parallel` and GPU-related marks if needed. Ref: [Markers for Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_markers/)
+  - To maximize code reuse, you may refer to
+    - `tests/conftest.py` for `omni_server` and `openai_client` fixtures, `generate_synthetic_image` and `assert_XXX_valid` helper.
+    - `tests/utils.py` for `@hardware_test(...)` and `hardware_marks`.
+    - [Parametrizing tests (pytest doc)](https://docs.pytest.org/en/stable/example/parametrize.html) to reuse test function implementation for different cases.
+  - Doc: add a concise dostring for each test function.
+  - Reference L4 test implementation: [tests/e2e/online_serving/test_qwen_image_edit_expansion.py](https://github.com/vllm-project/vllm-omni/blob/main/tests/e2e/online_serving/test_qwen_image_edit_expansion.py).
 
 ---
 

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 `codec_chunk_frames` (default=25) before sending. Set `initial_codec_chunk_frames` to emit smaller chunks during the initial phase for reduced TTFA
+- **Talker → Code2Wav**: Accumulated to `codec_chunk_frames` (default=25) before sending. During the initial phase, a dynamic initial chunk size (IC) is automatically selected based on server load to reduce TTFA. Use the per-request `initial_codec_chunk_frames` API field to override.
 - **Code2Wav**: Streaming decode with code2wav chunk_size
 
 With `async_chunk`:
@@ -67,6 +67,13 @@ Enabling **async_chunk** (False→True) sharply reduces time-to-first-audio (TTF
   </picture>
 </p>
 
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_e2e_performance.png">
+    <img alt="E2E Performance Data Comparison" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_e2e_performance.png" width=100%>
+  </picture>
+</p>
+
 ## Architecture
 ### Data Flow
 

docs/features/comfyui.md

@@ -47,7 +47,8 @@ If no, check your shell running the ComfyUI process. There may be some error mes
 This extension offers the following nodes based on the output modalities:
 
 - **Generate Image** for text-to-image and image-to-image tasks
-- **Multimodality Comprehension** for multimodality-to-text and multimodality-to-audio tasks
+- **Generate Video** for text-to-video and image-to-video tasks
+- **Multimodality Understanding** for multimodality-to-text and multimodality-to-audio tasks
 - **TTS** and **TTS Voice Clone** for TTS tasks
 
 This extension also offers example workflows (at **ComfyUI sidebar -> Templates -> vLLM-Omni**)
@@ -59,7 +60,7 @@ To build a simple workflow yourself,
 
 - Drag a generation node onto the canvas.
 - Depending on your need, grab built-in multimedia file loader nodes, such as **image->Load Image**, **image->video->Load Video**, **audio->Load Audio**
-- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**. For text output, you can install [ComfyUI-Custom-Scripts plugin](https://github.com/pythongosssss/ComfyUI-Custom-Scripts/) and grab its **utils->Show Text 🐍** node.
+- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**, **utils->Preview as Text**.
 - If you want to tune sampling parameters, grab corresponding nodes from **vLLM-Omni-> Sampling Params**.
     - For multi-stage models, you can connect multiple **AR Sampling Params** and **Diffusion Sampling Params** nodes to a **Multi-Stage Sampling Params List** node, and connect this node to the generation node.
     - For some multi-stage models like BAGEL, [only one stage's sampling parameters are exposed and tunable via vLLM-Omni's online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/). Thus, these models are treated as single-stage ones. Please check the vLLM-Omni documentation on how to correctly set each model's sampling parameters.

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/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.16.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

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
 ```

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/mkdocs/hooks/generate_examples.py

@@ -132,9 +132,19 @@ class Example:
             ".dat",
             ".db",
             ".sqlite",  # Data files
+            ".pyc",
+            ".pyo",
+            ".pyd",  # Python compiled
+            ".npy",
+            ".npz",
+            ".pkl",
+            ".pickle",  # Serialized data
         }
+        excluded_dirs = {"__pycache__", ".git", "node_modules", ".tox", ".mypy_cache"}
 
         def is_other_file(file: Path) -> bool:
+            if any(part in excluded_dirs for part in file.parts):
+                return False
             return file.is_file() and file != self.main_file and file.suffix.lower() not in binary_extensions
 
         return [file for file in self.path.rglob("*") if is_other_file(file)]

docs/models/supported_models.md

@@ -49,6 +49,7 @@ th {
 |`GlmImageForConditionalGeneration` | GLM-Image | `zai-org/GLM-Image` |
 |`NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` |
 |`MiMoAudioForConditionalGeneration` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` |
+|`Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` |
 
 
 ## List of Supported Models for NPU

docs/serving/speech_api.md

@@ -96,7 +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) |
+| `initial_codec_chunk_frames` | integer | null | Per-request initial chunk size override for TTFA tuning. When null, IC is computed dynamically based on server load. |
 
 **Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian
 

docs/user_guide/diffusion/parallelism_acceleration.md

@@ -18,23 +18,27 @@ The following parallelism methods are currently supported in vLLM-Omni:
 
 6. [HSDP](#hsdp): Hybrid Sharded Data Parallel shards model weights across GPUs using PyTorch FSDP2. This reduces per-GPU memory usage, enabling inference of large models on GPUs with limited memory.
 
+7. [Expert Parallel](#expert-parallelism): Expert Parallelism shards the Experts of a Mixture-of-Experts (MoE) layer across multiple devices. During the forward, a gating mechanism routes tokens to their designated experts, necessitating cross-cards communication(all-to-all) to dispatch tokens to the correct ranks and combine the results. This parallelism allows for massive scaling of model parameters without a proportional increase in the computational load per device.
+
 The following table shows which models are currently supported by parallelism method:
 
 ### ImageGen
 
-| Model                    | Model Identifier                     | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel |
-|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|:------------------:|
-| **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-Edit**      | `Qwen/Qwen-Image-Edit`               |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |
-| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509`          |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |
-| **Qwen-Image-Layered**   | `Qwen/Qwen-Image-Layered`            |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |
-| **Z-Image**              | `Tongyi-MAI/Z-Image-Turbo`           |     ✅      |    ✅    |      ❌       |  ✅ (TP=2 only)  |         ✅          |
-| **Stable-Diffusion3.5**  | `stabilityai/stable-diffusion-3.5`   |     ❌      |    ❌    |      ❌       |        ✅        |         ✅          |
-| **FLUX.2-klein**         | `black-forest-labs/FLUX.2-klein-4B`  |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |
-| **FLUX.1-dev**           | `black-forest-labs/FLUX.1-dev`       |     ❌      |    ❌    |      ✅       |        ✅        |         ❌          |
+| Model                    | Model Identifier                     | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel | Expert-Parallel |
+|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|:------------------:|:---------------:|
+| **LongCat-Image**        | `meituan-longcat/LongCat-Image`      |     ✅      |    ✅    |      ❌       |        ✅        |         ❌          |        N/A        |
+| **LongCat-Image-Edit**   | `meituan-longcat/LongCat-Image-Edit` |     ✅      |    ✅    |      ❌       |        ✅        |         ❌          |        N/A        |
+| **Ovis-Image**           | `OvisAI/Ovis-Image`                  |     ❌      |    ❌    |      ❌       |        ❌        |         ❌          |        N/A        |
+| **Qwen-Image**           | `Qwen/Qwen-Image`                    |     ✅      |    ✅    |      ✅       |        ✅        |         ✅          |        N/A        |
+| **Qwen-Image-Edit**      | `Qwen/Qwen-Image-Edit`               |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |        N/A        |
+| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509`          |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |        N/A        |
+| **Qwen-Image-Layered**   | `Qwen/Qwen-Image-Layered`            |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |        N/A        |
+| **Z-Image**              | `Tongyi-MAI/Z-Image-Turbo`           |     ✅      |    ✅    |      ❌       |  ✅ (TP=2 only)  |         ✅          |        N/A        |
+| **Stable-Diffusion3.5**  | `stabilityai/stable-diffusion-3.5`   |     ❌      |    ❌    |      ❌       |        ✅        |         ✅          |        N/A        |
+| **FLUX.2-klein**         | `black-forest-labs/FLUX.2-klein-4B`  |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |        N/A        |
+| **FLUX.1-dev**           | `black-forest-labs/FLUX.1-dev`       |     ❌      |    ❌    |      ✅       |        ✅        |         ❌          |        N/A        |
+| **FLUX.2-dev**           | `black-forest-labs/FLUX.2-dev`       |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |        N/A        |
+| **HunyuanImage3.0**      | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |        ✅        |
 
 !!! note "TP Limitations for Diffusion Models"
     We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP.
@@ -414,3 +418,28 @@ vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091 --use-hsdp --usp
 #### Adding HSDP Support to New Models
 
 For detailed instructions on adding HSDP support to new models, see the [HSDP Contributing Guide](../../design/feature/hsdp.md).
+
+### Expert Parallelism
+
+Unlike Tensor Parallelism which shards every layer's weights, EP only shards the MoE expert MLP blocks. This significantly reduces the memory footprint of MoE models (e.g., HunyuanImage3.0) while maintaining constant dense-equivalent compute efficiency. Expert Parallelism is enabled via `DiffusionParallelConfig.enable_expert_parallel`. And `self.ep = tp * sp * cfg * dp` for now, so at least one of TP/SP/CFG/DP should set when EP enabled.
+
+#### Offline Inference
+
+```python
+from vllm_omni import Omni
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+omni = Omni(
+    model="tencent/HunyuanImage-3.0",
+    parallel_config=DiffusionParallelConfig(tensor_parallel_size=8, enable_expert_parallel=True),
+)
+
+outputs = omni.generate(
+    "A brown and white dog is running on the grass",
+    OmniDiffusionSamplingParams(
+        num_inference_steps=50,
+        width=1024,
+        height=1024,
+    ),
+)
+```

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 \
-  --num_images_per_prompt 1 \
-  --num_inference_steps 4 \
+  --cfg-scale 4.0 \
+  --num-images-per-prompt 1 \
+  --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 \
-  --num_images_per_prompt 1 \
-  --num_inference_steps 4 \
+  --cfg-scale 4.0 \
+  --num-images-per-prompt 1 \
+  --num-inference-steps 4 \
   --output outputs/flux2_klein_4b_gguf.png
 ```
 

docs/user_guide/diffusion_acceleration.md

@@ -69,6 +69,7 @@ The following table shows which models are currently supported by each accelerat
 | **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
 | **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ |
 | **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
+| **FLUX.2-dev** | `black-forest-labs/FLUX.2-dev` | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ |
 
 ### VideoGen
 

docs/user_guide/examples/offline_inference/qwen3_tts.md

@@ -99,7 +99,8 @@ python end2end.py --query-type CustomVoice --streaming --output-dir /tmp/out_str
 ```
 
 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
+in the stage config). The initial chunk size is dynamically selected based on server load for reduced
+TTFA, and can be overridden per-request via the `initial_codec_chunk_frames` API field. 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

@@ -17,6 +17,31 @@ Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/
 | `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` | CustomVoice | Smaller/faster variant                                |
 | `Qwen/Qwen3-TTS-12Hz-0.6B-Base`        | Base        | Smaller/faster variant for voice cloning              |
 
+## Gradio Demo
+
+Two interactive Gradio demos are available, both supporting all 3 task types:
+
+| Demo | File | Transport | Streaming Quality |
+| ---- | ---- | --------- | ----------------- |
+| Standard | `gradio_demo.py` | HTTP chunked | May have small gaps between chunks |
+| FastRTC | `gradio_fastrtc_demo.py` | WebRTC | Gapless streaming (requires `pip install fastrtc`) |
+
+```bash
+# Option 1: Launch server + Standard Gradio together
+./run_gradio_demo.sh                                # CustomVoice (default)
+./run_gradio_demo.sh --task-type VoiceDesign        # VoiceDesign
+./run_gradio_demo.sh --task-type Base               # Voice cloning
+
+# Option 2: If server is already running
+python gradio_demo.py --api-base http://localhost:8000
+
+# Option 3: FastRTC demo (gapless streaming)
+pip install fastrtc
+python gradio_fastrtc_demo.py --api-base http://localhost:8000
+```
+
+Then open http://localhost:7860 in your browser.
+
 ## Run examples (Qwen3-TTS)
 
 ### Launch the Server
@@ -112,7 +137,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). Local paths are automatically base64-encoded by the client before sending to the server.
+- `--ref-audio`: Reference audio file path or URL for voice cloning (Base task)
 - `--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`)
@@ -187,6 +212,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
 
 ```
@@ -196,20 +274,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
@@ -220,7 +284,7 @@ Lists available voices for the loaded model:
     "task_type": "CustomVoice",
     "language": "Auto",
     "instructions": "Optional style instructions",
-    "ref_audio": "HTTP URL, base64 data URL, or file:// URI for voice cloning",
+    "ref_audio": "URL or base64 for voice cloning",
     "ref_text": "Reference audio transcript",
     "x_vector_only_mode": false,
     "max_new_tokens": 2048
@@ -253,6 +317,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 | Per-request initial chunk size override for TTFA tuning. When null, IC is computed dynamically based on server load. |
 | `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
@@ -261,7 +326,7 @@ Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/w
 
 | Parameter            | Type   | Required | Description                                  |
 | -------------------- | ------ | -------- | -------------------------------------------- |
-| `ref_audio`          | string | **Yes**  | Reference audio (HTTP URL, base64 data URL, or `file://` URI with `--allowed-local-media-path`) |
+| `ref_audio`          | string | **Yes**  | Reference audio (URL or base64 data URL)     |
 | `ref_text`           | string | No       | Transcript of reference audio (for ICL mode) |
 | `x_vector_only_mode` | bool   | No       | Use speaker embedding only (no ICL)          |
 
@@ -301,11 +366,27 @@ curl -X POST http://localhost:8091/v1/audio/speech \
 
 ## Example materials
 
+??? abstract "gradio_demo.py"
+    ``````py
+    --8<-- "examples/online_serving/qwen3_tts/gradio_demo.py"
+    ``````
+??? abstract "gradio_fastrtc_demo.py"
+    ``````py
+    --8<-- "examples/online_serving/qwen3_tts/gradio_fastrtc_demo.py"
+    ``````
 ??? abstract "openai_speech_client.py"
     ``````py
     --8<-- "examples/online_serving/qwen3_tts/openai_speech_client.py"
     ``````
+??? abstract "run_gradio_demo.sh"
+    ``````sh
+    --8<-- "examples/online_serving/qwen3_tts/run_gradio_demo.sh"
+    ``````
 ??? abstract "run_server.sh"
     ``````sh
     --8<-- "examples/online_serving/qwen3_tts/run_server.sh"
     ``````
+??? abstract "tts_common.py"
+    ``````py
+    --8<-- "examples/online_serving/qwen3_tts/tts_common.py"
+    ``````

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():
@@ -58,6 +57,7 @@ def parse_args():
         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
@@ -88,108 +88,80 @@ def main():
 
     from PIL import Image
 
-    if args.modality == "img2img":
-        from vllm_omni.entrypoints.omni_diffusion import OmniDiffusion
+    from vllm_omni.entrypoints.omni import Omni
 
-        print(f"[Info] Running in {args.modality} mode (Stage 1 only, cfg_parallel_size={args.cfg_parallel_size})")
+    omni_kwargs = {}
+    if args.stage_configs_path:
+        omni_kwargs["stage_configs_path"] = args.stage_configs_path
 
-        client = OmniDiffusion(
-            model=model_name,
-            parallel_config={"cfg_parallel_size": args.cfg_parallel_size},
-        )
+    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.image_path:
-            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.")
+    omni = Omni(model=model_name, **omni_kwargs)
 
-        result = client.generate(
-            prompts,
-            OmniDiffusionSamplingParams(
-                seed=52,
-                need_kv_receive=False,
-                num_inference_steps=args.steps,
-                extra_args={
-                    "cfg_text_scale": args.cfg_text_scale,
-                    "cfg_img_scale": args.cfg_img_scale,
-                },
-            ),
-        )
-
-        # 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,
+    formatted_prompts = []
+    for p in prompts:
+        if args.modality == "img2img":
+            if not args.image_path or not os.path.exists(args.image_path):
+                raise ValueError(f"img2img requires --image-path pointing to an existing file, got: {args.image_path}")
+            loaded_image = Image.open(args.image_path).convert("RGB")
+            final_prompt_text = f"<|fim_middle|><|im_start|>{p}<|im_end|>"
+            prompt_dict = {
+                "prompt": final_prompt_text,
+                "multi_modal_data": {"img2img": loaded_image},
+                "modalities": ["img2img"],
             }
-        )
-
-        omni = Omni(model=model_name, **omni_kwargs)
-
-        formatted_prompts = []
-        for p in args.prompts:
-            if args.modality == "img2text":
-                if args.image_path:
-                    loaded_image = Image.open(args.image_path).convert("RGB")
-                    final_prompt_text = f"<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n"
-                    prompt_dict = {
-                        "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"]}
-                if args.negative_prompt is not None:
-                    prompt_dict["negative_prompt"] = args.negative_prompt
-                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
-                extra = {
-                    "cfg_text_scale": args.cfg_text_scale,
-                    "cfg_img_scale": args.cfg_img_scale,
+            if args.negative_prompt is not None:
+                prompt_dict["negative_prompt"] = args.negative_prompt
+            formatted_prompts.append(prompt_dict)
+        elif args.modality == "img2text":
+            if args.image_path:
+                loaded_image = Image.open(args.image_path).convert("RGB")
+                final_prompt_text = f"<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n"
+                prompt_dict = {
+                    "prompt": final_prompt_text,
+                    "multi_modal_data": {"image": loaded_image},
+                    "modalities": ["text"],
                 }
-                if args.negative_prompt is not None:
-                    extra["negative_prompt"] = args.negative_prompt
-                diffusion_params.extra_args = extra  # type: ignore
+                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_dit_enable_taylorseer
+        --cache-backend cache_dit \
+        --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("--guidance_scale", type=float, default=1.0)
+    parser.add_argument("--cfg-scale", type=float, default=4.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_inference_steps", type=int, default=50)
-    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("--ring_degree", type=int, default=1)
-    parser.add_argument("--tensor_parallel_size", 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("--cache-backend", type=str, default=None, choices=["cache_dit", "tea_cache"])
+    parser.add_argument("--ulysses-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("--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_bn_compute_blocks", type=int, default=0)
-    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_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_taylorseer_order", 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-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-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-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("--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("--enforce_eager", action="store_true")
-    parser.add_argument("--vae_use_slicing", action="store_true")
-    parser.add_argument("--vae_use_tiling", action="store_true")
+    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("--cfg-parallel-size", type=int, default=1, choices=[1, 2])
+    parser.add_argument("--enforce-eager", action="store_true")
+    parser.add_argument("--vae-use-slicing", 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/image_to_image/image_edit.py

@@ -208,6 +208,11 @@ def parse_args() -> argparse.Namespace:
         default=1,
         help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
     )
+    parser.add_argument(
+        "--enable-expert-parallel",
+        action="store_true",
+        help="Enable expert parallelism for MoE layers.",
+    )
     parser.add_argument("--layers", type=int, default=4, help="Number of layers to decompose the input image into.")
     parser.add_argument(
         "--resolution",
@@ -348,6 +353,7 @@ def main():
         ring_degree=args.ring_degree,
         cfg_parallel_size=args.cfg_parallel_size,
         tensor_parallel_size=args.tensor_parallel_size,
+        enable_expert_parallel=args.enable_expert_parallel,
     )
 
     # Configure cache based on backend type
@@ -402,7 +408,7 @@ def main():
     else:
         print(f"  Input image size: {input_image.size}")
     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}, cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size}, enable_expert_parallel: {args.enable_expert_parallel}"
     )
     print(f"{'=' * 60}\n")
 

examples/offline_inference/image_to_video/image_to_video.py

@@ -22,7 +22,7 @@ Usage:
     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 \
+        --num-frames 121 --num-inference-steps 40 --guidance-scale 4.0 \
         --frame-rate 24 --fps 24 --output ltx2_i2v.mp4
 """
 

examples/offline_inference/qwen3_tts/README.md

@@ -96,7 +96,8 @@ python end2end.py --query-type CustomVoice --streaming --output-dir /tmp/out_str
 ```
 
 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
+in the stage config). The initial chunk size is dynamically selected based on server load for reduced
+TTFA, and can be overridden per-request via the `initial_codec_chunk_frames` API field. 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/text_to_image/README.md

@@ -106,6 +106,23 @@ python text_to_image.py \
   --seed 42
 ```
 
+### Flux.2-dev Models
+To start Flux.2-dev with a single GPU, cpu-offload must be enabled because the total size of its weights exceeds the 80GB memory capacity of the GPU.
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model black-forest-labs/FLUX.2-dev \
+  --prompt "a lovely bunny holding a sign that says 'vllm-omni'" \
+  --seed 42 \
+  --tensor-parallel-size 1 \
+  --num-images-per-prompt 1 \
+  --num-inference-steps 50 \
+  --guidance-scale 4.0 \
+  --height 1024 \
+  --width 1024 \
+  --enable-cpu-offload \
+  --output flux2-dev.png
+```
+
 ### Key Arguments
 
 **Common arguments:**

examples/offline_inference/text_to_image/text_to_image.py

@@ -170,6 +170,11 @@ def parse_args() -> argparse.Namespace:
         default=1,
         help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
     )
+    parser.add_argument(
+        "--enable-expert-parallel",
+        action="store_true",
+        help="Enable expert parallelism for MoE layers.",
+    )
     parser.add_argument(
         "--lora-path",
         type=str,
@@ -258,6 +263,7 @@ def main():
         cfg_parallel_size=args.cfg_parallel_size,
         tensor_parallel_size=args.tensor_parallel_size,
         vae_patch_parallel_size=args.vae_patch_parallel_size,
+        enable_expert_parallel=args.enable_expert_parallel,
     )
 
     # Check if profiling is requested via environment variable
@@ -323,7 +329,7 @@ def main():
     print(
         f"  Parallel configuration: tensor_parallel_size={args.tensor_parallel_size}, "
         f"ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, "
-        f"vae_patch_parallel_size={args.vae_patch_parallel_size}"
+        f"vae_patch_parallel_size={args.vae_patch_parallel_size}, enable_expert_parallel={args.enable_expert_parallel}."
     )
     print(f"  CPU offload: {args.enable_cpu_offload}")
     print(f"  Image size: {args.width}x{args.height}")

examples/offline_inference/text_to_video/text_to_video.py

@@ -127,6 +127,11 @@ def parse_args() -> argparse.Namespace:
         default=1,
         help="Number of GPUs used for VAE patch/tile parallelism (decode).",
     )
+    parser.add_argument(
+        "--enable-expert-parallel",
+        action="store_true",
+        help="Enable expert parallelism for MoE layers.",
+    )
     return parser.parse_args()
 
 
@@ -160,6 +165,7 @@ def main():
         cfg_parallel_size=args.cfg_parallel_size,
         tensor_parallel_size=args.tensor_parallel_size,
         vae_patch_parallel_size=args.vae_patch_parallel_size,
+        enable_expert_parallel=args.enable_expert_parallel,
     )
 
     # Check if profiling is requested via environment variable
@@ -193,7 +199,7 @@ def main():
     print(
         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}"
+        f" vae_patch_parallel_size={args.vae_patch_parallel_size}, enable_expert_parallel={args.enable_expert_parallel}"
     )
     print(f"  Video size: {args.width}x{args.height}")
     print(f"{'=' * 60}\n")

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/qwen3_tts/README.md

@@ -14,6 +14,31 @@ Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/
 | `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` | CustomVoice | Smaller/faster variant                                |
 | `Qwen/Qwen3-TTS-12Hz-0.6B-Base`        | Base        | Smaller/faster variant for voice cloning              |
 
+## Gradio Demo
+
+Two interactive Gradio demos are available, both supporting all 3 task types:
+
+| Demo | File | Transport | Streaming Quality |
+| ---- | ---- | --------- | ----------------- |
+| Standard | `gradio_demo.py` | HTTP chunked | May have small gaps between chunks |
+| FastRTC | `gradio_fastrtc_demo.py` | WebRTC | Gapless streaming (requires `pip install fastrtc`) |
+
+```bash
+# Option 1: Launch server + Standard Gradio together
+./run_gradio_demo.sh                                # CustomVoice (default)
+./run_gradio_demo.sh --task-type VoiceDesign        # VoiceDesign
+./run_gradio_demo.sh --task-type Base               # Voice cloning
+
+# Option 2: If server is already running
+python gradio_demo.py --api-base http://localhost:8000
+
+# Option 3: FastRTC demo (gapless streaming)
+pip install fastrtc
+python gradio_fastrtc_demo.py --api-base http://localhost:8000
+```
+
+Then open http://localhost:7860 in your browser.
+
 ## Run examples (Qwen3-TTS)
 
 ### Launch the Server
@@ -289,7 +314,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)                 |
+| `initial_codec_chunk_frames` | int | null | Per-request initial chunk size override for TTFA tuning. When null, IC is computed dynamically based on server load. |
 | `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

examples/online_serving/qwen3_tts/gradio_demo.py

@@ -0,0 +1,404 @@
+"""Gradio demo for Qwen3-TTS online serving via /v1/audio/speech API.
+
+Uses Gradio's built-in audio component for playback. Supports streaming
+(progressive PCM chunks) and non-streaming (full audio download) modes.
+
+For gapless streaming playback, see gradio_fastrtc_demo.py which uses WebRTC.
+
+Supports all 3 task types:
+  - CustomVoice: Predefined speaker with optional style instructions
+  - VoiceDesign: Natural language voice description
+  - Base: Voice cloning from reference audio (upload or URL)
+
+Usage:
+    # Start the server first (see run_server.sh), then:
+    python gradio_demo.py --api-base http://localhost:8000
+
+    # Or use run_gradio_demo.sh to start both server and demo together.
+"""
+
+import argparse
+import io
+
+import gradio as gr
+import httpx
+import numpy as np
+import soundfile as sf
+from tts_common import (
+    PCM_SAMPLE_RATE,
+    SUPPORTED_LANGUAGES,
+    TASK_TYPES,
+    add_common_args,
+    build_payload,
+    fetch_voices,
+    on_task_type_change,
+    stream_pcm_chunks,
+)
+
+
+def generate_speech(
+    api_base: str,
+    text: str,
+    task_type: str,
+    voice: str,
+    language: str,
+    instructions: str,
+    ref_audio: tuple | None,
+    ref_audio_url: str,
+    ref_text: str,
+    x_vector_only: bool,
+    response_format: str,
+    speed: float,
+):
+    """Non-streaming: call /v1/audio/speech and return audio."""
+    payload = build_payload(
+        text,
+        task_type,
+        voice,
+        language,
+        instructions,
+        ref_audio,
+        ref_audio_url,
+        ref_text,
+        x_vector_only,
+        response_format,
+        speed,
+        stream=False,
+    )
+
+    try:
+        with httpx.Client(timeout=300.0) as client:
+            resp = client.post(
+                f"{api_base}/v1/audio/speech",
+                json=payload,
+                headers={
+                    "Content-Type": "application/json",
+                    "Authorization": "Bearer EMPTY",
+                },
+            )
+    except httpx.TimeoutException:
+        raise gr.Error("Request timed out. The server may be busy.")
+    except httpx.ConnectError:
+        raise gr.Error(f"Cannot connect to server at {api_base}. Make sure the vLLM server is running.")
+
+    if resp.status_code != 200:
+        raise gr.Error(f"Server error ({resp.status_code}): {resp.text}")
+
+    content_type = resp.headers.get("content-type", "")
+    if "application/json" in content_type:
+        try:
+            error_data = resp.json()
+            raise gr.Error(f"Server error: {error_data}")
+        except ValueError:
+            pass
+
+    try:
+        if response_format == "pcm":
+            audio_np = np.frombuffer(resp.content, dtype=np.int16).astype(np.float32) / 32767.0
+            return (PCM_SAMPLE_RATE, audio_np)
+        audio_np, sample_rate = sf.read(io.BytesIO(resp.content))
+        if audio_np.ndim > 1:
+            audio_np = audio_np[:, 0]
+        return (sample_rate, audio_np.astype(np.float32))
+    except Exception as e:
+        raise gr.Error(f"Failed to decode audio response: {e}")
+
+
+def generate_speech_stream(
+    api_base: str,
+    text: str,
+    task_type: str,
+    voice: str,
+    language: str,
+    instructions: str,
+    ref_audio: tuple | None,
+    ref_audio_url: str,
+    ref_text: str,
+    x_vector_only: bool,
+    response_format: str,
+    speed: float,
+    chunk_seconds: float = 0.25,
+):
+    """Streaming: yield individual PCM chunks for progressive playback."""
+    payload = build_payload(
+        text,
+        task_type,
+        voice,
+        language,
+        instructions,
+        ref_audio,
+        ref_audio_url,
+        ref_text,
+        x_vector_only,
+        response_format,
+        speed,
+        stream=True,
+    )
+
+    min_chunk_samples = int(PCM_SAMPLE_RATE * chunk_seconds)
+    # Pre-buffer 2 chunks before first playback to avoid the gap
+    # between chunk 1 and chunk 2 during progressive streaming.
+    prebuffer_chunks = 2
+    pending: list[np.ndarray] = []
+    pending_len = 0
+    chunks_yielded = 0
+    prebuffer: list[np.ndarray] = []
+
+    try:
+        for samples in stream_pcm_chunks(api_base, payload):
+            float_samples = samples.astype(np.float32) / 32767.0
+            pending.append(float_samples)
+            pending_len += len(float_samples)
+            if pending_len >= min_chunk_samples:
+                audio_chunk = np.concatenate(pending)
+                pending.clear()
+                pending_len = 0
+
+                if chunks_yielded < prebuffer_chunks:
+                    prebuffer.append(audio_chunk)
+                    chunks_yielded += 1
+                    if chunks_yielded == prebuffer_chunks:
+                        yield (PCM_SAMPLE_RATE, np.concatenate(prebuffer))
+                        prebuffer.clear()
+                else:
+                    yield (PCM_SAMPLE_RATE, audio_chunk)
+        # Flush: prebuffer (if stream was shorter than 2 chunks) + remaining pending
+        remaining = prebuffer + pending
+        if remaining:
+            yield (PCM_SAMPLE_RATE, np.concatenate(remaining))
+    except httpx.TimeoutException:
+        raise gr.Error("Request timed out. The server may be busy.")
+    except httpx.ConnectError:
+        raise gr.Error(f"Cannot connect to server at {api_base}. Make sure the vLLM server is running.")
+
+
+def on_stream_change(stream: bool):
+    """When streaming is enabled, lock format to PCM and disable speed."""
+    if stream:
+        return (
+            gr.update(value="pcm", interactive=False),
+            gr.update(interactive=False),
+        )
+    return (
+        gr.update(value="wav", interactive=True),
+        gr.update(interactive=True),
+    )
+
+
+def build_interface(api_base: str, stream_chunk_seconds: float = 0.25):
+    """Build the Gradio interface."""
+    voices = fetch_voices(api_base)
+
+    css = """
+    #generate-btn button { width: 100%; }
+    .task-info { padding: 8px 12px; border-radius: 6px;
+                 background: #f0f4ff; margin-bottom: 8px; }
+    """
+
+    with gr.Blocks(css=css, title="Qwen3-TTS Demo") as demo:
+        gr.Markdown("# Qwen3-TTS Online Serving Demo")
+        gr.Markdown(f"**Server:** `{api_base}`")
+
+        with gr.Row():
+            # Left column: inputs
+            with gr.Column(scale=3):
+                text_input = gr.Textbox(
+                    label="Text to Synthesize",
+                    placeholder="Enter text here, e.g., Hello, how are you?",
+                    lines=4,
+                )
+
+                with gr.Row():
+                    task_type = gr.Radio(
+                        choices=TASK_TYPES,
+                        value="CustomVoice",
+                        label="Task Type",
+                        scale=2,
+                    )
+                    language = gr.Dropdown(
+                        choices=SUPPORTED_LANGUAGES,
+                        value="Auto",
+                        label="Language",
+                        scale=1,
+                    )
+
+                voice = gr.Dropdown(
+                    choices=voices,
+                    value=voices[0] if voices else None,
+                    label="Speaker",
+                    visible=True,
+                )
+
+                instructions = gr.Textbox(
+                    label="Instructions",
+                    placeholder=("e.g., Speak with excitement / A warm, friendly female voice"),
+                    lines=2,
+                    visible=True,
+                    info="Optional style/emotion instructions",
+                )
+
+                ref_audio = gr.Audio(
+                    label="Reference Audio (upload for voice cloning)",
+                    type="numpy",
+                    sources=["upload", "microphone"],
+                    visible=False,
+                )
+                ref_audio_url = gr.Textbox(
+                    label="Reference Audio URL",
+                    placeholder=("https://example.com/reference.wav (alternative to uploading)"),
+                    lines=1,
+                    visible=False,
+                )
+                ref_text = gr.Textbox(
+                    label="Reference Audio Transcript",
+                    placeholder=("Transcript of the reference audio (optional, improves quality)"),
+                    lines=2,
+                    visible=False,
+                )
+                x_vector_only = gr.Checkbox(
+                    label="Use x-vector only",
+                    value=False,
+                    visible=False,
+                    info=("Skip reference transcript, use speaker embedding only (lower quality)"),
+                )
+
+                with gr.Row():
+                    response_format = gr.Dropdown(
+                        choices=["wav", "mp3", "flac", "pcm", "aac", "opus"],
+                        value="wav",
+                        label="Audio Format",
+                        scale=1,
+                    )
+                    speed = gr.Slider(
+                        minimum=0.25,
+                        maximum=4.0,
+                        value=1.0,
+                        step=0.05,
+                        label="Speed",
+                        scale=1,
+                    )
+                    stream_checkbox = gr.Checkbox(
+                        label="Stream output",
+                        value=False,
+                        info=(
+                            "Enable streaming (PCM format, speed disabled). "
+                            "For gapless streaming, use gradio_fastrtc_demo.py"
+                        ),
+                        scale=1,
+                    )
+
+                generate_btn = gr.Button(
+                    "Generate Speech",
+                    variant="primary",
+                    size="lg",
+                    elem_id="generate-btn",
+                )
+
+            # Right column: output
+            with gr.Column(scale=2):
+                audio_output = gr.Audio(
+                    label="Generated Audio",
+                    interactive=False,
+                    streaming=True,
+                    autoplay=True,
+                )
+                gr.Markdown(
+                    "### Task Types\n"
+                    "- **CustomVoice**: Use a predefined speaker "
+                    "(Vivian, Ryan, etc.) with optional style instructions\n"
+                    "- **VoiceDesign**: Describe the desired voice in natural "
+                    "language (instructions required)\n"
+                    "- **Base**: Clone a voice from reference audio "
+                    "(upload a file or provide a URL)"
+                )
+
+        # Dynamic UI updates
+        task_type.change(
+            fn=on_task_type_change,
+            inputs=[task_type],
+            outputs=[
+                voice,
+                instructions,
+                ref_audio,
+                ref_audio_url,
+                ref_text,
+                x_vector_only,
+            ],
+        )
+
+        stream_checkbox.change(
+            fn=on_stream_change,
+            inputs=[stream_checkbox],
+            outputs=[response_format, speed],
+        )
+
+        all_inputs = [
+            text_input,
+            task_type,
+            voice,
+            language,
+            instructions,
+            ref_audio,
+            ref_audio_url,
+            ref_text,
+            x_vector_only,
+            response_format,
+            speed,
+        ]
+
+        def dispatch(stream_enabled, *args):
+            if stream_enabled:
+                yield from generate_speech_stream(
+                    api_base,
+                    *args,
+                    chunk_seconds=stream_chunk_seconds,
+                )
+            else:
+                yield generate_speech(api_base, *args)
+
+        generate_btn.click(
+            fn=dispatch,
+            inputs=[stream_checkbox] + all_inputs,
+            outputs=[audio_output],
+        )
+
+        demo.queue()
+    return demo
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Gradio demo for Qwen3-TTS online serving.")
+    add_common_args(parser)
+    parser.add_argument(
+        "--stream-chunk-seconds",
+        type=float,
+        default=0.25,
+        help=(
+            "Seconds of audio to buffer per streaming chunk (default: 0.25). "
+            "Lower = smoother but more overhead, higher = chunkier."
+        ),
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    print(f"Connecting to vLLM server at: {args.api_base}")
+    demo = build_interface(
+        args.api_base,
+        stream_chunk_seconds=args.stream_chunk_seconds,
+    )
+
+    try:
+        demo.launch(
+            server_name=args.host,
+            server_port=args.port,
+            share=args.share,
+        )
+    except KeyboardInterrupt:
+        print("\nShutting down...")
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/qwen3_tts/gradio_fastrtc_demo.py

@@ -0,0 +1,274 @@
+"""Gradio demo for Qwen3-TTS using FastRTC for gapless streaming audio.
+
+Uses WebRTC via FastRTC to deliver audio chunks as a continuous real-time
+stream, avoiding the inter-chunk gaps present in Gradio's built-in streaming.
+
+Requires: pip install fastrtc
+
+Supports all 3 task types:
+  - CustomVoice: Predefined speaker with optional style instructions
+  - VoiceDesign: Natural language voice description
+  - Base: Voice cloning from reference audio (URL only in this demo)
+
+Usage:
+    # Start the server first (see run_server.sh), then:
+    python gradio_fastrtc_demo.py --api-base http://localhost:8000
+
+    # With public sharing:
+    python gradio_fastrtc_demo.py --api-base http://localhost:8000 --share
+"""
+
+import argparse
+import logging
+from collections.abc import Generator
+
+import gradio as gr
+import numpy as np
+from fastrtc import WebRTC
+from tts_common import (
+    PCM_SAMPLE_RATE,
+    SUPPORTED_LANGUAGES,
+    TASK_TYPES,
+    add_common_args,
+    build_payload,
+    fetch_voices,
+    stream_pcm_chunks,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def make_tts_generator(api_base: str):
+    """Create a generator function for FastRTC's mode='receive'.
+
+    FastRTC calls this with the additional_inputs values whenever the user
+    clicks "Start Stream". It must be a synchronous generator yielding
+    (sample_rate, int16_ndarray) tuples.
+    """
+
+    def tts_generate(
+        text: str,
+        task_type: str,
+        voice: str,
+        language: str,
+        instructions: str,
+        ref_audio_url: str,
+        ref_text: str,
+        x_vector_only: bool,
+    ) -> Generator[tuple[int, np.ndarray], None, None]:
+        if not text or not text.strip():
+            return
+
+        payload = build_payload(
+            text,
+            task_type,
+            voice,
+            language,
+            instructions,
+            ref_audio=None,  # WebRTC UI doesn't support file upload
+            ref_audio_url=ref_audio_url,
+            ref_text=ref_text,
+            x_vector_only=x_vector_only,
+            stream=True,
+        )
+
+        try:
+            for samples in stream_pcm_chunks(api_base, payload):
+                yield (PCM_SAMPLE_RATE, samples)
+        except Exception:
+            logger.exception("TTS streaming error")
+
+    return tts_generate
+
+
+def build_interface(api_base: str):
+    """Build a custom Gradio Blocks UI matching gradio_demo.py layout."""
+    voices = fetch_voices(api_base)
+    handler = make_tts_generator(api_base)
+
+    css = """
+    #generate-btn button { width: 100%; }
+    .task-info { padding: 8px 12px; border-radius: 6px;
+                 background: #f0f4ff; margin-bottom: 8px; }
+    """
+
+    with gr.Blocks(css=css, title="Qwen3-TTS (FastRTC)") as demo:
+        gr.Markdown("# Qwen3-TTS Online Serving Demo (FastRTC)")
+        gr.Markdown(f"**Server:** `{api_base}` | **Transport:** WebRTC (gapless streaming)")
+
+        with gr.Row():
+            # Left column: inputs
+            with gr.Column(scale=3):
+                text_input = gr.Textbox(
+                    label="Text to Synthesize",
+                    placeholder="Enter text here, e.g., Hello, how are you?",
+                    lines=4,
+                )
+
+                with gr.Row():
+                    task_type = gr.Radio(
+                        choices=TASK_TYPES,
+                        value="CustomVoice",
+                        label="Task Type",
+                        scale=2,
+                    )
+                    language = gr.Dropdown(
+                        choices=SUPPORTED_LANGUAGES,
+                        value="Auto",
+                        label="Language",
+                        scale=1,
+                    )
+
+                voice = gr.Dropdown(
+                    choices=voices,
+                    value=voices[0] if voices else None,
+                    label="Speaker",
+                    visible=True,
+                )
+
+                instructions = gr.Textbox(
+                    label="Instructions",
+                    placeholder=("e.g., Speak with excitement / A warm, friendly female voice"),
+                    lines=2,
+                    visible=True,
+                    info="Optional style/emotion instructions",
+                )
+
+                # Base (voice clone) controls
+                # Note: file upload is not supported with WebRTC,
+                # use a URL instead.
+                ref_audio_url = gr.Textbox(
+                    label="Reference Audio URL",
+                    placeholder=("https://example.com/reference.wav (alternative to uploading)"),
+                    lines=1,
+                    visible=False,
+                )
+                ref_text = gr.Textbox(
+                    label="Reference Audio Transcript",
+                    placeholder=("Transcript of the reference audio (optional, improves quality)"),
+                    lines=2,
+                    visible=False,
+                )
+                x_vector_only = gr.Checkbox(
+                    label="Use x-vector only",
+                    value=False,
+                    visible=False,
+                    info=("Skip reference transcript, use speaker embedding only (lower quality)"),
+                )
+
+                generate_btn = gr.Button(
+                    "Generate Speech",
+                    variant="primary",
+                    size="lg",
+                    elem_id="generate-btn",
+                )
+
+            # Right column: output
+            with gr.Column(scale=2):
+                webrtc_output = WebRTC(
+                    label="Generated Audio (WebRTC)",
+                    mode="receive",
+                    modality="audio",
+                )
+                gr.Markdown(
+                    "### Task Types\n"
+                    "- **CustomVoice**: Use a predefined speaker "
+                    "(Vivian, Ryan, etc.) with optional style instructions\n"
+                    "- **VoiceDesign**: Describe the desired voice in natural "
+                    "language (instructions required)\n"
+                    "- **Base**: Clone a voice from reference audio "
+                    "(provide a URL)"
+                )
+
+        # Dynamic UI updates for task type
+        # FastRTC doesn't support file upload, so we only toggle
+        # 5 components (no ref_audio)
+        def on_task_type_change_fastrtc(task_type: str):
+            if task_type == "CustomVoice":
+                return (
+                    gr.update(visible=True),  # voice
+                    gr.update(visible=True, info="Optional style/emotion instructions"),
+                    gr.update(visible=False),  # ref_audio_url
+                    gr.update(visible=False),  # ref_text
+                    gr.update(visible=False),  # x_vector_only
+                )
+            elif task_type == "VoiceDesign":
+                return (
+                    gr.update(visible=False),
+                    gr.update(visible=True, info="Required: describe the voice style"),
+                    gr.update(visible=False),
+                    gr.update(visible=False),
+                    gr.update(visible=False),
+                )
+            elif task_type == "Base":
+                return (
+                    gr.update(visible=False),
+                    gr.update(visible=False),
+                    gr.update(visible=True),
+                    gr.update(visible=True),
+                    gr.update(visible=True),
+                )
+            return (
+                gr.update(visible=True),
+                gr.update(visible=True),
+                gr.update(visible=False),
+                gr.update(visible=False),
+                gr.update(visible=False),
+            )
+
+        task_type.change(
+            fn=on_task_type_change_fastrtc,
+            inputs=[task_type],
+            outputs=[
+                voice,
+                instructions,
+                ref_audio_url,
+                ref_text,
+                x_vector_only,
+            ],
+        )
+
+        # Wire up WebRTC streaming
+        all_inputs = [
+            text_input,
+            task_type,
+            voice,
+            language,
+            instructions,
+            ref_audio_url,
+            ref_text,
+            x_vector_only,
+        ]
+
+        webrtc_output.stream(
+            fn=handler,
+            inputs=all_inputs,
+            outputs=[webrtc_output],
+            trigger=generate_btn.click,
+        )
+
+        demo.queue()
+    return demo
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Gradio FastRTC demo for Qwen3-TTS (gapless streaming).")
+    add_common_args(parser)
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    logging.basicConfig(level=logging.INFO)
+    print(f"Connecting to vLLM server at: {args.api_base}")
+
+    demo = build_interface(args.api_base)
+    demo.launch(
+        server_name=args.host,
+        server_port=args.port,
+        share=args.share,
+    )
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/qwen3_tts/run_gradio_demo.sh

@@ -0,0 +1,203 @@
+#!/bin/bash
+# Launch both vLLM server and Gradio demo for Qwen3-TTS
+#
+# Usage:
+#   ./run_gradio_demo.sh                                    # Default: CustomVoice
+#   ./run_gradio_demo.sh --task-type VoiceDesign            # VoiceDesign model
+#   ./run_gradio_demo.sh --task-type Base --gradio-port 7861
+#
+# Options:
+#   --task-type TYPE        Task type: CustomVoice, VoiceDesign, Base (default: CustomVoice)
+#   --server-port PORT      Port for vLLM server (default: 8000)
+#   --gradio-port PORT      Port for Gradio demo (default: 7860)
+#   --server-host HOST      Host for vLLM server (default: 0.0.0.0)
+#   --gradio-ip IP          IP for Gradio demo (default: 127.0.0.1)
+#   --share                 Share Gradio demo publicly
+
+set -e
+
+# Default values
+TASK_TYPE="CustomVoice"
+SERVER_PORT=8000
+GRADIO_PORT=7860
+SERVER_HOST="0.0.0.0"
+GRADIO_IP="127.0.0.1"
+GRADIO_SHARE=false
+
+# Parse command line arguments
+while [[ $# -gt 0 ]]; do
+    case $1 in
+        --task-type)
+            TASK_TYPE="$2"
+            shift 2
+            ;;
+        --server-port)
+            SERVER_PORT="$2"
+            shift 2
+            ;;
+        --gradio-port)
+            GRADIO_PORT="$2"
+            shift 2
+            ;;
+        --server-host)
+            SERVER_HOST="$2"
+            shift 2
+            ;;
+        --gradio-ip)
+            GRADIO_IP="$2"
+            shift 2
+            ;;
+        --share)
+            GRADIO_SHARE=true
+            shift
+            ;;
+        --help)
+            echo "Usage: $0 [OPTIONS]"
+            echo ""
+            echo "Options:"
+            echo "  --task-type TYPE        Task type: CustomVoice, VoiceDesign, Base (default: CustomVoice)"
+            echo "  --server-port PORT      Port for vLLM server (default: 8000)"
+            echo "  --gradio-port PORT      Port for Gradio demo (default: 7860)"
+            echo "  --server-host HOST      Host for vLLM server (default: 0.0.0.0)"
+            echo "  --gradio-ip IP          IP for Gradio demo (default: 127.0.0.1)"
+            echo "  --share                 Share Gradio demo publicly"
+            echo ""
+            exit 0
+            ;;
+        *)
+            echo "Unknown option: $1"
+            echo "Use --help for usage information"
+            exit 1
+            ;;
+    esac
+done
+
+# Map task type to model
+case "$TASK_TYPE" in
+    CustomVoice)
+        MODEL="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
+        ;;
+    VoiceDesign)
+        MODEL="Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign"
+        ;;
+    Base)
+        MODEL="Qwen/Qwen3-TTS-12Hz-1.7B-Base"
+        ;;
+    *)
+        echo "Unknown task type: $TASK_TYPE"
+        echo "Supported: CustomVoice, VoiceDesign, Base"
+        exit 1
+        ;;
+esac
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+API_BASE="http://localhost:${SERVER_PORT}"
+
+echo "=========================================="
+echo "Qwen3-TTS Gradio Demo"
+echo "=========================================="
+echo "Task Type : $TASK_TYPE"
+echo "Model     : $MODEL"
+echo "Server    : http://${SERVER_HOST}:${SERVER_PORT}"
+echo "Gradio    : http://${GRADIO_IP}:${GRADIO_PORT}"
+echo "=========================================="
+
+# Cleanup on exit
+cleanup() {
+    echo ""
+    echo "Shutting down..."
+    if [ -n "$SERVER_PID" ]; then
+        echo "Stopping vLLM server (PID: $SERVER_PID)..."
+        kill "$SERVER_PID" 2>/dev/null || true
+        wait "$SERVER_PID" 2>/dev/null || true
+    fi
+    if [ -n "$GRADIO_PID" ]; then
+        echo "Stopping Gradio demo (PID: $GRADIO_PID)..."
+        kill "$GRADIO_PID" 2>/dev/null || true
+        wait "$GRADIO_PID" 2>/dev/null || true
+    fi
+    echo "Cleanup complete"
+    exit 0
+}
+trap cleanup SIGINT SIGTERM
+
+# Start vLLM server
+echo ""
+echo "Starting vLLM server..."
+LOG_FILE="/tmp/vllm_tts_server_${SERVER_PORT}.log"
+
+vllm-omni serve "$MODEL" \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
+    --host "$SERVER_HOST" \
+    --port "$SERVER_PORT" \
+    --gpu-memory-utilization 0.9 \
+    --trust-remote-code \
+    --enforce-eager \
+    --omni 2>&1 | tee "$LOG_FILE" &
+SERVER_PID=$!
+
+# Wait for server startup
+echo ""
+echo "Waiting for vLLM server to be ready..."
+STARTUP_FLAG="/tmp/vllm_tts_startup_flag_${SERVER_PORT}.tmp"
+rm -f "$STARTUP_FLAG"
+
+(
+    tail -f "$LOG_FILE" 2>/dev/null | grep -m 1 "Application startup complete" > /dev/null && touch "$STARTUP_FLAG"
+) &
+TAIL_PID=$!
+
+MAX_WAIT=300
+ELAPSED=0
+while [ $ELAPSED -lt $MAX_WAIT ]; do
+    if [ -f "$STARTUP_FLAG" ]; then
+        kill "$TAIL_PID" 2>/dev/null || true
+        wait "$TAIL_PID" 2>/dev/null || true
+        echo ""
+        echo "vLLM server is ready!"
+        break
+    fi
+    if ! kill -0 "$SERVER_PID" 2>/dev/null; then
+        kill "$TAIL_PID" 2>/dev/null || true
+        echo ""
+        echo "Error: vLLM server failed to start"
+        exit 1
+    fi
+    sleep 1
+    ELAPSED=$((ELAPSED + 1))
+done
+
+rm -f "$STARTUP_FLAG"
+
+if [ $ELAPSED -ge $MAX_WAIT ]; then
+    kill "$TAIL_PID" 2>/dev/null || true
+    echo "Error: Server startup timed out after ${MAX_WAIT}s"
+    kill "$SERVER_PID" 2>/dev/null || true
+    exit 1
+fi
+
+# Start Gradio demo
+echo ""
+echo "Starting Gradio demo..."
+cd "$SCRIPT_DIR"
+GRADIO_CMD=("python" "gradio_demo.py" "--api-base" "$API_BASE" "--ip" "$GRADIO_IP" "--port" "$GRADIO_PORT")
+if [ "$GRADIO_SHARE" = true ]; then
+    GRADIO_CMD+=("--share")
+fi
+
+"${GRADIO_CMD[@]}" &
+GRADIO_PID=$!
+
+echo ""
+echo "=========================================="
+echo "Both services are running!"
+echo "=========================================="
+echo "vLLM Server : http://${SERVER_HOST}:${SERVER_PORT}"
+echo "Gradio Demo : http://${GRADIO_IP}:${GRADIO_PORT}"
+echo ""
+echo "Press Ctrl+C to stop both services"
+echo "=========================================="
+echo ""
+
+wait $SERVER_PID $GRADIO_PID || true
+cleanup

examples/online_serving/qwen3_tts/tts_common.py

@@ -0,0 +1,218 @@
+"""Shared constants, helpers, and payload building for Qwen3-TTS Gradio demos."""
+
+import base64
+import io
+
+import gradio as gr
+import httpx
+import numpy as np
+import soundfile as sf
+
+SUPPORTED_LANGUAGES = [
+    "Auto",
+    "Chinese",
+    "English",
+    "Japanese",
+    "Korean",
+    "German",
+    "French",
+    "Russian",
+    "Portuguese",
+    "Spanish",
+    "Italian",
+]
+
+TASK_TYPES = ["CustomVoice", "VoiceDesign", "Base"]
+
+PCM_SAMPLE_RATE = 24000
+
+DEFAULT_API_BASE = "http://localhost:8000"
+
+
+def fetch_voices(api_base: str) -> list[str]:
+    """Fetch available voices from the server."""
+    try:
+        with httpx.Client(timeout=10.0) as client:
+            resp = client.get(
+                f"{api_base}/v1/audio/voices",
+                headers={"Authorization": "Bearer EMPTY"},
+            )
+        if resp.status_code == 200:
+            data = resp.json()
+            return data.get("voices", ["Vivian", "Ryan"])
+    except Exception:
+        pass
+    return ["Vivian", "Ryan"]
+
+
+def encode_audio_to_base64(audio_data: tuple) -> str:
+    """Encode Gradio audio input (sample_rate, numpy_array) to base64 data URL."""
+    sample_rate, audio_np = audio_data
+
+    if audio_np.dtype != np.int16:
+        if audio_np.dtype in (np.float32, np.float64):
+            audio_np = np.clip(audio_np, -1.0, 1.0)
+            audio_np = (audio_np * 32767).astype(np.int16)
+        else:
+            audio_np = audio_np.astype(np.int16)
+
+    buf = io.BytesIO()
+    sf.write(buf, audio_np, sample_rate, format="WAV")
+    wav_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
+    return f"data:audio/wav;base64,{wav_b64}"
+
+
+def build_payload(
+    text: str,
+    task_type: str,
+    voice: str,
+    language: str,
+    instructions: str,
+    ref_audio: tuple | None,
+    ref_audio_url: str,
+    ref_text: str,
+    x_vector_only: bool,
+    response_format: str = "pcm",
+    speed: float = 1.0,
+    stream: bool = True,
+) -> dict:
+    """Build the /v1/audio/speech request payload.
+
+    Raises gr.Error for invalid input so callers don't need to validate.
+    """
+    if not text or not text.strip():
+        raise gr.Error("Please enter text to synthesize.")
+
+    payload: dict = {
+        "input": text.strip(),
+        "response_format": "pcm" if stream else response_format,
+        "stream": stream,
+    }
+    if not stream:
+        payload["speed"] = speed
+
+    if task_type:
+        payload["task_type"] = task_type
+    if language:
+        payload["language"] = language
+
+    if task_type == "CustomVoice":
+        if voice:
+            payload["voice"] = voice
+        if instructions and instructions.strip():
+            payload["instructions"] = instructions.strip()
+
+    elif task_type == "VoiceDesign":
+        if not instructions or not instructions.strip():
+            raise gr.Error("VoiceDesign task requires voice style instructions.")
+        payload["instructions"] = instructions.strip()
+
+    elif task_type == "Base":
+        ref_audio_url_stripped = ref_audio_url.strip() if ref_audio_url else ""
+        if ref_audio_url_stripped:
+            payload["ref_audio"] = ref_audio_url_stripped
+        elif ref_audio is not None:
+            payload["ref_audio"] = encode_audio_to_base64(ref_audio)
+        else:
+            raise gr.Error("Base (voice clone) task requires reference audio. Upload a file or provide a URL.")
+        if ref_text and ref_text.strip():
+            payload["ref_text"] = ref_text.strip()
+        if x_vector_only:
+            payload["x_vector_only_mode"] = True
+
+    return payload
+
+
+def on_task_type_change(task_type: str):
+    """Update UI visibility based on selected task type."""
+    if task_type == "CustomVoice":
+        return (
+            gr.update(visible=True),  # voice dropdown
+            gr.update(visible=True, info="Optional style/emotion instructions"),
+            gr.update(visible=False),  # ref_audio
+            gr.update(visible=False),  # ref_audio_url
+            gr.update(visible=False),  # ref_text
+            gr.update(visible=False),  # x_vector_only
+        )
+    elif task_type == "VoiceDesign":
+        return (
+            gr.update(visible=False),
+            gr.update(visible=True, info="Required: describe the voice style"),
+            gr.update(visible=False),
+            gr.update(visible=False),
+            gr.update(visible=False),
+            gr.update(visible=False),
+        )
+    elif task_type == "Base":
+        return (
+            gr.update(visible=False),
+            gr.update(visible=False),
+            gr.update(visible=True),
+            gr.update(visible=True),
+            gr.update(visible=True),
+            gr.update(visible=True),
+        )
+    return (
+        gr.update(visible=True),
+        gr.update(visible=True),
+        gr.update(visible=False),
+        gr.update(visible=False),
+        gr.update(visible=False),
+        gr.update(visible=False),
+    )
+
+
+def stream_pcm_chunks(api_base: str, payload: dict):
+    """Stream raw PCM bytes from the server, yielding int16 numpy arrays.
+
+    Handles odd-byte boundaries between network chunks.
+    """
+    leftover = b""
+    with httpx.Client(timeout=300.0) as client:
+        with client.stream(
+            "POST",
+            f"{api_base}/v1/audio/speech",
+            json=payload,
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": "Bearer EMPTY",
+            },
+        ) as resp:
+            if resp.status_code != 200:
+                resp.read()
+                raise gr.Error(f"Server error ({resp.status_code}): {resp.text}")
+            for chunk in resp.iter_bytes():
+                if not chunk:
+                    continue
+                raw = leftover + chunk
+                usable = len(raw) - (len(raw) % 2)
+                leftover = raw[usable:]
+                if usable == 0:
+                    continue
+                yield np.frombuffer(raw[:usable], dtype=np.int16).copy()
+
+
+def add_common_args(parser):
+    """Add CLI arguments shared by both demos."""
+    parser.add_argument(
+        "--api-base",
+        default=DEFAULT_API_BASE,
+        help=f"Base URL for the vLLM API server (default: {DEFAULT_API_BASE}).",
+    )
+    parser.add_argument(
+        "--host",
+        default="0.0.0.0",
+        help="Host/IP for Gradio server (default: 0.0.0.0).",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=7860,
+        help="Port for Gradio server (default: 7860).",
+    )
+    parser.add_argument(
+        "--share",
+        action="store_true",
+        help="Share the Gradio demo publicly.",
+    )
+    return parser

pyproject.toml

@@ -51,6 +51,7 @@ dev = [
     "mooncake-transfer-engine==0.3.8.post1",
     "av", # for ComfyUI tests
     "openpyxl>=3.0.0", # for nightly CI
+    "pyttsx3>=2.99"
 ]
 
 docs = [

tests/benchmarks/test_serve_cli.py

@@ -3,13 +3,16 @@ from pathlib import Path
 
 import pytest
 
+from tests.conftest import OmniServerParams
 from tests.utils import hardware_test
 
 models = ["Qwen/Qwen2.5-Omni-7B"]
 stage_configs = [str(Path(__file__).parent.parent / "e2e" / "stage_configs" / "qwen2_5_omni_ci.yaml")]
 
 # Create parameter combinations for model and stage config
-test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
+test_params = [
+    OmniServerParams(model=model, stage_config_path=stage_config) for model in models for stage_config in stage_configs
+]
 
 
 @pytest.mark.core_model

tests/comfyui/conftest.py

@@ -63,6 +63,12 @@ def _setup_comfyui_test_environment():
     mock_comfy_api_input.AudioInput = AudioInput
     mock_comfy_api_input.VideoInput = VideoInput
     mock_comfy_api.input = mock_comfy_api_input
+    mock_comfy_api_latest = MagicMock()
+    mock_comfy_api_latest.Types.VideoComponents = MagicMock(side_effect=lambda **kwargs: kwargs)
+    mock_comfy_api_latest.InputImpl.VideoFromComponents = MagicMock(
+        side_effect=lambda _: VideoInput(b"mock_video_from_components")
+    )
+    mock_comfy_api.latest = mock_comfy_api_latest
 
     def mock_load(_: str | BinaryIO):
         """Mock nodes_audio.load that returns a waveform tensor (channels, samples) and sample rate."""
@@ -78,5 +84,6 @@ def _setup_comfyui_test_environment():
     # Install mock modules BEFORE importing any comfyui_vllm_omni code
     sys.modules["comfy_api"] = mock_comfy_api
     sys.modules["comfy_api.input"] = mock_comfy_api_input
+    sys.modules["comfy_api.latest"] = mock_comfy_api_latest
     sys.modules["comfy_extras"] = mock_comfy_extras
     sys.modules["comfy_extras.nodes_audio"] = mock_nodes_audio

tests/comfyui/test_comfyui_integration.py

@@ -20,11 +20,12 @@ import torch
 from comfy_api.input import AudioInput, VideoInput
 from comfyui_vllm_omni.nodes import (
     VLLMOmniGenerateImage,
+    VLLMOmniGenerateVideo,
     VLLMOmniTTS,
     VLLMOmniUnderstanding,
     VLLMOmniVoiceClone,
 )
-from comfyui_vllm_omni.utils.types import AutoregressionSamplingParams, DiffusionSamplingParams
+from comfyui_vllm_omni.utils.types import AutoregressionSamplingParams, DiffusionSamplingParams, WanModelSpecificParams
 from PIL import Image
 from vllm import SamplingParams
 from vllm.outputs import CompletionOutput, RequestOutput
@@ -34,6 +35,8 @@ from vllm_omni.entrypoints.cli.serve import OmniServeCommand
 from vllm_omni.inputs.data import OmniSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 
 class ServerCase(NamedTuple):
     """Parametrizing the model to serve."""
@@ -49,6 +52,7 @@ class SamplingCase(NamedTuple):
 
     kind: "SamplingKind"
     sampling_params: dict | list[dict] | None
+    lora: dict | None = None
 
 
 class SamplingKind(StrEnum):
@@ -58,11 +62,17 @@ class SamplingKind(StrEnum):
     UNDERSTANDING_AR_LIST = auto()
     TTS_NONE = auto()
     TTS_DIFFUSION_SINGLE = auto()
+    VIDEO_NONE = auto()
+    VIDEO_DIFFUSION_SINGLE = auto()
 
 
 # Pre-defined arguments to be used in function calls during the tests
 IMAGE_WIDTH = 64
 IMAGE_HEIGHT = 64
+VIDEO_WIDTH = 32
+VIDEO_HEIGHT = 32
+VIDEO_FPS = 8
+VIDEO_NUM_FRAMES = 5
 DIFFUSION_SINGLE_SAMPLING_PARAMS = DiffusionSamplingParams(
     {
         "n": 2,
@@ -102,6 +112,10 @@ AR_LIST_SAMPLING_PARAMS = [
     ),
 ]
 
+VIDEO_MODEL_PARAMS = WanModelSpecificParams({"guidance_scale_2": 5.0, "boundary_ratio": 0.98, "flow_shift": 12.0})
+
+LORA_PARAMS = {"local_path": "test_lora_path", "name": "test_name", "scale": 0.7, "int_id": 10}
+
 
 def _build_image_output(size: tuple[int, int] = (IMAGE_WIDTH, IMAGE_HEIGHT), color: str = "red") -> Image.Image:
     return Image.new("RGB", size, color=color)
@@ -181,6 +195,16 @@ def _build_diffusion_image_output_for_images_endpoint() -> OmniRequestOutput:
     )
 
 
+def _build_diffusion_video_output() -> OmniRequestOutput:
+    # Small video: VIDEO_NUM_FRAMES frames of (VIDEO_HEIGHT x VIDEO_WIDTH) RGB, shape (F, H, W, C)
+    video_frames = torch.zeros((VIDEO_NUM_FRAMES, VIDEO_HEIGHT, VIDEO_WIDTH, 3), dtype=torch.float32)
+    return OmniRequestOutput.from_diffusion(
+        request_id="test_req_video",
+        images=[video_frames],
+        final_output_type="video",
+    )
+
+
 def _build_diffusion_image_output_for_chat_endpoint() -> OmniRequestOutput:
     request_output = MagicMock()
     request_output.images = [_build_image_output(color="blue")]
@@ -193,12 +217,40 @@ def _build_diffusion_image_output_for_chat_endpoint() -> OmniRequestOutput:
     )
 
 
-def _assert_sampling_param_values(received: OmniSamplingParams, expected: dict[str, Any]):
+def _assert_sampling_param_values(
+    received: OmniSamplingParams, expected: dict[str, Any], expected_lora: dict | None = None
+):
     for key, expected_value in expected.items():
-        actual_value = getattr(received, key, None)
+        actual_value = getattr(received, key)
         assert actual_value == expected_value, (
             f"Expected sampling param '{key}'={expected_value}, got {actual_value}. The received sampling params: {received}"
         )
+    if expected_lora:
+        assert received.lora_request.lora_name == expected_lora["name"], (
+            f"Expected lora name={(expected_lora['name'])}, got {received.lora_request.lora_name}. The received sampling params: {received}"
+        )
+        assert received.lora_request.lora_int_id == expected_lora["int_id"], (
+            f"Expected lora int_id={expected_lora['int_id']}, got {received.lora_request.lora_int_id}. The received sampling params: {received}"
+        )
+        assert received.lora_request.lora_path == expected_lora["local_path"], (
+            f"Expected lora path={expected_lora['local_path']}, got {received.lora_request.lora_path}. The received sampling params: {received}"
+        )
+        assert received.lora_scale == expected_lora["scale"], (
+            f"Expected lora scale={expected_lora['scale']}, got {received.lora_scale}. The received sampling params: {received}"
+        )
+
+
+def _assert_model_param_values(received: OmniSamplingParams, expected: dict):
+    for key, expected_value in expected.items():
+        try:
+            actual_value = getattr(received, key)
+            expected_param_name = key
+        except AttributeError:
+            actual_value = received.extra_args.get(key, None)
+            expected_param_name = f'extra_args["{key}"]'
+        assert actual_value == expected_value, (
+            f"Expected model param '{expected_param_name}'={expected_value}, got {actual_value}. The received sampling params: {received}"
+        )
 
 
 def _build_mock_outputs(outputs: Iterable[OmniRequestOutput], sampling_case: SamplingCase, server_case: ServerCase):
@@ -235,6 +287,7 @@ def _build_mock_outputs(outputs: Iterable[OmniRequestOutput], sampling_case: Sam
                     "height": IMAGE_HEIGHT,
                     **expected,
                 },
+                LORA_PARAMS,
             )
         elif sampling_case.kind is SamplingKind.UNDERSTANDING_NONE:
             assert len(received_sampling_params_list) == 3
@@ -244,6 +297,33 @@ def _build_mock_outputs(outputs: Iterable[OmniRequestOutput], sampling_case: Sam
                 _assert_sampling_param_values(received_sampling_params_list[i], expected)
         elif sampling_case.kind in {SamplingKind.TTS_NONE, SamplingKind.TTS_DIFFUSION_SINGLE}:
             assert len(received_sampling_params_list) == 1
+        elif sampling_case.kind is SamplingKind.VIDEO_NONE:
+            assert len(received_sampling_params_list) == 1
+            _assert_sampling_param_values(
+                received_sampling_params_list[0],
+                {
+                    "width": VIDEO_WIDTH,
+                    "height": VIDEO_HEIGHT,
+                    "num_frames": VIDEO_NUM_FRAMES,
+                    "fps": VIDEO_FPS,
+                },
+            )
+        elif sampling_case.kind is SamplingKind.VIDEO_DIFFUSION_SINGLE:
+            assert len(received_sampling_params_list) == 1
+            expected = DIFFUSION_SINGLE_SAMPLING_PARAMS.copy()
+            expected["num_outputs_per_prompt"] = expected.pop("n")  # convert from n to num_outputs_per_prompt
+            _assert_sampling_param_values(
+                received_sampling_params_list[0],
+                {
+                    "width": VIDEO_WIDTH,
+                    "height": VIDEO_HEIGHT,
+                    "num_frames": VIDEO_NUM_FRAMES,
+                    "fps": VIDEO_FPS,
+                    **expected,
+                },
+                LORA_PARAMS,
+            )
+            _assert_model_param_values(received_sampling_params_list[0], VIDEO_MODEL_PARAMS)
         else:
             raise AssertionError(f"Unknown sampling case: {sampling_case.kind}")
 
@@ -401,9 +481,15 @@ def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni
 @pytest.mark.parametrize(
     "sampling_case",
     [
-        pytest.param(SamplingCase(kind=SamplingKind.IMAGE_NONE, sampling_params=None), id="no-sampling-params"),
         pytest.param(
-            SamplingCase(kind=SamplingKind.IMAGE_DIFFUSION_SINGLE, sampling_params=DIFFUSION_SINGLE_SAMPLING_PARAMS),
+            SamplingCase(kind=SamplingKind.IMAGE_NONE, sampling_params=None, lora=None), id="no-sampling-params"
+        ),
+        pytest.param(
+            SamplingCase(
+                kind=SamplingKind.IMAGE_DIFFUSION_SINGLE,
+                sampling_params=DIFFUSION_SINGLE_SAMPLING_PARAMS,
+                lora=LORA_PARAMS,
+            ),
             id="single-diffusion-sampling-params",
         ),
     ],
@@ -423,7 +509,8 @@ async def test_image_generation_node(api_server: str, model: str, image_input: b
         kwargs["image"] = torch.zeros((1, IMAGE_WIDTH, IMAGE_HEIGHT, 3), dtype=torch.float32)
     if sampling_case.sampling_params is not None:
         kwargs["sampling_params"] = sampling_case.sampling_params
-    print(f"!!!!!! Calling {model} node.generate with kwargs: {sampling_case.sampling_params}")
+    if sampling_case.lora:
+        kwargs["lora"] = sampling_case.lora
 
     result = await node.generate(**kwargs)
 
@@ -565,3 +652,75 @@ async def test_tts_nodes(api_server: str, node_cls, call_kwargs: dict, sampling_
     assert result[0]["sample_rate"] == 24000
     assert isinstance(result[0]["waveform"], torch.Tensor)
     assert result[0]["waveform"].shape == (1, 1, 24000)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "server_case,model,image_input",
+    [
+        pytest.param(
+            ServerCase(
+                served_model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+                stage_list=["diffusion"],
+                stage_configs=[{"stage_type": "diffusion"}],
+                outputs=[_build_diffusion_video_output()],
+            ),
+            "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+            False,
+            id="text-to-video",
+        ),
+        pytest.param(
+            ServerCase(
+                served_model="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
+                stage_list=["diffusion"],
+                stage_configs=[{"stage_type": "diffusion"}],
+                outputs=[_build_diffusion_video_output()],
+            ),
+            "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
+            True,
+            id="image-to-video",
+        ),
+    ],
+    indirect=["server_case"],
+)
+@pytest.mark.parametrize(
+    "sampling_case",
+    [
+        pytest.param(SamplingCase(kind=SamplingKind.VIDEO_NONE, sampling_params=None), id="no-sampling-params"),
+        pytest.param(
+            SamplingCase(
+                kind=SamplingKind.VIDEO_DIFFUSION_SINGLE,
+                sampling_params=DIFFUSION_SINGLE_SAMPLING_PARAMS,
+                lora=LORA_PARAMS,
+            ),
+            id="single-diffusion-sampling-params",
+        ),
+    ],
+    indirect=["sampling_case"],
+)
+async def test_video_generation_node(api_server: str, model: str, image_input: bool, sampling_case: SamplingCase):
+    node = VLLMOmniGenerateVideo()
+
+    kwargs = {
+        "url": api_server,
+        "model": model,
+        "prompt": "A beautiful sunset timelapse",
+        "negative_prompt": "",
+        "width": VIDEO_WIDTH,
+        "height": VIDEO_HEIGHT,
+        "fps": VIDEO_FPS,
+        "num_frames": VIDEO_NUM_FRAMES,
+        "model_params": VIDEO_MODEL_PARAMS,
+    }
+    if image_input:
+        kwargs["image"] = torch.zeros((1, VIDEO_HEIGHT, VIDEO_WIDTH, 3), dtype=torch.float32)
+    if sampling_case.sampling_params is not None:
+        kwargs["sampling_params"] = sampling_case.sampling_params
+    if sampling_case.lora:
+        kwargs["lora"] = sampling_case.lora
+
+    result = await node.generate(**kwargs)
+
+    assert isinstance(result, tuple)
+    assert len(result) == 1
+    assert isinstance(result[0], VideoInput)

tests/conftest.py

@@ -12,21 +12,27 @@ if "VLLM_TARGET_DEVICE" not in os.environ:
 
 import concurrent.futures
 import gc
+import multiprocessing
 import socket
 import subprocess
 import sys
 import threading
 import time
+from collections.abc import Generator
 from dataclasses import dataclass
+from io import BytesIO
 from pathlib import Path
-from typing import Any
+from typing import Any, NamedTuple
 
+import imageio.v3 as iio
 import numpy as np
 import psutil
 import pytest
+import soundfile as sf
 import torch
 import yaml
-from openai import OpenAI
+from openai import OpenAI, omit
+from PIL import Image
 from vllm import TextPrompt
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 from vllm.logger import init_logger
@@ -43,6 +49,65 @@ PromptImageInput = list[Any] | Any | None
 PromptVideoInput = list[Any] | Any | None
 
 
+class OmniServerParams(NamedTuple):
+    model: str
+    port: int | None = None
+    stage_config_path: str | None = None
+    server_args: list[str] | None = None
+
+
+def assert_image_valid(image: Path | Image.Image, *, width: int | None = None, height: int | None = None):
+    """Assert the file is a loadable image with optional exact dimensions."""
+    if isinstance(image, Path):
+        assert image.exists(), f"Image not found: {image}"
+        image = Image.open(image)
+        image.load()
+    assert image.width > 0 and image.height > 0
+    if width is not None:
+        assert image.width == width, f"Expected width={width}, got {image.width} in {image.name}"
+    if height is not None:
+        assert image.height == height, f"Expected height={height}, got {image.height} in {image.name}"
+    return image
+
+
+def assert_video_valid(frames: Path | np.ndarray, *, width: int, height: int, num_frames: int) -> None:
+    """Assert the MP4 has the expected resolution and exact frame count."""
+    if isinstance(frames, Path):
+        assert frames.exists(), f"Video not found: {frames}"
+        frames = iio.imread(str(frames), plugin="pyav", index=None)
+    assert frames.shape[0] == num_frames, f"Expected {num_frames} frames, got {frames.shape[0]}"
+    assert frames.shape[1] == height, f"Expected height={height}, got {frames.shape[1]}"
+    assert frames.shape[2] == width, f"Expected width={width}, got {frames.shape[2]}"
+
+
+def assert_audio_valid(path: Path, *, sample_rate: int, channels: int, duration_s: float) -> None:
+    """Assert the WAV has the expected sample rate, channel count, and duration."""
+
+    assert path.exists(), f"Audio not found: {path}"
+    info = sf.info(str(path))
+    assert info.samplerate == sample_rate, f"Expected sample_rate={sample_rate}, got {info.samplerate}"
+    assert info.channels == channels, f"Expected {channels} channel(s), got {info.channels}"
+    expected_frames = int(duration_s * sample_rate)
+    assert info.frames == expected_frames, (
+        f"Expected {expected_frames} frames ({duration_s}s @ {sample_rate} Hz), got {info.frames}"
+    )
+
+
+def decode_b64_image(b64: str):
+    img = Image.open(BytesIO(base64.b64decode(b64)))
+    img.load()
+    return img
+
+
+@pytest.fixture(scope="session")
+def model_prefix() -> str:
+    """Optional model-path prefix from MODEL_PREFIX env var.
+    Useful if models are downloaded to non-default local directories.
+    """
+    prefix = os.environ.get("MODEL_PREFIX", "")
+    return f"{prefix.rstrip('/')}/" if prefix else ""
+
+
 @pytest.fixture(autouse=True)
 def default_vllm_config():
     """Set a default VllmConfig for all tests.
@@ -217,102 +282,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
-    num_samples = int(sample_rate * duration)
+    def _pick_voice(engine: pyttsx3.Engine) -> str | None:
+        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
-    drop_density = 10  # Number of raindrops per second
-    drop_volume = 0.15  # Volume of individual raindrops
-    background_volume = 0.02  # Volume of background rain noise
+    engine = pyttsx3.init()
+    engine.setProperty("rate", 112)
+    engine.setProperty("volume", 1.0)
+    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
-    white_noise = np.random.randn(num_samples)
-    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
+    temp_wav = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
+    temp_wav.close()
 
-    # Add background noise to all channels
-    for ch in range(num_channels):
-        audio_data[:, ch] += bg_noise
+    try:
+        engine.save_to_file(phrase_text, temp_wav.name)
+        engine.runAndWait()
+        engine.stop()
 
-    # Total number of raindrops = density × duration × channels for stereo effect
-    total_drops = int(drop_density * duration * num_channels)
+        ready = False
+        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):
-        # Random timing for raindrop
-        drop_time = random.uniform(0, duration)
+        if not ready:
+            raise RuntimeError("pyttsx3 did not produce a WAV file in time.")
 
-        # Random duration of raindrop sound (0.01-0.05 seconds)
-        drop_duration = random.uniform(0.01, 0.05)
+        tts_audio, tts_sr = sf.read(temp_wav.name, dtype="float32", always_2d=True)
+    finally:
+        if os.path.exists(temp_wav.name):
+            os.unlink(temp_wav.name)
 
-        # Random frequency gives variation in raindrop pitch
-        drop_freq = random.uniform(500, 5000)  # Hz
+    if len(tts_audio) == 0:
+        raise RuntimeError("pyttsx3 produced an empty WAV file.")
 
-        # Random channel selection for stereo positioning
-        channel = random.randint(0, num_channels - 1)
+    tts_audio = _resample_audio(tts_audio, tts_sr, sample_rate)
+    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
-        start_sample = int(drop_time * sample_rate)
-        drop_samples = int(drop_duration * sample_rate)
-        end_sample = min(start_sample + drop_samples, num_samples)
+    lead_silence = min(int(sample_rate * 0.02), num_samples // 8)
+    pause_samples = int(sample_rate * 0.18)
+    start = lead_silence
+    phrase_len = tts_audio.shape[0]
 
-        if start_sample < end_sample:
-            # Generate the raindrop sound
-            num_drop_samples = end_sample - start_sample
-            t = np.arange(num_drop_samples) / sample_rate
+    while start < num_samples:
+        take = min(phrase_len, num_samples - start)
+        audio_data[start : start + take] = tts_audio[:take]
+        start += phrase_len + pause_samples
 
-            # Basic sine wave for raindrop sound
-            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))
+    max_amp = float(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()
-    result = {
-        "np_array": audio_array,
+    audio_bytes: bytes | None = None
+    output_path: str | None = None
+    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 +773,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")
-    text = model.transcribe(
-        output_path,
-        temperature=0.0,
-        word_timestamps=True,
-        condition_on_previous_text=False,
-    )["text"]
-    del model
-    if torch.cuda.is_available():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = whisper.load_model("small", device=device)
+    try:
+        text = model.transcribe(
+            output_path,
+            temperature=0.0,
+            word_timestamps=True,
+            condition_on_previous_text=False,
+        )["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 text:
-        return text
-    else:
-        return ""
+        if torch is not None and torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    return text or ""
+
+
+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):
@@ -951,15 +1121,16 @@ class OmniServer:
         max_wait = 1200  # 20 minutes
         start_time = time.time()
         while time.time() - start_time < max_wait:
-            try:
-                with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
-                    sock.settimeout(1)
-                    result = sock.connect_ex((self.host, self.port))
-                    if result == 0:
-                        print(f"Server ready on {self.host}:{self.port}")
-                        return
-            except Exception:
-                pass
+            # Check for process status
+            ret = self.proc.poll()
+            if ret is not None:
+                raise RuntimeError(f"Server processes exited with code {ret} before becoming ready.")
+            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+                sock.settimeout(1)
+                result = sock.connect_ex((self.host, self.port))
+                if result == 0:
+                    print(f"Server ready on {self.host}:{self.port}")
+                    return
             time.sleep(2)
 
         raise RuntimeError(f"Server failed to start within {max_wait} seconds")
@@ -1044,7 +1215,9 @@ def pytest_addoption(parser):
 
 
 @pytest.fixture(scope="session")
-def run_level(request):
+def run_level(request) -> str:
+    """A command-line argument that specifies the level of tests to run in this session.
+    See https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/CI_5levels/"""
     return request.config.getoption("--run-level")
 
 
@@ -1052,15 +1225,17 @@ _omni_server_lock = threading.Lock()
 
 
 @pytest.fixture(scope="module")
-def omni_server(request, run_level):
+def omni_server(request: pytest.FixtureRequest, run_level: str, model_prefix: str) -> Generator[OmniServer, Any, None]:
     """Start vLLM-Omni server as a subprocess with actual model weights.
     Uses session scope so the server starts only once for the entire test session.
     Multi-stage initialization can take 10-20+ minutes.
     """
     with _omni_server_lock:
-        *port, model, stage_config_path = request.param
-        port = port[0] if port else None
-        if run_level == "advanced_model":
+        params: OmniServerParams = request.param
+        model = model_prefix + params.model
+        port = params.port
+        stage_config_path = params.stage_config_path
+        if run_level == "advanced_model" and stage_config_path is not None:
             stage_config_path = modify_stage_config(
                 stage_config_path,
                 deletes={
@@ -1072,7 +1247,11 @@ def omni_server(request, run_level):
                 },
             )
 
-        server_args = ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"]
+        server_args = params.server_args or []
+        server_args = ["--stage-init-timeout", "120", *server_args]
+        if stage_config_path is not None:
+            server_args += ["--stage-configs-path", stage_config_path]
+
         with OmniServer(model, server_args, port=port) if port else OmniServer(model, server_args) as server:
             print("OmniServer started successfully")
             yield server
@@ -1092,6 +1271,17 @@ class OmniResponse:
     error_message: str | None = None
 
 
+@dataclass
+class DiffusionResponse:
+    text_content: str | None = None
+    images: list[Image.Image] | None = None
+    audios: list[Any] | None = None
+    videos: list[Any] | None = None
+    e2e_latency: float | None = None
+    success: bool = False
+    error_message: str | None = None
+
+
 def assert_omni_response(response: OmniResponse, request_config: dict[str, Any], run_level):
     """
     Validate response results.
@@ -1105,7 +1295,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"])
 
@@ -1140,6 +1330,51 @@ def assert_omni_response(response: OmniResponse, request_config: dict[str, Any],
             print(f"similarity is: {response.similarity}")
 
 
+def assert_diffusion_response(response: DiffusionResponse, request_config: dict[str, Any], run_level: str = None):
+    """
+    Validate diffusion response results.
+
+    Args:
+        response: DiffusionResponse object. Any not-None content will be validated based on the request_config.
+        request_config: Request configuration dictionary containing parameters like model, messages, extra_body.
+            When validating a certain modality, the corresponding params in request_config['extra_body'] must present.
+            It will be used to check against the multimedia file in the response.
+        run_level: Test run level (e.g., "core_model", "advanced_model")
+
+    Raises:
+        AssertionError: When the response does not meet validation criteria
+        KeyError: When the request_config does not contain necessary parameters for validation
+    """
+    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}")
+
+    extra_body = request_config.get("extra_body", {})
+
+    num_outputs_per_prompt = extra_body.get("num_outputs_per_prompt", 1)
+
+    if response.images is not None:
+        assert len(response.images) > 0, "No images in response"
+        assert len(response.images) == num_outputs_per_prompt, (
+            f"Expected {num_outputs_per_prompt} images, got {len(response.images)}"
+        )
+        if run_level == "advanced_model":
+            expected_width = extra_body["width"]  # intend to raise KeyError
+            expected_height = extra_body["height"]  # intend to raise KeyError
+            for img in response.images:
+                assert_image_valid(img, width=expected_width, height=expected_height)
+    if response.videos is not None:
+        raise NotImplementedError(
+            "Video validation is not implemented yet"
+        )  # consider using assert_video_valid defined above
+    if response.audios is not None:
+        raise NotImplementedError(
+            "Audio validation is not implemented yet"
+        )  # consider using assert_audio_valid defined above
+
+
 class OpenAIClientHandler:
     """
     OpenAI client handler class, encapsulating both streaming and non-streaming response processing logic.
@@ -1162,7 +1397,7 @@ class OpenAIClientHandler:
         self.client = OpenAI(base_url=f"http://{host}:{port}/v1", api_key=api_key)
         self.run_level = run_level
 
-    def _process_stream_response(self, chat_completion) -> OmniResponse:
+    def _process_stream_omni_response(self, chat_completion) -> OmniResponse:
         """
         Process streaming responses.
 
@@ -1223,7 +1458,7 @@ class OpenAIClientHandler:
 
         return result
 
-    def _process_non_stream_response(self, chat_completion) -> OmniResponse:
+    def _process_non_stream_omni_response(self, chat_completion) -> OmniResponse:
         """
         Process non-streaming responses.
 
@@ -1277,7 +1512,46 @@ class OpenAIClientHandler:
 
         return result
 
-    def send_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]:
+    def _process_diffusion_response(self, chat_completion) -> DiffusionResponse:
+        """
+        Process diffusion responses (image generation/editing).
+
+        Args:
+            chat_completion: OpenAI response object
+
+        Returns:
+            DiffusionResponse: Processed response object
+        """
+        result = DiffusionResponse()
+        start_time = time.perf_counter()
+
+        try:
+            images = []
+            # [TODO] reading video and audio output from API response for later validation
+
+            for choice in chat_completion.choices:
+                if hasattr(choice.message, "content") and choice.message.content is not None:
+                    content = choice.message.content
+                    if isinstance(content, list):
+                        for item in content:
+                            if hasattr(item, "image_url") and item.image_url is not None:
+                                image_url = item.image_url.url
+                                if image_url.startswith("data:image"):
+                                    b64_data = image_url.split(",", 1)[1]
+                                    img = decode_b64_image(b64_data)
+                                    images.append(img)
+
+            result.e2e_latency = time.perf_counter() - start_time
+            result.images = images if images else None
+            result.success = True
+
+        except Exception as e:
+            result.error_message = f"Diffusion response processing error: {str(e)}"
+            print(f"Error: {result.error_message}")
+
+        return result
+
+    def send_omni_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]:
         """
         Send OpenAI requests.
 
@@ -1303,9 +1577,9 @@ class OpenAIClientHandler:
             )
 
             if stream:
-                response = self._process_stream_response(chat_completion)
+                response = self._process_stream_omni_response(chat_completion)
             else:
-                response = self._process_non_stream_response(chat_completion)
+                response = self._process_non_stream_omni_response(chat_completion)
 
             assert_omni_response(response, request_config, run_level=self.run_level)
             responses.append(response)
@@ -1331,25 +1605,82 @@ class OpenAIClientHandler:
                     chat_completion = future.result()
 
                     if stream:
-                        response = self._process_stream_response(chat_completion)
+                        response = self._process_stream_omni_response(chat_completion)
                     else:
-                        response = self._process_non_stream_response(chat_completion)
+                        response = self._process_non_stream_omni_response(chat_completion)
 
                     assert_omni_response(response, request_config, run_level=self.run_level)
                     responses.append(response)
 
         return responses
 
+    def send_diffusion_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]:
+        """
+        Send OpenAI requests for diffusion models.
+
+        Args:
+            request_config: Request configuration dictionary containing parameters like model, messages
+            request_num: Number of requests to send concurrently, defaults to 1 (single request)
+        Returns:
+            List[OmniResponse]: List of response objects
+        """
+        responses = []
+        stream = request_config.get("stream", False)
+        modalities = request_config.get("modalities", omit)  # Most diffusion models don't require modalities param
+        extra_body = request_config.get("extra_body", None)
+
+        if stream:
+            raise NotImplementedError("Streaming is not currently implemented for diffusion model e2e test")
+
+        if request_num == 1:
+            # Send single request
+            chat_completion = self.client.chat.completions.create(
+                model=request_config.get("model"),
+                messages=request_config.get("messages"),
+                extra_body=extra_body,
+                modalities=modalities,
+            )
+
+            response = self._process_diffusion_response(chat_completion)
+            assert_diffusion_response(response, request_config, run_level=self.run_level)
+            responses.append(response)
+
+        else:
+            # Send concurrent requests
+            with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor:
+                futures = []
+
+                # Submit all request tasks
+                for _ in range(request_num):
+                    future = executor.submit(
+                        self.client.chat.completions.create,
+                        model=request_config.get("model"),
+                        messages=request_config.get("messages"),
+                        modalities=modalities,
+                        extra_body=extra_body,
+                    )
+                    futures.append(future)
+
+                # Process completed tasks
+                for future in concurrent.futures.as_completed(futures):
+                    chat_completion = future.result()
+                    response = self._process_diffusion_response(chat_completion)
+                    assert_diffusion_response(response, request_config, run_level=self.run_level)
+                    responses.append(response)
+
+        return responses
+
 
 @pytest.fixture
-def openai_client(omni_server, run_level):
-    """Create OpenAIClientHandler fixture"""
+def openai_client(omni_server: OmniServer, run_level: str):
+    """Create OpenAIClientHandler fixture to facilitate communication with OmniServer
+    with encapsulated request sending, concurrent requests, response handling, and validation."""
     return OpenAIClientHandler(host=omni_server.host, port=omni_server.port, api_key="EMPTY", run_level=run_level)
 
 
 class OmniRunner:
     """
-    Test runner for Omni models.
+    Offline test runner for Omni models.
     """
 
     def __init__(
@@ -1620,9 +1951,10 @@ class OmniRunner:
 
 
 @pytest.fixture(scope="module")
-def omni_runner(request):
+def omni_runner(request, model_prefix):
     with _omni_server_lock:
         model, stage_config_path = request.param
+        model = model_prefix + model
         with OmniRunner(model, seed=42, stage_configs_path=stage_config_path, stage_init_timeout=300) as runner:
             print("OmniRunner started successfully")
             yield runner

tests/diffusion/attention/test_attention_sp.py

@@ -1,5 +1,29 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Tests for Ulysses + Ring sequence-parallel attention correctness.
+
+What is tested
+--------------
+* ``test_sequence_parallel`` verifies that the ``Attention`` layer produces
+  numerically equivalent results whether the sequence is processed on a single
+  rank (baseline, SP=1) or sharded across multiple ranks via Ulysses/Ring SP.
+  The test spawns two separate multi-process runs with ``torch.multiprocessing.spawn``:
+  1. Baseline   – world_size=1, ulysses_degree=1, ring_degree=1.
+  2. SP run     – world_size=ulysses_degree*ring_degree, each rank holds a
+                  contiguous slice of the full sequence.
+  After both runs, rank-0 output tensors are compared element-wise with a
+  tolerance appropriate for the dtype (bfloat16).
+
+SP-plan hooks are NOT applied in this test
+------------------------------------------
+``ForwardContext.sp_plan_hooks_applied`` remains ``False``.  As a result,
+``ForwardContext.sp_active`` falls back to the "no hooks" branch: SP is
+considered active whenever ``parallel_config.sequence_parallel_size > 1``.
+This makes the test self-contained and suitable for standalone CI runs that
+do not exercise the full model-registry pipeline.
+"""
+
 import os
 import pickle
 import tempfile
@@ -27,6 +51,11 @@ def update_environment_variables(envs_dict: dict[str, str]):
         os.environ[k] = v
 
 
+def seed_everything(seed: int):
+    torch.manual_seed(seed)
+    current_omni_platform.manual_seed(seed)
+
+
 class TestAttentionModel(torch.nn.Module):
     """Test model using Attention layer."""
 
@@ -138,11 +167,10 @@ class TestMultiLayerAttentionModel(torch.nn.Module):
 @pytest.mark.parametrize("num_heads", [8])
 @pytest.mark.parametrize("head_size", [8])
 @pytest.mark.parametrize("causal", [False])
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])  # [torch.float16, torch.bfloat16]
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
 @pytest.mark.parametrize("use_sync", [False])
 @pytest.mark.parametrize("dynamic", [False])
 @pytest.mark.parametrize("use_compile", [False])
-@pytest.mark.parametrize("attn_backend", ["sdpa", "flash_attn"])
 def test_sequence_parallel(
     ulysses_degree: int,
     ring_degree: int,
@@ -156,7 +184,6 @@ def test_sequence_parallel(
     seq_len: int,
     num_heads: int,
     head_size: int,
-    attn_backend: str,
 ):
     """Test Ulysses attention by comparing with and without SP enabled."""
     sequence_parallel_size = ulysses_degree * ring_degree
@@ -200,7 +227,6 @@ def test_sequence_parallel(
                 model_state_file,
                 input_data_file,
                 True,  # is_baseline
-                attn_backend,
             ),
             nprocs=1,
         )
@@ -228,7 +254,6 @@ def test_sequence_parallel(
                 model_state_file,
                 input_data_file,
                 False,  # is_baseline
-                attn_backend,
             ),
             nprocs=sequence_parallel_size,
         )
@@ -324,12 +349,11 @@ def ulysses_attention_on_test_model(
     model_state_file: str,
     input_data_file: str,
     is_baseline: bool,
-    attn_backend: str,
 ):
     """Run Ulysses attention test on a test model and save results for comparison."""
     # Use fixed seed for reproducibility across baseline and SP runs
     RANDOM_SEED = 42
-    current_omni_platform.seed_everything(RANDOM_SEED)
+    seed_everything(RANDOM_SEED)
 
     mode_str = "Baseline (no SP)" if is_baseline else f"SP (ulysses={ulysses_degree}, ring={ring_degree})"
     print(f"\n[{mode_str}] Rank {local_rank}/{world_size} - Random seed set to {RANDOM_SEED}")
@@ -366,7 +390,6 @@ def ulysses_attention_on_test_model(
         model="test_model",
         dtype=dtype,
         parallel_config=parallel_config,
-        attention_backend=attn_backend,  # Set the attention backend here
     )
 
     # Initialize model parallel
@@ -411,11 +434,6 @@ def ulysses_attention_on_test_model(
             with open(model_state_file, "wb") as f:
                 pickle.dump(model_state, f)
 
-            # Generate and save full input data with fixed seed
-            # Reinitialize RNG to ensure reproducibility
-            torch.manual_seed(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_hsdp.py

@@ -11,7 +11,7 @@ import torch.nn as nn
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig
 
-pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.cpu]
+pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.cpu, pytest.mark.core_model]
 
 
 class TestHSDPInferenceConfig:

tests/diffusion/distributed/test_parallel_state_sp_groups.py

@@ -11,10 +11,7 @@ import torch
 
 from vllm_omni.diffusion.distributed.parallel_state import RankGenerator, set_seq_parallel_pg
 
-pytestmark = [
-    pytest.mark.diffusion,
-    pytest.mark.parallel,
-]
+pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.core_model, pytest.mark.cpu]
 
 
 def _fake_new_group_factory(created_groups: list[SimpleNamespace]):

tests/diffusion/distributed/test_sp_plan_hooks.py

@@ -47,6 +47,8 @@ requires_distributed = pytest.mark.skipif(
 pytestmark = [
     pytest.mark.diffusion,
     pytest.mark.parallel,
+    pytest.mark.core_model,
+    pytest.mark.cpu,
 ]
 
 # =============================================================================

tests/diffusion/distributed/test_vae_patch_parallel.py

@@ -8,6 +8,8 @@ import torch
 
 from vllm_omni.diffusion.distributed import vae_patch_parallel as vae_patch_parallel
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 
 class _DummyConfig:
     def __init__(self, **attrs):

tests/diffusion/models/flux2/test_flux2_transformer_tp.py

@@ -0,0 +1,244 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.flux2.flux2_transformer import (
+    Flux2PosEmbed,
+    Flux2Transformer2DModel,
+)
+
+
+# Initialize TP group before tests
+@pytest.fixture(scope="function", autouse=True)
+def setup_tp_group():
+    """Set up TP group for each test function"""
+    with patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=2):
+        with patch("vllm.distributed.parallel_state.get_tp_group") as mock_get_tp_group:
+            mock_tp_group = MagicMock()
+            mock_tp_group.world_size = 2
+            mock_get_tp_group.return_value = mock_tp_group
+            yield
+
+
+class TestFlux2TransformerWeightLoading:
+    """Test Flux2Transformer weight loading functionality"""
+
+    def test_weight_loading_tp2(self, setup_tp_group):
+        """Verify weights load correctly with TP=2"""
+        # Prepare test data
+        model = Flux2Transformer2DModel(
+            num_layers=2,
+            num_single_layers=1,
+            num_attention_heads=48,
+            attention_head_dim=128,
+            joint_attention_dim=15360,
+        )
+
+        # Mock TP=2 weight loading
+        mock_weights = []
+
+        # 1. Test regular weight loading
+        mock_weights.append(("x_embedder.weight", torch.randn(6144, 128)))
+        mock_weights.append(("context_embedder.weight", torch.randn(6144, 15360)))
+        mock_weights.append(("proj_out.weight", torch.randn(128, 6144)))
+
+        # 2. Test stacked_params_mapping weight loading
+        # Full weights - load_weights handles sharding internally
+        to_qkv_weight = torch.randn(18432, 6144)  # Full weights
+        mock_weights.append(("transformer_blocks.0.attn.to_qkv.weight", to_qkv_weight))
+
+        # Add_kv_proj weights
+        add_kv_proj_weight = torch.randn(18432, 6144)  # Full weights
+        mock_weights.append(("transformer_blocks.0.attn.add_kv_proj.weight", add_kv_proj_weight))
+
+        # 3. Test single block weight loading
+        single_block_qkv_mlp_weight = torch.randn(18432 + 18432 * 3, 6144)  # Full weights
+        mock_weights.append(("single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight", single_block_qkv_mlp_weight))
+
+        # Execute weight loading
+        loaded_params = model.load_weights(mock_weights)
+
+        # Verify
+        assert len(loaded_params) > 0, "Parameters should be loaded"
+
+        # Verify stacked_params_mapping is correctly set
+        assert model.stacked_params_mapping is not None
+        # Should have 6 mappings: 3 for to_qkv + 3 for add_kv_proj
+        assert len(model.stacked_params_mapping) == 6, "Should have 6 mappings"
+
+        # Verify weights are correctly loaded to corresponding layers
+        assert hasattr(model.transformer_blocks[0].attn.to_qkv, "weight")
+        # With TP=2, weight dimension on each GPU should be 18432/2 = 9216
+        assert model.transformer_blocks[0].attn.to_qkv.weight.shape == (9216, 6144), (
+            f"With TP=2, to_qkv weight dimension should be (9216, 6144), got {model.transformer_blocks[0].attn.to_qkv.weight.shape}"
+        )
+
+
+class TestFlux2RopePositionEmbedding:
+    """Test Flux2 RoPE position embedding functionality"""
+
+    def test_rope_position_embedding(self):
+        """Verify RoPE produces correct embeddings for 4D coordinates"""
+        # Prepare test data - use model default configuration
+        # axes_dims_rope default is (32, 32, 32, 32)
+        # get_1d_rotary_pos_embed outputs half the dimension for real/imag parts
+        # So actual output dimension should be (16+16+16+16) = 64
+        axes_dims = (32, 32, 32, 32)  # Use model default
+        rope_theta = 2000  # Model default is 2000
+        pos_embed = Flux2PosEmbed(theta=rope_theta, axes_dim=axes_dims)
+
+        # Create test IDs
+        seq_len = 10
+        ids = torch.randint(0, 100, (seq_len, 4))  # [S, 4]
+
+        # Forward pass
+        freqs_cos, freqs_sin = pos_embed(ids)
+
+        # Verify output shape - based on model config, expected dimension is 64
+        # Each axes_dim=32 outputs 16-dim real part, sum of 4 dimensions = 64
+        expected_dim = sum(axes_dims) // 2  # 128/2 = 64
+        assert freqs_cos.shape == (seq_len, expected_dim), (
+            f"Expected shape {(seq_len, expected_dim)}, got {freqs_cos.shape}"
+        )
+        assert freqs_sin.shape == (seq_len, expected_dim), (
+            f"Expected shape {(seq_len, expected_dim)}, got {freqs_sin.shape}"
+        )
+
+        # Verify output type - NPU may use float32, other devices use float64
+        assert freqs_cos.dtype in [torch.float32, torch.float64], "Should be float type"
+        assert freqs_sin.dtype in [torch.float32, torch.float64], "Should be float type"
+
+        # Verify value range
+        assert torch.all(freqs_cos >= -1) and torch.all(freqs_cos <= 1), "cos values should be in [-1, 1]"
+        assert torch.all(freqs_sin >= -1) and torch.all(freqs_sin <= 1), "sin values should be in [-1, 1]"
+
+        # Verify trigonometric relationship: cos^2 + sin^2 ≈ 1
+        cos_sq_sin_sq = freqs_cos**2 + freqs_sin**2
+        assert torch.allclose(cos_sq_sin_sq, torch.ones_like(cos_sq_sin_sq), atol=1e-6), "cos^2 + sin^2 should ≈ 1"
+
+        # Verify different positions produce different embeddings
+        ids_diff = torch.randint(0, 100, (seq_len, 4))
+        freqs_cos_diff, freqs_sin_diff = pos_embed(ids_diff)
+        assert not torch.allclose(freqs_cos, freqs_cos_diff), "Different positions should produce different embeddings"
+
+        # Verify same positions produce same embeddings
+        ids_same = ids.clone()
+        freqs_cos_same, freqs_sin_same = pos_embed(ids_same)
+        assert torch.allclose(freqs_cos, freqs_cos_same), "Same positions should produce same embeddings"
+        assert torch.allclose(freqs_sin, freqs_sin_same), "Same positions should produce same embeddings"
+
+
+class TestFlux2PackedModuleMapping:
+    """Test Flux2 packed module mapping functionality"""
+
+    def test_packed_module_mapping(self, setup_tp_group):
+        """Verify to_qkv packing matches HF checkpoint"""
+        model = Flux2Transformer2DModel(
+            num_layers=1,
+            num_single_layers=0,
+            num_attention_heads=48,
+            attention_head_dim=128,
+            joint_attention_dim=15360,
+        )
+
+        # Verify stacked_params_mapping is correctly initialized
+        model.load_weights([])  # Trigger stacked_params_mapping initialization
+        assert model.stacked_params_mapping is not None
+
+        # Verify mapping configuration
+        expected_mappings = [
+            (".to_qkv.", ".to_q.", "q"),
+            (".to_qkv.", ".to_k.", "k"),
+            (".to_qkv.", ".to_v.", "v"),
+            (".add_kv_proj", ".add_q_proj", "q"),
+            (".add_kv_proj", ".add_k_proj", "k"),
+            (".add_kv_proj", ".add_v_proj", "v"),
+        ]
+        assert model.stacked_params_mapping == expected_mappings
+
+        # Create mock HF checkpoint weights
+        hf_weights = []
+
+        # Mock HF separated Q/K/V weights
+        attn_block = model.transformer_blocks[0].attn
+        head_dim = 128
+        num_heads = 48
+        hidden_size = 6144
+        # Full weight dimension
+        full_shard_size = num_heads * head_dim  # 6144
+
+        # Q projection weights (full weights)
+        q_weight = torch.randn(full_shard_size, hidden_size)
+        hf_weights.append(("transformer_blocks.0.attn.to_q.weight", q_weight))
+
+        # K projection weights (full weights)
+        k_weight = torch.randn(full_shard_size, hidden_size)
+        hf_weights.append(("transformer_blocks.0.attn.to_k.weight", k_weight))
+
+        # V projection weights (full weights)
+        v_weight = torch.randn(full_shard_size, hidden_size)
+        hf_weights.append(("transformer_blocks.0.attn.to_v.weight", v_weight))
+
+        # Mock HF separated add_q/k/v projection weights
+        add_q_weight = torch.randn(full_shard_size, hidden_size)
+        hf_weights.append(("transformer_blocks.0.attn.add_q_proj.weight", add_q_weight))
+
+        add_k_weight = torch.randn(full_shard_size, hidden_size)
+        hf_weights.append(("transformer_blocks.0.attn.add_k_proj.weight", add_k_weight))
+
+        add_v_weight = torch.randn(full_shard_size, hidden_size)
+        hf_weights.append(("transformer_blocks.0.attn.add_v_proj.weight", add_v_weight))
+
+        # Execute weight loading
+        loaded_params = model.load_weights(hf_weights)
+
+        # Verify weights are loaded
+        assert len(loaded_params) > 0
+
+        # Verify final QKV weights are correctly combined (considering TP sharding)
+        # With TP=2, dimension on each GPU should be half of full dimension
+        expected_qkv_shape = (full_shard_size * 3 // 2, hidden_size)  # 18432/2 = 9216
+        assert attn_block.to_qkv.weight.shape == expected_qkv_shape, (
+            f"to_qkv weight dimension should be {expected_qkv_shape}, got {attn_block.to_qkv.weight.shape}"
+        )
+
+        expected_add_kv_shape = (full_shard_size * 3 // 2, hidden_size)
+        assert attn_block.add_kv_proj.weight.shape == expected_add_kv_shape, (
+            f"add_kv_proj weight dimension should be {expected_add_kv_shape}, got {attn_block.add_kv_proj.weight.shape}"
+        )
+
+    def test_packed_mapping_edge_cases(self, setup_tp_group):
+        """Test edge cases for packed mapping"""
+        model = Flux2Transformer2DModel(
+            num_layers=1,
+            num_single_layers=1,
+            num_attention_heads=48,
+            attention_head_dim=128,
+            joint_attention_dim=15360,
+        )
+        model.load_weights([])
+
+        # Test invalid weight names
+        invalid_weights = [("invalid.weight", torch.randn(10, 10))]
+        loaded_params = model.load_weights(invalid_weights)
+        assert len(loaded_params) == 0, "Should not load invalid weights"
+
+        # Test to_out weight renaming
+        to_out_weight = torch.randn(6144, 6144)
+        weights = [("transformer_blocks.0.attn.to_out.0.weight", to_out_weight)]
+        loaded_params = model.load_weights(weights)
+
+        # Check if to_out related weights are loaded
+        to_out_loaded = any("to_out" in p for p in loaded_params)
+        assert to_out_loaded, "to_out weights should be correctly renamed and loaded"
+
+        # Test partial weight loading
+        partial_weights = [
+            ("x_embedder.weight", torch.randn(6144, 128)),
+            ("transformer_blocks.0.attn.to_q.weight", torch.randn(6144, 6144)),  # Full weights
+        ]
+        loaded_params = model.load_weights(partial_weights)
+        assert len(loaded_params) == 2, "Should load two weights"
+        assert "x_embedder.weight" in loaded_params
+        assert any("to_qkv" in p for p in loaded_params), "to_q should be mapped to to_qkv"

tests/diffusion/models/nextstep_1_1/test_nextstep_cfg_parallel_layout.py

@@ -8,6 +8,8 @@ import vllm_omni.diffusion.models.nextstep_1_1.pipeline_nextstep_1_1 as nextstep
 from vllm_omni.diffusion.models.nextstep_1_1.modeling_nextstep_heads import FlowMatchingHead
 from vllm_omni.diffusion.models.nextstep_1_1.pipeline_nextstep_1_1 import NextStep11Pipeline
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 
 class _DummyImageHead:
     def __init__(self, token_dim: int):

tests/diffusion/models/z_image/test_zimage_tp_constraints.py

@@ -2,6 +2,8 @@ import pytest
 
 from vllm_omni.diffusion.models.z_image.z_image_transformer import validate_zimage_tp_constraints
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 
 def test_validate_zimage_tp_constraints_tp2_ok():
     ffn_hidden_dim, final_out_dims, supported_tp = validate_zimage_tp_constraints(

tests/diffusion/test_multiproc_executor_concurrency.py

@@ -12,7 +12,7 @@ from vllm_omni.diffusion.data import DiffusionOutput
 from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
 from vllm_omni.diffusion.scheduler import Scheduler
 
-pytestmark = [pytest.mark.diffusion]
+pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cpu]
 
 
 # ───────────────────────────────────────────── helpers ─────────────────────

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 = manager._normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=0)
+    normalized = normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=0)
     assert normalized is None
 
 

tests/distributed/omni_connectors/test_mooncake_transfer_engine_buffer.py

@@ -17,7 +17,7 @@ from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_c
 )
 
 # All tests in this file are pure-CPU unit tests for the memory allocator.
-pytestmark = [pytest.mark.cpu, pytest.mark.parallel]
+pytestmark = [pytest.mark.cpu, pytest.mark.parallel, pytest.mark.core_model]
 
 
 @pytest.mark.core_model

tests/distributed/omni_connectors/test_omni_connector_configs.py

@@ -8,6 +8,8 @@ import pytest
 # Use the new import path for initialization utilities
 from vllm_omni.distributed.omni_connectors.utils.initialization import load_omni_transfer_config
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 
 def get_config_files():
     """Helper to find config files."""

tests/distributed/omni_coordinator/test_load_balancer.py

@@ -0,0 +1,58 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from time import time
+
+from vllm_omni.distributed.omni_coordinator import (
+    InstanceInfo,
+    RandomBalancer,
+    StageStatus,
+)
+
+
+def test_load_balancer_select_returns_valid_index():
+    """Verify RandomBalancer.select() returns a valid index for instances."""
+    # Task structure mirrors async_omni; RandomBalancer ignores task contents.
+    task: dict = {
+        "request_id": "test",
+        "engine_inputs": None,
+        "sampling_params": None,
+    }
+
+    now = time()
+    instances = [
+        InstanceInfo(
+            input_addr="tcp://host:10001",
+            output_addr="tcp://host:10001-out",
+            stage_id=0,
+            status=StageStatus.UP,
+            queue_length=0,
+            last_heartbeat=now,
+            registered_at=now,
+        ),
+        InstanceInfo(
+            input_addr="tcp://host:10002",
+            output_addr="tcp://host:10002-out",
+            stage_id=0,
+            status=StageStatus.UP,
+            queue_length=1,
+            last_heartbeat=now,
+            registered_at=now,
+        ),
+        InstanceInfo(
+            input_addr="tcp://host:10003",
+            output_addr="tcp://host:10003-out",
+            stage_id=1,
+            status=StageStatus.UP,
+            queue_length=2,
+            last_heartbeat=now,
+            registered_at=now,
+        ),
+    ]
+
+    balancer = RandomBalancer()
+
+    index = balancer.select(task, instances)
+
+    assert isinstance(index, int)
+    assert 0 <= index < len(instances)

tests/distributed/omni_coordinator/test_omni_coord_client_for_hub.py

@@ -0,0 +1,119 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+import time
+
+import pytest
+import zmq
+
+from vllm_omni.distributed.omni_coordinator import (
+    InstanceList,
+    OmniCoordClientForHub,
+)
+
+
+def _bind_pub() -> tuple[zmq.Context, zmq.Socket, str]:
+    ctx = zmq.Context.instance()
+    pub = ctx.socket(zmq.PUB)
+    pub.bind("tcp://127.0.0.1:*")
+    endpoint = pub.getsockopt(zmq.LAST_ENDPOINT).decode("ascii")
+    return ctx, pub, endpoint
+
+
+def _wait_for_condition(cond, timeout: float = 2.0, interval: float = 0.01) -> bool:
+    start = time.time()
+    while time.time() - start < timeout:
+        if cond():
+            return True
+        time.sleep(interval)
+    return False
+
+
+def test_hub_client_caches_instance_list_from_pub():
+    """Verify OmniCoordClientForHub receives instance list updates from OmniCoordinator and caches for get_instance_list()."""
+    ctx, pub, endpoint = _bind_pub()
+
+    client = OmniCoordClientForHub(endpoint)
+    # ZMQ PUB/SUB slow-joiner: allow SUB to finish connecting before first send
+    time.sleep(0.2)
+
+    now = time.time()
+    instances_payload = [
+        {
+            "input_addr": "tcp://stage:10001",
+            "output_addr": "tcp://stage:10001-out",
+            "stage_id": 0,
+            "status": "up",
+            "queue_length": 0,
+            "last_heartbeat": now,
+            "registered_at": now,
+        },
+        {
+            "input_addr": "tcp://stage:10002",
+            "output_addr": "tcp://stage:10002-out",
+            "stage_id": 0,
+            "status": "up",
+            "queue_length": 1,
+            "last_heartbeat": now,
+            "registered_at": now,
+        },
+        {
+            "input_addr": "tcp://stage:10003",
+            "output_addr": "tcp://stage:10003-out",
+            "stage_id": 1,
+            "status": "error",
+            "queue_length": 5,
+            "last_heartbeat": now,
+            "registered_at": now,
+        },
+    ]
+
+    payload = {"instances": instances_payload, "timestamp": now}
+    pub.send(json.dumps(payload).encode("utf-8"))
+
+    assert _wait_for_condition(lambda: len(client.get_instance_list().instances) == 3)
+
+    inst_list = client.get_instance_list()
+    assert isinstance(inst_list, InstanceList)
+    assert len(inst_list.instances) == 3
+
+    for src, inst in zip(instances_payload, inst_list.instances, strict=True):
+        assert inst.input_addr == src["input_addr"]
+        assert inst.output_addr == src["output_addr"]
+        assert inst.stage_id == src["stage_id"]
+        assert inst.status.value == src["status"]
+
+    stage0 = client.get_instances_for_stage(0)
+    stage1 = client.get_instances_for_stage(1)
+
+    assert all(inst.stage_id == 0 for inst in stage0.instances)
+    assert all(inst.stage_id == 1 for inst in stage1.instances)
+
+    # Send an updated list with fewer instances and verify cache refresh.
+    updated_payload = {
+        "instances": instances_payload[:2],
+        "timestamp": now + 1.0,
+    }
+    pub.send(json.dumps(updated_payload).encode("utf-8"))
+
+    assert _wait_for_condition(lambda: len(client.get_instance_list().instances) == 2)
+    updated_list = client.get_instance_list()
+    assert len(updated_list.instances) == 2
+
+    client.close()
+    pub.close(0)
+    ctx.term()
+
+
+def test_hub_client_close_closes_sub_socket():
+    """Verify OmniCoordClientForHub.close() marks client as closed; second close raises."""
+    ctx, pub, endpoint = _bind_pub()
+    client = OmniCoordClientForHub(endpoint)
+    client.close()
+
+    with pytest.raises(RuntimeError, match="already closed"):
+        client.close()
+
+    pub.close(0)
+    ctx.term()

tests/distributed/omni_coordinator/test_omni_coord_client_for_stage.py

@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+
+import zmq
+
+from vllm_omni.distributed.omni_coordinator import (
+    OmniCoordClientForStage,
+    StageStatus,
+)
+
+
+def _bind_router() -> tuple[zmq.Context, zmq.Socket, str]:
+    ctx = zmq.Context.instance()
+    router = ctx.socket(zmq.ROUTER)
+    router.bind("tcp://127.0.0.1:*")
+    endpoint = router.getsockopt(zmq.LAST_ENDPOINT).decode("ascii")
+    return ctx, router, endpoint
+
+
+def _recv_event(router: zmq.Socket) -> dict:
+    frames = router.recv_multipart()
+    # ROUTER adds identity frame; the last frame is the payload.
+    payload = frames[-1]
+    return json.loads(payload.decode("utf-8"))
+
+
+def test_stage_client_auto_register_on_init():
+    """Verify OmniCoordClientForStage automatically sends initial registration/status-up event when created."""
+    ctx, router, endpoint = _bind_router()
+
+    input_addr = "tcp://stage:10001"
+    output_addr = "tcp://stage:10001-out"
+    stage_id = 0
+
+    client = OmniCoordClientForStage(endpoint, input_addr, output_addr, stage_id)
+
+    event = _recv_event(router)
+
+    assert event["event_type"] == "update"
+    assert event["status"] == StageStatus.UP.value
+    assert event["stage_id"] == stage_id
+    assert event["input_addr"] == input_addr
+    assert event["output_addr"] == output_addr
+
+    client.close()
+    router.close(0)
+    ctx.term()
+
+
+def test_stage_client_update_info_sends_correct_event():
+    """Verify OmniCoordClientForStage.update_info() sends status/load update events with expected fields."""
+    ctx, router, endpoint = _bind_router()
+
+    input_addr = "tcp://stage:10002"
+    output_addr = "tcp://stage:10002-out"
+    stage_id = 1
+
+    client = OmniCoordClientForStage(endpoint, input_addr, output_addr, stage_id)
+
+    # Discard initial registration event.
+    _recv_event(router)
+
+    client.update_info(status=StageStatus.ERROR)
+    client.update_info(queue_length=10)
+
+    first = _recv_event(router)
+    second = _recv_event(router)
+
+    assert first["status"] == StageStatus.ERROR.value
+    assert first["stage_id"] == stage_id
+    assert first["input_addr"] == input_addr
+    assert first["output_addr"] == output_addr
+
+    assert second["queue_length"] == 10
+    assert second["stage_id"] == stage_id
+    assert second["input_addr"] == input_addr
+    assert second["output_addr"] == output_addr
+
+    client.close()
+    router.close(0)
+    ctx.term()
+
+
+def test_stage_client_close_sends_down_status():
+    """Verify close() sends final status-down event before closing underlying socket."""
+    ctx, router, endpoint = _bind_router()
+
+    input_addr = "tcp://stage:10003"
+    output_addr = "tcp://stage:10003-out"
+    stage_id = 2
+
+    client = OmniCoordClientForStage(endpoint, input_addr, output_addr, stage_id)
+
+    # Discard initial registration event.
+    _recv_event(router)
+
+    client.close()
+
+    event = _recv_event(router)
+    assert event["status"] == StageStatus.DOWN.value
+    assert event["stage_id"] == stage_id
+    assert event["input_addr"] == input_addr
+    assert event["output_addr"] == output_addr
+
+    assert client._socket.closed  # DEALER socket no longer usable after close
+
+    router.close(0)
+    ctx.term()