[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

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

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/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/audio_generate_api.md

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

docs/serving/speech_api.md

@@ -96,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_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"
+    ``````

docs/user_guide/examples/online_serving/text_to_audio.md

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

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

examples/online_serving/stable_audio/README.md

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

examples/online_serving/stable_audio/curl_examples.sh

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

examples/online_serving/stable_audio/stable_audio_client.py

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

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

@@ -18,16 +18,21 @@ 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
@@ -44,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.
@@ -1057,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")
@@ -1150,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")
 
 
@@ -1158,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={
@@ -1178,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
@@ -1198,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.
@@ -1246,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.
@@ -1268,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.
 
@@ -1329,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.
 
@@ -1383,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.
 
@@ -1409,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)
@@ -1437,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__(
@@ -1726,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_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()

tests/distributed/omni_coordinator/test_omni_coordinator.py

@@ -0,0 +1,209 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+import time
+
+import zmq
+from vllm.v1.utils import get_engine_client_zmq_addr
+
+from vllm_omni.distributed.omni_coordinator import (
+    OmniCoordClientForStage,
+    OmniCoordinator,
+    StageStatus,
+)
+
+
+def _recv_instance_list(sub: zmq.Socket, timeout_ms: int = 2000) -> dict | None:
+    """Receive InstanceList JSON from SUB socket. Returns None on timeout."""
+    sub.setsockopt(zmq.RCVTIMEO, timeout_ms)
+    try:
+        data = sub.recv()
+        return json.loads(data.decode("utf-8"))
+    except zmq.Again:
+        return None
+
+
+def _wait_for_instance_list(
+    sub: zmq.Socket,
+    expected_count: int,
+    timeout: float = 3.0,
+) -> dict | None:
+    """Wait until received InstanceList with expected_count active instances."""
+    start = time.time()
+    while time.time() - start < timeout:
+        msg = _recv_instance_list(sub, timeout_ms=500)
+        if msg is not None and len(msg.get("instances", [])) == expected_count:
+            return msg
+    return None
+
+
+def test_omni_coordinator_registration_broadcast():
+    """Verify that after multiple OmniCoordClientForStage instances register,
+    OmniCoordinator publishes an InstanceList containing all registered instances.
+    """
+    router_addr = get_engine_client_zmq_addr(
+        local_only=False,
+        host="127.0.0.1",
+        port=0,
+    )
+    pub_addr = get_engine_client_zmq_addr(
+        local_only=False,
+        host="127.0.0.1",
+        port=0,
+    )
+    coordinator = OmniCoordinator(
+        router_zmq_addr=router_addr,
+        pub_zmq_addr=pub_addr,
+        heartbeat_timeout=1000.0,
+    )
+
+    sub_ctx = zmq.Context.instance()
+    sub = sub_ctx.socket(zmq.SUB)
+    sub.connect(pub_addr)
+    sub.setsockopt(zmq.SUBSCRIBE, b"")
+
+    # ZMQ PUB/SUB slow-joiner: allow SUB to connect before clients register.
+    time.sleep(0.3)
+
+    # Create 3 stage clients; each auto-registers on init.
+    clients = [
+        OmniCoordClientForStage(router_addr, "tcp://stage:10001", "tcp://stage:10001-out", 0),
+        OmniCoordClientForStage(router_addr, "tcp://stage:10002", "tcp://stage:10002-out", 0),
+        OmniCoordClientForStage(router_addr, "tcp://stage:10003", "tcp://stage:10003-out", 1),
+    ]
+
+    msg = _wait_for_instance_list(sub, expected_count=3)
+    assert msg is not None, "Expected InstanceList with 3 instances"
+    assert len(msg["instances"]) == 3
+    assert isinstance(msg["timestamp"], (int, float))
+
+    input_addrs = {inst["input_addr"] for inst in msg["instances"]}
+    assert "tcp://stage:10001" in input_addrs
+    assert "tcp://stage:10002" in input_addrs
+    assert "tcp://stage:10003" in input_addrs
+
+    for c in clients:
+        c.close()
+    coordinator.close()
+    sub.close(0)
+    sub_ctx.term()
+
+
+def test_omni_coordinator_heartbeat_timeout_handling():
+    """Verify that when a stage instance stops sending heartbeats,
+    OmniCoordinator marks it as unhealthy and excludes it from the active list.
+    """
+    router_addr = get_engine_client_zmq_addr(
+        local_only=False,
+        host="127.0.0.1",
+        port=0,
+    )
+    pub_addr = get_engine_client_zmq_addr(
+        local_only=False,
+        host="127.0.0.1",
+        port=0,
+    )
+    coordinator = OmniCoordinator(
+        router_zmq_addr=router_addr,
+        pub_zmq_addr=pub_addr,
+        heartbeat_timeout=5.0,
+    )
+
+    sub_ctx = zmq.Context.instance()
+    sub = sub_ctx.socket(zmq.SUB)
+    sub.connect(pub_addr)
+    sub.setsockopt(zmq.SUBSCRIBE, b"")
+
+    time.sleep(0.3)
+
+    # A and B: real clients that send heartbeats every 5s.
+    client_a = OmniCoordClientForStage(router_addr, "tcp://stage:a", "tcp://stage:a-out", 0)
+    client_b = OmniCoordClientForStage(router_addr, "tcp://stage:b", "tcp://stage:b-out", 0)
+
+    # C: raw DEALER that sends only registration, no heartbeat.
+    dealer_ctx = zmq.Context.instance()
+    dealer_c = dealer_ctx.socket(zmq.DEALER)
+    dealer_c.connect(router_addr)
+    reg_event = {
+        "input_addr": "tcp://stage:c",
+        "output_addr": "tcp://stage:c-out",
+        "stage_id": 0,
+        "event_type": "update",
+        "status": StageStatus.UP.value,
+        "queue_length": 0,
+    }
+    dealer_c.send(json.dumps(reg_event).encode("utf-8"))
+
+    msg = _wait_for_instance_list(sub, expected_count=3)
+    assert msg is not None, "Expected initial 3 instances"
+    assert len(msg["instances"]) == 3
+
+    # Wait for heartbeat timeout (timeout=5s, check interval ~2.5s).
+    time.sleep(8.0)
+
+    # Receive the update (C should be ERROR and excluded from active list).
+    msg_after_timeout = _wait_for_instance_list(sub, expected_count=2, timeout=5.0)
+    assert msg_after_timeout is not None, "Expected InstanceList with 2 instances after timeout"
+    instances = msg_after_timeout.get("instances", [])
+    input_addrs = {inst["input_addr"] for inst in instances}
+
+    assert "tcp://stage:a" in input_addrs
+    assert "tcp://stage:b" in input_addrs
+    assert "tcp://stage:c" not in input_addrs
+
+    client_a.close()
+    client_b.close()
+    dealer_c.close(0)
+    coordinator.close()
+    sub.close(0)
+    dealer_ctx.term()
+    sub_ctx.term()
+
+
+def test_omni_coordinator_instance_shutdown_handling():
+    """Verify that when a stage instance sends status='down',
+    OmniCoordinator removes it from the active list and broadcasts an updated list.
+    """
+    router_addr = get_engine_client_zmq_addr(
+        local_only=False,
+        host="127.0.0.1",
+        port=0,
+    )
+    pub_addr = get_engine_client_zmq_addr(
+        local_only=False,
+        host="127.0.0.1",
+        port=0,
+    )
+    coordinator = OmniCoordinator(
+        router_zmq_addr=router_addr,
+        pub_zmq_addr=pub_addr,
+        heartbeat_timeout=1000.0,
+    )
+
+    sub_ctx = zmq.Context.instance()
+    sub = sub_ctx.socket(zmq.SUB)
+    sub.connect(pub_addr)
+    sub.setsockopt(zmq.SUBSCRIBE, b"")
+
+    time.sleep(0.3)  # PUB/SUB slow-joiner
+
+    client = OmniCoordClientForStage(router_addr, "tcp://stage:shutdown", "tcp://stage:shutdown-out", 0)
+
+    msg = _wait_for_instance_list(sub, expected_count=1)
+    assert msg is not None
+    assert len(msg["instances"]) == 1
+    assert msg["instances"][0]["input_addr"] == "tcp://stage:shutdown"
+
+    # Send down status (simulating graceful shutdown).
+    client.update_info(status=StageStatus.DOWN)
+
+    # Receive updated list (should have 0 active instances).
+    msg = _wait_for_instance_list(sub, expected_count=0)
+    assert msg is not None
+    assert len(msg["instances"]) == 0
+
+    client.close()
+    coordinator.close()
+    sub.close(0)
+    sub_ctx.term()

tests/e2e/offline_inference/test_diffusion_cpu_offload.py

@@ -6,7 +6,7 @@ import pytest
 import torch
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
-from tests.utils import GPUMemoryMonitor, hardware_test
+from tests.utils import DeviceMemoryMonitor, hardware_test
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
 
@@ -23,9 +23,9 @@ models = ["riverclouds/qwen_image_random"]
 def inference(model_name: str, offload: bool = True):
     gc.collect()
     current_omni_platform.empty_cache()
-    device_index = torch.cuda.current_device()
-    torch.cuda.reset_peak_memory_stats(device=device_index)
-    monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02)
+    device_index = current_omni_platform.current_device()
+    current_omni_platform.reset_peak_memory_stats()
+    monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
     monitor.start()
     m = Omni(
         model=model_name,
@@ -33,7 +33,7 @@ def inference(model_name: str, offload: bool = True):
         # cache_backend="cache_dit",
         enable_cpu_offload=offload,
     )
-    torch.cuda.reset_peak_memory_stats(device=device_index)
+    current_omni_platform.reset_peak_memory_stats()
     height = 256
     width = 256
 
@@ -44,7 +44,7 @@ def inference(model_name: str, offload: bool = True):
             width=width,
             num_inference_steps=9,
             guidance_scale=0.0,
-            generator=torch.Generator("cuda").manual_seed(42),
+            generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
         ),
     )
     peak = monitor.peak_used_mb

tests/e2e/offline_inference/test_diffusion_layerwise_offload.py

@@ -5,7 +5,7 @@ import pytest
 import torch
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
-from tests.utils import GPUMemoryMonitor
+from tests.utils import DeviceMemoryMonitor
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
 
@@ -28,11 +28,9 @@ def run_inference(
     layerwise_offload: bool = False,
     num_inference_steps: int = 3,
 ) -> float:
-    # For now, only support on GPU, so apply torch.cuda operations here
-    # NPU / ROCm platforms are expected to be detected and skipped this test function
     current_omni_platform.empty_cache()
-    device_index = torch.cuda.current_device()
-    monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02)
+    device_index = current_omni_platform.current_device()
+    monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
     monitor.start()
 
     m = Omni(
@@ -44,7 +42,7 @@ def run_inference(
         flow_shift=5.0,
     )
 
-    torch.cuda.reset_peak_memory_stats(device=device_index)
+    current_omni_platform.reset_peak_memory_stats()
 
     # Refer to tests/e2e/offline_inference/test_t2v_model.py
     # Use minimal settings for testing
@@ -57,7 +55,7 @@ def run_inference(
         OmniDiffusionSamplingParams(
             height=height,
             width=width,
-            generator=torch.Generator("cuda").manual_seed(42),
+            generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
             guidance_scale=1.0,
             num_inference_steps=num_inference_steps,
             num_frames=num_frames,

tests/e2e/offline_inference/test_expert_parallel.py

@@ -0,0 +1,248 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+System test for Expert Parallel (EP).
+
+Tests verify that EP inference produces correct outputs compared to baseline.
+"""
+
+import gc
+import os
+import time
+from typing import NamedTuple
+
+import numpy as np
+import pytest
+import torch
+import torch.distributed as dist
+from PIL import Image
+
+from tests.utils import hardware_test
+from vllm_omni import Omni
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.platforms import current_omni_platform
+
+# Test configuration
+MODELS = ["tencent/HunyuanImage-3.0"]
+PROMPT = "A brown and white dog is running on the grass"
+DEFAULT_HEIGHT = 1024
+DEFAULT_WIDTH = 1024
+DEFAULT_SEED = 1234
+DEFAULT_STEPS = 50
+DEFAULT_GUIDANCE_SCALE = 0.0
+COS_SIM_MEAN_THRESHOLD = 1e-2
+COS_SIM_MAX_THRESHOLD = 1e-2
+MSE_THRESHOLD = 1e-2
+DIFF_MEAN_THRESHOLD = 5e-2
+DIFF_MAX_THRESHOLD = 1
+TOTAL_CARD_NUM = 8
+
+
+class InferenceResult(NamedTuple):
+    """Result of an inference run."""
+
+    images: list[Image.Image]
+    elapsed_ms: float
+
+
+def _cleanup_distributed():
+    """Clean up distributed environment and GPU resources."""
+    if dist.is_initialized():
+        dist.destroy_process_group()
+
+    for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]:
+        os.environ.pop(key, None)
+
+    gc.collect()
+    if current_omni_platform.is_available():
+        current_omni_platform.empty_cache()
+        current_omni_platform.synchronize()
+
+    time.sleep(5)
+
+
+def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float, float, float, float]:
+    """Return (mean_abs_diff, max_abs_diff) over RGB pixels in [0, 1]."""
+    ta = torch.from_numpy(np.asarray(a.convert("RGB"), dtype=np.float32) / 255.0)
+    tb = torch.from_numpy(np.asarray(b.convert("RGB"), dtype=np.float32) / 255.0)
+    assert ta.shape == tb.shape, f"Image shapes differ: {ta.shape} vs {tb.shape}"
+    abs_diff = torch.abs(ta - tb)
+
+    vec_a = ta.reshape(-1)
+    vec_b = tb.reshape(-1)
+    cos_sim = torch.nn.functional.cosine_similarity(vec_a, vec_b, dim=0, eps=1e-8)
+    mse = torch.nn.functional.mse_loss(ta, tb)
+    return abs_diff.mean().item(), abs_diff.max().item(), cos_sim.mean().item(), cos_sim.max().item(), mse.item()
+
+
+def _run_inference(
+    model_name: str,
+    tensor_parallel_size: int,
+    enable_expert_parallel: bool,
+    guidance_scale: float = DEFAULT_GUIDANCE_SCALE,
+    height: int = DEFAULT_HEIGHT,
+    width: int = DEFAULT_WIDTH,
+    seed: int = DEFAULT_SEED,
+    warmup: bool = True,
+) -> InferenceResult:
+    """Run inference with specified configuration.
+
+    Args:
+        warmup: If True, run one warmup iteration before the timed run.
+    """
+    parallel_config = DiffusionParallelConfig(
+        tensor_parallel_size=tensor_parallel_size,
+        enable_expert_parallel=enable_expert_parallel,
+    )
+    omni = Omni(model=model_name, parallel_config=parallel_config)
+
+    try:
+        # Warmup run (not timed)
+        if warmup:
+            _ = omni.generate(
+                PROMPT,
+                OmniDiffusionSamplingParams(
+                    height=height,
+                    width=width,
+                    num_inference_steps=DEFAULT_STEPS,
+                    guidance_scale=guidance_scale,
+                    generator=torch.Generator(current_omni_platform.device_type).manual_seed(seed),
+                    num_outputs_per_prompt=1,
+                ),
+            )
+
+        # Timed run
+        start = time.time()
+        outputs = omni.generate(
+            PROMPT,
+            OmniDiffusionSamplingParams(
+                height=height,
+                width=width,
+                num_inference_steps=DEFAULT_STEPS,
+                guidance_scale=guidance_scale,
+                generator=torch.Generator(current_omni_platform.device_type).manual_seed(seed),
+                num_outputs_per_prompt=1,
+            ),
+        )
+        elapsed_ms = (time.time() - start) * 1000
+
+        return InferenceResult(
+            images=outputs[0].request_output[0].images,
+            elapsed_ms=elapsed_ms,
+        )
+    finally:
+        omni.close()
+        _cleanup_distributed()
+
+
+# EP configurations: (tensor_parallel_size, enable_ep, height, width, is_perf_test)
+EP_TEST_CONFIG = [
+    (TOTAL_CARD_NUM, True, DEFAULT_HEIGHT, DEFAULT_WIDTH, True),
+]
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": TOTAL_CARD_NUM})
+@pytest.mark.parametrize("model_name", MODELS)
+def test_ep(model_name):
+    baseline_cache: dict[str, InferenceResult] = {}
+    results: list[dict] = []
+    # baseline_key = PROMPT + "_" + str(DEFAULT_HEIGHT) + "*" + str(DEFAULT_WIDTH) + "_" + str(DEFAULT_SEED) + "_" + str(DEFAULT_STEPS)
+    # baseline_cache[baseline_key] = _run_inference(
+    #     model_name=model_name,
+    #     tensor_parallel_size=TOTAL_CARD_NUM,
+    #     enable_expert_parallel=False,
+    #     warmup=True,
+    # )
+    device_count = current_omni_platform.get_device_count()
+
+    print("\n" + "=" * 90)
+    print(f"Sequence Parallel Test - Model: {model_name}")
+    print(f"Available GPUs: {device_count}")
+    print("=" * 90)
+
+    for tensor_parallel_size, enable_ep, height, width, is_perf_test in EP_TEST_CONFIG:
+        cache_key = PROMPT + "_" + str(height) + "*" + str(width) + "_" + str(DEFAULT_SEED) + "_" + str(DEFAULT_STEPS)
+        if cache_key not in baseline_cache:
+            print(
+                f"\n--- Running baseline{{ prompt: {PROMPT}, resolution: {height}x{width}, seed: {DEFAULT_SEED}, steps: {DEFAULT_STEPS} }}---"
+            )
+            baseline = _run_inference(
+                model_name=model_name,
+                tensor_parallel_size=TOTAL_CARD_NUM,
+                enable_expert_parallel=False,
+                warmup=True,
+            )
+            assert len(baseline.images) == 1
+            baseline_cache[cache_key] = baseline
+            print(f"[baseline] {height}x{width}: {baseline.elapsed_ms:.0f}ms")
+        else:
+            baseline = baseline_cache[cache_key]
+
+        print(f"\n--- Running test: {{ enable_ep: {enable_ep} }} ---")
+        infer_result = _run_inference(
+            model_name=model_name,
+            tensor_parallel_size=tensor_parallel_size,
+            enable_expert_parallel=enable_ep,
+            warmup=is_perf_test,
+        )
+        assert len(infer_result.images) == 1
+        mean_diff, max_diff, cos_sim_mean, cos_sim_max, mse = _diff_metrics(baseline.images[0], infer_result.images[0])
+        print(f"{cos_sim_mean}, {cos_sim_max}, {mse}")
+
+        # Build result entry
+        result = {
+            "prompt": PROMPT,
+            "enable_ep": enable_ep,
+            "height": height,
+            "width": width,
+            "baseline_ms": baseline.elapsed_ms,
+            "ep_ms": infer_result.elapsed_ms,
+            "mean_diff": mean_diff,
+            "max_diff": max_diff,
+            "cos_sim_mean": cos_sim_mean,
+            "cos_sim_max": cos_sim_max,
+            "mse": mse,
+            "is_perf_test": is_perf_test,
+        }
+        results.append(result)
+
+        # Output based on test type
+        if is_perf_test:
+            speedup = baseline.elapsed_ms / infer_result.elapsed_ms if infer_result.elapsed_ms > 0 else 0
+            result["speedup"] = speedup
+            print(
+                f"[enable_ep: {enable_ep}] {tensor_parallel_size} GPUs | "
+                f"baseline: {baseline.elapsed_ms:.0f}ms, ep: {infer_result.elapsed_ms:.0f}ms, "
+                f"speedup: {speedup:.2f}x"
+            )
+        else:
+            print(
+                f"[enable_ep: {enable_ep}] {tensor_parallel_size} GPUs | ep: {infer_result.elapsed_ms:.0f}ms (correctness only)"
+            )
+
+        print(
+            f"[enable_ep: {enable_ep}] diff: [mean={mean_diff:.6e}, max={max_diff:.6e}], cos_sim: [mean={cos_sim_mean:.6e}, max={cos_sim_max:.6e}], mse: {mse:.6e}"
+        )
+        assert mean_diff <= DIFF_MEAN_THRESHOLD and max_diff <= DIFF_MAX_THRESHOLD, (
+            f"[enable_ep: {enable_ep}] output differs from baseline: mean={mean_diff:.6e}, max={max_diff:.6e}"
+        )
+
+    print("\n" + "=" * 90)
+    print("SUMMARY")
+    print("=" * 90)
+    print(f"{'Mode':<15} {'GPUs':<6} {'Size':<10} {'Baseline':<12} {'EP':<12} {'Speedup':<10} {'Status'}")
+    print("-" * 90)
+    for r in results:
+        speedup_str = f"{r['speedup']:.2f}x" if r.get("speedup") else "N/A"
+        baseline_str = f"{r['baseline_ms']:.0f}ms" if r["is_perf_test"] else "N/A"
+        status = "PASS" if r["mean_diff"] <= DIFF_MEAN_THRESHOLD else "FAIL"
+        print(
+            f"{r['prompt']:<15.10} {r['enable_ep']:<6} {r['height']}x{r['width']:<5} "
+            f"{baseline_str:<12} {r['ep_ms']:.0f}ms{'':<7} {speedup_str:<10} {status}"
+        )
+    print("=" * 90)

tests/e2e/offline_inference/test_zimage_parallelism.py

@@ -22,7 +22,7 @@ import torch
 from PIL import Image
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
-from tests.utils import GPUMemoryMonitor, hardware_test
+from tests.utils import DeviceMemoryMonitor, hardware_test
 from vllm_omni import Omni
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -54,12 +54,13 @@ def _pil_to_float_rgb_tensor(img: Image.Image) -> torch.Tensor:
 
 
 def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]:
-    """Return (mean_abs_diff, max_abs_diff) over RGB pixels in [0, 1]."""
+    """Return (mean_abs_diff, p99_abs_diff) over RGB pixels in [0, 1]."""
     ta = _pil_to_float_rgb_tensor(a)
     tb = _pil_to_float_rgb_tensor(b)
     assert ta.shape == tb.shape, f"Image shapes differ: {ta.shape} vs {tb.shape}"
     abs_diff = torch.abs(ta - tb)
-    return abs_diff.mean().item(), abs_diff.max().item()
+    p99_abs_diff = torch.quantile(abs_diff.flatten(), 0.99).item()
+    return abs_diff.mean().item(), p99_abs_diff
 
 
 def _extract_single_image(outputs) -> Image.Image:
@@ -93,9 +94,9 @@ def _run_zimage_generate(
     if num_requests < 2:
         raise ValueError("num_requests must be >= 2 (1 warmup + >=1 timed)")
 
-    torch.cuda.empty_cache()
-    device_index = torch.cuda.current_device()
-    monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02)
+    current_omni_platform.empty_cache()
+    device_index = current_omni_platform.current_device()
+    monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
     monitor.start()
     m = Omni(
         model=_get_zimage_model(),
@@ -161,8 +162,8 @@ def _run_zimage_generate(
 def test_zimage_tensor_parallel_tp2(tmp_path: Path):
     if current_omni_platform.is_npu() or current_omni_platform.is_rocm():
         pytest.skip("Z-Image TP e2e test is only supported on CUDA for now.")
-    if not torch.cuda.is_available() or torch.cuda.device_count() < 2:
-        pytest.skip("Z-Image TP=2 requires >= 2 CUDA devices.")
+    if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2:
+        pytest.skip("Z-Image TP=2 requires >= 2 devices.")
 
     enforce_eager = False
 
@@ -196,18 +197,18 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path):
     assert tp1_img.width == width and tp1_img.height == height
     assert tp2_img.width == width and tp2_img.height == height
 
-    mean_abs_diff, max_abs_diff = _diff_metrics(tp1_img, tp2_img)
+    mean_abs_diff, p99_abs_diff = _diff_metrics(tp1_img, tp2_img)
     mean_threshold = 3e-2
-    max_threshold = 5e-1
+    p99_threshold = 2.5e-1
     print(
         "Z-Image TP image diff stats (TP=1 vs TP=2): "
-        f"mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e}; "
-        f"thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}; "
+        f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; "
+        f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; "
         f"tp1_img={tp1_path}, tp2_img={tp2_path}"
     )
-    assert mean_abs_diff <= mean_threshold and max_abs_diff <= max_threshold, (
-        f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e} "
-        f"(thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e})"
+    assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, (
+        f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e} "
+        f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})"
     )
 
     print(f"Z-Image TP perf (lower is better): tp1_time_s={tp1_time_s:.6f}, tp2_time_s={tp2_time_s:.6f}")
@@ -223,8 +224,8 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path):
 def test_zimage_vae_patch_parallel_tp2(tmp_path: Path):
     if current_omni_platform.is_npu() or current_omni_platform.is_rocm():
         pytest.skip("Z-Image VAE patch parallel e2e test is only supported on CUDA for now.")
-    if not torch.cuda.is_available() or torch.cuda.device_count() < 2:
-        pytest.skip("Z-Image VAE patch parallel TP=2 requires >= 2 CUDA devices.")
+    if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2:
+        pytest.skip("Z-Image VAE patch parallel TP=2 requires >= 2 devices.")
 
     enforce_eager = False
 
@@ -262,16 +263,16 @@ def test_zimage_vae_patch_parallel_tp2(tmp_path: Path):
     baseline_img.save(baseline_path)
     pp2_img.save(pp2_path)
 
-    mean_abs_diff, max_abs_diff = _diff_metrics(baseline_img, pp2_img)
+    mean_abs_diff, p99_abs_diff = _diff_metrics(baseline_img, pp2_img)
     mean_threshold = 5e-3
-    max_threshold = 1e-1
+    p99_threshold = 1e-1
     print(
         "Z-Image VAE patch parallel image diff stats (TP=2, pp=1 vs pp=2): "
-        f"mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e}; "
-        f"thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}; "
+        f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; "
+        f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; "
         f"pp1_img={baseline_path}, pp2_img={pp2_path}"
     )
-    assert mean_abs_diff <= mean_threshold and max_abs_diff <= max_threshold, (
-        f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e} "
-        f"(thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e})"
+    assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, (
+        f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e} "
+        f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})"
     )

tests/e2e/online_serving/test_qwen3_omni.py

@@ -8,6 +8,7 @@ from pathlib import Path
 import pytest
 
 from tests.conftest import (
+    OmniServerParams,
     dummy_messages_from_mix_data,
     generate_synthetic_audio,
     generate_synthetic_image,
@@ -51,7 +52,9 @@ else:
     stage_configs = [get_chunk_config()]
 
 # 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
+]
 
 
 def get_system_prompt():
@@ -120,7 +123,7 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
     }
 
     # Test single completion
-    openai_client.send_request(request_config)
+    openai_client.send_omni_request(request_config)
 
 
 @pytest.mark.advanced_model
@@ -146,4 +149,4 @@ def test_text_to_text_001(omni_server, openai_client) -> None:
         "key_words": {"text": ["beijing"]},
     }
 
-    openai_client.send_request(request_config, request_num=get_max_batch_size())
+    openai_client.send_omni_request(request_config, request_num=get_max_batch_size())

tests/e2e/online_serving/test_qwen3_omni_expansion.py

@@ -12,6 +12,7 @@ from pathlib import Path
 import pytest
 
 from tests.conftest import (
+    OmniServerParams,
     dummy_messages_from_mix_data,
     generate_synthetic_audio,
     generate_synthetic_image,
@@ -50,7 +51,9 @@ def get_chunk_config(default_path):
 default_path = str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml")
 stage_configs = [default_path, get_chunk_config(default_path)]
 # 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
+]
 
 
 def get_system_prompt():
@@ -106,7 +109,7 @@ def test_text_to_audio_001(omni_server, openai_client) -> None:
         "key_words": {"text": ["beijing"]},
     }
 
-    openai_client.send_request(request_config)
+    openai_client.send_omni_request(request_config)
 
 
 @pytest.mark.advanced_model
@@ -128,7 +131,7 @@ def test_text_to_text_audio_001(omni_server, openai_client) -> None:
         "key_words": {"text": ["beijing"]},
     }
 
-    openai_client.send_request(request_config, request_num=get_max_batch_size())
+    openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
 
 
 @pytest.mark.advanced_model
@@ -153,7 +156,7 @@ def test_image_to_text_001(omni_server, openai_client) -> None:
         "key_words": {"image": IMAGE_KEY},
     }
 
-    openai_client.send_request(request_config)
+    openai_client.send_omni_request(request_config)
 
 
 @pytest.mark.advanced_model
@@ -177,7 +180,7 @@ def test_image_to_audio_001(omni_server, openai_client) -> None:
         "key_words": {"image": IMAGE_KEY},
     }
 
-    openai_client.send_request(request_config)
+    openai_client.send_omni_request(request_config)
 
 
 @pytest.mark.advanced_model
@@ -201,7 +204,7 @@ def test_image_to_text_audio_001(omni_server, openai_client) -> None:
         "key_words": {"image": IMAGE_KEY},
     }
 
-    openai_client.send_request(request_config, request_num=get_max_batch_size())
+    openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
 
 
 @pytest.mark.advanced_model
@@ -225,7 +228,7 @@ def test_video_to_text_001(omni_server, openai_client) -> None:
         "key_words": {"video": VIDEO_KEY},
     }
 
-    openai_client.send_request(request_config)
+    openai_client.send_omni_request(request_config)
 
 
 @pytest.mark.advanced_model
@@ -249,7 +252,7 @@ def test_video_to_audio_001(omni_server, openai_client) -> None:
         "key_words": {"video": VIDEO_KEY},
     }
 
-    openai_client.send_request(request_config)
+    openai_client.send_omni_request(request_config)
 
 
 @pytest.mark.advanced_model
@@ -273,7 +276,7 @@ def test_video_to_text_audio_001(omni_server, openai_client) -> None:
         "key_words": {"video": VIDEO_KEY},
     }
 
-    openai_client.send_request(request_config, request_num=get_max_batch_size())
+    openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
 
 
 @pytest.mark.advanced_model
@@ -298,7 +301,7 @@ def test_text_audio_to_text_audio_001(omni_server, openai_client) -> None:
         "key_words": {"audio": AUDIO_KEY},
     }
 
-    openai_client.send_request(request_config)
+    openai_client.send_omni_request(request_config)
 
 
 @pytest.mark.advanced_model
@@ -324,7 +327,7 @@ def test_text_image_to_text_audio_001(omni_server, openai_client) -> None:
         "key_words": {"image": IMAGE_KEY},
     }
 
-    openai_client.send_request(request_config)
+    openai_client.send_omni_request(request_config)
 
 
 @pytest.mark.advanced_model
@@ -351,7 +354,7 @@ def test_text_video_to_text_audio_001(omni_server, openai_client) -> None:
         "key_words": {"video": VIDEO_KEY},
     }
 
-    openai_client.send_request(request_config)
+    openai_client.send_omni_request(request_config)
 
 
 @pytest.mark.skip(reason="There is a known issue with shape mismatch error.")
@@ -383,4 +386,4 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
         "stream": True,
         "key_words": {"audio": AUDIO_KEY, "image": IMAGE_KEY, "video": VIDEO_KEY},
     }
-    openai_client.send_request(request_config, request_num=get_max_batch_size())
+    openai_client.send_omni_request(request_config, request_num=get_max_batch_size())

tests/e2e/online_serving/test_qwen_image_edit_expansion.py

@@ -0,0 +1,164 @@
+"""
+Comprehensive tests of diffusion features that are available in online serving mode
+and are supported by the following models:
+- Qwen-Image-Edit: single image input
+- Qwen-Image-Edit-2509: two image inputs
+"""
+
+import pytest
+
+from tests.conftest import (
+    OmniServer,
+    OmniServerParams,
+    OpenAIClientHandler,
+    dummy_messages_from_mix_data,
+    generate_synthetic_image,
+)
+from tests.utils import hardware_marks
+
+EDIT_PROMPT = "Transform this modern, geometrist image into a Vincent van Gogh style impressionist painting."
+MULTI_EDIT_PROMPT = (
+    "Transform the first image into a Dadaism collage art. "
+    "Transform the second image into a Vincent van Gogh style painting. "
+    "Then juxtapose the two transformed images into a single artwork for visual contrast."
+)
+NEGATIVE_PROMPT = "blurry, low quality, modern, geometrist"
+SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
+PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
+
+
+# This test file targets two models, so I write a helper function.
+# If a similar test only involves one model, one can just define a global list variable.
+def _get_diffusion_feature_cases(model: str):
+    return [
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--cache-backend",
+                    "tea_cache",  # [TODO] may consider changing to cache_dit after #1779 is resolved. Currently cache_dit and layerwise offload cannot work together.
+                    "--enable-layerwise-offload",
+                ],
+            ),
+            id="single_card_001",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--cache-backend",
+                    "cache_dit",
+                    "--ulysses-degree",
+                    "2",
+                ],
+            ),
+            id="parallel_001",
+            marks=PARALLEL_FEATURE_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--cache-backend",
+                    "cache_dit",
+                    "--ring",
+                    "2",
+                ],
+            ),
+            id="parallel_002",
+            marks=PARALLEL_FEATURE_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--cache-backend",
+                    "tea_cache",
+                    "--cfg-parallel-size",
+                    "2",
+                ],
+            ),
+            id="parallel_003",
+            marks=PARALLEL_FEATURE_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--cache-backend",
+                    "cache_dit",
+                    "--tensor-parallel-size",
+                    "2",
+                    "--vae-patch-parallel-size",
+                    "2",
+                    "--vae-use-tiling",
+                ],
+            ),
+            id="parallel_004",
+            marks=PARALLEL_FEATURE_MARKS,
+        ),
+    ]
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parametrize(
+    "omni_server",
+    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit"),
+    indirect=True,
+)
+def test_qwen_image_edit(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+    """Test all diffusion features with Qwen-Image-Edit in regular end-user scenarios."""
+    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
+
+    messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=EDIT_PROMPT)
+
+    # CFG parallel is only activated when a negative prompt and true_cfg_scale > 1.0 are both present
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "extra_body": {
+            "height": 512,
+            "width": 512,
+            "num_inference_steps": 2,
+            "negative_prompt": NEGATIVE_PROMPT,
+            "true_cfg_scale": 4.0,
+            "seed": 42,
+        },
+    }
+
+    openai_client.send_diffusion_request(request_config)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parametrize(
+    "omni_server",
+    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2509"),
+    indirect=True,
+)
+def test_qwen_image_edit_2509(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+    """Test all diffusion features with Qwen-Image-Edit-2509 in regular end-user scenarios."""
+    image_data_url_1 = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
+    image_data_url_2 = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
+
+    messages = dummy_messages_from_mix_data(
+        image_data_url=[image_data_url_1, image_data_url_2], content_text=MULTI_EDIT_PROMPT
+    )
+
+    # CFG parallel is only activated when a negative prompt and true_cfg_scale > 1.0 are both present
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "extra_body": {
+            "height": 512,
+            "width": 512,
+            "num_inference_steps": 2,
+            "negative_prompt": NEGATIVE_PROMPT,
+            "true_cfg_scale": 4.0,
+            "seed": 42,
+        },
+    }
+
+    openai_client.send_diffusion_request(request_config)

tests/entrypoints/openai_api/test_image_server.py

@@ -136,7 +136,6 @@ def mock_async_diffusion(mocker: MockerFixture):
 
     async def generate(**kwargs):
         # Return n PIL images wrapped in result object
-        print("!!!!!!!!!!!!!!!!!!!!! kwargs", kwargs)
         n = kwargs["sampling_params_list"][0].num_outputs_per_prompt
         mock.captured_sampling_params_list = kwargs["sampling_params_list"]
         mock.captured_prompt = kwargs["prompt"]

tests/entrypoints/openai_api/test_serving_audio_generate.py

@@ -1,509 +0,0 @@
-# tests/entrypoints/openai_api/test_serving_audio_generate.py
-import logging
-from inspect import Signature, signature
-from unittest.mock import MagicMock, patch
-
-import pytest
-import torch
-from fastapi import FastAPI
-from fastapi.testclient import TestClient
-
-from vllm_omni.entrypoints.openai.protocol.audio import (
-    CreateAudio,
-    OpenAICreateAudioGenerateRequest,
-)
-from vllm_omni.entrypoints.openai.serving_audio_generate import (
-    OmniOpenAIServingAudioGenerate,
-)
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.outputs import OmniRequestOutput
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-logger = logging.getLogger(__name__)
-
-
-# Helper: create a mock audio output for endpoint tests
-def create_mock_audio_output(
-    request_id: str = "audiogen-mock-123",
-    sample_rate: int = 44100,
-    num_samples: int = 44100,
-    audio_key: str = "audio",
-) -> OmniRequestOutput:
-    """Return an OmniRequestOutput mimicking diffusion audio model output."""
-
-    audio_tensor = torch.sin(torch.linspace(0, 440 * 2 * torch.pi, num_samples))
-
-    return OmniRequestOutput.from_diffusion(
-        request_id=request_id,
-        images=[],
-        prompt=None,
-        metrics={},
-        multimodal_output={audio_key: audio_tensor, "sr": sample_rate},
-    )
-
-
-def _make_engine_client(*, audio_key: str = "audio", sample_rate: int = 44100):
-    """Build a mock engine client producing audio output."""
-    mock_engine_client = MagicMock()
-    mock_engine_client.errored = False
-    mock_engine_client.model_type = "StableAudioPipeline"
-    mock_engine_client.default_sampling_params_list = [{}]
-
-    async def mock_generate_fn(*args, **kwargs):
-        yield create_mock_audio_output(
-            request_id=kwargs.get("request_id", "audiogen-mock"),
-            sample_rate=sample_rate,
-            audio_key=audio_key,
-        )
-
-    mock_engine_client.generate = MagicMock(side_effect=mock_generate_fn)
-    return mock_engine_client
-
-
-def _make_server(engine_client=None):
-    """Build an OmniOpenAIServingAudioGenerate with mocks."""
-    if engine_client is None:
-        engine_client = _make_engine_client()
-
-    mock_models = MagicMock()
-    mock_models.is_base_model.return_value = True
-
-    return OmniOpenAIServingAudioGenerate(
-        engine_client=engine_client,
-        models=mock_models,
-        request_logger=MagicMock(),
-    )
-
-
-@pytest.fixture
-def test_app():
-    server = _make_server()
-
-    original_fn = server.create_audio_generate
-    sig = signature(original_fn)
-    new_params = [p for name, p in sig.parameters.items() if name != "raw_request"]
-    new_sig = Signature(parameters=new_params, return_annotation=sig.return_annotation)
-
-    async def patched_create_audio_generate(*args, **kwargs):
-        return await original_fn(*args, **kwargs)
-
-    patched_create_audio_generate.__signature__ = new_sig
-    server.create_audio_generate = patched_create_audio_generate
-
-    app = FastAPI()
-    app.add_api_route(
-        "/v1/audio/generate",
-        server.create_audio_generate,
-        methods=["POST"],
-        response_model=None,
-    )
-    return app
-
-
-@pytest.fixture
-def client(test_app):
-    return TestClient(test_app)
-
-
-# Request Validation (Pydantic model)
-class TestRequestValidation:
-    """Validate OpenAICreateAudioGenerateRequest pydantic constraints."""
-
-    def test_valid_minimal_request(self):
-        req = OpenAICreateAudioGenerateRequest(input="A calm piano melody")
-        assert req.input == "A calm piano melody"
-        assert req.response_format == "wav"
-        assert req.speed == 1.0
-
-    def test_fields_are_wired_correctly(self):
-        req = OpenAICreateAudioGenerateRequest(
-            input="rain sounds",
-            model="stable-audio",
-            response_format="flac",
-            speed=1.5,
-            audio_length=10.0,
-            audio_start=2.0,
-            negative_prompt="noise",
-            guidance_scale=7.5,
-            num_inference_steps=100,
-            seed=42,
-        )
-        assert req.input == "rain sounds"
-        assert req.model == "stable-audio"
-        assert req.response_format == "flac"
-        assert req.speed == 1.5
-        assert req.audio_length == 10.0
-        assert req.audio_start == 2.0
-        assert req.negative_prompt == "noise"
-        assert req.guidance_scale == 7.5
-        assert req.num_inference_steps == 100
-        assert req.seed == 42
-
-    def test_invalid_response_format(self):
-        with pytest.raises(Exception):
-            OpenAICreateAudioGenerateRequest(input="test", response_format="invalid_format")
-
-    def test_speed_lower_bound(self):
-        with pytest.raises(Exception):
-            OpenAICreateAudioGenerateRequest(input="test", speed=0.1)
-
-    def test_speed_upper_bound(self):
-        with pytest.raises(Exception):
-            OpenAICreateAudioGenerateRequest(input="test", speed=5.0)
-
-    def test_speed_at_boundaries(self):
-        req_low = OpenAICreateAudioGenerateRequest(input="test", speed=0.25)
-        assert req_low.speed == 0.25
-        req_high = OpenAICreateAudioGenerateRequest(input="test", speed=4.0)
-        assert req_high.speed == 4.0
-
-    def test_stream_format_sse_rejected(self):
-        with pytest.raises(Exception):
-            OpenAICreateAudioGenerateRequest(input="test", stream_format="sse")
-
-    def test_stream_format_audio_accepted(self):
-        req = OpenAICreateAudioGenerateRequest(input="test", stream_format="audio")
-        assert req.stream_format == "audio"
-
-
-# Constructor & Class Methods
-class TestConstructor:
-    def test_default_init(self):
-        server = _make_server()
-        assert server.diffusion_mode is False
-
-    def test_for_diffusion_factory(self):
-        engine_client = _make_engine_client()
-        mock_models = MagicMock()
-        mock_models.is_base_model.return_value = True
-
-        server = OmniOpenAIServingAudioGenerate.for_diffusion(
-            engine_client=engine_client,
-            models=mock_models,
-            request_logger=MagicMock(),
-        )
-        assert server.diffusion_mode is True
-
-    def test_is_stable_audio_model_true(self):
-        server = _make_server()
-        assert server._is_stable_audio_model() is True
-
-    def test_is_stable_audio_model_false(self):
-        engine = _make_engine_client()
-        engine.model_type = "SomeOtherModel"
-        server = _make_server(engine_client=engine)
-        assert server._is_stable_audio_model() is False
-
-
-# Parameter Wiring — verify request params reach the engine
-class TestParameterWiring:
-    """Ensure request parameters are correctly forwarded to the engine."""
-
-    @pytest.fixture
-    def server_and_engine(self):
-        engine = _make_engine_client()
-        server = _make_server(engine_client=engine)
-        return server, engine
-
-    @pytest.mark.asyncio
-    async def test_prompt_wiring(self, server_and_engine):
-        server, engine = server_and_engine
-        req = OpenAICreateAudioGenerateRequest(input="birds chirping")
-        await server.create_audio_generate(req)
-
-        engine.generate.assert_called_once()
-        call_kwargs = engine.generate.call_args[1]
-        assert call_kwargs["prompt"]["prompt"] == "birds chirping"
-        assert call_kwargs["output_modalities"] == ["audio"]
-
-    @pytest.mark.asyncio
-    async def test_negative_prompt_wiring(self, server_and_engine):
-        server, engine = server_and_engine
-        req = OpenAICreateAudioGenerateRequest(input="a calm ocean", negative_prompt="noise distortion")
-        await server.create_audio_generate(req)
-
-        call_kwargs = engine.generate.call_args[1]
-        assert call_kwargs["prompt"]["negative_prompt"] == "noise distortion"
-
-    @pytest.mark.asyncio
-    async def test_negative_prompt_absent(self, server_and_engine):
-        server, engine = server_and_engine
-        req = OpenAICreateAudioGenerateRequest(input="a calm ocean")
-        await server.create_audio_generate(req)
-
-        call_kwargs = engine.generate.call_args[1]
-        assert "negative_prompt" not in call_kwargs["prompt"]
-
-    @pytest.mark.asyncio
-    async def test_guidance_scale_wiring(self, server_and_engine):
-        server, engine = server_and_engine
-        req = OpenAICreateAudioGenerateRequest(input="test", guidance_scale=12.0)
-        await server.create_audio_generate(req)
-
-        call_kwargs = engine.generate.call_args[1]
-        sp = call_kwargs["sampling_params_list"][0]
-        assert isinstance(sp, OmniDiffusionSamplingParams)
-        assert sp.guidance_scale == 12.0
-
-    @pytest.mark.asyncio
-    async def test_num_inference_steps_wiring(self, server_and_engine):
-        server, engine = server_and_engine
-        req = OpenAICreateAudioGenerateRequest(input="test", num_inference_steps=200)
-        await server.create_audio_generate(req)
-
-        sp = engine.generate.call_args[1]["sampling_params_list"][0]
-        assert sp.num_inference_steps == 200
-
-    @pytest.mark.asyncio
-    async def test_seed_creates_generator(self, server_and_engine):
-        server, engine = server_and_engine
-        req = OpenAICreateAudioGenerateRequest(input="test", seed=42)
-
-        with patch("vllm_omni.entrypoints.openai.serving_audio_generate.torch") as mock_torch:
-            mock_gen = MagicMock()
-            mock_gen.manual_seed.return_value = mock_gen
-            mock_torch.Generator.return_value = mock_gen
-
-            await server.create_audio_generate(req)
-
-            mock_torch.Generator.assert_called_once()
-            mock_gen.manual_seed.assert_called_once_with(42)
-
-    @pytest.mark.asyncio
-    async def test_seed_none_skips_generator(self, server_and_engine):
-        server, engine = server_and_engine
-        req = OpenAICreateAudioGenerateRequest(input="test")
-
-        await server.create_audio_generate(req)
-
-        sp = engine.generate.call_args[1]["sampling_params_list"][0]
-        assert sp.generator is None
-
-    @pytest.mark.asyncio
-    async def test_audio_length_wiring(self, server_and_engine):
-        server, engine = server_and_engine
-        req = OpenAICreateAudioGenerateRequest(input="test", audio_length=10.0, audio_start=2.0)
-        await server.create_audio_generate(req)
-
-        sp = engine.generate.call_args[1]["sampling_params_list"][0]
-        assert sp.extra_args["audio_start_in_s"] == 2.0
-        assert sp.extra_args["audio_end_in_s"] == 12.0  # start + length
-
-    @pytest.mark.asyncio
-    async def test_audio_length_default_start(self, server_and_engine):
-        server, engine = server_and_engine
-        req = OpenAICreateAudioGenerateRequest(input="test", audio_length=5.0)
-        await server.create_audio_generate(req)
-
-        sp = engine.generate.call_args[1]["sampling_params_list"][0]
-        assert sp.extra_args["audio_start_in_s"] == 0.0
-        assert sp.extra_args["audio_end_in_s"] == 5.0
-
-    @pytest.mark.asyncio
-    async def test_no_audio_length_skips_extra_args(self, server_and_engine):
-        server, engine = server_and_engine
-        req = OpenAICreateAudioGenerateRequest(input="test")
-        await server.create_audio_generate(req)
-
-        sp = engine.generate.call_args[1]["sampling_params_list"][0]
-        assert sp.extra_args == {}
-
-    @pytest.mark.asyncio
-    async def test_defaults_not_set_when_omitted(self, server_and_engine):
-        """Guidance scale and num_inference_steps keep dataclass defaults when not in request."""
-        server, engine = server_and_engine
-        req = OpenAICreateAudioGenerateRequest(input="test")
-        await server.create_audio_generate(req)
-
-        sp = engine.generate.call_args[1]["sampling_params_list"][0]
-        defaults = OmniDiffusionSamplingParams()
-        assert sp.guidance_scale == defaults.guidance_scale
-        assert sp.num_inference_steps == defaults.num_inference_steps
-
-
-# Audio Response Format
-class TestAudioResponseFormat:
-    def test_wav_response(self, client):
-        payload = {"input": "a gentle rain", "response_format": "wav"}
-        response = client.post("/v1/audio/generate", json=payload)
-        assert response.status_code == 200
-        assert response.headers["content-type"] == "audio/wav"
-        assert len(response.content) > 0
-
-    def test_mp3_response(self, client):
-        payload = {"input": "a gentle rain", "response_format": "mp3"}
-        response = client.post("/v1/audio/generate", json=payload)
-        assert response.status_code == 200
-        assert response.headers["content-type"] == "audio/mpeg"
-        assert len(response.content) > 0
-
-    def test_flac_response(self, client):
-        payload = {"input": "a gentle rain", "response_format": "flac"}
-        response = client.post("/v1/audio/generate", json=payload)
-        assert response.status_code == 200
-        assert response.headers["content-type"] == "audio/flac"
-        assert len(response.content) > 0
-
-    def test_invalid_format_rejected(self, client):
-        payload = {"input": "test", "response_format": "banana"}
-        response = client.post("/v1/audio/generate", json=payload)
-        assert response.status_code == 422
-
-    @patch("vllm_omni.entrypoints.openai.serving_audio_generate.OmniOpenAIServingAudioGenerate.create_audio")
-    def test_speed_parameter_forwarded(self, mock_create_audio, test_app):
-        mock_audio_response = MagicMock()
-        mock_audio_response.audio_data = b"dummy_audio"
-        mock_audio_response.media_type = "audio/wav"
-        mock_create_audio.return_value = mock_audio_response
-
-        c = TestClient(test_app)
-        payload = {"input": "test", "response_format": "wav", "speed": 2.5}
-        c.post("/v1/audio/generate", json=payload)
-
-        mock_create_audio.assert_called_once()
-        audio_obj = mock_create_audio.call_args[0][0]
-        assert isinstance(audio_obj, CreateAudio)
-        assert audio_obj.speed == 2.5
-
-    @patch("vllm_omni.entrypoints.openai.serving_audio_generate.OmniOpenAIServingAudioGenerate.create_audio")
-    def test_sample_rate_from_output(self, mock_create_audio, test_app):
-        mock_audio_response = MagicMock()
-        mock_audio_response.audio_data = b"dummy"
-        mock_audio_response.media_type = "audio/wav"
-        mock_create_audio.return_value = mock_audio_response
-
-        c = TestClient(test_app)
-        payload = {"input": "test"}
-        c.post("/v1/audio/generate", json=payload)
-
-        audio_obj = mock_create_audio.call_args[0][0]
-        assert audio_obj.sample_rate == 44100  # Stable Audio default
-
-
-# Error Handling
-class TestErrorHandling:
-    @pytest.mark.asyncio
-    async def test_no_output_returns_error(self):
-        engine = _make_engine_client()
-
-        async def empty_gen(*args, **kwargs):
-            return
-            yield  # noqa: unreachable – makes this an async generator
-
-        engine.generate = MagicMock(side_effect=empty_gen)
-        server = _make_server(engine_client=engine)
-        req = OpenAICreateAudioGenerateRequest(input="test")
-        resp = await server.create_audio_generate(req)
-
-        # create_error_response returns an ErrorResponse with .error.message
-        assert "No output generated" in resp.error.message
-
-    @pytest.mark.asyncio
-    async def test_no_audio_in_output_returns_error(self):
-        engine = _make_engine_client()
-
-        async def gen_without_audio(*args, **kwargs):
-            yield OmniRequestOutput.from_diffusion(
-                request_id="test",
-                images=[],
-                prompt=None,
-                metrics={},
-                multimodal_output={},  # no audio key
-            )
-
-        engine.generate = MagicMock(side_effect=gen_without_audio)
-        server = _make_server(engine_client=engine)
-        req = OpenAICreateAudioGenerateRequest(input="test")
-        resp = await server.create_audio_generate(req)
-
-        assert "did not produce audio" in resp.error.message
-
-    @pytest.mark.asyncio
-    async def test_engine_errored_raises(self):
-        engine = _make_engine_client()
-        engine.errored = True
-        engine.dead_error = RuntimeError("engine is dead")
-        server = _make_server(engine_client=engine)
-
-        req = OpenAICreateAudioGenerateRequest(input="test")
-        with pytest.raises(RuntimeError, match="engine is dead"):
-            await server.create_audio_generate(req)
-
-    @pytest.mark.asyncio
-    async def test_model_outputs_key_fallback(self):
-        """Audio data under 'model_outputs' key should be accepted."""
-        engine = _make_engine_client(audio_key="model_outputs")
-        server = _make_server(engine_client=engine)
-        req = OpenAICreateAudioGenerateRequest(input="test")
-        resp = await server.create_audio_generate(req)
-
-        # Should succeed and return a Response with audio bytes
-        assert hasattr(resp, "body")
-        assert len(resp.body) > 0
-
-    @pytest.mark.asyncio
-    async def test_value_error_returns_error_response(self):
-        engine = _make_engine_client()
-
-        async def gen_value_error(*args, **kwargs):
-            raise ValueError("bad value")
-            yield  # noqa: unreachable
-
-        engine.generate = MagicMock(side_effect=gen_value_error)
-        server = _make_server(engine_client=engine)
-        req = OpenAICreateAudioGenerateRequest(input="test")
-        resp = await server.create_audio_generate(req)
-
-        assert "bad value" in resp.error.message
-
-    @pytest.mark.asyncio
-    async def test_generic_exception_returns_error_response(self):
-        engine = _make_engine_client()
-
-        async def gen_runtime_error(*args, **kwargs):
-            raise RuntimeError("something went wrong")
-            yield  # noqa: unreachable
-
-        engine.generate = MagicMock(side_effect=gen_runtime_error)
-        server = _make_server(engine_client=engine)
-        req = OpenAICreateAudioGenerateRequest(input="test")
-        resp = await server.create_audio_generate(req)
-
-        assert "Audio generation failed" in resp.error.message
-
-
-# End-to-End via TestClient
-class TestAudioGenerateAPI:
-    def test_basic_success(self, client):
-        payload = {"input": "ambient forest sounds"}
-        response = client.post("/v1/audio/generate", json=payload)
-        assert response.status_code == 200
-        assert len(response.content) > 0
-
-    def test_with_all_params(self, client):
-        payload = {
-            "input": "gentle piano",
-            "response_format": "wav",
-            "speed": 1.0,
-            "audio_length": 5.0,
-            "audio_start": 0.0,
-            "negative_prompt": "noise",
-            "guidance_scale": 7.0,
-            "num_inference_steps": 50,
-            "seed": 123,
-        }
-        response = client.post("/v1/audio/generate", json=payload)
-        assert response.status_code == 200
-        assert response.headers["content-type"] == "audio/wav"
-
-    def test_missing_input_rejected(self, client):
-        payload = {}
-        response = client.post("/v1/audio/generate", json=payload)
-        assert response.status_code == 422
-
-    def test_extra_unknown_fields_ignored(self, client):
-        payload = {"input": "test", "unknown_field": "value"}
-        response = client.post("/v1/audio/generate", json=payload)
-        # Pydantic v2 ignores extra fields by default
-        assert response.status_code == 200

tests/entrypoints/openai_api/test_serving_speech.py

@@ -1017,3 +1017,27 @@ class TestStreamingResponse:
         response = client.post("/v1/audio/speech", json={"input": "Hello", "response_format": "wav"})
         assert response.status_code == 200
         assert "audio/wav" in response.headers["content-type"]
+
+
+class TestAsyncOmniSupportedTasks:
+    """Test that AsyncOmni reports correct supported tasks based on output modalities."""
+
+    @pytest.mark.asyncio
+    async def test_tts_only_no_generate_task(self):
+        """TTS-only models (audio output, no text) should not include 'generate'."""
+        from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+        omni = AsyncOmni.__new__(AsyncOmni)
+        omni.output_modalities = [None, "audio"]
+        tasks = await omni.get_supported_tasks()
+        assert "generate" not in tasks
+
+    @pytest.mark.asyncio
+    async def test_omni_model_includes_generate(self):
+        """Models with text output (e.g. Qwen3-Omni) should include 'generate'."""
+        from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+        omni = AsyncOmni.__new__(AsyncOmni)
+        omni.output_modalities = ["text", None, "audio"]
+        tasks = await omni.get_supported_tasks()
+        assert "generate" in tasks

tests/entrypoints/test_omni_diffusion.py

@@ -612,7 +612,7 @@ def test_initialize_stage_configs_called_when_none(
     """Test that stage configs are auto-loaded when stage_configs_path is None."""
 
     def _fake_loader(
-        config_path: str,
+        model: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,
@@ -687,7 +687,7 @@ def test_generate_raises_on_length_mismatch(monkeypatch: pytest.MonkeyPatch, moc
     """Test that generate raises ValueError when sampling_params_list length doesn't match."""
 
     def _fake_loader(
-        config_path: str,
+        model: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,
@@ -742,7 +742,7 @@ def test_generate_pipeline_and_final_outputs(monkeypatch: pytest.MonkeyPatch, mo
     stage_cfg1["processed_input"] = ["processed-for-stage-1"]
 
     def _fake_loader(
-        config_path: str,
+        model: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,
@@ -848,7 +848,7 @@ def test_generate_pipeline_with_batch_input(monkeypatch: pytest.MonkeyPatch, moc
     stage_cfg1["stage_id"] = 1
 
     def _fake_loader(
-        config_path: str,
+        model: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,
@@ -969,7 +969,7 @@ def test_generate_no_final_output_returns_empty(
     stage_cfg1["final_output"] = False
 
     def _fake_loader(
-        config_path: str,
+        model: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,
@@ -1062,7 +1062,7 @@ def test_generate_sampling_params_none_use_default(
     stage_cfg1["final_output"] = False
 
     def _fake_loader(
-        config_path: str,
+        model: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,
@@ -1140,7 +1140,7 @@ def test_wait_for_stages_ready_timeout(monkeypatch: pytest.MonkeyPatch, mocker:
     """Test that _wait_for_stages_ready handles timeout correctly."""
 
     def _fake_loader(
-        config_path: str,
+        model: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,
@@ -1202,7 +1202,7 @@ def test_generate_handles_error_messages(monkeypatch: pytest.MonkeyPatch, mocker
     """Test that generate handles error messages from stages correctly."""
 
     def _fake_loader(
-        config_path: str,
+        model: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,
@@ -1283,7 +1283,7 @@ def test_close_sends_shutdown_signal(monkeypatch: pytest.MonkeyPatch, mocker: Mo
     """Test that close() sends shutdown signal to all input queues."""
 
     def _fake_loader(
-        config_path: str,
+        model: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,

tests/entrypoints/test_omni_llm.py

@@ -577,7 +577,7 @@ def test_initialize_stage_configs_called_when_none(
     """Test that stage configs are auto-loaded when stage_configs_path is None."""
 
     def _fake_loader(
-        config_path: str,
+        model: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,
@@ -654,7 +654,7 @@ def test_generate_raises_on_length_mismatch(monkeypatch: pytest.MonkeyPatch, moc
     """Test that generate raises ValueError when sampling_params_list length doesn't match."""
 
     def _fake_loader(
-        config_path: str,
+        model: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,
@@ -711,7 +711,7 @@ def test_generate_pipeline_and_final_outputs(monkeypatch: pytest.MonkeyPatch, mo
     stage_cfg1["processed_input"] = ["processed-for-stage-1"]
 
     def _fake_loader(
-        config_path: str,
+        model: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,
@@ -819,7 +819,7 @@ def test_generate_no_final_output_returns_empty(
     stage_cfg1["final_output"] = False
 
     def _fake_loader(
-        config_path: str,
+        model: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,
@@ -911,7 +911,7 @@ def test_generate_sampling_params_none_use_default(
     stage_cfg1["final_output"] = False
 
     def _fake_loader(
-        config_path: str,
+        model: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,
@@ -988,7 +988,7 @@ def test_wait_for_stages_ready_timeout(monkeypatch: pytest.MonkeyPatch, mocker:
     """Test that _wait_for_stages_ready handles timeout correctly."""
 
     def _fake_loader(
-        config_path: str,
+        model: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,
@@ -1052,7 +1052,7 @@ def test_generate_handles_error_messages(monkeypatch: pytest.MonkeyPatch, mocker
     """Test that generate handles error messages from stages correctly."""
 
     def _fake_loader(
-        config_path: str,
+        model: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,
@@ -1135,7 +1135,7 @@ def test_close_sends_shutdown_signal(monkeypatch: pytest.MonkeyPatch, mocker: Mo
     """Test that close() sends shutdown signal to all input queues."""
 
     def _fake_loader(
-        config_path: str,
+        model: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,

tests/examples/online_serving/test_qwen2_5_omni.py

@@ -12,7 +12,7 @@ from pathlib import Path
 
 import pytest
 
-from tests.conftest import convert_audio_file_to_text, cosine_similarity_text
+from tests.conftest import OmniServerParams, convert_audio_file_to_text, cosine_similarity_text
 from tests.utils import hardware_test
 
 models = ["Qwen/Qwen2.5-Omni-7B"]
@@ -22,7 +22,11 @@ stage_configs = [str(Path(__file__).parent.parent.parent / "e2e" / "stage_config
 
 example_dir = str(Path(__file__).parent.parent.parent.parent / "examples" / "online_serving" / "qwen2_5_omni")
 # Create parameter combinations for model and stage config
-test_params = [(8091, model, stage_config) for model in models for stage_config in stage_configs]
+test_params = [
+    OmniServerParams(model=model, port=8091, stage_config_path=stage_config)
+    for model in models
+    for stage_config in stage_configs
+]
 
 
 def run_cmd(command):

tests/examples/online_serving/test_qwen3_omni.py

@@ -12,7 +12,7 @@ from pathlib import Path
 
 import pytest
 
-from tests.conftest import convert_audio_file_to_text, cosine_similarity_text
+from tests.conftest import OmniServerParams, convert_audio_file_to_text, cosine_similarity_text
 from tests.utils import hardware_test
 
 models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
@@ -22,7 +22,11 @@ stage_configs = [str(Path(__file__).parent.parent.parent / "e2e" / "stage_config
 
 example_dir = str(Path(__file__).parent.parent.parent.parent / "examples" / "online_serving" / "qwen3_omni")
 # Create parameter combinations for model and stage config
-test_params = [(8091, model, stage_config) for model in models for stage_config in stage_configs]
+test_params = [
+    OmniServerParams(model=model, port=8091, stage_config_path=stage_config)
+    for model in models
+    for stage_config in stage_configs
+]
 
 
 def run_cmd(command):

tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_embed.py

@@ -19,6 +19,8 @@ from vllm.model_executor.models.qwen2_5_omni_thinker import (
     merge_interleaved_embeddings,
 )
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 # Fake token IDs
 AUDIO_TOKEN_ID = 1001
 IMAGE_TOKEN_ID = 1002

tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py

@@ -7,16 +7,20 @@ from types import SimpleNamespace
 import pytest
 import torch
 
+from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import (
+    compute_dynamic_initial_chunk_size,
+    max_ic_for_chunk_size,
+)
 from vllm_omni.model_executor.stage_input_processors.qwen3_tts import (
     talker2code2wav,
     talker2code2wav_async_chunk,
 )
 
-_FRAME = [1, 2, 3, 4]  # 4-codebook frame
-_Q = len(_FRAME)  # num quantizers
+_FRAME = [1, 2, 3, 4]
+_Q = len(_FRAME)
 
 
-def _req(rid: str, *, finished: bool, initial_codec_chunk_frames: int | None = None):
+def _req(rid, *, finished, initial_codec_chunk_frames=None):
     ai = None
     if initial_codec_chunk_frames is not None:
         entry = SimpleNamespace(list_data=[initial_codec_chunk_frames])
@@ -28,9 +32,10 @@ def _req(rid: str, *, finished: bool, initial_codec_chunk_frames: int | None = N
     )
 
 
-def _tm(*, chunk_frames=25, left_context=25, initial_chunk=0):
+def _tm(*, chunk_frames=25, left_context=25, max_num_seqs=1):
     return SimpleNamespace(
         code_prompt_token_ids=defaultdict(list),
+        scheduler_max_num_seqs=max_num_seqs,
         put_req_chunk=defaultdict(int),
         request_payload={},
         connector=SimpleNamespace(
@@ -38,16 +43,14 @@ def _tm(*, chunk_frames=25, left_context=25, initial_chunk=0):
                 "extra": {
                     "codec_chunk_frames": chunk_frames,
                     "codec_left_context_frames": left_context,
-                    "initial_codec_chunk_frames": initial_chunk,
                 }
             }
         ),
     )
 
 
-def _call(tm, rid, *, n_frames, put_req=0, finished=False, req_ic=None):
+def _call(tm, rid, *, n_frames, finished=False, req_ic=None):
     tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(n_frames)]
-    tm.put_req_chunk[rid] = put_req
     return talker2code2wav_async_chunk(
         transfer_manager=tm,
         pooling_output={"audio_codes": torch.zeros((0,))},
@@ -56,84 +59,68 @@ def _call(tm, rid, *, n_frames, put_req=0, finished=False, req_ic=None):
     )
 
 
-def test_does_not_emit_empty_chunk_when_not_finished():
+def test_empty_returns_none():
     tm = _tm()
-    request = _req("rid-empty", finished=False)
-
-    payload = talker2code2wav_async_chunk(
+    p = talker2code2wav_async_chunk(
         transfer_manager=tm,
         pooling_output={"audio_codes": torch.zeros((0,))},
-        request=request,
+        request=_req("r", finished=False),
     )
-
-    assert payload is None
+    assert p is None
 
 
-def test_flushes_tail_when_finished_without_pooler_output():
+def test_eof_marker_when_finished_empty():
     tm = _tm()
-    rid = "rid-tail"
-    tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(24)]
-    request = _req(rid, finished=True)
-
-    payload = talker2code2wav_async_chunk(
+    p = talker2code2wav_async_chunk(
         transfer_manager=tm,
         pooling_output=None,
-        request=request,
+        request=_req("r", finished=True),
+        is_finished=True,
     )
-
-    assert payload is not None
-    assert payload["finished"].item() is True
-    assert len(payload["code_predictor_codes"]) == _Q * 24
+    assert p == {"code_predictor_codes": [], "finished": torch.tensor(True, dtype=torch.bool)}
 
 
-def test_emits_eof_marker_when_finished_with_no_frames():
+def test_flush_on_finish():
     tm = _tm()
-    request = _req("rid-eof", finished=True)
-
-    payload = talker2code2wav_async_chunk(
+    tm.code_prompt_token_ids["r"] = [_FRAME[:] for _ in range(24)]
+    p = talker2code2wav_async_chunk(
         transfer_manager=tm,
         pooling_output=None,
-        request=request,
+        request=_req("r", finished=True),
+        is_finished=True,
     )
-
-    assert payload == {
-        "code_predictor_codes": [],
-        "finished": torch.tensor(True, dtype=torch.bool),
-    }
+    assert p is not None
+    assert p["finished"].item() is True
+    assert len(p["code_predictor_codes"]) == _Q * 24
 
 
 _CASES = [
-    # Normal path (initial=0): emit at chunk_size boundaries
-    ((25, 25, 0), (24, 0, False), None),
-    ((25, 25, 0), (25, 0, False), (0, 25)),
-    # Initial-chunk phase: hold, first emit, second emit
-    ((25, 25, 10), (9, 0, False), None),
-    ((25, 25, 10), (10, 0, False), (0, 10)),
-    ((25, 25, 10), (20, 1, False), (10, 20)),
-    # Non-divisible: holds at chunk boundary
-    ((25, 25, 12), (25, 2, False), None),
-    # Normal phase: offset by initial_coverage (chunk//initial * initial)
-    ((25, 25, 10), (45, 2, False), (20, 45)),
-    # Second normal emit (offset must stay stable)
-    ((25, 25, 10), (70, 3, False), (25, 50)),
-    # initial >= chunk clamps to chunk_size (behaves as normal)
-    ((25, 25, 30), (25, 0, False), (0, 25)),
-    # finished=True flushes IC tail
-    ((25, 25, 10), (5, 0, True), (0, 5)),
-    # finished=True flushes non-divisible IC residual
-    ((25, 25, 12), (25, 2, True), (24, 25)),
-    # finished=True flushes normal phase tail
-    ((25, 25, 10), (30, 2, True), (20, 30)),
+    # Dynamic IC=16, cs=25, initial_coverage=16
+    ((25, 25, 0), 24, False, None),  # IC phase: 24%16!=0 -> hold
+    ((25, 25, 0), 25, False, None),  # transition: adjusted=9, hold (no replay)
+    ((25, 25, 0), 41, False, (16, 41)),  # first normal emit, lc=16
+    # Per-request IC=10, cs=25, initial_coverage=20
+    ((25, 25, 10), 9, False, None),  # IC: hold
+    ((25, 25, 10), 10, False, (0, 10)),  # IC: emit at boundary
+    ((25, 25, 10), 25, False, None),  # transition: hold (no replay)
+    ((25, 25, 10), 45, False, (20, 45)),  # first normal emit, lc=20
+    ((25, 25, 10), 5, True, (0, 5)),  # finished flushes IC tail
+    ((25, 25, 10), 33, True, (20, 33)),  # finished flushes normal tail
+    # IC=8, cs=16: IC divides chunk_size evenly (edge case)
+    ((16, 25, 8), 8, False, (0, 8)),  # IC: emit
+    ((16, 25, 8), 16, False, None),  # transition: hold (no replay)
+    ((16, 25, 8), 24, False, (8, 24)),  # first normal emit, lc=8
+    # Per-request override: IC=15 at n_frames=10 -> 10%15!=0 -> hold
+    ((25, 25, 15), 10, False, None),
 ]
 
 
-@pytest.mark.parametrize("config, state, expected", _CASES)
-def test_streaming_decoding_with_variable_initial(config, state, expected):
-    chunk_frames, left_context, initial_chunk = config
-    n_frames, put_req, finished = state
-
-    tm = _tm(chunk_frames=chunk_frames, left_context=left_context, initial_chunk=initial_chunk)
-    payload = _call(tm, "r", n_frames=n_frames, put_req=put_req, finished=finished)
+@pytest.mark.parametrize("config, n_frames, finished, expected", _CASES)
+def test_streaming_phases(config, n_frames, finished, expected):
+    chunk_frames, left_context, req_ic_val = config
+    tm = _tm(chunk_frames=chunk_frames, left_context=left_context)
+    req_ic = req_ic_val if req_ic_val > 0 else None
+    payload = _call(tm, "r", n_frames=n_frames, finished=finished, req_ic=req_ic)
 
     if expected is None:
         assert payload is None
@@ -144,22 +131,86 @@ def test_streaming_decoding_with_variable_initial(config, state, expected):
         assert len(payload["code_predictor_codes"]) == _Q * exp_window
 
 
-def test_per_request_override_activates_initial_phase():
-    tm = _tm(initial_chunk=0)
-    payload = _call(tm, "r-override", n_frames=10, req_ic=10)
-    assert payload is not None
-    assert payload["left_context_size"] == 0
-    assert len(payload["code_predictor_codes"]) == _Q * 10
+def test_dynamic_ic_adapts_to_load():
+    # chunk_size=25 -> max_ic=16, steps=[2,4,8,16]
+    tm = _tm(max_num_seqs=8)
+
+    # Low load (1/8) -> IC=2 -> emit at 2
+    p1 = _call(tm, "r", n_frames=2)
+    assert p1 is not None
+    assert len(p1["code_predictor_codes"]) == _Q * 2
+
+    # High load: add 4 others -> active=5/8 -> IC=8 -> emit at 8
+    for i in range(4):
+        tm.code_prompt_token_ids[f"other-{i}"] = [[0]]
+    p2 = _call(tm, "r", n_frames=8)
+    assert p2 is not None
+    assert len(p2["code_predictor_codes"]) == _Q * 8
+
+    # Requests past initial phase still count in load factor
+    tm2 = _tm(max_num_seqs=4)
+    for i in range(3):
+        tm2.code_prompt_token_ids[f"long-{i}"] = [[0]] * 50  # well past cs=25
+    # active=4/4=1.0 -> IC=16
+    p3 = _call(tm2, "new", n_frames=16)
+    assert p3 is not None
+    assert len(p3["code_predictor_codes"]) == _Q * 16
 
 
-def test_per_request_override_wins_over_stage_config():
-    tm = _tm(initial_chunk=5)
-    payload = _call(tm, "r-override2", n_frames=10, put_req=0, req_ic=15)
-    assert payload is None
+def test_ic_load_change_mid_request():
+    """IC stateless: load spike mid-request shifts initial_coverage."""
+    tm = _tm(chunk_frames=25, left_context=25, max_num_seqs=8)
+
+    # Low load -> IC=2 -> emit at frame 2
+    p1 = _call(tm, "r", n_frames=2)
+    assert p1 is not None
+
+    # Spike load: 6 others -> IC=16 -> initial_coverage=16
+    for i in range(6):
+        tm.code_prompt_token_ids[f"other-{i}"] = [[0]] * 10
+
+    # adjusted=25-16=9, 9%25!=0 -> hold
+    assert _call(tm, "r", n_frames=25) is None
+
+    # First normal emit at 16+25=41
+    p3 = _call(tm, "r", n_frames=41)
+    assert p3 is not None
+    assert p3["left_context_size"] == 16
+
+
+@pytest.mark.parametrize(
+    "active,max_bs,max_ic,expected",
+    [
+        (0, 4, 32, 2),  # zero load -> min step
+        (2, 4, 32, 8),  # mid load
+        (4, 4, 32, 32),  # full load
+        (10, 4, 16, 16),  # over capacity, capped
+        (0, 4, 1, 1),  # max_ic below min step
+        (0, 0, 16, 2),  # zero capacity edge case
+    ],
+)
+def test_compute_dynamic_initial_chunk_size(active, max_bs, max_ic, expected):
+    assert compute_dynamic_initial_chunk_size(active, max_bs, max_ic) == expected
+
+
+@pytest.mark.parametrize(
+    "chunk_size,expected",
+    [
+        (25, 16),
+        (50, 32),
+        (70, 64),
+        (8, 4),
+        (4, 2),
+        (2, 1),
+        (1, 1),
+    ],
+)
+def test_max_ic_for_chunk_size(chunk_size, expected):
+    assert max_ic_for_chunk_size(chunk_size) == expected
 
 
 def test_first_streaming_chunk_prepends_ref_code_context():
-    tm = _tm(initial_chunk=10)
+    tm = _tm()
     rid = "r-ref"
     tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(10)]
     ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
@@ -177,7 +228,7 @@ def test_first_streaming_chunk_prepends_ref_code_context():
 
 
 def test_ref_code_context_only_applies_to_first_streaming_chunk():
-    tm = _tm(initial_chunk=10)
+    tm = _tm()
     rid = "r-ref2"
     tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(20)]
     tm.put_req_chunk[rid] = 1
@@ -196,7 +247,7 @@ def test_ref_code_context_only_applies_to_first_streaming_chunk():
 
 
 def test_ref_code_context_can_be_buffered_before_first_emit():
-    tm = _tm(initial_chunk=10)
+    tm = _tm()
     rid = "r-ref-buffered"
     ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
 

tests/perf/scripts/run_benchmark.py

@@ -107,8 +107,14 @@ def omni_server(request):
         print("OmniServer stopped")
 
 
-def run_benchmark(args: list, test_name: str, flow, dataset_name: str, num_prompt) -> Any:
-    """Generate synthetic image with random values."""
+def run_benchmark(
+    args: list,
+    test_name: str,
+    flow,
+    dataset_name: str,
+    num_prompt,
+) -> Any:
+    """Run a single benchmark iteration and return the parsed result JSON."""
     current_dt = datetime.now().strftime("%Y%m%d-%H%M%S")
     result_filename = f"result_{test_name}_{dataset_name}_{flow}_{num_prompt}_{current_dt}.json"
     if "--result-filename" in args:
@@ -117,10 +123,6 @@ def run_benchmark(args: list, test_name: str, flow, dataset_name: str, num_promp
         ["vllm", "bench", "serve", "--omni"]
         + args
         + [
-            "--backend",
-            "openai-chat-omni",
-            "--endpoint",
-            "/v1/chat/completions",
             "--save-result",
             "--result-dir",
             os.environ.get("BENCHMARK_DIR", "tests"),
@@ -196,7 +198,10 @@ def benchmark_params(request, omni_server):
     total = len(all_params)
     print(f"\n  Running benchmark {current}/{total} for {test_name}")
 
-    return {"test_name": test_name, "params": all_params[param_index]}
+    return {
+        "test_name": test_name,
+        "params": all_params[param_index],
+    }
 
 
 def assert_result(result, params, num_prompt):
@@ -266,7 +271,11 @@ def test_performance_benchmark(omni_server, benchmark_params):
     for qps, num_prompt in zip(qps_list, num_prompt_list):
         args = args + ["--request-rate", str(qps), "--num-prompts", str(num_prompt)]
         result = run_benchmark(
-            args=args, test_name=test_name, flow=qps, dataset_name=dataset_name, num_prompt=num_prompt
+            args=args,
+            test_name=test_name,
+            flow=qps,
+            dataset_name=dataset_name,
+            num_prompt=num_prompt,
         )
         assert_result(result, params, num_prompt=num_prompt)
 
@@ -274,6 +283,10 @@ def test_performance_benchmark(omni_server, benchmark_params):
     for concurrency, num_prompt in zip(max_concurrency_list, num_prompt_list):
         args = args + ["--max-concurrency", str(concurrency), "--num-prompts", str(num_prompt), "--request-rate", "inf"]
         result = run_benchmark(
-            args=args, test_name=test_name, flow=concurrency, dataset_name=dataset_name, num_prompt=num_prompt
+            args=args,
+            test_name=test_name,
+            flow=concurrency,
+            dataset_name=dataset_name,
+            num_prompt=num_prompt,
         )
         assert_result(result, params, num_prompt=num_prompt)

tests/perf/stage_configs/qwen3_tts.yaml

@@ -0,0 +1,99 @@
+# Stage config for running Qwen3-TTS with 2-stage architecture
+# Stage 0: Talker (text -> 8-layer RVQ codec codes)
+# Stage 1: Code2Wav (codec codes -> audio waveform)
+#
+# The following config has been verified on 1x H100-80G GPU.
+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

tests/perf/tests/test.json

@@ -8,6 +8,8 @@
         "benchmark_params": [
             {
                 "dataset_name": "random",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
                 "num_prompts": [
                     10,
                     40
@@ -28,6 +30,8 @@
             },
             {
                 "dataset_name": "random-mm",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
                 "num_prompts": [
                     10,
                     40
@@ -88,6 +92,8 @@
         "benchmark_params": [
             {
                 "dataset_name": "random",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
                 "num_prompts": [
                     10,
                     40
@@ -108,6 +114,8 @@
             },
             {
                 "dataset_name": "random-mm",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
                 "num_prompts": [
                     10,
                     40
@@ -140,5 +148,38 @@
                 }
             }
         ]
+    },
+    {
+        "test_name": "test_qwen3_tts",
+        "server_params": {
+            "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
+            "stage_config_name": "qwen3_tts.yaml"
+        },
+        "benchmark_params": [
+            {
+                "dataset_name": "random",
+                "backend": "openai-audio-speech",
+                "endpoint": "/v1/audio/speech",
+                "num_prompts": [
+                    10,
+                    40
+                ],
+                "max_concurrency": [
+                    1,
+                    4
+                ],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "extra_body": {
+                    "voice": "Vivian",
+                    "language": "English"
+                },
+                "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
+                "baseline": {
+                    "mean_audio_ttfp_ms": 100000,
+                    "mean_audio_rtf": 100000
+                }
+            }
+        ]
     }
 ]

tests/utils.py

@@ -20,6 +20,8 @@ from typing_extensions import ParamSpec
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import cuda_device_count_stateless
 
+from vllm_omni.platforms import current_omni_platform
+
 _P = ParamSpec("_P")
 
 if current_platform.is_rocm():
@@ -441,27 +443,11 @@ def npu_marks(*, res: str, num_cards: int):
         return [mark for mark in [test_platform, test_resource, test_distributed] if mark is not None]
 
 
-def hardware_test(*, res: dict[str, str], num_cards: int | dict[str, int] = 1):
+def hardware_marks(*, res: dict[str, str], num_cards: int | dict[str, int] = 1):
     """
-    Decorate a test for multiple hardware platforms with a single call.
-    Automatically wraps the test with @create_new_process_for_each_test() for distributed tests.
-
-    Args:
-        res: Mapping from platform to resource type. Supported platforms/resources:
-            - cuda: L4, H100
-            - rocm: MI325
-            - npu: A2, A3
-        num_cards: Number of cards required. Can be:
-            - int: same card count for all platforms (default: 1)
-            - dict: per-platform card count, e.g., {"cuda": 2, "rocm": 2}
-
-    Example:
-        @hardware_test(
-            res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
-            num_cards={"cuda": 2, "rocm": 2, "npu": 2},
-        )
-        def test_multi_platform():
-            ...
+    Get a collection of pytest marks to apply for `@hardware_test`,
+    including CUDA, ROCm, and NPU,
+    based on the specified platforms and resources.
     """
     # Validate platforms
     # Don't validate platform details in this decorator
@@ -484,7 +470,7 @@ def hardware_test(*, res: dict[str, str], num_cards: int | dict[str, int] = 1):
                 num_cards_dict[platform] = 1
 
     # Collect marks from all platforms
-    all_marks: list[Callable[[Callable[_P, None]], Callable[_P, None]]] = []
+    all_marks: list[pytest.MarkDecorator] = []
     for platform, resource in res.items():
         cards = num_cards_dict[platform]
         if platform == "cuda" or platform == "rocm":
@@ -494,6 +480,32 @@ def hardware_test(*, res: dict[str, str], num_cards: int | dict[str, int] = 1):
         else:
             raise ValueError(f"Unsupported platform: {platform}")
         all_marks.extend(marks)
+    return all_marks
+
+
+def hardware_test(*, res: dict[str, str], num_cards: int | dict[str, int] = 1):
+    """
+    Decorate a test for multiple hardware platforms with a single call.
+    Automatically wraps the test with @create_new_process_for_each_test() for distributed tests.
+
+    Args:
+        res: Mapping from platform to resource type. Supported platforms/resources:
+            - cuda: L4, H100
+            - rocm: MI325
+            - npu: A2, A3
+        num_cards: Number of cards required. Can be:
+            - int: same card count for all platforms (default: 1)
+            - dict: per-platform card count, e.g., {"cuda": 2, "rocm": 2}
+
+    Example:
+        @hardware_test(
+            res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
+            num_cards={"cuda": 2, "rocm": 2, "npu": 2},
+        )
+        def test_multi_platform():
+            ...
+    """
+    all_marks = hardware_marks(res=res, num_cards=num_cards)
 
     def wrapper(f: Callable[_P, None]) -> Callable[_P, None]:
         func = f
@@ -504,8 +516,8 @@ def hardware_test(*, res: dict[str, str], num_cards: int | dict[str, int] = 1):
     return wrapper
 
 
-class GPUMemoryMonitor:
-    """Poll global device memory usage via CUDA APIs."""
+class DeviceMemoryMonitor:
+    """Poll global device memory usage."""
 
     def __init__(self, device_index: int, interval: float = 0.05):
         self.device_index = device_index
@@ -518,8 +530,8 @@ class GPUMemoryMonitor:
         def monitor_loop() -> None:
             while not self._stop_event.is_set():
                 try:
-                    with torch.cuda.device(self.device_index):
-                        free_bytes, total_bytes = torch.cuda.mem_get_info()
+                    with current_omni_platform.device(self.device_index):
+                        free_bytes, total_bytes = current_omni_platform.mem_get_info()
                     used_mb = (total_bytes - free_bytes) / (1024**2)
                     self._peak_used_mb = max(self._peak_used_mb, used_mb)
                 except Exception:
@@ -537,8 +549,8 @@ class GPUMemoryMonitor:
 
     @property
     def peak_used_mb(self) -> float:
-        fallback_alloc = torch.cuda.max_memory_allocated(device=self.device_index) / (1024**2)
-        fallback_reserved = torch.cuda.max_memory_reserved(device=self.device_index) / (1024**2)
+        fallback_alloc = current_omni_platform.max_memory_allocated(device=self.device_index) / (1024**2)
+        fallback_reserved = current_omni_platform.max_memory_reserved(device=self.device_index) / (1024**2)
         return max(self._peak_used_mb, fallback_alloc, fallback_reserved)
 
     def __del__(self):

vllm_omni/benchmarks/patch/patch.py

@@ -208,10 +208,11 @@ async def async_request_openai_chat_omni_completions(
 async def async_request_openai_audio_speech(
     request_func_input: RequestFuncInput, session: aiohttp.ClientSession, pbar: tqdm | None = None
 ) -> MixRequestFuncOutput:
-    """Non-streaming request to /v1/audio/speech endpoint.
+    """Streaming request to /v1/audio/speech endpoint.
 
-    The endpoint returns raw audio bytes (e.g. WAV). Pass voice, instructions,
-    and other TTS-specific fields via ``extra_body``.
+    Sends ``stream=true`` with ``response_format=pcm`` so the server returns
+    raw PCM chunks as they are decoded. This allows measuring TTFP (time to
+    first audio packet) separately from E2EL.
     """
     api_url = request_func_input.api_url
     _validate_api_url(api_url, "OpenAI Audio Speech API", "audio/speech")
@@ -219,16 +220,11 @@ async def async_request_openai_audio_speech(
     payload = {
         "model": request_func_input.model_name if request_func_input.model_name else request_func_input.model,
         "input": request_func_input.prompt,
+        "stream": True,
+        "response_format": "pcm",
     }
     _update_payload_common(payload, request_func_input)
 
-    response_format = payload.get("response_format", "wav")
-    if response_format == "pcm":
-        raise ValueError(
-            "pcm response format is not supported yet. \
-        Please use other formats like wav, mp3, etc. instead."
-        )
-
     headers = {
         "Content-Type": "application/json",
         "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
@@ -238,41 +234,38 @@ async def async_request_openai_audio_speech(
     output = MixRequestFuncOutput()
     output.prompt_len = request_func_input.prompt_len
 
+    # PCM format: 16-bit signed, 24 kHz, mono
+    sample_rate = 24000
+    sample_width = 2  # 16-bit = 2 bytes
+    channels = 1
+
     st = time.perf_counter()
     output.start_time = st
+    total_pcm_bytes = 0
     try:
         async with session.post(url=api_url, json=payload, headers=headers) as response:
             if response.status == 200:
-                audio_bytes = await response.read()
+                async for chunk in response.content.iter_any():
+                    if not chunk:
+                        continue
+                    timestamp = time.perf_counter()
+                    if output.audio_ttfp == 0.0:
+                        output.audio_ttfp = timestamp - st
+                        output.ttft = output.audio_ttfp
+                    total_pcm_bytes += len(chunk)
+
                 end_time = time.perf_counter()
                 output.latency = end_time - st
-                # ttft = latency since this is a non-streaming request
-                # hence there is no distinction between first and last token/audio
-                output.ttft = output.latency
-                output.audio_ttfp = output.latency
 
-                try:
-                    audio_segment = AudioSegment.from_file(io.BytesIO(audio_bytes))
-                    output.audio_duration = len(audio_segment) / 1000.0
-                    frame_width = audio_segment.frame_width
-                    if frame_width > 0:
-                        output.audio_frames = len(audio_segment.raw_data) // frame_width
-                    else:
-                        output.audio_frames = 0
-                        logger.warning("Audio frame width is zero")
-                    if output.audio_duration > 0:
-                        # rtf = audio_generate_time / audio_duration and
-                        # audio_generate_time = latency since this is a non-streaming request
-                        # so the time to receive last portion of audio is the latency
-                        output.audio_rtf = output.latency / output.audio_duration
-                    else:
-                        output.audio_rtf = 0
-                        logger.warning("Audio duration is zero")
-                    output.success = True
-                except Exception as e:
-                    output.success = False
-                    output.error = f"Failed to parse audio response: {e}"
-                    logger.error(f"ERROR: Failed to parse audio response: {e}")
+                total_samples = total_pcm_bytes // (sample_width * channels)
+                output.audio_duration = total_samples / sample_rate
+                output.audio_frames = total_samples
+                if output.audio_duration > 0:
+                    output.audio_rtf = output.latency / output.audio_duration
+                else:
+                    output.audio_rtf = 0
+                    logger.warning("Audio duration is zero")
+                output.success = True
             else:
                 output.error = response.reason or ""
                 output.success = False

vllm_omni/diffusion/data.py

@@ -43,6 +43,9 @@ class DiffusionParallelConfig:
     tensor_parallel_size: int = 1
     """Number of tensor parallel groups."""
 
+    enable_expert_parallel: bool = False
+    """Enable expert parallelism for MoE layers (TP is still used for non-MoE layers)."""
+
     sequence_parallel_size: int | None = None
     """Number of sequence parallel groups. sequence_parallel_size = ring_degree * ulysses_degree"""
 
@@ -457,6 +460,17 @@ class OmniDiffusionConfig:
     quantization: str | None = None
     quantization_config: "DiffusionQuantizationConfig | dict[str, Any] | None" = None
 
+    @property
+    def is_moe(self) -> bool:
+        num_experts = self.tf_model_config.get("num_experts", None)
+        if isinstance(num_experts, int):
+            return num_experts > 1
+
+        if isinstance(num_experts, (list, tuple)):
+            return any(isinstance(n, int) and n > 1 for n in num_experts)
+
+        return False
+
     def settle_port(self, port: int, port_inc: int = 42, max_attempts: int = 100) -> int:
         """
         Find an available port with retry logic.

vllm_omni/diffusion/distributed/parallel_state.py

@@ -35,6 +35,8 @@ from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion import envs
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.forward_context import get_forward_context
 from vllm_omni.platforms import current_omni_platform
 
 from .group_coordinator import (
@@ -51,7 +53,7 @@ logger = init_logger(__name__)
 
 
 _WORLD: GroupCoordinator | None = None
-# get _TP from vllm.distributed.parallel_state
+# get _TP&_EP from vllm.distributed.parallel_state
 _SP: SequenceParallelGroupCoordinator | None = None
 _PP: PipelineGroupCoordinator | None = None
 _CFG: GroupCoordinator | None = None
@@ -188,6 +190,7 @@ class RankGenerator:
         self.fs = fs
         self.rank_offset = rank_offset
         self.world_size = tp * sp * pp * cfg * dp
+        self.ep = tp * sp * cfg * dp  # EP level exclude PP
 
         self.name_to_size = {
             "tp": self.tp,
@@ -250,6 +253,15 @@ class RankGenerator:
                 ranks.append(group)
             return ranks
 
+        if token == "ep":
+            ranks = []
+            num_pp_stages = self.pp
+            for i in range(num_pp_stages):
+                start = i * self.ep + self.rank_offset
+                end = start + self.ep
+                ranks.append(list(range(start, end)))
+            return ranks
+
         mask = self.get_mask(self.order, token)
         ranks = generate_masked_orthogonal_rank_groups(self.world_size, self.ordered_size, mask)
         if self.rank_offset > 0:
@@ -474,6 +486,7 @@ def init_model_parallel_group(
         "data",
         "pipeline",
         "tensor",
+        "expert",
         "sequence",
         "classifier_free_guidance",
         "fully_shard",
@@ -670,6 +683,7 @@ def initialize_model_parallel(
     pipeline_parallel_size: int = 1,
     fully_shard_degree: int = 1,
     hsdp_replicate_size: int = 1,
+    enable_expert_parallel: bool = False,
     backend: str | None = None,
 ) -> None:
     if backend is None:
@@ -837,6 +851,18 @@ def initialize_model_parallel(
         parallel_mode="fully_shard",
     )
 
+    if enable_expert_parallel:
+        od_config: OmniDiffusionConfig | None = get_forward_context().omni_diffusion_config
+        if od_config and od_config.is_moe:
+            vllm_parallel_state._EP = init_model_parallel_group(
+                group_ranks=rank_generator.get_ranks("ep"),
+                local_rank=get_world_group().local_rank,
+                backend=backend,
+                parallel_mode="expert",
+            )
+        else:
+            raise RuntimeError("Expert parallelism enabled for a non-MoE model ")
+
     init_dit_group(dit_parallel_size, backend)
 
 
@@ -861,6 +887,10 @@ def destroy_model_parallel():
         vllm_parallel_state._TP.destroy()
     vllm_parallel_state._TP = None
 
+    if vllm_parallel_state._EP:
+        vllm_parallel_state._EP.destroy()
+    vllm_parallel_state._EP = None
+
     global _PP
     if _PP:
         _PP.destroy()

vllm_omni/diffusion/forward_context.py

@@ -31,15 +31,35 @@ class ForwardContext:
     # Original sequence length before padding (for removing padding in gather)
     sp_original_seq_len: int | None = None
 
-    # SP active scope tracking
+    # Set by registry when _sp_plan hooks are applied.
+    # When True, sp_active is determined by _sp_shard_depth (for _sp_plan hooks)
+    # When False, sp_active defaults to True when sequence_parallel_size > 1 (for manual SP, standalone tests, etc.)
+    sp_plan_hooks_applied: bool = False
+    # SP active scope tracking within the _sp_plan hook mechanism.
     # Tracks the depth of SP sharding - incremented on shard, decremented on gather
     # Used by attention layers to determine if SP communication should be enabled
     _sp_shard_depth: int = 0
 
     @property
     def sp_active(self) -> bool:
-        """Returns True when inside an SP sharded region (between shard and gather)."""
-        return self._sp_shard_depth > 0
+        """Returns True when SP attention parallelism should be enabled.
+
+        - If _sp_plan hooks are applied: use _sp_shard_depth (0 = outside sharded region).
+        - If _sp_plan hooks are NOT applied: default to True when sequence_parallel_size > 1,
+          since _sp_shard_depth is only meaningful within the _sp_plan hook mechanism.
+        """
+        if self.sp_plan_hooks_applied:
+            return self._sp_shard_depth > 0
+        # No _sp_plan: assume SP active when configured (manual SP, standalone tests)
+        if self.omni_diffusion_config is None:
+            raise ValueError(
+                "omni_diffusion_config is not set when checking sp_active! "
+                "This usually means set_forward_context() was not called. "
+                "Please call with set_forward_context(omni_diffusion_config=...)."
+            )
+
+        sp_size = self.omni_diffusion_config.parallel_config.sequence_parallel_size
+        return sp_size is not None and sp_size > 1
 
     def __post_init__(self):
         pass

vllm_omni/diffusion/models/flux2/__init__.py

@@ -0,0 +1,17 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Flux2 diffusion model components."""
+
+from vllm_omni.diffusion.models.flux2.flux2_transformer import (
+    Flux2Transformer2DModel,
+)
+from vllm_omni.diffusion.models.flux2.pipeline_flux2 import (
+    Flux2Pipeline,
+    get_flux2_post_process_func,
+)
+
+__all__ = [
+    "Flux2Pipeline",
+    "Flux2Transformer2DModel",
+    "get_flux2_post_process_func",
+]

vllm_omni/diffusion/models/flux2/flux2_transformer.py

@@ -0,0 +1,770 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Iterable
+from types import SimpleNamespace
+from typing import Any
+
+import torch
+from diffusers.models.embeddings import (
+    TimestepEmbedding,
+    Timesteps,
+    get_1d_rotary_pos_embed,
+)
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.normalization import AdaLayerNormContinuous
+from diffusers.utils import is_torch_npu_available
+from torch import nn
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.layers.rope import RotaryEmbedding
+
+
+class Flux2SwiGLU(nn.Module):
+    """SwiGLU activation used by Flux2."""
+
+    def __init__(self):
+        super().__init__()
+        self.gate_fn = nn.SiLU()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x1, x2 = x.chunk(2, dim=-1)
+        return self.gate_fn(x1) * x2
+
+
+class Flux2FeedForward(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        dim_out: int | None = None,
+        mult: float = 3.0,
+        inner_dim: int | None = None,
+        bias: bool = False,
+    ):
+        super().__init__()
+        if inner_dim is None:
+            inner_dim = int(dim * mult)
+        dim_out = dim_out or dim
+
+        self.linear_in = MergedColumnParallelLinear(
+            dim,
+            [inner_dim, inner_dim],
+            bias=bias,
+            return_bias=False,
+        )
+        self.act_fn = Flux2SwiGLU()
+        self.linear_out = RowParallelLinear(
+            inner_dim,
+            dim_out,
+            bias=bias,
+            input_is_parallel=True,
+            return_bias=False,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.linear_in(x)
+        x = self.act_fn(x)
+        return self.linear_out(x)
+
+
+class Flux2Attention(nn.Module):
+    def __init__(
+        self,
+        query_dim: int,
+        heads: int = 8,
+        dim_head: int = 64,
+        dropout: float = 0.0,
+        bias: bool = False,
+        added_kv_proj_dim: int | None = None,
+        added_proj_bias: bool | None = True,
+        out_bias: bool = True,
+        eps: float = 1e-5,
+        out_dim: int = None,
+        elementwise_affine: bool = True,
+    ):
+        super().__init__()
+        self.head_dim = dim_head
+        self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+        self.query_dim = query_dim
+        self.out_dim = out_dim if out_dim is not None else query_dim
+        self.heads = out_dim // dim_head if out_dim is not None else heads
+        self.dropout = dropout
+        self.added_kv_proj_dim = added_kv_proj_dim
+
+        self.to_qkv = QKVParallelLinear(
+            hidden_size=query_dim,
+            head_size=self.head_dim,
+            total_num_heads=self.heads,
+            bias=bias,
+        )
+        self.query_num_heads = self.to_qkv.num_heads
+        self.kv_num_heads = self.to_qkv.num_kv_heads
+
+        self.norm_q = RMSNorm(dim_head, eps=eps)
+        self.norm_k = RMSNorm(dim_head, eps=eps)
+
+        self.to_out = nn.ModuleList(
+            [
+                RowParallelLinear(
+                    self.inner_dim,
+                    self.out_dim,
+                    bias=out_bias,
+                    input_is_parallel=True,
+                    return_bias=False,
+                ),
+                nn.Dropout(dropout),
+            ]
+        )
+
+        if added_kv_proj_dim is not None:
+            self.norm_added_q = RMSNorm(dim_head, eps=eps)
+            self.norm_added_k = RMSNorm(dim_head, eps=eps)
+            self.add_kv_proj = QKVParallelLinear(
+                hidden_size=added_kv_proj_dim,
+                head_size=self.head_dim,
+                total_num_heads=self.heads,
+                bias=added_proj_bias,
+            )
+            self.add_query_num_heads = self.add_kv_proj.num_heads
+            self.add_kv_num_heads = self.add_kv_proj.num_kv_heads
+            self.to_add_out = RowParallelLinear(
+                self.inner_dim,
+                query_dim,
+                bias=out_bias,
+                input_is_parallel=True,
+                return_bias=False,
+            )
+
+        self.rope = RotaryEmbedding(is_neox_style=False)
+        self.attn = Attention(
+            num_heads=self.query_num_heads,
+            head_size=self.head_dim,
+            softmax_scale=1.0 / (self.head_dim**0.5),
+            causal=False,
+            num_kv_heads=self.kv_num_heads,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+        **kwargs,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        qkv, _ = self.to_qkv(hidden_states)
+        query, key, value = qkv.chunk(3, dim=-1)
+
+        encoder_query = encoder_key = encoder_value = None
+        if encoder_hidden_states is not None and self.added_kv_proj_dim is not None:
+            encoder_qkv, _ = self.add_kv_proj(encoder_hidden_states)
+            encoder_query, encoder_key, encoder_value = encoder_qkv.chunk(3, dim=-1)
+
+        query = query.unflatten(-1, (self.query_num_heads, -1))
+        key = key.unflatten(-1, (self.kv_num_heads, -1))
+        value = value.unflatten(-1, (self.kv_num_heads, -1))
+
+        query = self.norm_q(query)
+        key = self.norm_k(key)
+
+        if encoder_hidden_states is not None and self.added_kv_proj_dim is not None:
+            encoder_query = encoder_query.unflatten(-1, (self.add_query_num_heads, -1))
+            encoder_key = encoder_key.unflatten(-1, (self.add_kv_num_heads, -1))
+            encoder_value = encoder_value.unflatten(-1, (self.add_kv_num_heads, -1))
+
+            encoder_query = self.norm_added_q(encoder_query)
+            encoder_key = self.norm_added_k(encoder_key)
+
+            query = torch.cat([encoder_query, query], dim=1)
+            key = torch.cat([encoder_key, key], dim=1)
+            value = torch.cat([encoder_value, value], dim=1)
+
+        if image_rotary_emb is not None:
+            cos, sin = image_rotary_emb
+            cos = cos.to(query.dtype)
+            sin = sin.to(query.dtype)
+            query = self.rope(query, cos, sin)
+            key = self.rope(key, cos, sin)
+
+        attn_metadata = None
+        if attention_mask is not None:
+            if attention_mask.dim() == 3:
+                attention_mask = attention_mask.unsqueeze(1)
+            attn_metadata = AttentionMetadata(attn_mask=attention_mask)
+
+        hidden_states = self.attn(query, key, value, attn_metadata)
+        hidden_states = hidden_states.flatten(2, 3).to(query.dtype)
+
+        if encoder_hidden_states is not None:
+            context_len = encoder_hidden_states.shape[1]
+            encoder_hidden_states, hidden_states = hidden_states.split_with_sizes(
+                [context_len, hidden_states.shape[1] - context_len],
+                dim=1,
+            )
+            encoder_hidden_states = self.to_add_out(encoder_hidden_states)
+
+        hidden_states = self.to_out[0](hidden_states)
+        hidden_states = self.to_out[1](hidden_states)
+
+        if encoder_hidden_states is not None:
+            return hidden_states, encoder_hidden_states
+        return hidden_states
+
+
+class Flux2ParallelSelfAttention(nn.Module):
+    """
+    Parallel attention block that fuses QKV projections with MLP input projections.
+    """
+
+    def __init__(
+        self,
+        query_dim: int,
+        heads: int = 8,
+        dim_head: int = 64,
+        dropout: float = 0.0,
+        bias: bool = False,
+        out_bias: bool = True,
+        eps: float = 1e-5,
+        out_dim: int = None,
+        elementwise_affine: bool = True,
+        mlp_ratio: float = 4.0,
+        mlp_mult_factor: int = 2,
+    ):
+        super().__init__()
+        self.head_dim = dim_head
+        self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+        self.query_dim = query_dim
+        self.out_dim = out_dim if out_dim is not None else query_dim
+        self.heads = out_dim // dim_head if out_dim is not None else heads
+        self.dropout = dropout
+
+        self.mlp_ratio = mlp_ratio
+        self.mlp_hidden_dim = int(query_dim * self.mlp_ratio)
+        self.mlp_mult_factor = mlp_mult_factor
+
+        self.to_qkv_mlp_proj = ColumnParallelLinear(
+            self.query_dim,
+            self.inner_dim * 3 + self.mlp_hidden_dim * self.mlp_mult_factor,
+            bias=bias,
+            gather_output=True,
+        )
+        self.mlp_act_fn = Flux2SwiGLU()
+
+        self.norm_q = RMSNorm(dim_head, eps=eps)
+        self.norm_k = RMSNorm(dim_head, eps=eps)
+
+        self.to_out = ColumnParallelLinear(
+            self.inner_dim + self.mlp_hidden_dim,
+            self.out_dim,
+            bias=out_bias,
+            gather_output=True,
+        )
+        self.rope = RotaryEmbedding(is_neox_style=False)
+        self.attn = Attention(
+            num_heads=self.heads,
+            head_size=self.head_dim,
+            softmax_scale=1.0 / (self.head_dim**0.5),
+            causal=False,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        hidden_states, _ = self.to_qkv_mlp_proj(hidden_states)
+        qkv, mlp_hidden_states = torch.split(
+            hidden_states,
+            [3 * self.inner_dim, self.mlp_hidden_dim * self.mlp_mult_factor],
+            dim=-1,
+        )
+
+        query, key, value = qkv.chunk(3, dim=-1)
+        query = query.unflatten(-1, (self.heads, -1))
+        key = key.unflatten(-1, (self.heads, -1))
+        value = value.unflatten(-1, (self.heads, -1))
+
+        query = self.norm_q(query)
+        key = self.norm_k(key)
+
+        if image_rotary_emb is not None:
+            cos, sin = image_rotary_emb
+            cos = cos.to(query.dtype)
+            sin = sin.to(query.dtype)
+            query = self.rope(query, cos, sin)
+            key = self.rope(key, cos, sin)
+
+        attn_metadata = None
+        if attention_mask is not None:
+            if attention_mask.dim() == 3:
+                attention_mask = attention_mask.unsqueeze(1)
+            attn_metadata = AttentionMetadata(attn_mask=attention_mask)
+
+        attn_output = self.attn(query, key, value, attn_metadata)
+        attn_output = attn_output.flatten(2, 3).to(query.dtype)
+
+        mlp_hidden_states = self.mlp_act_fn(mlp_hidden_states)
+        hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=-1)
+        hidden_states, _ = self.to_out(hidden_states)
+        return hidden_states
+
+
+class Flux2SingleTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        mlp_ratio: float = 3.0,
+        eps: float = 1e-6,
+        bias: bool = False,
+    ):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.attn = Flux2ParallelSelfAttention(
+            query_dim=dim,
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            out_dim=dim,
+            bias=bias,
+            out_bias=bias,
+            eps=eps,
+            mlp_ratio=mlp_ratio,
+            mlp_mult_factor=2,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | None,
+        temb_mod_params: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+        joint_attention_kwargs: dict[str, Any] | None = None,
+        split_hidden_states: bool = False,
+        text_seq_len: int | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        if encoder_hidden_states is not None:
+            text_seq_len = encoder_hidden_states.shape[1]
+            hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        mod_shift, mod_scale, mod_gate = temb_mod_params
+
+        norm_hidden_states = self.norm(hidden_states)
+        norm_hidden_states = (1 + mod_scale) * norm_hidden_states + mod_shift
+
+        joint_attention_kwargs = joint_attention_kwargs or {}
+        attn_output = self.attn(
+            hidden_states=norm_hidden_states,
+            image_rotary_emb=image_rotary_emb,
+            **joint_attention_kwargs,
+        )
+
+        hidden_states = hidden_states + mod_gate * attn_output
+        if hidden_states.dtype == torch.float16:
+            hidden_states = hidden_states.clip(-65504, 65504)
+
+        if split_hidden_states:
+            encoder_hidden_states, hidden_states = hidden_states[:, :text_seq_len], hidden_states[:, text_seq_len:]
+            return encoder_hidden_states, hidden_states
+        return hidden_states
+
+
+class Flux2TransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        mlp_ratio: float = 3.0,
+        eps: float = 1e-6,
+        bias: bool = False,
+    ):
+        super().__init__()
+        self.mlp_hidden_dim = int(dim * mlp_ratio)
+
+        self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.norm1_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+
+        self.attn = Flux2Attention(
+            query_dim=dim,
+            added_kv_proj_dim=dim,
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            out_dim=dim,
+            bias=bias,
+            added_proj_bias=bias,
+            out_bias=bias,
+            eps=eps,
+        )
+
+        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.ff = Flux2FeedForward(dim=dim, dim_out=dim, mult=mlp_ratio, bias=bias)
+
+        self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.ff_context = Flux2FeedForward(dim=dim, dim_out=dim, mult=mlp_ratio, bias=bias)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb_mod_params_img: tuple[tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...],
+        temb_mod_params_txt: tuple[tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...],
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+        joint_attention_kwargs: dict[str, Any] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        joint_attention_kwargs = joint_attention_kwargs or {}
+
+        (shift_msa, scale_msa, gate_msa), (shift_mlp, scale_mlp, gate_mlp) = temb_mod_params_img
+        (c_shift_msa, c_scale_msa, c_gate_msa), (c_shift_mlp, c_scale_mlp, c_gate_mlp) = temb_mod_params_txt
+
+        norm_hidden_states = self.norm1(hidden_states)
+        norm_hidden_states = (1 + scale_msa) * norm_hidden_states + shift_msa
+
+        norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states)
+        norm_encoder_hidden_states = (1 + c_scale_msa) * norm_encoder_hidden_states + c_shift_msa
+
+        attn_output, context_attn_output = self.attn(
+            hidden_states=norm_hidden_states,
+            encoder_hidden_states=norm_encoder_hidden_states,
+            image_rotary_emb=image_rotary_emb,
+            **joint_attention_kwargs,
+        )
+
+        attn_output = gate_msa * attn_output
+        hidden_states = hidden_states + attn_output
+
+        norm_hidden_states = self.norm2(hidden_states)
+        norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp
+        ff_output = self.ff(norm_hidden_states)
+        hidden_states = hidden_states + gate_mlp * ff_output
+
+        context_attn_output = c_gate_msa * context_attn_output
+        encoder_hidden_states = encoder_hidden_states + context_attn_output
+
+        norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
+        norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp) + c_shift_mlp
+        context_ff_output = self.ff_context(norm_encoder_hidden_states)
+        encoder_hidden_states = encoder_hidden_states + c_gate_mlp * context_ff_output
+        if encoder_hidden_states.dtype == torch.float16:
+            encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
+
+        return encoder_hidden_states, hidden_states
+
+
+class Flux2PosEmbed(nn.Module):
+    def __init__(self, theta: int, axes_dim: list[int]):
+        super().__init__()
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        # Expected ids shape: [S, len(self.axes_dim)]
+        cos_out = []
+        sin_out = []
+        pos = ids.float()
+        is_mps = ids.device.type == "mps"
+        is_npu = ids.device.type == "npu"
+        freqs_dtype = torch.float32 if (is_mps or is_npu) else torch.float64
+        # Unlike Flux 1, loop over len(self.axes_dim) rather than ids.shape[-1]
+        for i in range(len(self.axes_dim)):
+            freqs_cis = get_1d_rotary_pos_embed(
+                self.axes_dim[i],
+                pos[..., i],
+                theta=self.theta,
+                use_real=False,
+                freqs_dtype=freqs_dtype,
+            )
+            cos_out.append(freqs_cis.real)
+            sin_out.append(freqs_cis.imag)
+        freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device)
+        freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device)
+        return freqs_cos, freqs_sin
+
+
+class Flux2TimestepGuidanceEmbeddings(nn.Module):
+    def __init__(
+        self,
+        in_channels: int = 256,
+        embedding_dim: int = 6144,
+        bias: bool = False,
+        guidance_embeds: bool = True,
+    ):
+        super().__init__()
+        self.time_proj = Timesteps(num_channels=in_channels, flip_sin_to_cos=True, downscale_freq_shift=0)
+        self.timestep_embedder = TimestepEmbedding(
+            in_channels=in_channels,
+            time_embed_dim=embedding_dim,
+            sample_proj_bias=bias,
+        )
+
+        if guidance_embeds:
+            self.guidance_embedder = TimestepEmbedding(
+                in_channels=in_channels,
+                time_embed_dim=embedding_dim,
+                sample_proj_bias=bias,
+            )
+        else:
+            self.guidance_embedder = None
+
+    def forward(self, timestep: torch.Tensor, guidance: torch.Tensor | None) -> torch.Tensor:
+        timesteps_proj = self.time_proj(timestep)
+        timesteps_emb = self.timestep_embedder(timesteps_proj.to(timestep.dtype))
+
+        if guidance is not None and self.guidance_embedder is not None:
+            guidance_proj = self.time_proj(guidance)
+            guidance_emb = self.guidance_embedder(guidance_proj.to(guidance.dtype))
+            return timesteps_emb + guidance_emb
+        return timesteps_emb
+
+
+class Flux2Modulation(nn.Module):
+    def __init__(self, dim: int, mod_param_sets: int = 2, bias: bool = False):
+        super().__init__()
+        self.mod_param_sets = mod_param_sets
+
+        self.linear = nn.Linear(dim, dim * 3 * self.mod_param_sets, bias=bias)
+        self.act_fn = nn.SiLU()
+
+    def forward(self, temb: torch.Tensor) -> tuple[tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...]:
+        mod = self.act_fn(temb)
+        mod = self.linear(mod)
+
+        if mod.ndim == 2:
+            mod = mod.unsqueeze(1)
+        mod_params = torch.chunk(mod, 3 * self.mod_param_sets, dim=-1)
+        # Return tuple of 3-tuples of modulation params shift/scale/gate
+        return tuple(mod_params[3 * i : 3 * (i + 1)] for i in range(self.mod_param_sets))
+
+
+class Flux2Transformer2DModel(nn.Module):
+    """
+    The Transformer model introduced in Flux 2.
+    """
+
+    _repeated_blocks = ["Flux2TransformerBlock", "Flux2SingleTransformerBlock"]
+
+    def __init__(
+        self,
+        patch_size: int = 1,
+        in_channels: int = 128,
+        out_channels: int | None = None,
+        num_layers: int = 8,
+        num_single_layers: int = 48,
+        attention_head_dim: int = 128,
+        num_attention_heads: int = 48,
+        joint_attention_dim: int = 15360,
+        timestep_guidance_channels: int = 256,
+        mlp_ratio: float = 3.0,
+        axes_dims_rope: tuple[int, ...] = (32, 32, 32, 32),
+        rope_theta: int = 2000,
+        eps: float = 1e-6,
+        guidance_embeds: bool = True,
+    ):
+        super().__init__()
+        self.stacked_params_mapping = None
+        self.out_channels = out_channels or in_channels
+        self.inner_dim = num_attention_heads * attention_head_dim
+        self.config = SimpleNamespace(
+            patch_size=patch_size,
+            in_channels=in_channels,
+            out_channels=self.out_channels,
+            num_layers=num_layers,
+            num_single_layers=num_single_layers,
+            attention_head_dim=attention_head_dim,
+            num_attention_heads=num_attention_heads,
+            joint_attention_dim=joint_attention_dim,
+            timestep_guidance_channels=timestep_guidance_channels,
+            mlp_ratio=mlp_ratio,
+            axes_dims_rope=axes_dims_rope,
+            rope_theta=rope_theta,
+            eps=eps,
+            guidance_embeds=guidance_embeds,
+        )
+
+        self.pos_embed = Flux2PosEmbed(theta=rope_theta, axes_dim=axes_dims_rope)
+        self.time_guidance_embed = Flux2TimestepGuidanceEmbeddings(
+            in_channels=timestep_guidance_channels,
+            embedding_dim=self.inner_dim,
+            bias=False,
+            guidance_embeds=guidance_embeds,
+        )
+
+        self.double_stream_modulation_img = Flux2Modulation(self.inner_dim, mod_param_sets=2, bias=False)
+        self.double_stream_modulation_txt = Flux2Modulation(self.inner_dim, mod_param_sets=2, bias=False)
+        self.single_stream_modulation = Flux2Modulation(self.inner_dim, mod_param_sets=1, bias=False)
+
+        self.x_embedder = nn.Linear(in_channels, self.inner_dim, bias=False)
+        self.context_embedder = nn.Linear(joint_attention_dim, self.inner_dim, bias=False)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                Flux2TransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                    mlp_ratio=mlp_ratio,
+                    eps=eps,
+                    bias=False,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+        self.single_transformer_blocks = nn.ModuleList(
+            [
+                Flux2SingleTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                    mlp_ratio=mlp_ratio,
+                    eps=eps,
+                    bias=False,
+                )
+                for _ in range(num_single_layers)
+            ]
+        )
+
+        self.norm_out = AdaLayerNormContinuous(
+            self.inner_dim, self.inner_dim, elementwise_affine=False, eps=eps, bias=False
+        )
+        self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=False)
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return next(self.parameters()).dtype
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor = None,
+        timestep: torch.LongTensor = None,
+        img_ids: torch.Tensor = None,
+        txt_ids: torch.Tensor = None,
+        guidance: torch.Tensor | None = None,
+        joint_attention_kwargs: dict[str, Any] | None = None,
+        return_dict: bool = True,
+    ) -> torch.Tensor | Transformer2DModelOutput:
+        joint_attention_kwargs = joint_attention_kwargs or {}
+
+        num_txt_tokens = encoder_hidden_states.shape[1]
+
+        timestep = timestep.to(hidden_states.dtype) * 1000
+        if guidance is not None:
+            guidance = guidance.to(hidden_states.dtype) * 1000
+
+        temb = self.time_guidance_embed(timestep, guidance)
+
+        double_stream_mod_img = self.double_stream_modulation_img(temb)
+        double_stream_mod_txt = self.double_stream_modulation_txt(temb)
+        single_stream_mod = self.single_stream_modulation(temb)[0]
+
+        hidden_states = self.x_embedder(hidden_states)
+        encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+
+        if img_ids.ndim == 3:
+            img_ids = img_ids[0]
+        if txt_ids.ndim == 3:
+            txt_ids = txt_ids[0]
+
+        if is_torch_npu_available():
+            freqs_cos_image, freqs_sin_image = self.pos_embed(img_ids.cpu())
+            image_rotary_emb = (freqs_cos_image.npu(), freqs_sin_image.npu())
+            freqs_cos_text, freqs_sin_text = self.pos_embed(txt_ids.cpu())
+            text_rotary_emb = (freqs_cos_text.npu(), freqs_sin_text.npu())
+        else:
+            image_rotary_emb = self.pos_embed(img_ids)
+            text_rotary_emb = self.pos_embed(txt_ids)
+        concat_rotary_emb = (
+            torch.cat([text_rotary_emb[0], image_rotary_emb[0]], dim=0),
+            torch.cat([text_rotary_emb[1], image_rotary_emb[1]], dim=0),
+        )
+
+        for index_block, block in enumerate(self.transformer_blocks):
+            encoder_hidden_states, hidden_states = block(
+                hidden_states=hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                temb_mod_params_img=double_stream_mod_img,
+                temb_mod_params_txt=double_stream_mod_txt,
+                image_rotary_emb=concat_rotary_emb,
+                joint_attention_kwargs=joint_attention_kwargs,
+            )
+
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        for index_block, block in enumerate(self.single_transformer_blocks):
+            hidden_states = block(
+                hidden_states=hidden_states,
+                encoder_hidden_states=None,
+                temb_mod_params=single_stream_mod,
+                image_rotary_emb=concat_rotary_emb,
+                joint_attention_kwargs=joint_attention_kwargs,
+            )
+
+        hidden_states = hidden_states[:, num_txt_tokens:, ...]
+        hidden_states = self.norm_out(hidden_states, temb)
+        output = self.proj_out(hidden_states)
+
+        if not return_dict:
+            return (output,)
+
+        return Transformer2DModelOutput(sample=output)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+            (".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"),
+        ]
+        # Expose packed shard mappings for LoRA handling of fused projections.
+        self.stacked_params_mapping = stacked_params_mapping
+
+        params_dict = dict(self.named_parameters())
+
+        for name, buffer in self.named_buffers():
+            if name.endswith(".beta") or name.endswith(".eps"):
+                params_dict[name] = buffer
+
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            original_name = name
+            mapped = False
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in original_name:
+                    continue
+                name = original_name.replace(weight_name, param_name)
+                param = params_dict.get(name)
+                if param is None:
+                    break
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                loaded_params.add(name)
+                mapped = True
+                break
+            if mapped:
+                continue
+
+            name = original_name
+            if name not in params_dict and ".to_out.0." in name:
+                name = name.replace(".to_out.0.", ".to_out.")
+            # Some GGUF checkpoints include quantized tensors for modules that
+            # are intentionally left unquantized in this model.
+            param = params_dict.get(name)
+            if param is None:
+                continue
+            weight_loader = getattr(param, "weight_loader", default_weight_loader)
+            weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params