Initial plan

Signed-off-by: Chenguang ZHENG <645327136@qq.com>

.buildkite/pipeline.yml

@@ -20,12 +20,19 @@ steps:
   - label: "Simple Unit Test"
     depends_on: image-build
     commands:
-      - pytest -v -s tests/entrypoints/
-      - pytest -v -s tests/diffusion/cache/
-      - pytest -v -s tests/diffusion/lora/
-      - pytest -v -s tests/model_executor/models/qwen2_5_omni/test_audio_length.py
-      - pytest -v -s tests/worker/
-      - pytest -v -s tests/distributed/omni_connectors/test_kv_flow.py
+    - |
+      pytest -v -s \
+        tests/entrypoints/ \
+        tests/diffusion/cache/ \
+        tests/diffusion/lora/ \
+        tests/model_executor/models/qwen2_5_omni/test_audio_length.py \
+        tests/worker/ \
+        tests/distributed/omni_connectors/test_kv_flow.py \
+        --cov=vllm_omni \
+        --cov-branch \
+        --cov-report=term-missing \
+        --cov-report=html \
+        --cov-report=xml
     agents:
       queue: "gpu_1_queue"
     plugins:

.buildkite/test-nightly.yaml

@@ -0,0 +1,45 @@
+steps:
+  - label: ":docker: Build image"
+    key: image-build
+    commands:
+      - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
+      - "docker build --file docker/Dockerfile.ci -t vllm-omni-ci ."
+      - "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
+      - "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
+    agents:
+      queue: "cpu_queue_premerge"
+
+  - label: "Omni Model Test with H100"
+    timeout_in_minutes: 180
+    depends_on: image-build
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_expansion.py
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: public.ecr.aws/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

.github/workflows/build_wheel.yml

@@ -7,6 +7,7 @@ on:
     paths-ignore:
       - 'docs/**'
       - 'examples/**'
+      - 'apps/**'
       - '.github/**'
       - '.buildkite/**'
       - '**.md'
@@ -20,6 +21,7 @@ on:
     paths-ignore:
       - 'docs/**'
       - 'examples/**'
+      - 'apps/**'
       - '.github/**'
       - '.buildkite/**'
       - '*.yml'

.github/workflows/comfyui-publish.yml

@@ -0,0 +1,27 @@
+name: ComfyUI Integration -- Publish to Comfy registry
+on:
+  workflow_dispatch:
+  push:
+    tags:
+      - "*"
+    paths:
+      - "apps/ComfyUI-vLLM-Omni/**"
+
+permissions:
+  issues: write
+
+defaults:
+  run:
+    working-directory: apps/ComfyUI-vLLM-Omni
+
+jobs:
+  publish-node:
+    name: Publish Custom Node to registry
+    runs-on: ubuntu-latest
+    steps:
+      - name: ♻️ Check out code
+        uses: actions/checkout@v4
+      - name: 📦 Publish Custom Node
+        uses: Comfy-Org/publish-node-action@main
+        with:
+          personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}

.github/workflows/comfyui-validate.yml

@@ -0,0 +1,40 @@
+name: ComfyUI Integration -- Validate backwards compatibility
+
+on:
+  pull_request:
+    branches:
+      - master
+      - main
+    paths:
+      - "apps/ComfyUI-vLLM-Omni/**"
+
+defaults:
+  run:
+    working-directory: apps/ComfyUI-vLLM-Omni
+
+jobs:
+  check-base-path:
+    runs-on: ubuntu-latest
+    outputs:
+      exists: ${{ steps.check.outputs.exists }}
+    defaults:
+      run:
+        working-directory: .
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event.pull_request.base.sha }}
+          fetch-depth: 1
+      - id: check
+        run: |
+          if [ -d "apps/ComfyUI-vLLM-Omni" ]; then
+            echo "exists=true" >> "$GITHUB_OUTPUT"
+          else
+            echo "exists=false" >> "$GITHUB_OUTPUT"
+          fi
+  validate:
+    needs: check-base-path
+    if: needs.check-base-path.outputs.exists == 'true'
+    runs-on: ubuntu-latest
+    steps:
+      - uses: comfy-org/node-diff@main

.gitignore

@@ -153,6 +153,14 @@ dmypy.json
 # Cython debug symbols
 cython_debug/
 
+# Claude
+CLAUDE.md
+.claude/
+
+# Codex
+AGENTS.md
+.codex/
+
 # cursor
 .cursor/
 
@@ -214,6 +222,7 @@ data/
 datasets/
 *.csv
 *.json
+!apps/ComfyUI-vLLM-Omni/example_workflows/*.json
 *.jsonl
 *.parquet
 

README.md

@@ -9,15 +9,16 @@ 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://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true"><b>Slides</b></a> |
 </p>
 
+
 ---
 
 *Latest News* 🔥
 
 - [2026/02] We released [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0) - This is the first **stable release** of vLLM-Omni that expands Omni’s diffusion / image-video generation and audio / TTS stack, improves distributed execution and memory efficiency, and broadens platform/backend coverage (GPU/ROCm/NPU/XPU). It also brings meaningful upgrades to serving APIs, profiling & benchmarking, and overall stability. Please check our latest [paper](https://arxiv.org/abs/2602.02204) for architecture design and performance results.
-- [2026/01] We released [0.12.0rc1](https://github.com/vllm-project/vllm-omni/releases/tag/v0.12.0rc1) - a major RC milestone focused on maturing the diffusion stack, strengthening OpenAI-compatible serving, expanding omni-model coverage, and improving stability across platforms (GPU/NPU/ROCm), please check our latest [design](https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true).
+- [2026/01] We released [0.12.0rc1](https://github.com/vllm-project/vllm-omni/releases/tag/v0.12.0rc1) - a major RC milestone focused on maturing the diffusion stack, strengthening OpenAI-compatible serving, expanding omni-model coverage, and improving stability across platforms (GPU/NPU/ROCm).
 - [2025/11] vLLM community officially released [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni) in order to support omni-modality models serving.
 
 ---

apps/ComfyUI-vLLM-Omni/.gitignore

@@ -0,0 +1,115 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# OSX useful to ignore
+*.DS_Store
+.AppleDouble
+.LSOverride
+
+# Thumbnails
+._*
+
+# Files that might appear in the root of a volume
+.DocumentRevisions-V100
+.fseventsd
+.Spotlight-V100
+.TemporaryItems
+.Trashes
+.VolumeIcon.icns
+.com.apple.timemachine.donotpresent
+
+# Directories potentially created on remote AFP share
+.AppleDB
+.AppleDesktop
+Network Trash Folder
+Temporary Items
+.apdisk
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+env/
+venv/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+
+# Sphinx documentation
+docs/_build/
+
+# IntelliJ Idea
+.idea
+*.iml
+*.ipr
+*.iws
+
+# PyBuilder
+target/
+
+# Cookiecutter
+output/
+python_boilerplate/
+cookiecutter-pypackage-env/
+
+# vscode settings
+.history/
+*.code-workspace
+
+# Frontend extension
+node_modules/
+.env
+.env.local
+.env.development.local
+.env.test.local
+.env.production.local
+npm-debug.log*
+yarn-debug.log*
+yarn-error.log*
+node.zip
+.vscode/
+.claude/
+.codemate/

apps/ComfyUI-vLLM-Omni/LICENSE

@@ -0,0 +1,15 @@
+Apache Software License 2.0
+
+Copyright (c) 2026, Zeyu Huang
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.

apps/ComfyUI-vLLM-Omni/README.md

@@ -0,0 +1,170 @@
+# vLLM-Omni
+
+vLLM-Omni offers a ComfyUI integration on top of its online serving API.
+It can send model inference requests to either a locally running vLLM-Omni service or a remote one.
+
+## Requirement
+
+- Python 3.12 or above
+- [ComfyUI installed](https://docs.comfy.org/installation/system_requirements)
+- [vLLM-Omni installed](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/) on either the same device or another device discoverable via the internet.
+- No need to install additional packages apart from those already required by ComfyUI.
+
+> [!TIP]
+> If you run both ComfyUI and vLLM-Omni on the same device, you can create separate virtual environments and use different Python versions for them.
+
+
+## Installation
+
+Copy this folder to the `custom_nodes` subfolder of your ComfyUI installation. Your directory should look like `ComfyUI/custom_nodes/ComfyUI-vLLM-Omni`.
+
+If you are running ComfyUI during copying, you should restart ComfyUI to load this extension.
+
+> [!TIP]
+> You can use utility websites such as https://download-directory.github.io/ to download a subdirectory of a repo. Also checkout community discussions (e.g., https://stackoverflow.com/questions/7106012/download-a-single-folder-or-directory-from-a-github-repository) for more info.
+
+On the device and virtual environment you run ComfyUI, launch ComfyUI with
+```bash
+cd ComfyUI
+
+# The regular way
+python main.py
+
+# If you are mainly using this node, launch it faster with
+python main.py --cpu
+```
+
+On the device and virtual environment you run vLLM-Omni, start a model service with
+```bash
+vllm serve The_Model_ID_to_Serve --omni --port 8000
+```
+
+Check **ComfyUI's sidebar -> Node Library**. There should be a new folder named **vLLM-Omni**.
+If no, check your shell running the ComfyUI process. There may be some error messages before the line `Import times for custom nodes:` and the line `To see the GUI go to: http://127.0.0.1:8188`.
+
+## Quickstart
+
+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
+- **Multimodality Comprehension** 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**)
+
+> [!INFO]
+> The node UI and feature designs are intended to match vLLM-Omni online serving interfaces. It cannot offer more than what the interfaces support.
+
+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.
+- 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.
+    - For multi-stage models where all stages are either autoregression or diffusion, you can also connect only a single Sampling Params node, indicating that this set of sampling parameters will be used for all stages.
+
+**The following features are tested**:
+
+- Single-node workflows for
+    - Multimodal Comprehension (e.g., Qwen Omni, BAGEL)
+    - Text-to-Image Generation (e.g., Qwen-Image)
+    - Image-to-Image Generation (e.g., Qwen-Image-Edit)
+    - TTS (e.g., Qwen TTS, including VoiceDesign, VoiceClone, CustomVoice)
+
+**The following features are not currently tested**. They may work or break. You are welcomed to test it out and offer comments.
+
+- Multi-node workflow that connects multiple model services together.
+
+## Screenshots and Examples
+
+### Multimodal comprehension (e.g., Qwen Omni series, BAGEL)
+
+(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Annotated Example**)
+
+<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-comprehension.jpg">
+    <img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-comprehension.jpg" width=55%>
+  </picture>
+</p>
+
+> [!TIP]
+> Although this node enables all-modality input, you should check whether the specific model you host and request for supports the modalities you connect to the node.
+
+You can configure per-stage sampling parameters for multi-stage models.
+
+<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-multi-stage.jpg">
+    <img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-multi-stage.jpg" width=55%>
+  </picture>
+</p>
+
+### Text-to-image and image-to-image generation (e.g., Z-Image-Turbo, Qwen-Image-Edit, BAGEL)
+
+(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Image 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-image-generation.jpg">
+    <img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-image-generation.jpg" width=55%>
+  </picture>
+</p>
+
+> [!TIP]
+> The node automatically choose text-to-image or image-to-image 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**)
+
+<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-tts.jpg">
+    <img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-tts.jpg" width=55%>
+  </picture>
+</p>
+
+> [!TIP]
+> There is a dedicated node for VoiceClone tasks with reference audio input. Other simple text-to-speech tasks should use the regular TTS node.
+
+## Develop
+
+Follow the [development convention and rules of vLLM-Omni](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/).
+
+## Limitation and Non-Goals
+
+- Single server mode only. No automatic load balancing or failover.
+- Features set is bounded to vLLM-Omni's online service capability, including
+    - The types of models supported in online mode,
+    - The types of sampling parameters supported in the online mode,
+    - The ways to send files (primarily through full-length base64 in JSON payload),
+    - Figuring out errors in the payload (such as unsupported fields by a specific model) if the endpoint does not explicitly return an error,
+    - (The lack of) Authentication
+    - (The lack of) Progress indicator
+
+## Support
+
+If you are new to ComfyUI, please check out [its documentation](https://docs.comfy.org/) for usage instructions.
+
+If you are new to vLLM-Omni, please also check out [its documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/) for usage instructions.
+
+Whenever you find an issue or problem, please
+
+- First find out if this is an upstream limitation of vLLM-Omni's online serving mode, by [checking their documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/examples/).
+- [Open an issue](https://github.com/vllm-project/vllm-omni/issues) that clearly describes this ComfyUI or online service problem.
+
+## Acknowledgements
+
+Features
+
+- https://github.com/dougbtv/comfyui-vllm-omni/ The official reference implementation for ComfyUI integration with vLLM-Omni's DALL-E compatible image generation API.
+- https://github.com/Comfy-Org/ComfyUI/tree/master/comfy_extras ComfyUI's built-in node implementations.
+
+UI/UX design references
+
+- https://github.com/sgl-project/sglang/pull/15271 SGLang Diffusion's official ComfyUI integration for image and video generation.
+- https://github.com/SXQBW/ComfyUI-Qwen-Omni A third party ComfyUI integration for Qwen Omni series.
+- https://github.com/flybirdxx/ComfyUI-Qwen-TTS https://github.com/DarioFT/ComfyUI-Qwen3-TTS Third  party ComfyUI integrations for Qwen TTS series.

apps/ComfyUI-vLLM-Omni/__init__.py

@@ -0,0 +1,52 @@
+"""Top-level package for comfyui_vllm_omni."""  # noqa: N999  # This is not a python library intended to be imported
+
+__all__ = [
+    "NODE_CLASS_MAPPINGS",
+    "NODE_DISPLAY_NAME_MAPPINGS",
+    "WEB_DIRECTORY",
+]
+
+__author__ = """Zeyu Huang"""
+__email__ = "11222265+fhfuih@users.noreply.github.com"
+__version__ = "0.0.1"
+
+from .comfyui_vllm_omni.nodes import (
+    VLLMOmniARSampling,
+    VLLMOmniComprehension,
+    VLLMOmniDiffusionSampling,
+    VLLMOmniGenerateImage,
+    VLLMOmniQwenTTSParams,
+    VLLMOmniSamplingParamsList,
+    VLLMOmniTTS,
+    VLLMOmniVoiceClone,
+)
+
+# A dictionary that contains all nodes you want to export with their names
+NODE_CLASS_MAPPINGS = {
+    # === Generation ===
+    "VLLMOmniGenerateImage": VLLMOmniGenerateImage,
+    "VLLMOmniComprehension": VLLMOmniComprehension,
+    "VLLMOmniTTS": VLLMOmniTTS,
+    "VLLMOmniVoiceClone": VLLMOmniVoiceClone,
+    # === Params ===
+    "VLLMOmniARSampling": VLLMOmniARSampling,
+    "VLLMOmniDiffusionSampling": VLLMOmniDiffusionSampling,
+    "VLLMOmniSamplingParamsList": VLLMOmniSamplingParamsList,
+    "VLLMOmniQwenTTSParams": VLLMOmniQwenTTSParams,
+}
+
+# A dictionary that contains the friendly/humanly readable titles for the nodes
+NODE_DISPLAY_NAME_MAPPINGS = {
+    # === Generation ===
+    "VLLMOmniGenerateImage": "Generate Image",
+    "VLLMOmniComprehension": "Multimodality Comprehension",
+    "VLLMOmniTTS": "TTS (Text to Speech)",
+    "VLLMOmniVoiceClone": "TTS Voice Cloning",
+    # === Params ===
+    "VLLMOmniARSampling": "AR Sampling Params",
+    "VLLMOmniDiffusionSampling": "Diffusion Sampling Params",
+    "VLLMOmniSamplingParamsList": "Multi-Stage Sampling Params List",
+    "VLLMOmniQwenTTSParams": "Qwen TTS Params",
+}
+
+WEB_DIRECTORY = "./web"

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

@@ -0,0 +1 @@
+# noqa: N999   # This is not a python library intended to be imported

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

@@ -0,0 +1,569 @@
+from typing import Literal, cast
+
+import torch
+from comfy_api.input import AudioInput, VideoInput
+
+from .utils.api_client import VLLMOmniClient
+from .utils.logger import get_logger
+from .utils.models import lookup_model_spec
+from .utils.types import AudioFormat
+from .utils.validators import (
+    add_sampling_parameters_to_stage,
+    validate_model_and_sampling_params_types,
+)
+
+logger = get_logger(__name__)
+
+
+class _VLLMOmniGenerateBase:
+    """Base class for vLLM-Omni generation nodes with shared functionality."""
+
+    CATEGORY = "vLLM-Omni"
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, url, model) -> str | Literal[True]:
+        """
+        Can only validate this model's own input. Cannot check inputs from other nodes.
+        See: https://docs.comfy.org/custom-nodes/backend/server_overview#validate_inputs
+        """
+        if not url:
+            return "URL must not be empty"
+        if not model:
+            return "Model must not be empty"
+        return True
+
+
+class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "url": ("STRING", {"default": "http://localhost:8000/v1"}),
+                "model": ("STRING", {"default": "Tongyi-MAI/Z-Image-Turbo"}),
+                "prompt": ("STRING", {"multiline": True}),
+                "negative_prompt": ("STRING", {"multiline": True, "default": ""}),
+                "width": ("INT", {"default": 512, "min": 64, "max": 2048}),
+                "height": ("INT", {"default": 512, "min": 64, "max": 2048}),
+            },
+            "optional": {
+                "image": ("IMAGE",),
+                "mask": ("MASK",),
+                # "video": ("VIDEO",),
+                # "audio": ("AUDIO",),
+                "sampling_params": ("SAMPLING_PARAMS",),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("image",)
+    FUNCTION = "generate"
+
+    async def generate(
+        self,
+        url: str,
+        model: str,
+        prompt: str,
+        width: int,
+        height: int,
+        negative_prompt: str | None = None,
+        image: torch.Tensor | None = None,
+        mask: torch.Tensor | None = None,
+        audio: AudioInput | None = None,  # Hidden & unused
+        video: VideoInput | None = None,  # Hidden & unused
+        sampling_params: dict | list[dict] | None = None,
+        **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:
+            raise ValueError("Mask input provided without an image input.")
+
+        client = VLLMOmniClient(url)
+
+        spec, pattern = lookup_model_spec(model)
+        is_bagel = pattern is not None and "bagel" in pattern.lower()
+
+        # Prefer DALL-E compatible API for simple (one-stage) diffusion models
+        if (spec is None or spec["stages"] == ["diffusion"]) and not is_bagel:
+            sampling_params = cast(dict | None, sampling_params)
+            if audio is None and image is None and video is None:
+                # No multimodal input --- use DALL-E image generation
+                logger.info("Using DALL-E image generation endpoint")
+                output = await client.generate_image(
+                    model=model,
+                    prompt=prompt,
+                    width=width,
+                    height=height,
+                    negative_prompt=negative_prompt,
+                    sampling_params=sampling_params,
+                )
+                return (output,)
+            elif image is not None and audio is None and video is None:
+                # Image and text input --- use DALL-E image edit
+                logger.info("Using DALL-E image edit endpoint")
+                output = await client.edit_image(
+                    model=model,
+                    prompt=prompt,
+                    image=image,
+                    width=width,
+                    height=height,
+                    negative_prompt=negative_prompt,
+                    mask=mask,
+                    sampling_params=sampling_params,
+                )
+                return (output,)
+
+        logger.info("Using chat completion endpoint")
+        sampling_params = add_sampling_parameters_to_stage(
+            model, sampling_params, "diffusion", width=width, height=height
+        )
+        logger.debug("Edited sampling params: %s", sampling_params)
+
+        output = await client.generate_image_chat_completion(
+            model=model,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image=image,
+            audio=audio,
+            video=video,
+            sampling_params=sampling_params,
+        )
+
+        return (output,)
+
+
+class VLLMOmniComprehension(_VLLMOmniGenerateBase):
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "url": ("STRING", {"default": "http://localhost:8000/v1"}),
+                "model": ("STRING", {"default": "Qwen/Qwen2.5-Omni-7B"}),
+                "prompt": ("STRING", {"multiline": True}),
+                "output_text": ("BOOLEAN", {"default": True}),
+                "output_audio": ("BOOLEAN", {"default": True}),
+                "use_audio_in_video": ("BOOLEAN", {"default": True}),
+            },
+            "optional": {
+                "image": ("IMAGE",),
+                "video": ("VIDEO",),
+                "audio": ("AUDIO",),
+                "sampling_params": ("SAMPLING_PARAMS",),
+            },
+        }
+
+    RETURN_TYPES = ("STRING", "AUDIO")
+    RETURN_NAMES = ("text_response", "audio_response")
+    FUNCTION = "generate"
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, url, model, output_text, output_audio) -> str | Literal[True]:  # type: ignore[reportIncompatibleMethodOverride]
+        super_validation = super().VALIDATE_INPUTS(url, model)
+        if isinstance(super_validation, str):
+            return super_validation
+        if not output_text and not output_audio:
+            return "At least one of output_text or output_audio must be True."
+        return True
+
+    async def generate(
+        self,
+        url: str,
+        model: str,
+        prompt: str,
+        image: torch.Tensor | None = None,
+        audio: AudioInput | None = None,
+        video: VideoInput | None = None,
+        sampling_params: dict | list[dict] | None = None,
+        output_text: bool = True,
+        output_audio: bool = True,
+        use_audio_in_video: bool = True,
+        **kwargs,
+    ) -> tuple[str, AudioInput]:
+        logger.info("Uncaught kwargs: %s", kwargs)
+        logger.debug("Got sampling params: %s", sampling_params)
+        validate_model_and_sampling_params_types(model, sampling_params)
+
+        client = VLLMOmniClient(url)
+        spec, pattern = lookup_model_spec(model)
+        is_bagel = pattern is not None and "bagel" in pattern.lower()
+
+        if is_bagel:
+            # A lot of special handlings here...
+            if output_audio:
+                raise ValueError("BAGEL models do not support audio output.")
+            if audio is not None or video is not None:
+                raise ValueError("BAGEL models do not support audio or video input.")
+            (
+                text_response,
+                _,
+            ) = await client.generate_comprehension_chat_completion(
+                model=model,
+                prompt=prompt,
+                image=image,
+                audio=None,
+                video=None,
+                sampling_params=sampling_params,
+                modalities=["text"],
+            )
+        else:
+            modalities = []
+            if output_text:
+                modalities.append("text")
+            if output_audio:
+                modalities.append("audio")
+
+            if use_audio_in_video and video is not None:
+                use_audio_in_video = True
+            else:
+                use_audio_in_video = False
+
+            (
+                text_response,
+                audio,
+            ) = await client.generate_comprehension_chat_completion(
+                model=model,
+                prompt=prompt,
+                image=image,
+                audio=audio,
+                video=video,
+                sampling_params=sampling_params,
+                modalities=modalities,
+                # == extra kwargs ==
+                mm_processor_kwargs={"use_audio_in_video": use_audio_in_video},
+            )
+
+        if text_response is None:
+            text_response = ""
+        if audio is None:
+            channels = 1
+            duration = 1
+            sample_rate = 44100
+            num_samples = int(round(duration * sample_rate))
+            waveform = torch.zeros((1, channels, num_samples), dtype=torch.float32)
+            audio = {"waveform": waveform, "sample_rate": sample_rate}
+
+        return (text_response, audio)
+
+
+class VLLMOmniTTS(_VLLMOmniGenerateBase):
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "url": ("STRING", {"default": "http://localhost:8000/v1"}),
+                "model": (
+                    "STRING",
+                    {"default": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"},
+                ),
+                "input": ("STRING", {"multiline": True}),
+                "voice": ("STRING", {"default": "Vivian"}),
+                "response_format": (["mp3", "opus", "aac", "flac", "wav", "pcm"],),
+                "speed": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.25, "max": 4.0, "step": 0.01},
+                ),
+            },
+            "optional": {
+                "model_specific_params": ("TTS_PARAMS",),
+            },
+        }
+
+    RETURN_TYPES = ("AUDIO",)
+    RETURN_NAMES = ("audio",)
+    FUNCTION = "generate"
+
+    async def generate(
+        self,
+        url: str,
+        model: str,
+        input: str,
+        voice: str,
+        response_format: AudioFormat,
+        speed: float,
+        model_specific_params: dict | None,
+        **kwargs,
+    ) -> tuple[AudioInput]:
+        logger.info("Got extra kwargs in TTS: %s", kwargs)
+
+        is_qwen_tts = "qwen3-tts" in model.lower()
+        extra_params_type = None if model_specific_params is None else model_specific_params["type"]
+        if not is_qwen_tts and extra_params_type == "qwen-tts":
+            raise ValueError(
+                "You have provided Qwen-specific TTS params."
+                "However, the model appears to not be a Qwen TTS model (no 'Qwen3-TTS' in model name)."
+            )
+
+        combined_params = {**kwargs, **(model_specific_params or {})}
+        combined_params.pop("type", None)  # Internal fields in model_specific_params
+
+        client = VLLMOmniClient(url)
+
+        audio = await client.generate_speech(
+            model=model,
+            input=input,
+            voice=voice,
+            response_format=response_format,
+            speed=speed,
+            **combined_params,
+        )
+        return (audio,)
+
+
+class VLLMOmniVoiceClone(_VLLMOmniGenerateBase):
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "url": ("STRING", {"default": "http://localhost:8000/v1"}),
+                "model": ("STRING", {"default": "Qwen/Qwen3-TTS-12Hz-1.7B-Base"}),
+                "input": ("STRING", {"multiline": True}),
+                "voice": ("STRING", {"default": "Vivian"}),
+                "response_format": (["mp3", "opus", "aac", "flac", "wav", "pcm"],),
+                "speed": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.25, "max": 4.0, "step": 0.01},
+                ),
+                "ref_audio": ("AUDIO",),
+                "ref_text": ("STRING", {"multiline": True}),
+                "x_vector_only_mode": ("BOOLEAN", {"default": False}),
+            },
+            "optional": {
+                "model_specific_params": ("TTS_PARAMS",),
+            },
+        }
+
+    RETURN_TYPES = ("AUDIO",)
+    RETURN_NAMES = ("audio",)
+    FUNCTION = "generate"
+
+    async def generate(
+        self,
+        url: str,
+        model: str,
+        input: str,
+        voice: str,
+        response_format: AudioFormat,
+        speed: float,
+        ref_audio: AudioInput,
+        ref_text: str,
+        x_vector_only_mode: bool,
+        model_specific_params: dict | None,
+        **kwargs,
+    ):
+        is_qwen_tts = "qwen3-tts" in model.lower()
+        extra_params_type = None if model_specific_params is None else model_specific_params["type"]
+        if not is_qwen_tts and extra_params_type == "qwen-tts":
+            raise ValueError(
+                "You have provided Qwen-specific TTS params."
+                "However, the model appears to not be a Qwen TTS model (no 'Qwen3-TTS' in model name)."
+            )
+
+        combined_params = {
+            "ref_audio": ref_audio,
+            "ref_text": ref_text,
+            "x_vector_only_mode": x_vector_only_mode,
+            **kwargs,
+            **(model_specific_params or {}),
+        }
+        combined_params.pop("type", None)  # Internal fields in model_specific_params
+
+        client = VLLMOmniClient(url)
+
+        audio = await client.generate_speech(
+            model=model,
+            input=input,
+            voice=voice,
+            response_format=response_format,
+            speed=speed,
+            **combined_params,
+        )
+        return (audio,)
+
+
+class VLLMOmniARSampling:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "max_tokens": ("INT", {"default": 100, "min": 1, "max": 10000}),
+                "temperature": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01},
+                ),
+                "top_p": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01},
+                ),
+                "repetition_penalty": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.01},
+                ),
+                # === Put seed at last. ===
+                # Whenever a field named "seed" is present, ComfyUI adds another field called "control after generate"
+                "seed": (
+                    "INT",
+                    {
+                        "default": -1,
+                        "min": -1,
+                        "step": 1,
+                        "tooltip": "-1 means to not provide a seed.",
+                    },
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("SAMPLING_PARAMS",)
+    RETURN_NAMES = ("AR sampling params",)
+    FUNCTION = "get_params"
+    CATEGORY = "vLLM-Omni/Sampling Params"
+
+    def get_params(self, seed, **kwargs):
+        params = {
+            "type": "autoregression",  # for internal use, removed before sending the request
+            **kwargs,
+        }
+        if seed >= 0:
+            params["seed"] = seed
+        return (params,)
+
+
+class VLLMOmniDiffusionSampling:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "n": (
+                    "INT",
+                    {
+                        "default": 1,
+                        "min": 0,
+                        "max": 10,
+                        "step": 1,
+                        "tooltip": "Number of images to generate",
+                    },
+                ),
+                "num_inference_steps": (
+                    "INT",
+                    {
+                        "default": 50,
+                        "min": 1,
+                        "max": 1000,
+                        "tooltip": "Number of denoising steps (higher = better quality, slower).",
+                    },
+                ),
+                "guidance_scale": (
+                    "FLOAT",
+                    {
+                        "default": 7.5,
+                        "min": 0.0,
+                        "max": 20.0,
+                        "step": 0.1,
+                        "tooltip": "Classifier-free guidance scale (higher = more prompt adherence).",
+                    },
+                ),
+                "true_cfg_scale": (
+                    "FLOAT",
+                    {
+                        "default": 1.0,
+                        "min": 0.0,
+                        "max": 20.0,
+                        "step": 0.5,
+                        "tooltip": "True CFG scale for advanced control (model-specific).",
+                    },
+                ),
+                "vae_use_slicing": (
+                    "BOOLEAN",
+                    {
+                        "default": False,
+                        "tooltip": "Enable VAE slicing for reduced memory usage (slight quality trade-off)",
+                    },
+                ),
+                # === Put seed at last. ===
+                # Whenever a field named "seed" is present, ComfyUI adds another field called "control after generate"
+                "seed": (
+                    "INT",
+                    {
+                        "default": -1,
+                        "min": -1,
+                        "step": 1,
+                        "tooltip": "-1 means to not provide a seed.",
+                    },
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("SAMPLING_PARAMS",)
+    RETURN_NAMES = ("diffusion sampling params",)
+    FUNCTION = "get_params"
+    CATEGORY = "vLLM-Omni/Sampling Params"
+
+    def get_params(self, seed, **kwargs):
+        params = {
+            "type": "diffusion",  # for internal use, removed before sending the request
+            **kwargs,
+        }
+        if seed >= 0:
+            params["seed"] = seed
+        return (params,)
+
+
+class VLLMOmniSamplingParamsList:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "param1": ("SAMPLING_PARAMS",),
+            },
+            "optional": {
+                "param2": ("SAMPLING_PARAMS",),
+                "param3": ("SAMPLING_PARAMS",),
+            },
+        }
+
+    RETURN_TYPES = ("SAMPLING_PARAMS",)
+    RETURN_NAMES = ("param list",)
+    FUNCTION = "aggregate"
+    CATEGORY = "vLLM-Omni/Sampling Params"
+
+    def aggregate(self, param1: dict, param2: dict | None = None, param3: dict | None = None):
+        for i, p in enumerate((param1, param2, param3)):
+            if isinstance(p, list):
+                raise ValueError(
+                    f"Input {i} is a Multi-Stage Sampling Params List. "
+                    f"Expected a single sampling parameters node (either AR or Diffusion)."
+                )
+
+        params = [param1]
+        if param2 is not None:
+            params.append(param2)
+        if param3 is not None:
+            params.append(param3)
+        return (params,)
+
+
+class VLLMOmniQwenTTSParams:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "task_type": (
+                    ["CustomVoice", "VoiceDesign", "Base"],
+                    {"default": "CustomVoice"},
+                ),
+                "language": (
+                    ["Auto", "Chinese", "English", "Japanese", "Korean"],
+                    {"default": "Auto"},
+                ),
+                "instructions": ("STRING", {"multiline": True}),
+                "max_new_tokens": ("INT", {"default": 2048, "min": 1}),
+            }
+        }
+
+    RETURN_TYPES = ("TTS_PARAMS",)
+    RETURN_NAMES = ("Qwen TTS params",)
+    FUNCTION = "get_params"
+    CATEGORY = "vLLM-Omni/TTS Params"
+
+    def get_params(self, **kwargs):
+        return ({"type": "qwen-tts", **kwargs},)

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

@@ -0,0 +1,473 @@
+"""
+An high-level API client adapter that forwards ComfyUI inputs to vLLM-Omni's REST API,
+and transforms the API responses back to ComfyUI formats.
+
+The image generation part is derived from dougbtv/comfyui-vllm-omni by Doug (@dougbtv).
+Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed under the MIT License.
+"""
+
+from typing import Any
+
+import aiohttp
+import av.error
+import torch
+from comfy_api.input import AudioInput, VideoInput
+
+from .format import (
+    audio_to_base64,
+    base64_to_audio,
+    base64_to_image_tensor,
+    bytes_to_audio,
+    image_tensor_to_base64,
+    image_tensor_to_png_bytes,
+    video_to_base64,
+)
+from .logger import get_logger, pretty_printer
+from .models import lookup_model_spec
+from .types import AudioFormat
+
+logger = get_logger(__name__)
+
+
+class VLLMOmniClient:
+    def __init__(self, base_url: str, timeout: float = 300.0):
+        self.base_url = base_url
+        self.timeout = aiohttp.ClientTimeout(total=timeout)
+
+    async def generate_image(
+        self,
+        *,
+        model: str,
+        prompt: str,
+        width: int,
+        height: int,
+        negative_prompt: str | None = None,
+        sampling_params: 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 = {
+            "model": model,
+            "prompt": prompt,
+            "size": size,
+            "response_format": "b64_json",
+        }
+        if negative_prompt:
+            payload["negative_prompt"] = negative_prompt
+        if sampling_params is not None:
+            # Only select specific sampling params
+            for k in (
+                "n",
+                "num_inference_steps",
+                "guidance_scale",
+                "true_cfg_scale",
+                "vae_use_slicing",
+            ):
+                if k in sampling_params and sampling_params[k] is not None:
+                    payload[k] = sampling_params[k]
+        logger.debug("img gen payload: %s", payload)
+
+        url = self.base_url + "/images/generations"
+        async with aiohttp.ClientSession(timeout=self.timeout) as session:
+            try:
+                async with session.post(
+                    url,
+                    json=payload,
+                    headers={"Content-Type": "application/json"},
+                ) 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}")
+                    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")
+
+                    image_tensors = []
+                    for idx, img in enumerate(data["data"]):
+                        if "b64_json" not in img:
+                            raise RuntimeError(f"API returned image #{idx} without 'b64_json' field")
+                        base64_str = img["b64_json"]
+                        tensor = base64_to_image_tensor(base64_str)
+                        image_tensors.append(tensor)
+                        logger.debug("Image #%d has shape %s", idx, tensor.shape)
+
+                    batch_tensor = torch.stack(image_tensors, dim=0)
+                    logger.debug("batch_tensor output has shape: %s", batch_tensor.shape)
+                    return batch_tensor
+
+            except aiohttp.ClientError as e:
+                raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
+
+    async def edit_image(
+        self,
+        *,
+        model: str,
+        prompt: str,
+        image: torch.Tensor,
+        width: int,
+        height: int,
+        negative_prompt: str | None = None,
+        mask: torch.Tensor | None = None,
+        sampling_params: dict | None = None,
+    ) -> torch.Tensor:
+        """Run image editing via DALLE API"""
+        await self._check_model_exist(model)
+
+        size = f"{width}x{height}"
+        image_filename = "image.png"  # Required for multipart form
+        form = aiohttp.FormData()
+        form.add_field("model", model)
+        form.add_field(
+            "image",
+            image_tensor_to_png_bytes(image, image_filename),
+            filename=image_filename,
+            content_type="image/png",
+        )
+        form.add_field("prompt", prompt)
+        form.add_field("size", size)
+        if negative_prompt:
+            form.add_field("negative_prompt", negative_prompt)
+        if sampling_params is not None:
+            # Only select specific sampling params
+            for k in ("n", "num_inference_steps", "guidance_scale", "true_cfg_scale"):
+                if k in sampling_params and sampling_params[k] is not None:
+                    form.add_field(k, str(sampling_params[k]))
+        if mask is not None:
+            mask_filename = "mask.png"
+            form.add_field(
+                "mask",
+                image_tensor_to_png_bytes(mask, mask_filename),
+                filename=mask_filename,
+                content_type="image/png",
+            )
+
+        url = self.base_url + "/images/edits"
+        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}")
+
+                    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")
+
+                    image_tensors = []
+                    for idx, img in enumerate(data["data"]):
+                        if "b64_json" not in img:
+                            raise RuntimeError(f"API returned image #{idx} without 'b64_json' field")
+                        base64_str = img["b64_json"]
+                        tensor = base64_to_image_tensor(base64_str)
+                        image_tensors.append(tensor)
+
+                    return torch.stack(image_tensors, dim=0)
+
+            except aiohttp.ClientError as e:
+                raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
+
+    async def generate_image_chat_completion(
+        self,
+        *,
+        model: str,
+        prompt: str,
+        negative_prompt: str | None = None,
+        image: torch.Tensor | None = None,
+        audio: AudioInput | None = None,
+        video: VideoInput | None = None,
+        sampling_params: dict | list[dict] | None = None,
+    ) -> torch.Tensor:
+        payload = VLLMOmniClient._prepare_chat_completion_messages(
+            model=model,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image=image,
+            audio=audio,
+            video=video,
+            sampling_params=sampling_params,
+            modalities=["image"],
+        )
+        choices = await self._generate_base_chat_completion(model, payload)
+
+        image_tensors = []
+        for idx, img_content in enumerate(choices[0]["message"]["content"]):
+            base64_str = img_content.get("image_url", {}).get("url", "")
+            if not base64_str:
+                raise RuntimeError(f"API returned image #{idx} without image url")
+            tensor = base64_to_image_tensor(base64_str)
+            image_tensors.append(tensor)
+
+        return torch.stack(image_tensors, dim=0)
+
+    async def generate_comprehension_chat_completion(
+        self,
+        *,
+        model: str,
+        prompt: str,
+        image: torch.Tensor | None = None,
+        audio: AudioInput | None = None,
+        video: VideoInput | None = None,
+        sampling_params: dict | list[dict] | None = None,
+        modalities: list[str] = ["text", "audio"],
+        **extra_body,
+    ) -> tuple[str | None, AudioInput | None]:
+        # Response may contain two choices: one with text, one with audio
+        payload = VLLMOmniClient._prepare_chat_completion_messages(
+            model=model,
+            prompt=prompt,
+            negative_prompt=None,
+            image=image,
+            audio=audio,
+            video=video,
+            sampling_params=sampling_params,
+            modalities=modalities,
+            **extra_body,
+        )
+
+        choices = await self._generate_base_chat_completion(model, payload)
+        text_response = None
+        audio_base64 = None
+        for choice in choices:
+            try:
+                text_response = choice["message"]["content"]
+            except (KeyError, TypeError):
+                # Either this case (text response) or the audio response case will be hit. Checking None's later.
+                pass
+            try:
+                audio_base64 = choice["message"]["audio"]["data"]
+            except (KeyError, TypeError):
+                # Either this case (text response) or the audio response case will be hit. Checking None's later.
+                pass
+        if audio_base64 is None and text_response is None:
+            raise RuntimeError(
+                "API response missing both '.message.audio' and 'message.content' fields."
+                f"The choices object is {choices}"
+            )
+        if audio_base64 is not None:
+            audio = base64_to_audio(audio_base64)
+            logger.debug(
+                "audio sample rate %d, audio shape %s, duration in second %f",
+                audio["sample_rate"],
+                audio["waveform"].shape,
+                audio["waveform"].shape[2] / audio["sample_rate"],
+            )
+        else:
+            audio = None
+        return text_response, audio
+
+    async def generate_speech(
+        self,
+        *,
+        model: str,
+        input: str,
+        voice: str,
+        response_format: AudioFormat,
+        speed: float,
+        **extra_params,
+    ) -> AudioInput:
+        await self._check_model_exist(model)
+
+        ref_audio: AudioInput | None = extra_params.pop("ref_audio", None)
+
+        payload = {
+            "model": model,
+            "input": input,
+            "voice": voice,
+            "response_format": response_format,
+            "speed": speed,
+            **extra_params,
+        }
+
+        if ref_audio is not None:
+            audio_base64 = audio_to_base64(ref_audio)
+            payload["ref_audio"] = audio_base64
+
+        logger.debug("Omni TTS payload: %s", pretty_printer.pformat(payload))
+
+        url = self.base_url + "/audio/speech"
+        async with aiohttp.ClientSession(timeout=self.timeout) as session:
+            try:
+                async with session.post(
+                    url,
+                    json=payload,
+                    headers={"Content-Type": "application/json"},
+                ) 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:
+                        audio_bytes = await response.read()
+                    except aiohttp.ContentTypeError as e:
+                        raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}")
+
+                    try:
+                        audio = bytes_to_audio(audio_bytes)
+                    except av.error.InvalidDataError as e:
+                        raise ValueError(
+                            f"Invalid audio data received from vLLM-Omni: {e}"
+                            "Check if you have input unsupported arguments (such as 'voice')"
+                        )
+                    return audio
+
+            except aiohttp.ClientError as e:
+                raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
+
+    async def _generate_base_chat_completion(self, model: str, payload: dict[str, Any]) -> list[dict[str, Any]]:
+        logger.debug("Omni payload: %s", pretty_printer.pformat(payload))
+        await self._check_model_exist(model)
+
+        url = self.base_url + "/chat/completions"
+        async with aiohttp.ClientSession(timeout=self.timeout) as session:
+            try:
+                async with session.post(
+                    url,
+                    json=payload,
+                    headers={"Content-Type": "application/json"},
+                ) 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}")
+
+                    logger.debug(
+                        "chat completion response: %s",
+                        pretty_printer.pformat(data),
+                    )
+
+                    try:
+                        return data["choices"]
+                    except (KeyError, TypeError):
+                        raise RuntimeError("Invalid JSON response from vLLM-Omni: missing 'choices' field")
+
+            except aiohttp.ClientError as e:
+                raise RuntimeError(f"Network error connecting to vLLM-Omni at {self.base_url}: {e}")
+
+    async def _check_model_exist(self, model: str):
+        url = self.base_url + "/models"
+        async with aiohttp.ClientSession(timeout=self.timeout) as session:
+            try:
+                async with session.get(
+                    url,
+                    headers={"Content-Type": "application/json"},
+                ) 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} "
+                            f"when getting hosted model list: {error_text}"
+                        )
+
+                    try:
+                        data = await response.json()
+                    except aiohttp.ContentTypeError as e:
+                        raise RuntimeError(f"Invalid JSON response when getting hosted model list from vLLM-Omni: {e}")
+
+            except aiohttp.ClientError as e:
+                raise RuntimeError(f"Network error connecting to vLLM-Omni at {self.base_url}: {e}")
+        try:
+            model_list = data["data"]
+            model_found = next((True for m in model_list if m["id"] == model), False)
+        except (KeyError, TypeError):
+            raise RuntimeError(f"Invalid JSON response of the hosted model list: {data}")
+
+        if not model_found:
+            raise ValueError(f"Model {model} not served at {self.base_url}.")
+
+    @staticmethod
+    def _prepare_chat_completion_messages(
+        *,
+        model: str,
+        prompt: str,
+        negative_prompt: str | None,
+        image: torch.Tensor | None = None,
+        audio: AudioInput | None = None,
+        video: VideoInput | None = None,
+        sampling_params: dict | list[dict] | None = None,
+        modalities: list[str] | None = None,  # diffusion don't have this field
+        **extra_body,
+    ):
+        message_content: list[dict] = [{"type": "text", "text": prompt}]
+        if image is not None:
+            message_content.append(
+                {
+                    "type": "image_url",
+                    "image_url": {"url": image_tensor_to_base64(image)},
+                }
+            )
+        if audio is not None:
+            message_content.append({"type": "audio_url", "audio_url": {"url": audio_to_base64(audio)}})
+        if video is not None:
+            message_content.append({"type": "video_url", "video_url": {"url": video_to_base64(video)}})
+        messages = [{"role": "user", "content": message_content}]
+
+        combined_extra_body: dict[str, Any] = {}
+        if sampling_params is not None:
+            spec, _ = lookup_model_spec(model)
+            is_single_sampling_param = isinstance(sampling_params, dict) or len(sampling_params) == 1
+
+            # Exclude internal key
+            if isinstance(sampling_params, dict):
+                sampling_params = {k: v for k, v in sampling_params.items() if k != "type"}
+            else:
+                sampling_params = [{k: v for k, v in sp.items() if k != "type"} for sp in sampling_params]
+
+            if (spec is None and is_single_sampling_param) or (spec is not None and spec["stages"] == ["diffusion"]):
+                # Diffusion format: extra_body directly contains sampling params.
+                # Validation should have taken care of matching sampling params' types.
+                # * Use this mode if the model is a simple one-stage diffusion model.
+                # * Fallback to this mode if model is not registered and a single sampling param is provided.
+                sampling_params = sampling_params if isinstance(sampling_params, dict) else sampling_params[0]
+                combined_extra_body: dict[str, Any] = {**sampling_params}
+                if "n" in combined_extra_body:
+                    combined_extra_body["num_outputs_per_prompt"] = combined_extra_body["n"]
+                    del combined_extra_body["n"]
+            else:
+                # Use AR style payload, extra_body has a sampling_params_list field
+                combined_extra_body: dict[str, Any] = {"sampling_params_list": sampling_params}
+
+        if negative_prompt:
+            combined_extra_body["negative_prompt"] = negative_prompt
+
+        if extra_body:
+            combined_extra_body.update(extra_body)
+
+        payload: dict[str, Any] = {"messages": messages, "model": model}
+        if combined_extra_body:
+            payload["extra_body"] = combined_extra_body
+        if modalities:
+            payload["modalities"] = modalities
+
+        spec, _ = lookup_model_spec(model)
+        if spec:
+            preprocessor = spec.get("payload_preprocessor", None)
+            if preprocessor is not None:
+                payload = preprocessor(payload)
+
+        return payload

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

@@ -0,0 +1,230 @@
+"""Image/tensor format helpers.
+
+The image generation part is derived from dougbtv/comfyui-vllm-omni by Doug (@dougbtv).
+Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed under the MIT License.
+"""
+
+import base64
+import mimetypes
+from io import BytesIO
+
+import av
+import numpy as np
+import torch
+from comfy_api.input import AudioInput, VideoInput
+from comfy_extras import nodes_audio
+from PIL import Image
+
+from .logger import get_logger
+
+logger = get_logger(__name__)
+
+
+def base64_to_image_tensor(base64_str: str, mode: str = "RGB") -> torch.Tensor:
+    """
+    Convert base64-encoded image to ComfyUI image tensor.
+
+    Args:
+        base64_str: Base64-encoded image string
+        mode: PIL image mode (default RGB for transparency support)
+
+    Returns:
+        torch.Tensor with shape (1, H, W, C) in float32 [0, 1] range
+
+    Raises:
+        ValueError: If base64 string is invalid or image cannot be decoded
+    """
+    if base64_str.startswith("data:image"):
+        _, base64_str = base64_str.split(",", 1)
+
+    try:
+        # Decode base64 to bytes
+        image_bytes = base64.b64decode(base64_str)
+    except Exception as e:
+        raise ValueError(f"Invalid base64 string: {e}")
+
+    # Create BytesIO object for PIL
+    image_bytesio = BytesIO(image_bytes)
+
+    # Open with PIL and convert to desired mode
+    try:
+        pil_image = Image.open(image_bytesio)
+        pil_image = pil_image.convert(mode)
+    except Exception as e:
+        raise RuntimeError(f"Failed to open image: {e}")
+
+    image_array = np.asarray(pil_image).astype(np.float32) / 255.0
+    image_tensor = torch.from_numpy(image_array)
+    return image_tensor
+
+
+def image_tensor_to_png_bytes(tensor: torch.Tensor, filename: str = "image.png") -> BytesIO:
+    """
+    Convert ComfyUI image tensor to PNG BytesIO for multipart upload.
+
+    This function converts a ComfyUI IMAGE tensor to a PNG-encoded BytesIO object
+    suitable for multipart/form-data upload. The BytesIO object has its .name
+    attribute set, which is required by aiohttp for file uploads.
+
+    Args:
+        tensor: ComfyUI IMAGE tensor with shape (B, H, W, C), dtype float32, range [0, 1]
+        filename: Name attribute to set on BytesIO (default: "image.png")
+
+    Returns:
+        BytesIO object containing PNG-encoded image with .name attribute set
+
+    Raises:
+        ValueError: If tensor format is invalid (not 4D, wrong dtype, etc.)
+    """
+    if tensor.ndim != 4:
+        raise ValueError(f"Expected 4D tensor with shape (B, H, W, C), got {tensor.ndim}D tensor")
+
+    image_tensor = tensor[0]  # Shape: (H, W, C)
+    image_np = (image_tensor.cpu().numpy() * 255.0).astype(np.uint8)
+    pil_image = Image.fromarray(image_np)
+
+    # Save to BytesIO as image file
+    img_bytes = BytesIO()
+    # Set name attribute (required for multipart upload and mimetype detection)
+    img_bytes.name = filename
+    try:
+        pil_image.save(img_bytes)
+    except Exception as e:
+        raise RuntimeError(f"Failed to save image as file: {e}")
+
+    # Reset position to beginning
+    img_bytes.seek(0)
+
+    return img_bytes
+
+
+def image_tensor_to_base64(tensor: torch.Tensor, filename: str = "image.png") -> str:
+    """
+    Convert ComfyUI image tensor to base64-encoded image string.
+
+    Args:
+        tensor: ComfyUI IMAGE tensor with shape (B, H, W, C), dtype float32, range [0, 1]
+        filename: Name attribute to set on BytesIO (default: "image.png")
+        format: File format of the output image file buffer (default: "PNG")
+
+    Returns:
+        Base64-encoded image string
+
+    Raises:
+        ValueError: If tensor format is invalid (not 4D, wrong dtype, etc.)
+    """
+    img_bytes = image_tensor_to_png_bytes(tensor, filename)
+    img_bytes.seek(0)
+    byte_data = img_bytes.read()
+    base64_str = base64.b64encode(byte_data).decode("utf-8")
+    mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
+    return f"data:{mime_type};base64,{base64_str}"
+
+
+def video_to_bytes(video: VideoInput, filename: str = "video.mp4") -> BytesIO:
+    output_buffer = BytesIO()
+    output_buffer.name = filename
+    video.save_to(output_buffer)
+    output_buffer.seek(0)
+    return output_buffer
+
+
+def video_to_base64(video: VideoInput, filename: str = "video.mp4") -> str:
+    video_buffer = video_to_bytes(video, filename)
+    video_buffer.seek(0)
+    byte_data = video_buffer.read()
+    base64_str = base64.b64encode(byte_data).decode("utf-8")
+    mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
+    return f"data:{mime_type};base64,{base64_str}"
+
+
+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"]
+    format = filename.rsplit(".", maxsplit=1)[1]
+    layout = "mono" if waveform.shape[0] == 1 else "stereo"
+
+    output_buffer = BytesIO()
+    output_buffer.name = filename
+    output_container = av.open(output_buffer, mode="w", format=format)
+    if format == "opus":
+        out_stream = output_container.add_stream("libopus", rate=sample_rate, layout=layout)
+        if quality == "64k":
+            out_stream.bit_rate = 64000  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+        elif quality == "96k":
+            out_stream.bit_rate = 96000  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+        elif quality == "128k":
+            out_stream.bit_rate = 128000  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+        elif quality == "192k":
+            out_stream.bit_rate = 192000  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+        elif quality == "320k":
+            out_stream.bit_rate = 320000  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+    elif format == "mp3":
+        out_stream = output_container.add_stream("libmp3lame", rate=sample_rate, layout=layout)
+        if quality == "V0":
+            out_stream.codec_context.qscale = 1  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+        elif quality == "128k":
+            out_stream.bit_rate = 128000  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+        elif quality == "320k":
+            out_stream.bit_rate = 320000  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+    else:  # format == "flac":
+        out_stream = output_container.add_stream("flac", rate=sample_rate, layout=layout)
+
+    frame = av.AudioFrame.from_ndarray(
+        waveform.movedim(0, 1).reshape(1, -1).float().numpy(),
+        format="flt",
+        layout=layout,
+    )
+    frame.sample_rate = sample_rate
+    frame.pts = 0
+    output_container.mux(out_stream.encode(frame))  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+    # Flush encoder
+    output_container.mux(out_stream.encode(None))  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+    output_container.close()
+    output_buffer.seek(0)
+
+    return output_buffer
+
+
+def audio_to_base64(audio: AudioInput, filename: str = "audio.mp3", quality: str = "128k") -> str:
+    audio_buffer = audio_to_bytes(audio, filename, quality)
+    audio_buffer.seek(0)
+    byte_data = audio_buffer.read()
+    base64_str = base64.b64encode(byte_data).decode("utf-8")
+    mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
+    return f"data:{mime_type};base64,{base64_str}"
+
+
+def bytes_to_audio(audio_bytes: bytes) -> AudioInput:
+    """
+    Convert audio bytes to ComfyUI audio tensor.
+
+    Args:
+        audio_bytes: Audio file bytes
+    Returns:
+        torch.Tensor with shape (B, C, T) in float32 range [-1, 1]
+    """
+    audio_buffer = BytesIO(audio_bytes)
+    waveform, sample_rate = nodes_audio.load(audio_buffer)  # type: ignore # Although expect string argument, it calls av.open underneath, which supports BytesIO (file-like)
+    return {"waveform": waveform.unsqueeze(0), "sample_rate": sample_rate}
+
+
+def base64_to_audio(base64_str: str) -> AudioInput:
+    """
+    Convert base64-encoded audio to ComfyUI audio tensor.
+
+    Args:
+        base64_str: Base64-encoded audio string
+    Returns:
+        torch.Tensor with shape (B, C, T) in float32 range [-1, 1]
+    """
+    if base64_str.startswith("data:audio"):
+        _, base64_str = base64_str.split(",", 1)
+
+    try:
+        # Decode base64 to bytes
+        audio_bytes = base64.b64decode(base64_str)
+    except Exception as e:
+        raise ValueError(f"Invalid base64 string: {e}")
+
+    return bytes_to_audio(audio_bytes)

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

@@ -0,0 +1,125 @@
+"""Centralized logger configuration for vLLM-Omni ComfyUI."""
+
+import logging
+import pprint
+import sys
+from typing import Any
+
+
+def get_logger(name: str) -> logging.Logger:
+    """
+    Get or create a logger with proper formatting.
+
+    Args:
+        name: Logger name (typically __name__ of the calling module)
+
+    Returns:
+        Configured logger instance
+    """
+    logger = logging.getLogger(name)
+
+    # Only configure if not already configured
+    if not logger.handlers:
+        logger.setLevel(logging.DEBUG)
+
+        # Create console handler
+        handler = logging.StreamHandler(sys.stdout)
+        handler.setLevel(logging.INFO)
+
+        # Create formatter
+        formatter = logging.Formatter(
+            fmt="(ComfyUI-vLLM-Omni) [%(levelname)s] %(asctime)s [%(filename)s:%(lineno)s] %(message)s",
+            datefmt="%Y-%m-%d %H:%M:%S",
+        )
+        handler.setFormatter(formatter)
+
+        # Add handler to logger
+        logger.addHandler(handler)
+
+        # Prevent propagation to root logger
+        logger.propagate = False
+
+    return logger
+
+
+class OmitBase64PrettyPrinter(pprint.PrettyPrinter):
+    """
+    A PrettyPrinter that redacts specific field names with '...'
+    wherever they appear in nested structures.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def _format(self, obj: Any, stream, indent: int, allowance: int, context, level: int) -> None:
+        # Check if this is a dict with redacted keys
+        if isinstance(obj, dict):
+            # Create a copy with redacted values
+            display_obj = {}
+            for key, value in obj.items():
+                if key == "data" or key == "url":
+                    if value.startswith("data:"):
+                        base64_header = value.split(",", 1)[0]
+                        display_obj[key] = f"{base64_header},***"
+                    elif value.startswith("http://") or value.startswith("https://"):
+                        display_obj[key] = value
+                    elif len(value) > 10:
+                        display_obj[key] = f"{value[:10]}***"
+                    else:
+                        display_obj[key] = value
+                else:
+                    display_obj[key] = value
+            obj = display_obj
+
+        # Handle list/tuple/set containing dicts that might have redacted keys
+        # (pprint will recursively call _format on nested items, so this
+        # handles arbitrary nesting automatically)
+
+        super()._format(obj, stream, indent, allowance, context, level)
+
+
+pretty_printer = OmitBase64PrettyPrinter()
+
+
+# ========== PPrint EXAMPLES ==========
+
+if __name__ == "__main__":
+    data = {
+        "messages": [
+            {
+                "role": "system",
+                "content": (
+                    "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group,"
+                    "capable of perceiving auditory and visual inputs, as well as generating text and speech."
+                ),
+            },
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": "What sound is it, and what is the drawing about?",
+                    },
+                    {
+                        "type": "text",
+                        "text": "What sound is it, and what is the drawing about?",
+                    },
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAwAAAA"},
+                    },
+                    {
+                        "type": "audio_url",
+                        "audio_url": {"url": "data:audio/mpeg;base64,SUQzBAAAAAAAIlRTU0UAAAAOAAADT="},
+                    },
+                ],
+            },
+        ],
+        "extra_body": {"mm_processor_kwargs": {"use_audio_in_video": False}},
+        "modalities": ["text"],
+    }
+
+    # Create printer that redacts 'password' and 'token' fields
+
+    print("\nRedactingPrettyPrinter (hides secrets):")
+    pretty_printer.pprint(data)

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

@@ -0,0 +1,119 @@
+import re
+
+from .types import Modality, ModelMode, Spec
+
+
+def _bagel_payload_preprocessor(payload: dict) -> dict:
+    try:
+        for message in payload["messages"]:
+            for content in message["content"]:
+                if content["type"] == "text":
+                    content["text"] = "<|im_start|>" + content["text"] + "<|im_end|>"
+    except (KeyError, TypeError):
+        raise RuntimeError("Internal Error: malformatted BAGEL payload")
+    extra_body = payload.pop("extra_body", {})
+    return {**payload, **extra_body}
+
+
+def _qwen25_payload_preprocessor(payload: dict) -> dict:
+    if payload["messages"][0]["role"] != "system":
+        payload["messages"] = [
+            {
+                "role": "system",
+                "content": (
+                    "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group,"
+                    "capable of perceiving auditory and visual inputs, as well as generating text and speech."
+                ),
+            },
+            *payload["messages"],
+        ]
+    return payload
+
+
+_MODEL_PIPELINE_SPECS: dict[str, Spec] = {
+    r"BAGEL-7B-MoT": {
+        "stages": [
+            "diffusion"  # The vLLM-Omni interface treats it as a single-stage diffusion model
+        ],
+        "modes": [
+            {
+                "mode": ModelMode.COMPREHENSION,
+                "input_modalities": [Modality.TEXT, Modality.IMAGE],
+            }
+        ],
+        "payload_preprocessor": _bagel_payload_preprocessor,
+    },
+    r"Qwen2.5-Omni*": {
+        "stages": ["autoregression", "autoregression", "autoregression"],
+        "payload_preprocessor": _qwen25_payload_preprocessor,
+        "modes": [
+            {
+                "mode": ModelMode.COMPREHENSION,
+                "input_modalities": [
+                    Modality.TEXT,
+                    Modality.IMAGE,
+                    Modality.VIDEO,
+                    Modality.AUDIO,
+                ],
+            }
+        ],
+    },
+    r"Qwen3-Omni*": {
+        "stages": ["autoregression", "autoregression", "autoregression"],
+        "modes": [
+            {
+                "mode": ModelMode.COMPREHENSION,
+                "input_modalities": [
+                    Modality.TEXT,
+                    Modality.IMAGE,
+                    Modality.VIDEO,
+                    Modality.AUDIO,
+                ],
+            }
+        ],
+    },
+}
+# Convert dict keys to regex patterns
+MODEL_PIPELINE_SPECS: dict[re.Pattern, Spec] = {}
+for k, v in _MODEL_PIPELINE_SPECS.items():
+    MODEL_PIPELINE_SPECS[re.compile(k)] = v
+del _MODEL_PIPELINE_SPECS
+
+
+def lookup_model_spec(model: str) -> tuple[Spec | None, str | None]:
+    try:
+        last_component = model.rstrip("/").rsplit("/", 1)[-1]
+    except IndexError:
+        last_component = model
+    for pattern, spec in MODEL_PIPELINE_SPECS.items():
+        if pattern.search(last_component):
+            return spec, pattern.pattern
+    return None, None
+
+
+# ============== DEMONSTRATION ==============
+
+if __name__ == "__main__":
+    test_paths = [
+        "Qwen/Qwen2.5-Omni-7B",
+        "MyModels/Qwen2.5-Omni-3B",
+        "/root/home/Qwen2.5-Omni-7B",
+        "Qwen/Qwen3-Omni",
+        "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+        "Custom/Path/UnknownModel-Instruct",
+        "Not/Matching/Anything",
+    ]
+
+    test_payload = {"messages": [{"role": "user", "content": "prompt"}]}
+
+    print("Testing registry lookups:\n")
+    for path in test_paths:
+        spec, _ = lookup_model_spec(path)
+        if spec:
+            if preprocessor := spec.get("payload_preprocessor"):
+                result = preprocessor(test_payload)
+                print(f"✓ {path:<40} → {result}")
+            else:
+                print(f"✓ {path:<40} → No preprocessor")
+        else:
+            print(f"✗ {path:<40} → No match")

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

@@ -0,0 +1,39 @@
+from collections.abc import Callable
+from enum import Enum, auto
+from typing import (
+    Any,
+    Literal,
+    NotRequired,
+    TypeAlias,
+    TypedDict,
+)
+
+AudioFormat: TypeAlias = Literal["mp3", "opus", "aac", "flac", "wav", "pcm"]
+
+
+class ModelMode(Enum):
+    IMAGE_GENERATION = auto()
+    VIDEO_GENERATION = auto()
+    AUDIO_GENERATION = auto()
+    COMPREHENSION = auto()
+
+
+class Modality(Enum):
+    TEXT = auto()  # maybe not useful. Prompt is always required
+    IMAGE = auto()
+    VIDEO = auto()
+    AUDIO = auto()
+
+
+class ModelModeSpec(TypedDict):
+    mode: ModelMode
+    input_modalities: list[Modality]
+
+
+PayloadPreprocessor: TypeAlias = Callable[[dict[str, Any]], dict[str, Any]]
+
+
+class Spec(TypedDict):
+    stages: list[Literal["diffusion", "autoregression"]]
+    modes: list[ModelModeSpec]
+    payload_preprocessor: NotRequired[PayloadPreprocessor]

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

@@ -0,0 +1,84 @@
+from .logger import get_logger
+from .models import lookup_model_spec
+
+logger = get_logger(__name__)
+
+
+def validate_model_and_sampling_params_types(
+    model_name: str,
+    sampling_param_list: dict | list[dict] | None = None,
+):
+    # Check if model name exists
+    if not model_name:
+        raise ValueError("Model name must not be empty.")
+
+    # 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")
+        return
+    if sampling_param_list is None:
+        return
+
+    # Check the number of stages and their data types
+    stages = pipeline_spec["stages"]
+    if isinstance(sampling_param_list, list):
+        # Check that the lengths match
+        if len(stages) != len(sampling_param_list):
+            raise ValueError(
+                f"Sampling parameter list length {len(sampling_param_list)} does not match "
+                f"number of stages {len(stages)} for model {model_name}."
+            )
+        # Check that each stage's type match
+        for i, sp in enumerate(sampling_param_list):
+            if "type" not in sp:
+                raise RuntimeError("Internal error: unknown sampling parameter type")
+            if sp["type"] != stages[i]:
+                raise ValueError(
+                    f"Sampling parameter type ({sp['type']}) does not match "
+                    f"stage type ({stages[i]}) at index {i} for model {model_name}."
+                )
+    elif isinstance(sampling_param_list, dict):
+        if "type" not in sampling_param_list:
+            raise RuntimeError("Internal error: unknown sampling parameter type")
+        # Check that the provided single sampling param matches all stages
+        elif any(stage != sampling_param_list["type"] for stage in stages):
+            raise ValueError(
+                f"When passing a single sampling parameter node, all stages of the model must match "
+                f"the provided sampling parameter's type. "
+                f"However, the stages of model {model_name} are: {stages}. "
+                f"The provided sampling parameter is {sampling_param_list['type']}"
+            )
+
+
+def add_sampling_parameters_to_stage(
+    model_name: str,
+    sampling_param_list: dict | list[dict] | None,
+    stage_type: str,
+    /,
+    **params_to_add,
+) -> dict | list[dict]:
+    """
+    Given a model's name and the sampling parameter list to query this model,
+    add arbitrary additional parameters to the sampling parameters of all stages of the given type.
+    """
+    pipeline_spec, _ = lookup_model_spec(model_name)
+    if not pipeline_spec:
+        logger.warning(
+            f"Since the model {model_name} is not in our list, we cannot ensure if "
+            f"the fields ({tuple(params_to_add.keys())}) are added to the correct stage's sampling params. "
+            f"We will do it heuristically."
+        )
+        pipeline_spec = {"stages": ["diffusion"]}
+
+    stages = pipeline_spec["stages"]
+    if isinstance(sampling_param_list, dict):
+        sampling_param_list.update(params_to_add)
+    elif sampling_param_list is None:
+        sampling_param_list = params_to_add.copy()
+    else:
+        for i, stage in enumerate(stages):
+            if stage == stage_type:
+                sampling_param_list[i].update(params_to_add)
+
+    return sampling_param_list

apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Annotated Example.json

@@ -0,0 +1 @@
+{"id":"1c99f525-0a37-45ba-a28a-7df7c3af66b4","revision":0,"last_node_id":12,"last_link_id":14,"nodes":[{"id":1,"type":"VLLMOmniComprehension","pos":[1191.2177053682556,144.66829928181377],"size":[400,268],"flags":{},"order":8,"mode":0,"inputs":[{"localized_name":"image","name":"image","shape":7,"type":"IMAGE","link":1},{"localized_name":"video","name":"video","shape":7,"type":"VIDEO","link":2},{"localized_name":"audio","name":"audio","shape":7,"type":"AUDIO","link":3},{"localized_name":"sampling_params","name":"sampling_params","shape":7,"type":"SAMPLING_PARAMS","link":14},{"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":"output_text","name":"output_text","type":"BOOLEAN","widget":{"name":"output_text"},"link":null},{"localized_name":"output_audio","name":"output_audio","type":"BOOLEAN","widget":{"name":"output_audio"},"link":null},{"localized_name":"use_audio_in_video","name":"use_audio_in_video","type":"BOOLEAN","widget":{"name":"use_audio_in_video"},"link":null}],"outputs":[{"localized_name":"text_response","name":"text_response","type":"STRING","links":[8]},{"localized_name":"audio_response","name":"audio_response","type":"AUDIO","links":[9]}],"properties":{"Node name for S&R":"VLLMOmniComprehension"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen2.5-Omni-7B","",true,true,true]},{"id":3,"type":"LoadVideo","pos":[729.5984141255855,-198.631920454299],"size":[282.798828125,233.0743408203125],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"file","name":"file","type":"COMBO","widget":{"name":"file"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"VIDEO","name":"VIDEO","type":"VIDEO","links":[2]}],"properties":{"Node name for S&R":"LoadVideo"},"widgets_values":["draw.mp4","image"]},{"id":4,"type":"LoadAudio","pos":[729.8037086965753,99.86963519703949],"size":[282.798828125,136],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"COMBO","widget":{"name":"audio"},"link":null},{"localized_name":"audioUI","name":"audioUI","type":"AUDIO_UI","widget":{"name":"audioUI"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"AUDIOUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"AUDIO","name":"AUDIO","type":"AUDIO","links":[3]}],"properties":{"Node name for S&R":"LoadAudio"},"widgets_values":["Megan-Fox.mp3",null,null]},{"id":5,"type":"VLLMOmniARSampling","pos":[510.3517536642828,658.073751009259],"size":[270,178],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"max_tokens","name":"max_tokens","type":"INT","widget":{"name":"max_tokens"},"link":null},{"localized_name":"temperature","name":"temperature","type":"FLOAT","widget":{"name":"temperature"},"link":null},{"localized_name":"top_p","name":"top_p","type":"FLOAT","widget":{"name":"top_p"},"link":null},{"localized_name":"repetition_penalty","name":"repetition_penalty","type":"FLOAT","widget":{"name":"repetition_penalty"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"AR sampling params","name":"AR sampling params","type":"SAMPLING_PARAMS","links":[12]}],"properties":{"Node name for S&R":"VLLMOmniARSampling"},"widgets_values":[100,1,1,1,-1,"randomize"]},{"id":7,"type":"VLLMOmniARSampling","pos":[503.33235181647115,419.34158016181806],"size":[270,178],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"max_tokens","name":"max_tokens","type":"INT","widget":{"name":"max_tokens"},"link":null},{"localized_name":"temperature","name":"temperature","type":"FLOAT","widget":{"name":"temperature"},"link":null},{"localized_name":"top_p","name":"top_p","type":"FLOAT","widget":{"name":"top_p"},"link":null},{"localized_name":"repetition_penalty","name":"repetition_penalty","type":"FLOAT","widget":{"name":"repetition_penalty"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"AR sampling params","name":"AR sampling params","type":"SAMPLING_PARAMS","links":[5,13]}],"properties":{"Node name for S&R":"VLLMOmniARSampling"},"widgets_values":[100,1,1,1,-1,"randomize"]},{"id":8,"type":"VLLMOmniSamplingParamsList","pos":[820.6056617389042,426.38372037182273],"size":[263.066015625,66],"flags":{},"order":7,"mode":0,"inputs":[{"localized_name":"param1","name":"param1","type":"SAMPLING_PARAMS","link":5},{"localized_name":"param2","name":"param2","shape":7,"type":"SAMPLING_PARAMS","link":12},{"localized_name":"param3","name":"param3","shape":7,"type":"SAMPLING_PARAMS","link":13}],"outputs":[{"localized_name":"param list","name":"param list"Line truncated

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

@@ -0,0 +1 @@
+{"id":"91f75acc-8040-40f6-865a-2e8a7cfd6672","revision":0,"last_node_id":11,"last_link_id":23,"nodes":[{"id":3,"type":"PreviewImage","pos":[1281.8455167767304,-69.02638461454333],"size":[305.6628685610174,307.0336884172169],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":22}],"outputs":[],"properties":{"Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":11,"type":"VLLMOmniGenerateImage","pos":[816.0962515873642,-137.51112584854445],"size":[400,278],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"image","name":"image","shape":7,"type":"IMAGE","link":21},{"localized_name":"mask","name":"mask","shape":7,"type":"MASK","link":null},{"localized_name":"sampling_params","name":"sampling_params","shape":7,"type":"SAMPLING_PARAMS","link":23},{"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}],"outputs":[{"localized_name":"image","name":"image","type":"IMAGE","links":[22]}],"properties":{"Node name for S&R":"VLLMOmniGenerateImage"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen-Image-Edit","Put this figure in a realistic mountain view","",512,512]},{"id":8,"type":"VLLMOmniDiffusionSampling","pos":[478.59266934006774,183.67711984955648],"size":[284.205078125,202],"flags":{},"order":0,"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":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"diffusion sampling params","name":"diffusion sampling params","type":"SAMPLING_PARAMS","links":[23]}],"properties":{"Node name for S&R":"VLLMOmniDiffusionSampling"},"widgets_values":[4,50,1,1,false,42,"fixed"]},{"id":10,"type":"MarkdownNote","pos":[227.99306819575278,-231.24143306069843],"size":[240.3326478922474,136.3505820791881],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[],"title":"Note: Task and Input","properties":{},"widgets_values":["vLLM-Omni nodes are categorized based on the output modality. The \"Generate Image\" node supports both text-to-image generation or image-to-image generation (a.k.a. image editing). The node will route to the correct endpoint 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":[21]},{"localized_name":"MASK","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["example.png","image"]}],"links":[[21,4,0,11,0,"IMAGE"],[22,11,0,3,0,"IMAGE"],[23,8,0,11,2,"SAMPLING_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":[276.8427259980847,549.4885226647582]}},"version":0.4}

apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni TTS.json

@@ -0,0 +1 @@
+{"id":"c0a11241-f0e6-4f5c-8488-b3d344a9df3d","revision":0,"last_node_id":7,"last_link_id":3,"nodes":[{"id":4,"type":"PreviewAudio","pos":[690.1759063763212,226.76444957981545],"size":[270,88],"flags":{},"order":5,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"AUDIO","link":2},{"localized_name":"audioUI","name":"audioUI","type":"AUDIO_UI","widget":{"name":"audioUI"},"link":null}],"outputs":[],"properties":{"Node name for S&R":"PreviewAudio"},"widgets_values":[]},{"id":5,"type":"LoadAudio","pos":[-83.61093919682321,103.80510898465587],"size":[282.798828125,136],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"COMBO","widget":{"name":"audio"},"link":null},{"localized_name":"audioUI","name":"audioUI","type":"AUDIO_UI","widget":{"name":"audioUI"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"AUDIOUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"AUDIO","name":"AUDIO","type":"AUDIO","links":[3]}],"properties":{"Node name for S&R":"LoadAudio"},"widgets_values":["Megan-Fox.mp3",null,null]},{"id":3,"type":"VLLMOmniVoiceClone","pos":[248.63466280064262,168.13112893706375],"size":[400,306],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"ref_audio","name":"ref_audio","type":"AUDIO","link":3},{"localized_name":"model_specific_params","name":"model_specific_params","shape":7,"type":"TTS_PARAMS","link":1},{"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":"input","name":"input","type":"STRING","widget":{"name":"input"},"link":null},{"localized_name":"voice","name":"voice","type":"STRING","widget":{"name":"voice"},"link":null},{"localized_name":"response_format","name":"response_format","type":"COMBO","widget":{"name":"response_format"},"link":null},{"localized_name":"speed","name":"speed","type":"FLOAT","widget":{"name":"speed"},"link":null},{"localized_name":"ref_text","name":"ref_text","type":"STRING","widget":{"name":"ref_text"},"link":null},{"localized_name":"x_vector_only_mode","name":"x_vector_only_mode","type":"BOOLEAN","widget":{"name":"x_vector_only_mode"},"link":null}],"outputs":[{"localized_name":"audio","name":"audio","type":"AUDIO","links":[2]}],"properties":{"Node name for S&R":"VLLMOmniVoiceClone"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen3-TTS-12Hz-1.7B-Base","Someone just spilled a cup of coffee on my jacket this morning in the subway. Now I have to wear this stained jacket for a whole day!","","mp3",1,"",true]},{"id":6,"type":"MarkdownNote","pos":[255.52667408769662,27.26842775454861],"size":[381.58204108637847,88],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[],"title":"Note: TTS Nodes","properties":{},"widgets_values":["Apart from the Voice Cloning node, there is also another TTS node for simple text-to-speech tasks (without extra inputs for reference audio)."],"color":"#432","bgcolor":"#000"},{"id":7,"type":"MarkdownNote","pos":[-196.99453003423977,598.3693693727904],"size":[382.4539509209102,88],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[],"title":"Note: Model-Specific Parameters","properties":{},"widgets_values":["TTS models often require some tailor-made parameters. If you need to customize these parameters, grab one that matches the model you are requesting from \"TTS Params\" subfolder."],"color":"#432","bgcolor":"#000"},{"id":2,"type":"VLLMOmniQwenTTSParams","pos":[-207.8925310560497,342.7942867650348],"size":[400,200],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"task_type","name":"task_type","type":"COMBO","widget":{"name":"task_type"},"link":null},{"localized_name":"language","name":"language","type":"COMBO","widget":{"name":"language"},"link":null},{"localized_name":"instructions","name":"instructions","type":"STRING","widget":{"name":"instructions"},"link":null},{"localized_name":"max_new_tokens","name":"max_new_tokens","type":"INT","widget":{"name":"max_new_tokens"},"link":null}],"outputs":[{"localized_name":"Qwen TTS params","name":"Qwen TTS params","type":"TTS_PARAMS","links":[1]}],"properties":{"Node name for S&R":"VLLMOmniQwenTTSParams"},"widgets_values":["Base","Auto","Super angry",2048]}],"links":[[1,2,0,3,1,"TTS_PARAMS"],[2,3,0,4,0,"AUDIO"],[3,5,0,3,0,"AUDIO"]],"groups":[{"id":1,"title":"Model-Specific Parameters","bounding":[-228.38328407738007,264.42790274716276,443.9887293164586,439.6853171248897],"color":"#3f789e","font_size":24,"flags":{}}],"config":{},"extra":{"workflowRendererVersion":"LG","ds":{"scale":1.1469075819486911,"offset":[898.8819468322333,258.7179979718389]}},"version":0.4}

apps/ComfyUI-vLLM-Omni/web/main.js

@@ -0,0 +1,21 @@
+/**
+ * @file This file is intended to add dynamic fields to vLLM-Omni nodes
+ * based on widget (in-node form fields) and input (connection link) values and changes.
+ * However, this functionality is currently disabled/commented out.
+ * Because it introduces too much complexity,
+ * and it may even conflict with the current backend (Python) validation for unknown reasons (pending ComfyUI upstream fixes).
+ */
+
+import { app } from "../../scripts/app.js";
+app.registerExtension({
+    name: "vllm.vllm_omni",
+    async beforeRegisterNodeDef(nodeType, nodeData, app) {
+        if (!nodeData.name.startsWith("VLLMOmni")) {
+            return
+        }
+        // Stub frontend plugin for now
+    },
+    async setup() {
+        console.info("vLLM-Omni Setup complete!")
+    },
+})

docker/Dockerfile.npu

@@ -1,5 +1,5 @@
-ARG VLLM_ASCEND_IMAGE=quay.nju.edu.cn/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.11.0rc2
+ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
+ARG VLLM_ASCEND_TAG=v0.14.0rc1
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 ARG APP_DIR=/vllm-workspace/vllm-omni
@@ -8,7 +8,7 @@ WORKDIR ${APP_DIR}
 COPY . .
 
 # Install vllm-omni with dev dependencies
-RUN pip install --no-cache-dir -e ".[dev]"
+RUN pip install --no-cache-dir -e . --no-build-isolation
 
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
 

docker/Dockerfile.npu.a3

@@ -0,0 +1,15 @@
+ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
+ARG VLLM_ASCEND_TAG=v0.14.0rc1-a3
+FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
+
+ARG APP_DIR=/vllm-workspace/vllm-omni
+WORKDIR ${APP_DIR}
+
+COPY . .
+
+# Install vllm-omni with dev dependencies
+RUN pip install --no-cache-dir -e . --no-build-isolation
+
+ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
+
+ENTRYPOINT []

docker/Dockerfile.rocm

@@ -15,11 +15,7 @@ RUN mkdir -p ${COMMON_WORKDIR}/vllm-omni
 
 # Step 2: Copy vllm-omni code and install without uv
 COPY . ${COMMON_WORKDIR}/vllm-omni
-RUN cd ${COMMON_WORKDIR}/vllm-omni && uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]"
-
-# When we are installing onnxruntime-rocm, we need to uninstall the system-installed onnxruntime first.
-# These are the dependencies of Qwen3-TTS.
-RUN uv pip uninstall onnxruntime --system && uv pip install --no-cache-dir onnxruntime-rocm sox --system
+RUN cd ${COMMON_WORKDIR}/vllm-omni && uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]" --no-build-isolation
 
 RUN ln -sf /usr/bin/python3 /usr/bin/python
 

docs/.nav.yml

@@ -11,17 +11,23 @@ nav:
   - Examples:
     - examples/README.md
     - Offline Inference:
+      - BAGEL-7B-MoT: user_guide/examples/offline_inference/bagel.md
       - Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
       - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
+      - LoRA Inference Examples: user_guide/examples/offline_inference/lora_inference.md
       - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
       - Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
       - Qwen3-TTS Offline Inference: user_guide/examples/offline_inference/qwen3_tts.md
+      - Text-To-Audio: user_guide/examples/offline_inference/text_to_audio.md
       - Text-To-Image: user_guide/examples/offline_inference/text_to_image.md
       - Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
     - Online Serving:
+      - BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
       - Image-To-Image: user_guide/examples/online_serving/image_to_image.md
+      - Online LoRA Inference (Diffusion): user_guide/examples/online_serving/lora_inference.md
       - Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
       - Qwen3-Omni: user_guide/examples/online_serving/qwen3_omni.md
+      - Qwen3-TTS Online Serving: user_guide/examples/online_serving/qwen3_tts.md
       - Text-To-Image: user_guide/examples/online_serving/text_to_image.md
   - General:
     - usage/*
@@ -38,6 +44,7 @@ nav:
       - Cache-DiT: user_guide/diffusion/cache_dit_acceleration.md
       - Parallelism Acceleration: user_guide/diffusion/parallelism_acceleration.md
       - CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md
+    - ComfyUI: features/comfyui.md
 - Developer Guide:
   - General:
     - contributing/README.md
@@ -54,6 +61,7 @@ nav:
     - Feature Design:
       - design/feature/disaggregated_inference.md
       - design/feature/ray_based_execution.md
+      - design/feature/omni_connectors/
     - Module Design:
       - design/module/ar_module.md
       - design/module/dit_module.md

docs/api/README.md

@@ -10,6 +10,9 @@ Main entry points for vLLM-Omni inference and serving.
 - [vllm_omni.entrypoints.chat_utils.OmniAsyncMultiModalContentParser][]
 - [vllm_omni.entrypoints.chat_utils.OmniAsyncMultiModalItemTracker][]
 - [vllm_omni.entrypoints.chat_utils.parse_chat_messages_futures][]
+- [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][]
+- [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][]
+- [vllm_omni.entrypoints.cli.benchmark.serve.OmniBenchmarkServingSubcommand][]
 - [vllm_omni.entrypoints.cli.serve.OmniServeCommand][]
 - [vllm_omni.entrypoints.client_request_state.ClientRequestState][]
 - [vllm_omni.entrypoints.log_utils.OrchestratorMetrics][]
@@ -26,7 +29,9 @@ Main entry points for vLLM-Omni inference and serving.
 
 Input data structures for multi-modal inputs.
 
+- [vllm_omni.inputs.data.OmniDiffusionSamplingParams][]
 - [vllm_omni.inputs.data.OmniEmbedsPrompt][]
+- [vllm_omni.inputs.data.OmniTextPrompt][]
 - [vllm_omni.inputs.data.OmniTokenInputs][]
 - [vllm_omni.inputs.data.OmniTokensPrompt][]
 - [vllm_omni.inputs.parse.parse_singleton_prompt_omni][]
@@ -58,6 +63,7 @@ Core scheduling and caching components.
 - [vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler][]
 - [vllm_omni.core.sched.output.OmniCachedRequestData][]
 - [vllm_omni.core.sched.output.OmniNewRequestData][]
+- [vllm_omni.core.sched.output.OmniSchedulerOutput][]
 - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.DistributedGroupResidualVectorQuantization][]
 - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.DistributedResidualVectorQuantization][]
 - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.EuclideanCodebook][]
@@ -88,20 +94,23 @@ Configuration classes.
 
 Worker classes and model runners for distributed inference.
 
-- [vllm_omni.diffusion.worker.gpu_diffusion_model_runner.GPUDiffusionModelRunner][]
-- [vllm_omni.diffusion.worker.gpu_diffusion_worker.GPUDiffusionWorker][]
-- [vllm_omni.diffusion.worker.gpu_diffusion_worker.WorkerProc][]
-- [vllm_omni.diffusion.worker.npu.npu_worker.NPUWorker][]
-- [vllm_omni.diffusion.worker.npu.npu_worker.NPUWorkerProc][]
+- [vllm_omni.diffusion.worker.diffusion_model_runner.DiffusionModelRunner][]
+- [vllm_omni.diffusion.worker.diffusion_worker.DiffusionWorker][]
+- [vllm_omni.diffusion.worker.diffusion_worker.WorkerProc][]
+- [vllm_omni.platforms.npu.worker.npu_ar_model_runner.ExecuteModelState][]
+- [vllm_omni.platforms.npu.worker.npu_ar_model_runner.NPUARModelRunner][]
+- [vllm_omni.platforms.npu.worker.npu_ar_worker.NPUARWorker][]
+- [vllm_omni.platforms.npu.worker.npu_generation_model_runner.NPUGenerationModelRunner][]
+- [vllm_omni.platforms.npu.worker.npu_generation_worker.NPUGenerationWorker][]
+- [vllm_omni.platforms.npu.worker.npu_model_runner.OmniNPUModelRunner][]
+- [vllm_omni.platforms.xpu.worker.xpu_ar_model_runner.XPUARModelRunner][]
+- [vllm_omni.platforms.xpu.worker.xpu_ar_worker.XPUARWorker][]
+- [vllm_omni.platforms.xpu.worker.xpu_generation_model_runner.XPUGenerationModelRunner][]
+- [vllm_omni.platforms.xpu.worker.xpu_generation_worker.XPUGenerationWorker][]
 - [vllm_omni.worker.gpu_ar_model_runner.ExecuteModelState][]
 - [vllm_omni.worker.gpu_ar_model_runner.GPUARModelRunner][]
 - [vllm_omni.worker.gpu_ar_worker.GPUARWorker][]
 - [vllm_omni.worker.gpu_generation_model_runner.GPUGenerationModelRunner][]
 - [vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker][]
 - [vllm_omni.worker.gpu_model_runner.OmniGPUModelRunner][]
-- [vllm_omni.worker.npu.npu_ar_model_runner.ExecuteModelState][]
-- [vllm_omni.worker.npu.npu_ar_model_runner.NPUARModelRunner][]
-- [vllm_omni.worker.npu.npu_ar_worker.NPUARWorker][]
-- [vllm_omni.worker.npu.npu_generation_model_runner.NPUGenerationModelRunner][]
-- [vllm_omni.worker.npu.npu_generation_worker.NPUGenerationWorker][]
-- [vllm_omni.worker.npu.npu_model_runner.OmniNPUModelRunner][]
+- [vllm_omni.worker.mixins.OmniWorkerMixin][]

docs/cli/bench/serve.md

@@ -52,7 +52,7 @@ You can use `vllm bench serve --omni --help=all` to get descriptions of all para
 - `--percentile-metrics`  
         "Comma-separated list of selected metrics to report percentiles."
                     "This argument specifies the metrics to report percentiles."
-                    'Allowed metric names are "ttft", "tpot", "itl", "e2el", "audio_ttfp", "audio_rtf". '
+                    'Allowed metric names are "ttft", "tpot", "itl", "e2el", "audio_ttfp", "audio_rtf", "audio_duration". '
 
 - `--save-result`  
 Specify to save benchmark results to a json file
@@ -220,48 +220,48 @@ If successful, you will see the following output:
 ============ Serving Benchmark Result ============
 Successful requests:                     2
 Failed requests:                         0
-Benchmark duration (s):                  24.35
-Request throughput (req/s):              0.08
+Benchmark duration (s):                  3.89
+Request throughput (req/s):              0.51
 Peak concurrent requests:                2.00
 ----------------End-to-end Latency----------------
-Mean E2EL (ms):                          22576.23
-Median E2EL (ms):                        22576.23
-P99 E2EL (ms):                           24205.72
+Mean E2EL (ms):                          3824.76
+Median E2EL (ms):                        3824.76
+P99 E2EL (ms):                           3888.54
 ================== Text Result ===================
 Total input tokens:                      30
-Total generated tokens:                  8973
-Output token throughput (tok/s):         368.52
-Peak output token throughput (tok/s):    81.00
+Total generated tokens:                  10101
+Output token throughput (tok/s):         2595.57
+Peak output token throughput (tok/s):    111.00
 Peak concurrent requests:                2.00
-Total Token throughput (tok/s):          369.76
+Total Token throughput (tok/s):          2603.28
 ---------------Time to First Token----------------
-Mean TTFT (ms):                          125.16
-Median TTFT (ms):                        125.16
-P99 TTFT (ms):                           155.88
+Mean TTFT (ms):                          117.15
+Median TTFT (ms):                        117.15
+P99 TTFT (ms):                           142.69
 -----Time per Output Token (excl. 1st token)------
-Mean TPOT (ms):                          5.01
-Median TPOT (ms):                        5.01
-P99 TPOT (ms):                           5.42
+Mean TPOT (ms):                          0.73
+Median TPOT (ms):                        0.73
+P99 TPOT (ms):                           0.74
 ---------------Inter-token Latency----------------
-Mean ITL (ms):                           34.15
-Median ITL (ms):                         0.01
-P99 ITL (ms):                            376.19
+Mean ITL (ms):                           16.47
+Median ITL (ms):                         16.19
+P99 ITL (ms):                            52.55
 ================== Audio Result ==================
-Total audio duration generated(s):       3.95
-Total audio frames generated:            94890
-Audio throughput(audio duration/s):      0.16
+Total audio duration generated(s):       15.79
+Total audio frames generated:            379050
+Audio throughput(audio duration/s):      4.06
 ---------------Time to First Packet---------------
-Mean AUDIO_TTFP (ms):                    11756.89
-Median AUDIO_TTFP (ms):                  11756.89
-P99 AUDIO_TTFP (ms):                     20854.25
+Mean AUDIO_TTFP (ms):                    3701.37
+Median AUDIO_TTFP (ms):                  3701.37
+P99 AUDIO_TTFP (ms):                     3762.25
 -----------------Real Time Factor-----------------
-Mean AUDIO_RTF:                          3.75
-Median AUDIO_RTF:                        3.75
-P99 AUDIO_RTF:                           7.39
+Mean AUDIO_RTF:                          0.47
+Median AUDIO_RTF:                        0.47
+P99 AUDIO_RTF:                           0.48
 ==================================================
 ```
 Notes:
-We use (audio generation time - first packet latency) / audio duration to calculate RTF.
+We use audio generation time / audio duration to calculate RTF.
 
 </details>
 

docs/features/comfyui.md

@@ -0,0 +1,70 @@
+# vLLM-Omni ComfyUI Integration
+
+vLLM-Omni offers a ComfyUI integration on top of its online serving API.
+It can send model inference requests to either a locally running vLLM-Omni service or a remote one.
+
+## Requirement
+
+- Python 3.12 or above
+- [ComfyUI installed](https://docs.comfy.org/installation/system_requirements)
+- [vLLM-Omni installed](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/) on either the same device or another device discoverable via the internet.
+- No need to install additional packages apart from those already required by ComfyUI.
+
+!!! tip
+    If you run both ComfyUI and vLLM-Omni on the same device, you can create separate virtual environments and use different Python versions for them.
+
+
+## Installation
+
+Copy the `apps/ComfyUI-vLLM-Omni` folder to the `custom_nodes` subfolder of your ComfyUI installation. Your directory should look like `ComfyUI/custom_nodes/ComfyUI-vLLM-Omni`.
+
+If you are running ComfyUI during copying, you should restart ComfyUI to load this extension.
+
+!!! tip
+    You can use utility websites such as https://download-directory.github.io/ to download a subdirectory of a repo. Also checkout community discussions (e.g., https://stackoverflow.com/questions/7106012/download-a-single-folder-or-directory-from-a-github-repository) for more info.
+
+On the device and virtual environment you run ComfyUI, launch ComfyUI with
+```bash
+cd ComfyUI
+
+# The regular way
+python main.py
+
+# If you are mainly using this node, launch it faster with
+python main.py --cpu
+```
+
+On the device and virtual environment you run vLLM-Omni, start a model service with
+```bash
+vllm serve The_Model_ID_to_Serve --omni --port 8000
+```
+
+Check **ComfyUI's sidebar -> Node Library**. There should be a new folder named **vLLM-Omni**.
+If no, check your shell running the ComfyUI process. There may be some error messages before the line `Import times for custom nodes:` and the line `To see the GUI go to: http://127.0.0.1:8188`.
+
+## Quickstart
+
+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
+- **TTS** and **TTS Voice Clone** for TTS tasks
+
+This extension also offers example workflows (at **ComfyUI sidebar -> Templates -> vLLM-Omni**)
+
+!!! info
+    The node UI and feature designs are intended to match vLLM-Omni online serving interfaces. It cannot offer more than what the interfaces support.
+
+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.
+- 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.
+    - For multi-stage models where all stages are either autoregression or diffusion, you can also connect only a single Sampling Params node, indicating that this set of sampling parameters will be used for all stages.
+
+## Examples & Screenshots
+
+Please read the [ComfyUI integration's README](https://github.com/vllm-project/vllm-omni/tree/main/apps/ComfyUI-vLLM-Omni) for more info.

docs/getting_started/installation/gpu.md

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

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

@@ -17,12 +17,10 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
 # --8<-- [start:pre-built-wheels]
 
 #### Installation of vLLM
-Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
-
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.14.0 --torch-backend=auto
+uv pip install vllm==0.15.0 --torch-backend=auto
 ```
 
 #### Installation of vLLM-Omni
@@ -36,10 +34,10 @@ 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.14.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.15.0 release version of the library
 
 ```bash
-uv pip install vllm==0.14.0 --torch-backend=auto
+uv pip install vllm==0.15.0 --torch-backend=auto
 ```
 
 The release 0.14.0 of vLLM is based on PyTorch 2.9.0 which requires CUDA 12.9 environment.
@@ -58,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.14.0
+git checkout v0.15.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.14.0/vllm-0.14.0-cp38-abi3-manylinux_2_31_x86_64.whl
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.15.0/vllm-0.15.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

@@ -9,10 +9,60 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
 
 # --8<-- [start:pre-built-wheels]
 
+#### Installation of vLLM
+
+vLLM-Omni is built based on vLLM. Please install it with command below.
+```bash
+uv pip install vllm==0.14.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.14.0/rocm700
+```
+
+#### Installation of vLLM-Omni
+
+```bash
+# we need to add --no-build-isolation as the torch
+# is not obtained from pypi, we have to install using the
+# torch installed in our environment
+uv pip install vllm-omni
+
+# Optional if want to run Qwen3 TTS
+uv pip uninstall onnxruntime # should be removed before we can install onnxruntime-rocm
+uv pip install onnxruntime-rocm sox
+```
+
 # --8<-- [end:pre-built-wheels]
 
 # --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.14.0 release version of the library
+
+```bash
+uv pip install vllm==0.14.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.14.0/rocm700
+```
+
+The release 0.14.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
+```bash
+git clone https://github.com/vllm-project/vllm-omni.git
+cd vllm-omni
+VLLM_OMNI_TARGET_DEVICE=rocm uv pip install -e .
+# OR
+uv pip install -e . --no-build-isolation
+```
+
+<details><summary>(Optional) Installation of vLLM from source</summary>
+If you want to check, modify or debug with source code of vLLM, install the library from source with the following instructions:
+
+```bash
+git clone https://github.com/vllm-project/vllm.git
+cd vllm
+git checkout v0.14.0
+python3 -m pip install -r requirements/rocm.txt
+python3 setup.py develop
+```
+
 # --8<-- [end:build-wheel-from-source]
 
 # --8<-- [start:build-docker]

docs/getting_started/installation/npu.md

@@ -16,7 +16,13 @@ vLLM-Omni supports NPU through the vLLM Ascend Plugin (vllm-ascend). This is a c
 
 ## Installation
 
-### Recommended
+### Set up using Docker
+
+=== "NPU"
+
+    --8<-- "docs/getting_started/installation/npu/npu.inc.md:pre-built-images"
+
+### Build wheel from source
 
 === "NPU"
 

docs/getting_started/installation/npu/npu.inc.md

@@ -13,10 +13,10 @@ export DEVICE0=/dev/davinci0
 export DEVICE1=/dev/davinci1
 # Update the vllm-ascend image
 # Atlas A2:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1
 # Atlas A3:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0-a3
-export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1-a3
+export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1
 docker run --rm \
     --name vllm-omni-npu \
     --shm-size=1g \
@@ -34,26 +34,79 @@ docker run --rm \
     -p 8000:8000 \
     -it $IMAGE bash
 
-# Install the missing dependency of mooncake in the origin image.
-apt update
-apt install libjemalloc2
-echo "export LD_PRELOAD=/usr/lib/$(uname -m)-linux-gnu/libjemalloc.so.2:$LD_PRELOAD" >> ~/.bashrc
-source ~/.bashrc
-
 # Inside the container, install vLLM-Omni from source
 cd /vllm-workspace
 git clone -b v0.14.0 https://github.com/vllm-project/vllm-omni.git
-cd vllm-omni
-pip install -v -e .
-export VLLM_WORKER_MULTIPROC_METHOD=spawn
 
-# (Optional) Disable mooncake for stable capability
-mv /usr/local/Ascend/ascend-toolkit/latest/python/site-packages/mooncake \
-   /usr/local/Ascend/ascend-toolkit/latest/python/site-packages/mooncake.disabled
+cd vllm-omni
+VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
+# OR pip install -v -e . --no-build-isolation
+export VLLM_WORKER_MULTIPROC_METHOD=spawn
 ```
 
 The default workdir is `/workspace`, with vLLM, vLLM-Ascend and vLLM-Omni code placed in `/vllm-workspace` installed in development mode.
 
 For other installation methods (pip installation, building from source, custom Docker builds), please refer to the [vllm-ascend installation guide](https://docs.vllm.ai/projects/ascend/en/latest/installation.html).
 
+We are keeping [issue #997](https://github.com/vllm-project/vllm-omni/issues/997) up to date with the aligned versions of vLLM, vLLM-Ascend, and vLLM-Omni, and also outlining the Q1 roadmap there.
+
 # --8<-- [end:installation]
+
+# --8<-- [start:pre-built-images]
+
+`vllm-ascend` offers Docker images for deployment. You can just pull the **prebuilt image** from the image repository [ascend/vllm-ascend](https://quay.io/repository/ascend/vllm-ascend?tab=tags) and run it with bash.
+
+Supported images as following.
+
+| image name | Hardware | OS |
+|-|-|-|
+| image-tag | Atlas A2 | Ubuntu |
+| image-tag-a3 | Atlas A3 | Ubuntu |
+
+vLLM-Omni offers Docker images for Ascend NPU deployment. You can just pull the **prebuilt image** from the image repository [ascend/vllm-ascend](https://quay.io/repository/ascend/vllm-ascend?tab=tags) and run it with bash.
+
+Here's an example deployment command that has been verified on 6 x NPUs:
+
+```bash
+# Atlas A2:
+# export IMAGE=quay.io/ascend/vllm-omni:v0.14.0
+# Atlas A3:
+# export IMAGE=quay.io/ascend/vllm-omni:v0.14.0-a3
+export IMAGE=quay.io/ascend/vllm-omni:v0.14.0
+docker run --rm \
+    --name vllm-omni-npu \
+    --shm-size=1g \
+    --device /dev/davinci0 \
+    --device /dev/davinci1 \
+    --device /dev/davinci2 \
+    --device /dev/davinci3 \
+    --device /dev/davinci4 \
+    --device /dev/davinci5 \
+    --device /dev/davinci_manager \
+    --device /dev/devmm_svm \
+    --device /dev/hisi_hdc \
+    -v /usr/local/dcmi:/usr/local/dcmi \
+    -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
+    -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
+    -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
+    -v /etc/ascend_install.info:/etc/ascend_install.info \
+    -v /root/.cache:/root/.cache \
+    -p 8000:8000 \
+    -it $IMAGE bash
+```
+
+!!! tip
+    You can use this docker image to serve models the same way you would with in vLLM! To do so, make sure you overwrite the default entrypoint (`vllm serve --omni`) which works only for models supported in the vLLM-Omni project.
+
+Or build IMAGE from **source code**:
+
+```bash
+git clone https://github.com/vllm-project/vllm-omni.git
+cd vllm-omni
+# A2
+docker build -t vllm-omni-dev-image:latest -f ./docker/Dockerfile.npu .
+# A3
+# docker build -t vllm-omni-dev-image:latest -f ./docker/Dockerfile.npu.a3 .
+```
+
+# --8<-- [end:pre-built-images]

docs/getting_started/quickstart.md

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

docs/models/supported_models.md

@@ -25,6 +25,7 @@ th {
 | `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` |
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` |
 | `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
+| `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` |
 |`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |
 | `WanPipeline` | Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` |
 | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` |
@@ -51,6 +52,7 @@ th {
 
 | Architecture | Models | Example HF Models |
 |--------------|--------|-------------------|
+| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` |
 | `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B`|
 | `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |
 | `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` |
@@ -59,3 +61,8 @@ th {
 | `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` |
 |`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |
+|`LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` |
+|`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` |
+|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` |
+|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` |
+|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` |

docs/usage/faq.md

@@ -22,7 +22,7 @@ A: At first, you can check current [issues](https://github.com/vllm-project/vllm
 
 > Q: Does vLLM-Omni support AWQ or any other quantization?
 
-A: vLLM-Omni partitions model into several stages. For AR stages, it will reuse main logic of LLMEngine in vLLM. So current quantization supported in vLLM should be also supported in vLLM-Omni for them. But systematic verification is ongoing. For quantization for DiffusionEngine, we are working on it. Please stay tuned and welcome contribution!
+A: We plan to introduce GGUF FP8 prequantized models and online FP8 quantization in version 0.16.0. Support for other quantization types will follow in future releases. For details, please see our [Q1 quantization roadmap](https://github.com/vllm-project/vllm-omni/issues/1057).
 
 > Q: Does vLLM-Omni support multimodal streaming input and output?
 

docs/user_guide/diffusion/parallelism_acceleration.md

@@ -49,7 +49,7 @@ The following table shows which models are currently supported by parallelism me
 
 | Model | Model Identifier | Ulysses-SP | Ring-SP | Tensor-Parallel |
 |-------|------------------|------------|---------|--------------------------|
-| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ❌ |
+| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ |
 
 ### Tensor Parallelism
 

docs/user_guide/diffusion_acceleration.md

@@ -50,6 +50,7 @@ The following table shows which models are currently supported by each accelerat
 | **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ❌ | ❌ |
 | **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ✅ | ❌ | ❌ | ✅ |
 | **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ❌ | ❌ | ❌ |
+| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ❌ |
 
 ### VideoGen
 

docs/user_guide/examples/offline_inference/bagel.md

@@ -2,6 +2,7 @@
 
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/bagel>.
 
+
 ## Set up
 
 Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
@@ -99,7 +100,7 @@ python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
 
 BAGEL-7B-MoT supports **multiple modality modes** for different use cases.
 
-The default yaml configuration deploys Thinker and DiT on the same GPU. You can use the default configuration file: [`bagel.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel.yaml)
+The default yaml configuration deploys Thinker and DiT on the same GPU. You can use the default configuration file: [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel.yaml)
 
 #### 📌 Command Line Arguments (end2end.py)
 
@@ -177,3 +178,10 @@ sudo apt install ffmpeg
 | Stage-0 (Thinker)   | **15.04 GiB** **+ KV Cache** |
 | Stage-1 (DiT)       | **26.50 GiB**                |
 | Total               | **~42 GiB + KV Cache**       |
+
+## Example materials
+
+??? abstract "end2end.py"
+    ``````py
+    --8<-- "examples/offline_inference/bagel/end2end.py"
+    ``````

docs/user_guide/examples/offline_inference/image_to_image.md

@@ -47,10 +47,15 @@ Key arguments:
 - `--image`: path(s) to the source image(s) (PNG/JPG, converted to RGB). Can specify multiple images.
 - `--prompt` / `--negative_prompt`: text description (string).
 - `--cfg_scale`: true classifier-free guidance scale (default: 4.0). Classifier-free guidance is enabled by setting cfg_scale > 1 and providing a negative_prompt. Higher guidance scale encourages images closely linked to the text prompt, usually at the expense of lower image quality.
-- `--cfg_parallel_size`: the number of devices to run CFG Parallel. CFG Parallel is valid only if classifier-free guidance is enabled and `cfg_parallel_size` is set to 2.
 - `--guidance_scale`: guidance scale for guidance-distilled models (default: 1.0, disabled). Unlike classifier-free guidance (--cfg_scale), guidance-distilled models take the guidance scale directly as an input parameter. Enabled when guidance_scale > 1. Ignored when not using guidance-distilled models.
 - `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower).
 - `--output`: path to save the generated PNG.
+- `--vae_use_slicing`: enable VAE slicing for memory optimization.
+- `--vae_use_tiling`: enable VAE tiling for memory optimization.
+- `--cfg_parallel_size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+
+> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
 
 ## Example materials
 

docs/user_guide/examples/offline_inference/image_to_video.md

@@ -7,12 +7,19 @@ This example demonstrates how to generate videos from images using Wan2.2 Image-
 
 ## Local CLI Usage
 
+Download the example image:
+
+```bash
+wget https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg
+```
+
 ### Wan2.2-I2V-A14B-Diffusers (MoE)
+
 ```bash
 python image_to_video.py \
   --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
-  --image input.png \
-  --prompt "A cat playing with yarn, smooth motion" \
+  --image cherry_blossom.jpg \
+  --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
   --negative_prompt "<optional quality filter>" \
   --height 480 \
   --width 832 \
@@ -27,11 +34,12 @@ python image_to_video.py \
 ```
 
 ### Wan2.2-TI2V-5B-Diffusers (Unified)
+
 ```bash
 python image_to_video.py \
   --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
-  --image input.png \
-  --prompt "A cat playing with yarn, smooth motion" \
+  --image cherry_blossom.jpg \
+  --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
   --negative_prompt "<optional quality filter>" \
   --height 480 \
   --width 832 \
@@ -52,12 +60,17 @@ Key arguments:
 - `--num_frames`: Number of frames (default 81).
 - `--guidance_scale` and `--guidance_scale_high`: CFG scale (applied to low/high-noise stages for MoE).
 - `--negative_prompt`: Optional list of artifacts to suppress.
-- `--cfg_parallel_size`: the number of devices to run CFG Parallel. CFG Parallel is valid only if classifier-free guidance is enabled and `cfg_parallel_size` is set to 2.
 - `--boundary_ratio`: Boundary split ratio for two-stage MoE models.
 - `--flow_shift`: Scheduler flow shift (5.0 for 720p, 12.0 for 480p).
 - `--num_inference_steps`: Number of denoising steps (default 50).
 - `--fps`: Frames per second for the saved MP4 (requires `diffusers` export_to_video).
 - `--output`: Path to save the generated video.
+- `--vae_use_slicing`: Enable VAE slicing for memory optimization.
+- `--vae_use_tiling`: Enable VAE tiling for memory optimization.
+- `--cfg_parallel_size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+
+> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
 
 ## Example materials
 

docs/user_guide/examples/offline_inference/lora_inference.md

@@ -1,8 +1,9 @@
-# LoRA-Inference
+# LoRA Inference(Diffusion)
 
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/lora_inference>.
 
-This contains examples for using LoRA (Low-Rank Adaptation) adapters with vLLM-omni diffusion models for offline inference.
+
+This directory contains examples for using LoRA (Low-Rank Adaptation) adapters with vLLM-omni diffusion models for offline inference.
 The example uses the  `stabilityai/stable-diffusion-3.5-medium` as the default model, but you can replace it with other models in vLLM-omni.
 
 ## Overview

docs/user_guide/examples/offline_inference/qwen3_tts.md

@@ -1,4 +1,4 @@
-# Qwen3-TTS Offline Inference
+# Qwen3-TTS
 
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen3_tts>.
 
@@ -16,6 +16,15 @@ Qwen3 TTS provides multiple task variants for speech generation:
 ## Setup
 Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
 
+### ROCm Dependencies
+
+You will need to install these two dependencies `onnxruntime-rocm` and `sox`.
+
+```
+pip uninstall onnxruntime # should be removed before we can install onnxruntime-rocm
+pip install onnxruntime-rocm sox
+```
+
 ## Quick Start
 
 Run a single sample for a task:

docs/user_guide/examples/offline_inference/text_to_audio.md

@@ -0,0 +1,47 @@
+# Text-To-Audio
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_audio>.
+
+
+The `stabilityai/stable-audio-open-1.0` pipeline generates audio from text prompts.
+
+## Prerequisites
+
+If you use a gated model (e.g., `stabilityai/stable-audio-open-1.0`), ensure you have access:
+
+1. **Accept Model License**: Visit the model page on Hugging Face (e.g., [stabilityai/stable-audio-open-1.0]) and accept the user agreement.
+2. **Authenticate**: Log in to Hugging Face locally to access the gated model.
+   ```bash
+   huggingface-cli login
+   ```
+
+## Local CLI Usage
+
+```bash
+python text_to_audio.py \
+  --model stabilityai/stable-audio-open-1.0 \
+  --prompt "The sound of a hammer hitting a wooden surface" \
+  --negative_prompt "Low quality" \
+  --seed 42 \
+  --guidance_scale 7.0 \
+  --audio_length 10.0 \
+  --num_inference_steps 100 \
+  --output stable_audio_output.wav
+```
+
+Key arguments:
+
+- `--prompt`: text description (string).
+- `--negative_prompt`: negative prompt for classifier-free guidance.
+- `--seed`: integer seed for deterministic generation.
+- `--guidance_scale`: classifier-free guidance scale.
+- `--audio_length`: audio duration in seconds.
+- `--num_inference_steps`: diffusion sampling steps.(more steps = higher quality, slower).
+- `--output`: path to save the generated WAV file.
+
+## Example materials
+
+??? abstract "text_to_audio.py"
+    ``````py
+    --8<-- "examples/offline_inference/text_to_audio/text_to_audio.py"
+    ``````

docs/user_guide/examples/offline_inference/text_to_image.md

@@ -51,7 +51,8 @@ if __name__ == "__main__":
 
     For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input
     list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
-    batched inputs, you can [modify this configuration](../../../configuration/stage_configs.md) to let the model accept a longer batch of prompts.
+    batched inputs, you can [modify this configuration](../../../configuration/stage_configs.md
+    ) to let the model accept a longer batch of prompts.
 
 Apart from string prompt, vLLM-Omni also supports dictionary prompts in the same style as vLLM.
 This is useful for models that support negative prompts.
@@ -95,11 +96,16 @@ Key arguments:
 - `--prompt`: text description (string).
 - `--seed`: integer seed for deterministic sampling.
 - `--cfg_scale`: true CFG scale (model-specific guidance strength).
-- `--cfg_parallel_size`: the number of devices to run CFG Parallel. CFG Parallel is valid only if classifier-free guidance is enabled and `cfg_parallel_size` is set to 2.
 - `--num_images_per_prompt`: number of images to generate per prompt (saves as `output`, `output_1`, ...).
 - `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower).
 - `--height/--width`: output resolution (defaults 1024x1024).
 - `--output`: path to save the generated PNG.
+- `--vae_use_slicing`: enable VAE slicing for memory optimization.
+- `--vae_use_tiling`: enable VAE tiling for memory optimization.
+- `--cfg_parallel_size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+
+> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
 
 > ℹ️ Qwen-Image currently publishes best-effort presets at `1328x1328`, `1664x928`, `928x1664`, `1472x1140`, `1140x1472`, `1584x1056`, and `1056x1584`. Adjust `--height/--width` accordingly for the most reliable outcomes.
 

docs/user_guide/examples/offline_inference/text_to_video.md

@@ -12,10 +12,11 @@ python text_to_video.py \
   --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
   --negative_prompt "<optional quality filter>" \
   --height 480 \
-  --width 640 \
-  --num_frames 32 \
+  --width 832 \
+  --num_frames 33 \
   --guidance_scale 4.0 \
   --guidance_scale_high 3.0 \
+  --flow_shift 12.0 \
   --num_inference_steps 40 \
   --fps 16 \
   --output t2v_out.mp4
@@ -24,14 +25,19 @@ python text_to_video.py \
 Key arguments:
 
 - `--prompt`: text description (string).
-- `--height/--width`: output resolution (defaults 720x1280). Dimensions should align with Wan VAE downsampling (multiples of 8).
+- `--height/--width`: output resolution (defaults 480x832, i.e. 480P). Dimensions should align with Wan VAE downsampling (multiples of 8).
 - `--num_frames`: Number of frames (Wan default is 81).
-- `--guidance_scale` and `--guidance_scale_high`: CFG scale (applied to low/high)..
+- `--guidance_scale` and `--guidance_scale_high`: CFG scale (applied to low/high).
 - `--negative_prompt`: optional list of artifacts to suppress (the PR demo used a long Chinese string).
-- `--cfg_parallel_size`: the number of devices to run CFG Parallel. CFG Parallel is valid only if classifier-free guidance is enabled and `cfg_parallel_size` is set to 2.
-- `--boundary_ratio`: Boundary split ratio for low/high DiT.
+- `--boundary_ratio`: Boundary split ratio for low/high DiT. Default `0.875` uses both transformers for best quality. Set to `1.0` to load only the low-noise transformer (saves noticeable memory with good quality, recommended if memory is limited). Set to `0.0` loads only the high-noise transformer (not recommended, lower quality).
 - `--fps`: frames per second for the saved MP4 (requires `diffusers` export_to_video).
 - `--output`: path to save the generated video.
+- `--vae_use_slicing`: enable VAE slicing for memory optimization.
+- `--vae_use_tiling`: enable VAE tiling for memory optimization.
+- `--cfg_parallel_size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+
+> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
 
 ## Example materials
 

docs/user_guide/examples/online_serving/bagel.md

@@ -2,9 +2,10 @@
 
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/bagel>.
 
+
 ## 🛠️ Installation
 
-Please refer to [README.md](../../../README.md)
+Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
 
 ## Run examples (BAGEL-7B-MoT)
 
@@ -74,7 +75,7 @@ python openai_chat_client.py \
 
 BAGEL-7B-MoT supports **multiple modality modes** for different use cases.
 
-The default yaml configuration deploys Thinker and DiT on the same GPU. You can use the default configuration file: [`bagel.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel.yaml)
+The default yaml configuration deploys Thinker and DiT on the same GPU. You can use the default configuration file: [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel.yaml)
 
 | Modality    | Input        | Output | Description                            |
 | ----------- | ------------ | ------ | -------------------------------------- |
@@ -230,3 +231,14 @@ sudo apt install ffmpeg
 | Stage-0 (Thinker)   | **15.04 GiB** **+ KV Cache** |
 | Stage-1 (DiT)       | **26.50 GiB**                |
 | Total               | **~42 GiB + KV Cache**       |
+
+## Example materials
+
+??? abstract "openai_chat_client.py"
+    ``````py
+    --8<-- "examples/online_serving/bagel/openai_chat_client.py"
+    ``````
+??? abstract "run_server.sh"
+    ``````sh
+    --8<-- "examples/online_serving/bagel/run_server.sh"
+    ``````

docs/user_guide/examples/online_serving/lora_inference.md

@@ -1,7 +1,8 @@
-# LoRA-Inference
+# LoRA Inference(Diffusion)
 
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/lora_inference>.
 
+
 This example shows how to use **per-request LoRA** with vLLM-Omni diffusion models via the OpenAI-compatible Chat Completions API.
 
 > Note: The LoRA adapter path must be readable on the **server** machine (usually a local path or a mounted directory).
@@ -55,15 +56,17 @@ lora_adapter/
 └── adapter_model.safetensors
 ```
 
+## Example materials
+
 ??? abstract "openai_chat_client.py"
     ``````py
     --8<-- "examples/online_serving/lora_inference/openai_chat_client.py"
     ``````
 ??? abstract "run_curl_lora_inference.sh"
-    ``````py
+    ``````sh
     --8<-- "examples/online_serving/lora_inference/run_curl_lora_inference.sh"
     ``````
 ??? abstract "run_server.sh"
-    ``````py
+    ``````sh
     --8<-- "examples/online_serving/lora_inference/run_server.sh"
     ``````

docs/user_guide/examples/online_serving/qwen3_tts.md

@@ -0,0 +1,189 @@
+# Qwen3-TTS Online Serving
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/qwen3_tts>.
+
+
+This directory contains examples for running Qwen3-TTS models with vLLM-Omni's online serving API.
+
+## Supported Models
+
+| Model | Task Type | Description |
+|-------|-----------|-------------|
+| `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | CustomVoice | Predefined speaker voices with optional style control |
+| `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | VoiceDesign | Natural language voice style description |
+| `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Base | Voice cloning from reference audio |
+
+## Quick Start
+
+### 1. Start the Server
+
+```bash
+# CustomVoice model (default)
+./run_server.sh
+
+# Or specify task type
+./run_server.sh CustomVoice
+./run_server.sh VoiceDesign
+./run_server.sh Base
+```
+
+### 2. Run the Client
+
+```bash
+# CustomVoice: Use predefined speaker
+python openai_speech_client.py \
+    --text "你好，我是通义千问" \
+    --voice Vivian \
+    --language Chinese
+
+# CustomVoice with style instruction
+python openai_speech_client.py \
+    --text "今天天气真好" \
+    --voice Ryan \
+    --instructions "用开心的语气说"
+
+# VoiceDesign: Describe the voice style
+python openai_speech_client.py \
+    --model Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign \
+    --task-type VoiceDesign \
+    --text "哥哥，你回来啦" \
+    --instructions "体现撒娇稚嫩的萝莉女声，音调偏高"
+
+# Base: Voice cloning
+python openai_speech_client.py \
+    --model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
+    --task-type Base \
+    --text "Hello, this is a cloned voice" \
+    --ref-audio /path/to/reference.wav \
+    --ref-text "Original transcript of the reference audio"
+```
+
+### 3. Using curl
+
+```bash
+# Simple TTS request
+curl -X POST http://localhost:8000/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Hello, how are you?",
+        "voice": "Vivian",
+        "language": "English"
+    }' --output output.wav
+
+# With style instruction
+curl -X POST http://localhost:8000/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "I am so excited!",
+        "voice": "Vivian",
+        "instructions": "Speak with great enthusiasm"
+    }' --output excited.wav
+
+# List available voices in CustomVoice models
+curl http://localhost:8000/v1/audio/voices
+```
+
+## API Reference
+
+### Endpoint
+
+```
+POST /v1/audio/speech
+```
+
+This endpoint follows the [OpenAI Audio Speech API](https://platform.openai.com/docs/api-reference/audio/createSpeech) format with additional Qwen3-TTS parameters.
+
+### Request Body
+
+```json
+{
+    "input": "Text to synthesize",
+    "voice": "Vivian",
+    "response_format": "wav",
+    "task_type": "CustomVoice",
+    "language": "Auto",
+    "instructions": "Optional style instructions",
+    "ref_audio": "URL or base64 for voice cloning",
+    "ref_text": "Reference audio transcript",
+    "x_vector_only_mode": false,
+    "max_new_tokens": 2048
+}
+```
+
+> **Note:** The `model` field is optional when serving a single model, as the server already knows which model is loaded.
+
+### Response
+
+Returns audio data in the requested format (default: WAV).
+
+## Parameters
+
+### Standard OpenAI Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `input` | string | required | Text to synthesize |
+| `voice` | string | "Vivian" | Speaker/voice name |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `speed` | float | 1.0 | Playback speed (0.25-4.0) |
+| `model` | string | optional | Model name (optional when serving single model) |
+
+### Qwen3-TTS Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `task_type` | string | "CustomVoice" | Task: CustomVoice, VoiceDesign, or Base |
+| `language` | string | "Auto" | Language: Auto, Chinese, English, Japanese, Korean |
+| `instructions` | string | "" | Voice style/emotion instructions |
+| `max_new_tokens` | int | 2048 | Maximum tokens to generate |
+
+### Voice Clone Parameters (Base task)
+
+| Parameter | Type | Required | Description |
+|-----------|------|----------|-------------|
+| `ref_audio` | string | Yes* | Reference audio (file path, URL, or base64) |
+| `ref_text` | string | No | Transcript of reference audio (for ICL mode) |
+| `x_vector_only_mode` | bool | false | Use speaker embedding only (no ICL) |
+
+## Python Usage
+
+```python
+import httpx
+
+# Simple request
+response = httpx.post(
+    "http://localhost:8000/v1/audio/speech",
+    json={
+        "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
+        "input": "Hello world",
+        "voice": "Vivian",
+    },
+    timeout=300.0,
+)
+
+with open("output.wav", "wb") as f:
+    f.write(response.content)
+```
+
+## Limitations
+
+- **No streaming**: Audio is generated completely before being returned. Streaming will be supported after the pipeline is disaggregated (see RFC #938).
+- **Single request**: Batch processing is not yet optimized for online serving.
+
+## Troubleshooting
+
+1. **Connection refused**: Make sure the server is running on the correct port
+2. **Out of memory**: Reduce `--gpu-memory-utilization` in run_server.sh
+3. **Unsupported speaker**: Check supported speakers via model documentation
+4. **Voice clone fails**: Ensure you're using the Base model variant for voice cloning
+
+## Example materials
+
+??? abstract "openai_speech_client.py"
+    ``````py
+    --8<-- "examples/online_serving/qwen3_tts/openai_speech_client.py"
+    ``````
+??? abstract "run_server.sh"
+    ``````sh
+    --8<-- "examples/online_serving/qwen3_tts/run_server.sh"
+    ``````

examples/offline_inference/image_to_video/README.md

@@ -4,12 +4,19 @@ This example demonstrates how to generate videos from images using Wan2.2 Image-
 
 ## Local CLI Usage
 
+Download the example image:
+
+```bash
+wget https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg
+```
+
 ### Wan2.2-I2V-A14B-Diffusers (MoE)
+
 ```bash
 python image_to_video.py \
   --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
-  --image input.png \
-  --prompt "A cat playing with yarn, smooth motion" \
+  --image cherry_blossom.jpg \
+  --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
   --negative_prompt "<optional quality filter>" \
   --height 480 \
   --width 832 \
@@ -24,11 +31,12 @@ python image_to_video.py \
 ```
 
 ### Wan2.2-TI2V-5B-Diffusers (Unified)
+
 ```bash
 python image_to_video.py \
   --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
-  --image input.png \
-  --prompt "A cat playing with yarn, smooth motion" \
+  --image cherry_blossom.jpg \
+  --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
   --negative_prompt "<optional quality filter>" \
   --height 480 \
   --width 832 \

examples/offline_inference/lora_inference/README.md

@@ -1,4 +1,4 @@
-# LoRA Inference Examples
+# LoRA Inference(Diffusion)
 
 This directory contains examples for using LoRA (Low-Rank Adaptation) adapters with vLLM-omni diffusion models for offline inference.
 The example uses the  `stabilityai/stable-diffusion-3.5-medium` as the default model, but you can replace it with other models in vLLM-omni.

examples/offline_inference/qwen3_tts/README.md

@@ -1,4 +1,4 @@
-# Qwen3-TTS Offline Inference
+# Qwen3-TTS
 
 This directory contains an offline demo for running Qwen3 TTS models with vLLM Omni. It builds task-specific inputs and generates WAV files locally.
 

examples/offline_inference/text_to_video/text_to_video.py

@@ -109,7 +109,12 @@ def parse_args() -> argparse.Namespace:
         choices=[1, 2],
         help="Number of GPUs used for classifier free guidance parallel size.",
     )
-
+    parser.add_argument(
+        "--tensor_parallel_size",
+        type=int,
+        default=1,
+        help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
+    )
     return parser.parse_args()
 
 
@@ -141,6 +146,7 @@ def main():
         ulysses_degree=args.ulysses_degree,
         ring_degree=args.ring_degree,
         cfg_parallel_size=args.cfg_parallel_size,
+        tensor_parallel_size=args.tensor_parallel_size,
     )
 
     # Check if profiling is requested via environment variable
@@ -173,7 +179,7 @@ def main():
     print(f"  Inference steps: {args.num_inference_steps}")
     print(f"  Frames: {args.num_frames}")
     print(
-        f"  Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}"
+        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}"
     )
     print(f"  Video size: {args.width}x{args.height}")
     print(f"{'=' * 60}\n")

examples/online_serving/image_to_image/README.md

@@ -12,6 +12,9 @@ For **multi-image** input editing, use **Qwen-Image-Edit-2509** (QwenImageEditPl
 vllm serve Qwen/Qwen-Image-Edit --omni --port 8092
 ```
 
+!!! note
+    If you encounter Out-of-Memory (OOM) issues or have limited GPU memory, you can enable VAE slicing and tiling to reduce memory usage, --vae-use-slicing --vae-use-tiling
+
 ### Multi-Image Edit (Qwen-Image-Edit-2509)
 
 ```bash

examples/online_serving/lora_inference/README.md

@@ -1,4 +1,4 @@
-# Online LoRA Inference (Diffusion)
+# LoRA Inference(Diffusion)
 
 This example shows how to use **per-request LoRA** with vLLM-Omni diffusion models via the OpenAI-compatible Chat Completions API.
 

examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py

@@ -308,19 +308,24 @@ def get_multi_audios_query(custom_prompt: str | None = None):
 
 def get_use_audio_in_video_query(
     video_path: str | None = None,
-    audio_path: str | None = None,
     custom_prompt: str | None = None,
 ):
+    """Query for use_audio_in_video mode.
+
+    When use_audio_in_video=True, audio is automatically extracted from the video
+    by the server. Do NOT send a separate audio_url - this would cause a mismatch
+    between the number of audio and video items.
+    """
     question = custom_prompt or (
         "Describe the content of the video in details, then convert what the baby say into text."
     )
     video_url = get_video_url_from_path(video_path)
-    audio_url = get_audio_url_from_path(audio_path)
+    # Note: audio is extracted from video automatically when use_audio_in_video=True
+    # Do not include a separate audio_url here
     return {
         "role": "user",
         "content": [
             {"type": "video_url", "video_url": {"url": video_url}},
-            {"type": "audio_url", "audio_url": {"url": audio_url}},
             {"type": "text", "text": question},
         ],
     }
@@ -397,7 +402,6 @@ def run_multimodal_generation(args) -> None:
     elif args.query_type == "use_audio_in_video":
         prompt = query_func(
             video_path=video_path,
-            audio_path=audio_path,
             custom_prompt=custom_prompt,
         )
     else:
@@ -484,7 +488,7 @@ def parse_args():
         "--query-type",
         "-q",
         type=str,
-        default="use_mixed_modalities",
+        default="use_audio_in_video",
         choices=query_map.keys(),
         help="Query type.",
     )

examples/online_serving/qwen3_tts/README.md

@@ -1,4 +1,4 @@
-# Qwen3-TTS Online Serving
+# Qwen3-TTS
 
 This directory contains examples for running Qwen3-TTS models with vLLM-Omni's online serving API.
 

examples/online_serving/qwen3_tts/openai_speech_client.py

@@ -111,9 +111,14 @@ def run_tts_generation(args) -> None:
         print(response.text)
         return
 
-    if response.content.decode("utf-8").startswith('{"error"'):
-        print(f"Error: {response.content.decode('utf-8')}")
-        return
+    # Check for JSON error response (only if content is valid UTF-8 text)
+    try:
+        text = response.content.decode("utf-8")
+        if text.startswith('{"error"'):
+            print(f"Error: {text}")
+            return
+    except UnicodeDecodeError:
+        pass  # Binary audio data, not an error
 
     # Save audio response
     output_path = args.output or "tts_output.wav"

pyproject.toml

@@ -1,5 +1,9 @@
 [build-system]
-requires = ["setuptools>=61.0", "wheel"]
+requires = [
+    "setuptools>=77.0.3,<81.0.0",
+    "wheel",
+    "setuptools-scm>=8.0",
+]
 build-backend = "setuptools.build_meta"
 
 [project]
@@ -8,16 +12,16 @@ version = "0.14.0"
 description = "A framework for efficient model inference with omni-modality models"
 readme = "README.md"
 requires-python = ">=3.10,<3.14"
-license = {text = "Apache-2.0"}
+license = "Apache-2.0"
 authors = [
     {name = "vLLM-Omni Team"}
 ]
 keywords = ["vllm", "multimodal", "diffusion", "transformer", "inference", "serving"]
+dynamic = ["dependencies"]
 classifiers = [
     "Development Status :: 3 - Alpha",
     "Intended Audience :: Developers",
     "Intended Audience :: Science/Research",
-    "License :: OSI Approved :: Apache Software License",
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
@@ -26,26 +30,10 @@ classifiers = [
     "Topic :: Software Development :: Libraries :: Python Modules",
 ]
 
-
-dependencies = [
-    # Core runtime dependencies (required for actual usage)
-    "omegaconf>=2.3.0",
-    "librosa>=0.11.0",
-    "resampy>=0.4.3",
-    "diffusers>=0.36.0",
-    "accelerate==1.12.0",
-    "gradio==5.50",
-    "soundfile>=0.13.1",
-    "cache-dit==1.2.0",
-    "tqdm>=4.66.0",
-    "torchsde>=0.2.6",  # Required for Stable Audio scheduler
-    "fa3-fwd==0.0.1", # flash attention 3, maintained by @ZJY0516
-    "openai-whisper>=20250625",
-    "imageio[ffmpeg]>=2.37.2",
-    "onnxruntime>=1.19.0",
-    "sox>=1.5.0",
-    # "vllm==0.14.0",  # TODO: fix the entrypoints overwrite problem
-]
+# Dependencies are now managed dynamically via setup.py based on detected hardware platform.
+# This allows automatic installation of the correct platform-specific dependencies (CUDA/ROCm/CPU/XPU/NPU)
+# without requiring extras like [cuda]. See requirements/ directory for platform-specific dependencies.
+# Note: vllm is intentionally excluded due to entrypoints overwrite issue.
 
 [project.optional-dependencies]
 
@@ -155,11 +143,7 @@ python_classes = ["Test*"]
 python_functions = ["test_*"]
 addopts = [
     "--strict-markers",
-    "--strict-config",
-    "--cov=vllm_omni",
-    "--cov-report=term-missing",
-    "--cov-report=html",
-    "--cov-report=xml",
+    "--strict-config"
 ]
 markers = [
     # ci/cd required
@@ -204,4 +188,5 @@ extend-ignore-identifiers-re = [
     ".*NOTHINK.*",
     ".*nin.*",
     "Ono_Anna",
+    ".*cann.*",
 ]

requirements/common.txt

@@ -0,0 +1,14 @@
+# Common dependencies for all platforms
+omegaconf>=2.3.0
+librosa>=0.11.0
+resampy>=0.4.3
+diffusers>=0.36.0
+accelerate==1.12.0
+gradio==5.50
+soundfile>=0.13.1
+cache-dit==1.2.0
+tqdm>=4.66.0
+torchsde>=0.2.6
+openai-whisper>=20250625
+imageio[ffmpeg]>=2.37.2
+sox>=1.5.0

requirements/cpu.txt

@@ -0,0 +1,2 @@
+-r common.txt
+onnxruntime>=1.23.2

requirements/cuda.txt

@@ -0,0 +1,3 @@
+-r common.txt
+onnxruntime>=1.23.2
+fa3-fwd==0.0.1

requirements/npu.txt

@@ -0,0 +1,3 @@
+-r common.txt
+onnxruntime-cann>=1.23.2
+torchaudio==2.9.0

requirements/rocm.txt

@@ -0,0 +1,2 @@
+-r common.txt
+onnxruntime-rocm>=1.22.2

requirements/xpu.txt

@@ -0,0 +1,2 @@
+-r common.txt
+onnxruntime>=1.23.2

setup.py

@@ -0,0 +1,171 @@
+"""
+Setup script for vLLM-Omni with hardware-dependent installation.
+
+This setup.py implements platform-aware dependency routing so users can run
+`pip install vllm-omni` and automatically receive the correct platform-specific
+dependencies (CUDA/ROCm/CPU/XPU/NPU) without requiring extras like `[cuda]`.
+"""
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+from setuptools import setup
+
+
+def uninstall_onnxruntime() -> None:
+    """
+    Uninstall onnxruntime package if it exists.
+
+    This is necessary for ROCm environments where onnxruntime may conflict
+    with ROCm-specific dependencies.
+    """
+    try:
+        import pkg_resources
+
+        try:
+            pkg_resources.get_distribution("onnxruntime")
+            print("Found onnxruntime installed, uninstalling for ROCm compatibility...")
+            subprocess.check_call(
+                [sys.executable, "-m", "pip", "uninstall", "-y", "onnxruntime"],
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+            )
+            print("Successfully uninstalled onnxruntime")
+        except pkg_resources.DistributionNotFound:
+            print("onnxruntime not installed, skipping uninstall")
+    except Exception as e:
+        print(f"Warning: Failed to uninstall onnxruntime: {e}")
+
+
+def detect_target_device() -> str:
+    """
+    Detect the target device for installation following RFC priority rules.
+
+    Priority order:
+    1. VLLM_OMNI_TARGET_DEVICE environment variable (highest priority)
+    2. Torch backend detection (cuda, rocm, npu, xpu)
+    3. CPU fallback (default)
+
+    Returns:
+        str: Device name ('cuda', 'rocm', 'npu', 'xpu', or 'cpu')
+    """
+    # Priority 1: Explicit override via environment variable
+    target_device = os.environ.get("VLLM_OMNI_TARGET_DEVICE")
+    if target_device:
+        valid_devices = ["cuda", "rocm", "npu", "xpu", "cpu"]
+        if target_device.lower() in valid_devices:
+            print(f"Using target device from VLLM_OMNI_TARGET_DEVICE: {target_device.lower()}")
+            return target_device.lower()
+        else:
+            print(f"Warning: Invalid VLLM_OMNI_TARGET_DEVICE '{target_device}', falling back to auto-detection")
+
+    # Priority 2: Torch backend detection
+    # This is a code path for when user is using
+    # --no-build-isolation flag
+    try:
+        import torch
+
+        # Check for CUDA
+        if torch.version.cuda is not None:
+            print("Detected CUDA backend from torch")
+            return "cuda"
+
+        # Check for ROCm (AMD)
+        if torch.version.hip is not None:
+            print("Detected ROCm backend from torch")
+            uninstall_onnxruntime()
+            return "rocm"
+
+        # Check for NPU (Ascend)
+        if hasattr(torch, "npu"):
+            try:
+                if torch.npu.is_available():
+                    print("Detected NPU backend from torch")
+                    return "npu"
+            except Exception:
+                pass
+
+        # Check for XPU (Intel)
+        if hasattr(torch, "xpu"):
+            try:
+                if torch.xpu.is_available():
+                    print("Detected XPU backend from torch")
+                    return "xpu"
+            except Exception:
+                pass
+
+        print("No GPU backend detected in torch, defaulting to CPU")
+        return "cpu"
+
+    except ImportError:
+        print("PyTorch not found, defaulting to CUDA installation")
+        return "cuda"
+
+
+def load_requirements(file_path: Path) -> list[str]:
+    """
+    Load requirements from a file, supporting -r directive for recursive loading.
+
+    Args:
+        file_path: Path to the requirements file
+
+    Returns:
+        List of requirement strings
+    """
+    requirements = []
+
+    if not file_path.exists():
+        print(f"Warning: Requirements file not found: {file_path}")
+        return requirements
+
+    with open(file_path) as f:
+        for line in f:
+            line = line.strip()
+
+            # Skip empty lines and comments
+            if not line or line.startswith("#"):
+                continue
+
+            # Handle -r directive for recursive loading
+            if line.startswith("-r "):
+                nested_file = line[3:].strip()
+                nested_path = file_path.parent / nested_file
+                requirements.extend(load_requirements(nested_path))
+            else:
+                requirements.append(line)
+
+    return requirements
+
+
+def get_install_requires() -> list[str]:
+    """
+    Get the list of dependencies based on detected platform.
+
+    Returns:
+        List of requirement strings for the detected platform
+    """
+    device = detect_target_device()
+    requirements_dir = Path(__file__).parent / "requirements"
+    requirements_file = requirements_dir / f"{device}.txt"
+
+    print(f"Loading requirements from: {requirements_file}")
+    requirements = load_requirements(requirements_file)
+
+    if not requirements:
+        print(f"Warning: No requirements loaded for device '{device}'")
+    else:
+        print(f"Loaded {len(requirements)} requirements for {device}")
+
+    return requirements
+
+
+if __name__ == "__main__":
+    # Get platform-specific dependencies
+    install_requires = get_install_requires()
+
+    # Setup configuration
+    setup(
+        install_requires=install_requires,
+    )

tests/conftest.py

@@ -60,6 +60,12 @@ def clean_gpu_memory_between_tests():
     _run_post_test_cleanup()
 
 
+@pytest.fixture(autouse=True)
+def log_test_name_before_test(request):
+    print(f"--- Running test: {request.node.name}")
+    yield
+
+
 def _run_pre_test_cleanup(enable_force=False):
     if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1" and not enable_force:
         print("GPU cleanup disabled")

tests/diffusion/distributed/test_parallel_state_sp_groups.py

@@ -0,0 +1,71 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for SP subgroup construction in parallel_state.py."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.distributed.parallel_state import RankGenerator, set_seq_parallel_pg
+
+pytestmark = [
+    pytest.mark.diffusion,
+    pytest.mark.parallel,
+]
+
+
+def _fake_new_group_factory(created_groups: list[SimpleNamespace]):
+    def _fake_new_group(ranks, *args, **kwargs):
+        group = SimpleNamespace(ranks=list(ranks))
+        created_groups.append(group)
+        return group
+
+    return _fake_new_group
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    "rank, expected_ulysses, expected_ring",
+    [
+        (0, [0, 2], [0]),
+        (1, [1, 3], [1]),
+        (2, [0, 2], [2]),
+        (3, [1, 3], [3]),
+    ],
+)
+def test_set_seq_parallel_pg_uses_explicit_sp_groups(rank, expected_ulysses, expected_ring, monkeypatch):
+    created_groups: list[SimpleNamespace] = []
+    monkeypatch.setattr(torch.distributed, "new_group", _fake_new_group_factory(created_groups))
+
+    # tp=2, sp=2 -> SP groups are non-contiguous: [0,2] and [1,3]
+    sp_group_ranks = RankGenerator(2, 2, 1, 1, 1, "tp-sp-pp-cfg-dp").get_ranks("sp")
+
+    ulysses_pg, ring_pg = set_seq_parallel_pg(
+        sp_ulysses_degree=2,
+        sp_ring_degree=1,
+        rank=rank,
+        world_size=4,
+        sp_group_ranks=sp_group_ranks,
+    )
+
+    assert ulysses_pg.ranks == expected_ulysses
+    assert ring_pg.ranks == expected_ring
+
+
+@pytest.mark.cpu
+def test_set_seq_parallel_pg_validates_sp_group_ranks(monkeypatch):
+    created_groups: list[SimpleNamespace] = []
+    monkeypatch.setattr(torch.distributed, "new_group", _fake_new_group_factory(created_groups))
+
+    # world_size=4, sp_size=2 -> expect 2 groups, provide 1 to trigger validation
+    with pytest.raises(ValueError, match="Invalid sp_group_ranks"):
+        set_seq_parallel_pg(
+            sp_ulysses_degree=2,
+            sp_ring_degree=1,
+            rank=0,
+            world_size=4,
+            sp_group_ranks=[[0, 2]],
+        )

tests/e2e/offline_inference/test_cache_dit.py

@@ -24,6 +24,7 @@ if str(REPO_ROOT) not in sys.path:
 
 from vllm_omni import Omni
 from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.platforms import current_omni_platform
 
 os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
 
@@ -62,7 +63,7 @@ def test_cache_dit(model_name: str):
                 width=width,
                 num_inference_steps=num_inference_steps,
                 guidance_scale=0.0,
-                generator=torch.Generator("cuda").manual_seed(42),
+                generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
                 num_outputs_per_prompt=1,  # Single output for speed
             ),
         )

tests/e2e/offline_inference/test_t2i_model.py

@@ -47,7 +47,7 @@ def test_diffusion_model(model_name: str):
                 width=width,
                 num_inference_steps=2,
                 guidance_scale=0.0,
-                generator=torch.Generator("cuda").manual_seed(42),
+                generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
                 num_outputs_per_prompt=2,
             ),
         )

tests/e2e/offline_inference/test_t2v_model.py

@@ -14,6 +14,7 @@ if str(REPO_ROOT) not in sys.path:
 
 from vllm_omni import Omni
 from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.platforms import current_omni_platform
 
 # os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
@@ -42,7 +43,7 @@ def test_video_diffusion_model(model_name: str):
             num_frames=num_frames,
             num_inference_steps=2,
             guidance_scale=1.0,
-            generator=torch.Generator("cuda").manual_seed(42),
+            generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
         ),
     )
     first_output = outputs[0]

tests/e2e/online_serving/test_qwen3_omni_expansion.py

@@ -8,6 +8,7 @@ import concurrent.futures
 import os
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+import threading
 import time
 from pathlib import Path
 
@@ -21,7 +22,6 @@ from tests.conftest import (
     dummy_messages_from_mix_data,
     generate_synthetic_audio,
     generate_synthetic_image,
-    modify_stage_config,
 )
 
 models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
@@ -33,6 +33,29 @@ stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omn
 test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
 
 
+_omni_server_lock = threading.Lock()
+
+
+@pytest.fixture(scope="module")
+def omni_server(request):
+    """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:
+        model, stage_config_path = request.param
+
+        print(f"Starting OmniServer with model: {model}")
+
+        with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"]) as server:
+            print("OmniServer started successfully")
+            yield server
+            print("OmniServer stopping...")
+
+        print("OmniServer stopped")
+
+
+@pytest.fixture
 def client(omni_server):
     """OpenAI client for the running vLLM-Omni server."""
     return openai.OpenAI(
@@ -86,227 +109,180 @@ def get_deploy_config(deploy_type="TP1"):
     return result.get(deploy_type, result["TP1"])
 
 
-@pytest.mark.parametrize("test_config", test_params)
-def test_text_to_text_001(test_config: tuple[str, str]) -> None:
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_text_to_text_001(client: openai.OpenAI, omni_server) -> None:
     """Test processing text, generating text output via OpenAI API."""
-    model, stage_config_path = test_config
-    with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "90"]) as server:
-        messages = dummy_messages_from_mix_data(system_prompt=get_system_prompt(), content_text=get_prompt())
+    messages = dummy_messages_from_mix_data(system_prompt=get_system_prompt(), content_text=get_prompt())
 
-        # Test single completion
-        api_client = client(server)
-        start_time = time.perf_counter()
-        chat_completion = api_client.chat.completions.create(
-            model=server.model, messages=messages, max_tokens=20, modalities=["text"]
-        )
-        # Verify E2E
-        print(f"the request e2e is: {time.perf_counter() - start_time}")
-        # TODO: Verify the E2E latency after confirmation baseline.
+    # Test single completion
+    start_time = time.perf_counter()
+    chat_completion = client.chat.completions.create(
+        model=omni_server.model, messages=messages, max_tokens=20, modalities=["text"]
+    )
+    # Verify E2E
+    print(f"the request e2e is: {time.perf_counter() - start_time}")
+    # TODO: Verify the E2E latency after confirmation baseline.
 
-        # Verify only output text
-        assert len(chat_completion.choices) == 1, "The generated content includes more than just text."
+    # Verify only output text
+    assert len(chat_completion.choices) == 1, "The generated content includes more than just text."
 
-        # Verify text output success
-        text_choice = chat_completion.choices[0]
-        assert text_choice.message.content is not None, "No text output is generated"
-        assert chat_completion.usage.completion_tokens <= 20, "The output length more than the requested max_tokens."
-        assert "beijing" in text_choice.message.content.lower(), "The output do not contain keywords."
+    # Verify text output success
+    text_choice = chat_completion.choices[0]
+    assert text_choice.message.content is not None, "No text output is generated"
+    assert chat_completion.usage.completion_tokens <= 20, "The output length more than the requested max_tokens."
+    assert "beijing" in text_choice.message.content.lower(), "The output do not contain keywords."
 
 
-@pytest.mark.parametrize("test_config", test_params)
-def test_audio_to_text_001(test_config: tuple[str, str]) -> None:
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_audio_to_text_001(client: openai.OpenAI, omni_server) -> None:
     """Test processing text, generating text output via OpenAI API."""
-    model, stage_config_path = test_config
-    deploy_config = get_deploy_config()
-    deploy_config[0]["default_sampling_params.ignore_eos"] = True
-    stage_config_path = modify_stage_config(stage_config_path, deploy_config)
-    with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "90"]) as server:
-        audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(1, 1)['base64']}"
-        messages = dummy_messages_from_mix_data(audio_data_url=audio_data_url)
-        # Test single completion
-        api_client = client(server)
-        start_time = time.perf_counter()
-        chat_completion = api_client.chat.completions.create(
-            model=server.model, messages=messages, max_tokens=200, modalities=["text"]
-        )
-        # Verify only output text
-        assert len(chat_completion.choices) == 1, "The generated content includes more than just text."
+    audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(1, 1)['base64']}"
+    messages = dummy_messages_from_mix_data(audio_data_url=audio_data_url)
+    # Test single completion
+    start_time = time.perf_counter()
+    chat_completion = client.chat.completions.create(model=omni_server.model, messages=messages, modalities=["text"])
+    # Verify only output text
+    assert len(chat_completion.choices) == 1, "The generated content includes more than just text."
 
-        # Verify text output success
-        text_choice = chat_completion.choices[0]
-        assert text_choice.message.content is not None, "No text output is generated"
-        assert chat_completion.usage.completion_tokens == 200, (
-            "The output length differs from the requested max_tokens."
-        )
+    # Verify text output success
+    text_choice = chat_completion.choices[0]
+    assert text_choice.message.content is not None, "No text output is generated"
+    assert "water" in text_choice.message.content.lower(), "The output do not contain keywords."
 
-        # Verify E2E
-        print(f"the request e2e is: {time.perf_counter() - start_time}")
-        # TODO: Verify the E2E latency after confirmation baseline.
+    # Verify E2E
+    print(f"the request e2e is: {time.perf_counter() - start_time}")
+    # TODO: Verify the E2E latency after confirmation baseline.
 
 
-@pytest.mark.parametrize("test_config", test_params)
-def test_audio_to_text_audio_001(test_config: tuple[str, str]) -> None:
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_audio_to_text_audio_001(client: openai.OpenAI, omni_server) -> None:
     """Test processing text, generating audio output via OpenAI API."""
 
-    model, stage_config_path = test_config
+    audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(1, 5)['base64']}"
+
+    messages = dummy_messages_from_mix_data(audio_data_url=audio_data_url)
     num_concurrent_requests = get_max_batch_size()
-    stage_config_path = modify_stage_config(
-        stage_config_path,
-        {
-            "stage_args": {
-                0: {"runtime.max_batch_size": num_concurrent_requests},
-                1: {"runtime.max_batch_size": num_concurrent_requests},
-            }
-        },
-    )
-    with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "90"]) as server:
-        audio_data_url = []
-        for _ in range(5):
-            audio_data_url.append(f"data:audio/wav;base64,{generate_synthetic_audio(1, 5)['base64']}")
-
-        messages = dummy_messages_from_mix_data(audio_data_url=audio_data_url)
-
-        # Test single completion
-        api_client = client(server)
-        e2e_list = list()
-        with concurrent.futures.ThreadPoolExecutor(max_workers=num_concurrent_requests) as executor:
-            # Submit multiple completion requests concurrently
-            futures = [
-                executor.submit(api_client.chat.completions.create, model=server.model, messages=messages)
-                for _ in range(num_concurrent_requests)
-            ]
-            start_time = time.perf_counter()
-            # Wait for all requests to complete and collect results
-            chat_completions = list()
-            for future in concurrent.futures.as_completed(futures):
-                chat_completions.append(future.result())
-                # Verify E2E
-                current_e2e = time.perf_counter() - start_time
-                print(f"the request e2e is: {current_e2e}")
-                # TODO: Verify the E2E latency after confirmation baseline.
-                e2e_list.append(current_e2e)
-
-        print(f"the avg e2e is: {sum(e2e_list) / len(e2e_list)}")
-        # Verify all completions succeeded
-        assert len(chat_completions) == num_concurrent_requests, "Not all requests succeeded."
-        for chat_completion in chat_completions:
-            # Verify audio output success
-            audio_message = chat_completion.choices[1].message
-            audio_data = audio_message.audio.data
-            assert audio_data is not None, "No audio output is generated"
-            assert audio_message.audio.expires_at > time.time(), "The generated audio has expired."
-
-            # Verify text output success
-            text_choice = chat_completion.choices[0]
-            text_content = text_choice.message.content
-            assert text_choice.message.content is not None, "No text output is generated"
-
-            # Verify text output same as audio output
-            audio_content = convert_audio_to_text(audio_data)
-            print(f"text content is: {text_content}")
-            print(f"audio content is: {audio_content}")
-            similarity = cosine_similarity_text(audio_content, text_content)
-            print(f"similarity between audio and text is: {similarity}")
-            assert similarity > 0.9, "The audio content is not same as the text"
-
-
-@pytest.mark.parametrize("test_config", test_params)
-def test_image_to_text_001(test_config: tuple[str, str]) -> None:
-    """Test processing text, generating text output via OpenAI API."""
-    model, stage_config_path = test_config
-    deploy_config = get_deploy_config()
-    stage_config_path = modify_stage_config(stage_config_path, deploy_config)
-
-    with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "90"]) as server:
-        image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
-        messages = dummy_messages_from_mix_data(image_data_url=image_data_url)
-        # Test single completion
-        api_client = client(server)
+    # Test single completion
+    e2e_list = list()
+    with concurrent.futures.ThreadPoolExecutor(max_workers=num_concurrent_requests) as executor:
+        # Submit multiple completion requests concurrently
+        futures = [
+            executor.submit(client.chat.completions.create, model=omni_server.model, messages=messages)
+            for _ in range(num_concurrent_requests)
+        ]
         start_time = time.perf_counter()
-        chat_completion = api_client.chat.completions.create(
-            model=server.model, messages=messages, max_tokens=100, modalities=["text"]
-        )
-        # Verify E2E
-        print(f"the request e2e is: {time.perf_counter() - start_time}")
-        # TODO: Verify the E2E latency after confirmation baseline.
+        # Wait for all requests to complete and collect results
+        chat_completions = list()
+        for future in concurrent.futures.as_completed(futures):
+            chat_completions.append(future.result())
+            # Verify E2E
+            current_e2e = time.perf_counter() - start_time
+            print(f"the request e2e is: {current_e2e}")
+            # TODO: Verify the E2E latency after confirmation baseline.
+            e2e_list.append(current_e2e)
 
-        # Verify only output text
-        assert len(chat_completion.choices) == 1, "The generated content includes more than just text."
+    print(f"the avg e2e is: {sum(e2e_list) / len(e2e_list)}")
+    # Verify all completions succeeded
+    assert len(chat_completions) == num_concurrent_requests, "Not all requests succeeded."
+    for chat_completion in chat_completions:
+        # Verify audio output success
+        audio_message = chat_completion.choices[1].message
+        audio_data = audio_message.audio.data
+        assert audio_data is not None, "No audio output is generated"
+        assert audio_message.audio.expires_at > time.time(), "The generated audio has expired."
+
+        # Verify text output success
+        text_choice = chat_completion.choices[0]
+        text_content = text_choice.message.content
+        assert text_choice.message.content is not None, "No text output is generated"
+
+        # Verify text output same as audio output
+        audio_content = convert_audio_to_text(audio_data)
+        print(f"text content is: {text_content}")
+        print(f"audio content is: {audio_content}")
+        similarity = cosine_similarity_text(audio_content, text_content)
+        print(f"similarity between audio and text is: {similarity}")
+        assert similarity > 0.9, "The audio content is not same as the text"
+
+
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_image_to_text_001(client: openai.OpenAI, omni_server) -> None:
+    """Test processing text, generating text output via OpenAI API."""
+    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
+    messages = dummy_messages_from_mix_data(image_data_url=image_data_url)
+    # Test single completion
+    start_time = time.perf_counter()
+    chat_completion = client.chat.completions.create(
+        model=omni_server.model, messages=messages, max_tokens=100, modalities=["text"]
+    )
+    # Verify E2E
+    print(f"the request e2e is: {time.perf_counter() - start_time}")
+    # TODO: Verify the E2E latency after confirmation baseline.
+
+    # Verify only output text
+    assert len(chat_completion.choices) == 1, "The generated content includes more than just text."
+
+    # Verify text output success
+    text_choice = chat_completion.choices[0]
+    text_content = text_choice.message.content
+    assert text_content is not None, "No text output is generated"
+    assert chat_completion.usage.completion_tokens <= 100, "The output length more than the requested max_tokens."
+    assert "square" in text_content.lower(), "The output do not contain keywords."
+
+
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_image_to_text_audio_001(client: openai.OpenAI, omni_server) -> None:
+    """Test processing text, generating audio output via OpenAI API."""
+
+    num_concurrent_requests = get_max_batch_size()
+
+    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(1280, 720)['base64']}"
+
+    messages = dummy_messages_from_mix_data(image_data_url=image_data_url)
+
+    # Test single completion
+    e2e_list = list()
+    with concurrent.futures.ThreadPoolExecutor(max_workers=num_concurrent_requests) as executor:
+        # Submit multiple completion requests concurrently
+        futures = [
+            executor.submit(
+                client.chat.completions.create,
+                model=omni_server.model,
+                messages=messages,
+            )
+            for _ in range(num_concurrent_requests)
+        ]
+        start_time = time.perf_counter()
+        # Wait for all requests to complete and collect results
+        chat_completions = list()
+        for future in concurrent.futures.as_completed(futures):
+            chat_completions.append(future.result())
+            # Verify E2E
+            current_e2e = time.perf_counter() - start_time
+            print(f"the request e2e is: {current_e2e}")
+            # TODO: Verify the E2E latency after confirmation baseline.
+            e2e_list.append(current_e2e)
+
+    print(f"the avg e2e is: {sum(e2e_list) / len(e2e_list)}")
+    # Verify all completions succeeded
+    assert len(chat_completions) == num_concurrent_requests, "Not all requests succeeded."
+    for chat_completion in chat_completions:
+        # Verify audio output success
+        audio_message = chat_completion.choices[1].message
+        audio_data = audio_message.audio.data
+        assert audio_data is not None, "No audio output is generated"
+        assert audio_message.audio.expires_at > time.time(), "The generated audio has expired."
 
         # Verify text output success
         text_choice = chat_completion.choices[0]
         text_content = text_choice.message.content
         assert text_content is not None, "No text output is generated"
-        assert chat_completion.usage.completion_tokens <= 100, "The output length more than the requested max_tokens."
         assert "square" in text_content.lower(), "The output do not contain keywords."
 
-
-@pytest.mark.parametrize("test_config", test_params)
-def test_image_to_text_audio_001(test_config: tuple[str, str]) -> None:
-    """Test processing text, generating audio output via OpenAI API."""
-
-    model, stage_config_path = test_config
-    num_concurrent_requests = 5
-    stage_config_path = modify_stage_config(
-        stage_config_path,
-        {
-            "stage_args": {
-                0: {"runtime.max_batch_size": num_concurrent_requests},
-                1: {"runtime.max_batch_size": num_concurrent_requests},
-            }
-        },
-    )
-    with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "90"]) as server:
-        image_data_url = []
-        for _ in range(4):
-            image_data_url.append(f"data:image/jpeg;base64,{generate_synthetic_image(1280, 720)['base64']}")
-
-        messages = dummy_messages_from_mix_data(image_data_url=image_data_url)
-
-        # Test single completion
-        api_client = client(server)
-        e2e_list = list()
-        with concurrent.futures.ThreadPoolExecutor(max_workers=num_concurrent_requests) as executor:
-            # Submit multiple completion requests concurrently
-            futures = [
-                executor.submit(
-                    api_client.chat.completions.create,
-                    model=server.model,
-                    messages=messages,
-                )
-                for _ in range(num_concurrent_requests)
-            ]
-            start_time = time.perf_counter()
-            # Wait for all requests to complete and collect results
-            chat_completions = list()
-            for future in concurrent.futures.as_completed(futures):
-                chat_completions.append(future.result())
-                # Verify E2E
-                current_e2e = time.perf_counter() - start_time
-                print(f"the request e2e is: {current_e2e}")
-                # TODO: Verify the E2E latency after confirmation baseline.
-                e2e_list.append(current_e2e)
-
-        print(f"the avg e2e is: {sum(e2e_list) / len(e2e_list)}")
-        # Verify all completions succeeded
-        assert len(chat_completions) == num_concurrent_requests, "Not all requests succeeded."
-        for chat_completion in chat_completions:
-            # Verify audio output success
-            audio_message = chat_completion.choices[1].message
-            audio_data = audio_message.audio.data
-            assert audio_data is not None, "No audio output is generated"
-            assert audio_message.audio.expires_at > time.time(), "The generated audio has expired."
-
-            # Verify text output success
-            text_choice = chat_completion.choices[0]
-            text_content = text_choice.message.content
-            assert text_content is not None, "No text output is generated"
-            assert "square" in text_content.lower(), "The output do not contain keywords."
-
-            # Verify text output same as audio output
-            audio_content = convert_audio_to_text(audio_data)
-            print(f"text content is: {text_content}")
-            print(f"audio content is: {audio_content}")
-            similarity = cosine_similarity_text(audio_content, text_content)
-            print(f"similarity between audio and text is: {similarity}")
-            assert similarity > 0.9, "The audio content is not same as the text"
+        # Verify text output same as audio output
+        audio_content = convert_audio_to_text(audio_data)
+        print(f"text content is: {text_content}")
+        print(f"audio content is: {audio_content}")
+        similarity = cosine_similarity_text(audio_content, text_content)
+        print(f"similarity between audio and text is: {similarity}")
+        assert similarity > 0.9, "The audio content is not same as the text"

tests/entrypoints/openai_api/test_image_server.py

@@ -338,6 +338,25 @@ def test_with_seed(test_client):
     assert response.status_code == 200
 
 
+def test_with_seed_zero(test_client):
+    """Test with seed=0 for reproducibility"""
+    response = test_client.post(
+        "/v1/images/generations",
+        json={
+            "prompt": "a tree",
+            "seed": 0,
+            "size": "1024x1024",
+        },
+    )
+    assert response.status_code == 200
+    engine = test_client.app.state.engine_client
+    captured = engine.captured_sampling_params_list[0]
+    # Verify that seed=0 is correctly passed
+    assert captured.seed == 0, (
+        f"Expected seed=0, but got seed={captured.seed}. This indicates the bug where seed=0 is treated as falsy."
+    )
+
+
 def test_with_custom_parameters(test_client):
     """Test with custom diffusion parameters"""
     response = test_client.post(
@@ -814,3 +833,55 @@ def test_image_edit_compression_png(async_omni_test_client):
 
     assert len(img_bytes_10) < len(img_bytes_50)
     assert len(img_bytes_50) < len(img_bytes_100)
+
+
+def test_image_edit_with_seed_zero(async_omni_test_client):
+    """Test that seed=0 is correctly handled in image editing.
+
+    Previously, seed=0 was incorrectly replaced by a random seed due to the
+    falsy value check using `or` operator. This test ensures seed=0 is
+    properly passed through to the sampling parameters in image editing.
+    """
+    img_bytes_1 = make_test_image_bytes((16, 16))
+
+    response = async_omni_test_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes_1)],
+        data={
+            "prompt": "edit this image",
+            "seed": 0,
+        },
+    )
+    assert response.status_code == 200
+    engine = async_omni_test_client.app.state.engine_client
+    captured_sampling_params = engine.captured_sampling_params_list[-1]
+    # Verify that seed=0 is correctly passed
+    assert captured_sampling_params.seed == 0, (
+        f"Expected seed=0, but got seed={captured_sampling_params.seed}. "
+        "This indicates the bug where seed=0 is treated as falsy."
+    )
+
+
+def test_image_edit_with_seed_zero_single_stage(test_client):
+    """Test that seed=0 is correctly handled in image editing (single stage).
+
+    Test seed=0 handling in image editing with single stage path.
+    """
+    img_bytes_1 = make_test_image_bytes((16, 16))
+
+    response = test_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes_1)],
+        data={
+            "prompt": "edit this image",
+            "seed": 0,
+        },
+    )
+    assert response.status_code == 200
+    engine = test_client.app.state.engine_client
+    captured_sampling_params = engine.captured_sampling_params_list[0]
+    # Verify that seed=0 is correctly passed
+    assert captured_sampling_params.seed == 0, (
+        f"Expected seed=0, but got seed={captured_sampling_params.seed}. "
+        "This indicates the bug where seed=0 is treated as falsy."
+    )

tests/worker/test_process_gpu_memory.py

@@ -0,0 +1,154 @@
+"""Tests for process-scoped GPU memory accounting."""
+
+import os
+from unittest import mock
+
+import pytest
+
+
+class TestParseCudaVisibleDevices:
+    def test_empty(self):
+        from vllm_omni.worker.gpu_memory_utils import parse_cuda_visible_devices
+
+        with mock.patch.dict(os.environ, {}, clear=True):
+            os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+            assert parse_cuda_visible_devices() == []
+
+        with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": ""}):
+            assert parse_cuda_visible_devices() == []
+
+    def test_integer_indices(self):
+        from vllm_omni.worker.gpu_memory_utils import parse_cuda_visible_devices
+
+        with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "2,3,5"}):
+            assert parse_cuda_visible_devices() == [2, 3, 5]
+
+        with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0"}):
+            assert parse_cuda_visible_devices() == [0]
+
+    def test_uuids(self):
+        from vllm_omni.worker.gpu_memory_utils import parse_cuda_visible_devices
+
+        uuid1 = "GPU-12345678-1234-1234-1234-123456789abc"
+        uuid2 = "GPU-87654321-4321-4321-4321-cba987654321"
+        with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": f"{uuid1},{uuid2}"}):
+            assert parse_cuda_visible_devices() == [uuid1, uuid2]
+
+    def test_mig_ids(self):
+        from vllm_omni.worker.gpu_memory_utils import parse_cuda_visible_devices
+
+        mig1 = "MIG-GPU-12345678-1234-1234-1234-123456789abc/0/0"
+        mig2 = "MIG-GPU-12345678-1234-1234-1234-123456789abc/1/0"
+        with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": f"{mig1},{mig2}"}):
+            assert parse_cuda_visible_devices() == [mig1, mig2]
+
+    def test_spaces(self):
+        from vllm_omni.worker.gpu_memory_utils import parse_cuda_visible_devices
+
+        with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": " 2 , 3 , 5 "}):
+            assert parse_cuda_visible_devices() == [2, 3, 5]
+
+
+class TestGetProcessGpuMemory:
+    @pytest.mark.skipif(not os.path.exists("/dev/nvidia0"), reason="No GPU")
+    def test_returns_memory_for_current_process(self):
+        import torch
+
+        from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory
+
+        if not torch.cuda.is_available():
+            pytest.skip("CUDA not available")
+
+        device = torch.device("cuda:0")
+        tensor = torch.zeros(1000, 1000, device=device)
+
+        memory = get_process_gpu_memory(0)
+        assert memory >= 0
+
+        del tensor
+        torch.cuda.empty_cache()
+
+    def test_raises_on_invalid_device(self):
+        from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory
+
+        with (
+            mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit"),
+            mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"),
+            mock.patch("vllm.third_party.pynvml.nvmlDeviceGetCount", return_value=1),
+        ):
+            with pytest.raises(RuntimeError, match="Invalid GPU device"):
+                get_process_gpu_memory(5)
+
+    def test_returns_zero_when_process_not_found(self):
+        from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory
+
+        with (
+            mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit"),
+            mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"),
+            mock.patch("vllm.third_party.pynvml.nvmlDeviceGetCount", return_value=8),
+            mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlDeviceGetHandleByIndex"),
+            mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlDeviceGetComputeRunningProcesses", return_value=[]),
+        ):
+            memory = get_process_gpu_memory(0)
+            assert memory == 0
+
+    def test_uses_uuid_when_provided(self):
+        from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory
+
+        uuid = "GPU-12345678-1234-1234-1234-123456789abc"
+        mock_handle = mock.MagicMock()
+
+        with (
+            mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": uuid}),
+            mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit"),
+            mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"),
+            mock.patch("vllm.third_party.pynvml.nvmlDeviceGetHandleByUUID", return_value=mock_handle) as mock_by_uuid,
+            mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlDeviceGetComputeRunningProcesses", return_value=[]),
+        ):
+            memory = get_process_gpu_memory(0)
+            assert memory == 0
+            mock_by_uuid.assert_called_once_with(uuid)
+
+    def test_raises_on_invalid_uuid(self):
+        from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory
+
+        uuid = "GPU-invalid-uuid"
+
+        with (
+            mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": uuid}),
+            mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit"),
+            mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"),
+            mock.patch("vllm.third_party.pynvml.nvmlDeviceGetHandleByUUID", side_effect=Exception("Invalid UUID")),
+        ):
+            with pytest.raises(RuntimeError, match="Failed to get NVML handle"):
+                get_process_gpu_memory(0)
+
+    def test_returns_none_on_nvml_init_failure(self):
+        from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory
+
+        with (
+            mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit", side_effect=Exception("NVML unavailable")),
+            mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"),
+        ):
+            result = get_process_gpu_memory(0)
+            assert result is None
+
+
+class TestIsProcessScopedMemoryAvailable:
+    def test_returns_true_when_nvml_works(self):
+        from vllm_omni.worker.gpu_memory_utils import is_process_scoped_memory_available
+
+        with (
+            mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit"),
+            mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"),
+        ):
+            assert is_process_scoped_memory_available() is True
+
+    def test_returns_false_when_nvml_fails(self):
+        from vllm_omni.worker.gpu_memory_utils import is_process_scoped_memory_available
+
+        with mock.patch(
+            "vllm_omni.worker.gpu_memory_utils.nvmlInit",
+            side_effect=Exception("NVML unavailable"),
+        ):
+            assert is_process_scoped_memory_available() is False

vllm_omni/benchmarks/metrics/metrics.py

@@ -14,13 +14,17 @@ class MultiModalsBenchmarkMetrics(BenchmarkMetrics):
     median_audio_ttfp_ms: float = 0.0
     std_audio_ttfp_ms: float = 0.0
     percentiles_audio_ttfp_ms: list[tuple[float, float]] = None
-    total_audio_duration_ms: float = 0.0
+    total_audio_duration_s: float = 0.0
     total_audio_frames: int = 0
     audio_throughput: float = 0.0
     mean_audio_rtf: float = 0.0
     median_audio_rtf: float = 0.0
     std_audio_rtf: float = 0.0
     percentiles_audio_rtf: list[tuple[float, float]] = None
+    mean_audio_duration_s: float = 0.0
+    median_audio_duration_s: float = 0.0
+    std_audio_duration_s: float = 0.0
+    percentiles_audio_duration_s: list[tuple[float, float]] = None
 
 
 def print_metrics(
@@ -73,7 +77,7 @@ def print_text_metrics(task_type, selected_percentile_metrics, metrics: MultiMod
 
 def print_audio_metrics(selected_percentile_metrics, metrics: MultiModalsBenchmarkMetrics):
     print("{s:{c}^{n}}".format(s=" Audio Result ", n=50, c="="))
-    print("{:<40} {:<10.2f}".format("Total audio duration generated(s):", metrics.total_audio_duration_ms))
+    print("{:<40} {:<10.2f}".format("Total audio duration generated(s):", metrics.total_audio_duration_s))
     print("{:<40} {:<10}".format("Total audio frames generated:", metrics.total_audio_frames))
     print("{:<40} {:<10.2f}".format("Audio throughput(audio duration/s):", metrics.audio_throughput))
     for metric in selected_percentile_metrics:
@@ -92,16 +96,23 @@ def process_one_metric(
         "e2el": "End-to-end Latency",
         "audio_ttfp": "Time to First Packet",
         "audio_rtf": "Real Time Factor",
+        "audio_duration": "Audio Duration",
     }
 
     header = metric_header_map.get(metric_attribute_name, metric_attribute_name)
     print("{s:{c}^{n}}".format(s=header, n=50, c="-"))
 
     is_audio_rtf = metric_attribute_name == "audio_rtf"
+    is_audio_duration = metric_attribute_name == "audio_duration"
 
-    suffix = "" if is_audio_rtf else "_ms"
-    unit_suffix = "" if is_audio_rtf else " (ms)"
-
+    suffix = "_ms"
+    unit_suffix = " (ms)"
+    if is_audio_duration:
+        suffix = "_s"
+        unit_suffix = " (s)"
+    elif is_audio_rtf:
+        suffix = ""
+        unit_suffix = ""
     mean_attr_name = f"mean_{metric_attribute_name}{suffix}"
     mean_value = getattr(metrics, mean_attr_name, 0.0)
     print(f"{f'Mean {metric_attribute_name.upper()}{unit_suffix}:':<40} {mean_value:<10.2f}")
@@ -212,6 +223,9 @@ def calculate_metrics(
                 good_completed += 1
 
     if completed == 0:
+        warnings.formatwarning = (
+            lambda msg, category, filename, lineno, line=None: f"{filename}:{lineno}: {category.__name__}: {msg}\n"
+        )
         warnings.warn(
             "All requests failed. This is likely due to a misconfiguration on the benchmark arguments.",
             stacklevel=2,
@@ -296,7 +310,11 @@ def calculate_metrics(
         std_audio_ttfp_ms=np.std(audio_ttfps or 0) * 1000,
         median_audio_ttfp_ms=np.median(audio_ttfps or 0) * 1000,
         percentiles_audio_ttfp_ms=[(p, np.percentile(audio_ttfps or 0, p) * 1000) for p in selected_percentiles],
-        total_audio_duration_ms=sum(audio_duration),
+        mean_audio_duration_s=np.mean(audio_duration or 0),
+        std_audio_duration_s=np.std(audio_duration or 0),
+        median_audio_duration_s=np.median(audio_duration or 0),
+        percentiles_audio_duration_s=[(p, np.percentile(audio_duration or 0, p)) for p in selected_percentiles],
+        total_audio_duration_s=sum(audio_duration),
         total_audio_frames=sum(audio_frames),
         audio_throughput=sum(audio_duration) / dur_s,
         mean_audio_rtf=np.mean(audio_rtfs or 0),

vllm_omni/benchmarks/patch/patch.py

@@ -113,8 +113,8 @@ async def async_request_openai_chat_omni_completions(
     st = time.perf_counter()
     output.start_time = st
     most_recent_timestamp = st
+    timestamp = st
     audio_generate_time = 0.0
-    audio_first_timestamp = st
     try:
         async with session.post(url=api_url, json=payload, headers=headers) as response:
             if response.status == 200:
@@ -133,7 +133,6 @@ async def async_request_openai_chat_omni_completions(
                             continue
 
                         chunk = message.removeprefix("data: ")
-
                         if chunk != "[DONE]":
                             timestamp = time.perf_counter()
                             data = json.loads(chunk)
@@ -148,11 +147,11 @@ async def async_request_openai_chat_omni_completions(
                                     else:
                                         output.itl.append(timestamp - most_recent_timestamp)
                                     generated_text += content or ""
+                                    most_recent_timestamp = timestamp
                                 elif modality == "audio":
                                     if output.audio_ttfp == 0.0:
-                                        audio_first_timestamp = timestamp
                                         output.audio_ttfp = timestamp - st
-                                    audio_generate_time = timestamp - audio_first_timestamp
+                                    audio_generate_time = timestamp - st
                                     if content != "":
                                         audio_bytes = base64.b64decode(content)
                                         seg = AudioSegment.from_file(io.BytesIO(audio_bytes))
@@ -164,8 +163,8 @@ async def async_request_openai_chat_omni_completions(
 
                             elif usage := data.get("usage"):
                                 output.output_tokens = usage.get("completion_tokens")
-                            most_recent_timestamp = timestamp
 
+                output.latency = timestamp - st
                 output.generated_text = generated_text
                 if generated_audio is not None:
                     output.audio_duration = len(generated_audio) / 1000.0
@@ -181,9 +180,7 @@ async def async_request_openai_chat_omni_completions(
                     else:
                         output.audio_rtf = 0
                         logger.warning("Audio duration is zero")
-
                 output.success = True
-                output.latency = most_recent_timestamp - st
             else:
                 output.error = response.reason or ""
                 output.success = False
@@ -473,6 +470,9 @@ async def benchmark(
             "request_goodput": metrics.request_goodput if goodput_config_dict else None,
             "output_throughput": metrics.output_throughput,
             "total_token_throughput": metrics.total_token_throughput,
+            "total_audio_duration_s": metrics.total_audio_duration_s,
+            "total_audio_frames": metrics.total_audio_frames,
+            "audio_throughput": metrics.audio_throughput,
             "input_lens": [output.prompt_len for output in outputs],
             "output_lens": actual_output_lens,
             "ttfts": [output.ttft for output in outputs],
@@ -505,8 +505,20 @@ async def benchmark(
         if metric_attribute_name not in selected_percentile_metrics:
             return
         is_audio_rtf = metric_attribute_name == "audio_rtf"
+        is_audio_duration = metric_attribute_name == "audio_duration"
 
-        suffix = "" if is_audio_rtf else "_ms"
+        suffix = "_ms"
+        if is_audio_duration:
+            suffix = "_s"
+        elif is_audio_rtf:
+            suffix = ""
+        mean_attr_name = f"mean_{metric_attribute_name}{suffix}"
+        mean_value = getattr(metrics, mean_attr_name, 0.0)
+        result[mean_attr_name] = mean_value
+
+        median_attr_name = f"median_{metric_attribute_name}{suffix}"
+        median_value = getattr(metrics, median_attr_name, 0.0)
+        result[median_attr_name] = median_value
         for p, value in getattr(metrics, f"percentiles_{metric_attribute_name}{suffix}"):
             p_word = str(int(p)) if int(p) == p else str(p)
             result[f"p{p_word}_{metric_attribute_name}{suffix}"] = value

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -67,9 +67,11 @@ class OmniGenerationScheduler(VLLMScheduler):
             if self.omni_connector is not None:
                 get_chunk_for_generation(self.omni_connector, request)
 
-            # OMNI: Skip requests that are already finished or not in self.requests
-            # This can happen when connector marks request as finished
-            if request.status == RequestStatus.FINISHED_STOPPED or request.request_id not in self.requests:
+            # OMNI: Skip requests that are not in self.requests
+            # This can happen when connector marks request as finished and it's removed from requests
+            if request.request_id not in self.requests or (
+                self.omni_connector is None and request.status == RequestStatus.FINISHED_STOPPED
+            ):
                 already_finished_reqs.add(request)
                 req_index += 1
                 continue
@@ -108,9 +110,10 @@ class OmniGenerationScheduler(VLLMScheduler):
             if self.omni_connector is not None:
                 get_chunk_for_generation(self.omni_connector, request)
 
-            # OMNI: Skip requests that are already finished or not in self.requests
-            # This can happen when connector marks request as finished
-            if request.status == RequestStatus.FINISHED_STOPPED or request.request_id not in self.requests:
+            # OMNI: Skip requests that are not in self.requests
+            if request.request_id not in self.requests or (
+                self.omni_connector is None and request.status == RequestStatus.FINISHED_STOPPED
+            ):
                 # Pop the finished request from waiting queue and don't schedule it
                 self.waiting.pop_request()
                 continue

vllm_omni/diffusion/cache/cache_dit_backend.py

@@ -273,19 +273,77 @@ def enable_cache_for_longcat_image(pipeline: Any, cache_config: Any) -> Callable
 
 
 def enable_cache_for_flux(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
-    """Enable cache-dit for Flux dual-transformer architecture.
-
-    Flux uses two transformers (transformer and transformer_2) that need
-    to be enabled together using BlockAdapter.
+    """Enable cache-dit for Flux.1-dev pipeline.
 
     Args:
         pipeline: The Flux pipeline instance.
         cache_config: DiffusionCacheConfig instance with cache configuration.
-
     Returns:
-        A refresh function that can be called to update cache context with new num_inference_steps.
+        A refresh function that can be called with a new ``num_inference_steps``
+        to update the cache context for the pipeline.
     """
-    raise NotImplementedError("cache-dit is not implemented for Flux pipeline.")
+    # Build DBCacheConfig for transformer
+    db_cache_config = _build_db_cache_config(cache_config)
+
+    calibrator = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
+
+    # Build ParamsModifier for transformer
+    modifier = ParamsModifier(
+        cache_config=db_cache_config,
+        calibrator_config=calibrator,
+    )
+
+    logger.info(
+        f"Enabling cache-dit on Flux transformer with BlockAdapter: "
+        f"Fn={db_cache_config.Fn_compute_blocks}, "
+        f"Bn={db_cache_config.Bn_compute_blocks}, "
+        f"W={db_cache_config.max_warmup_steps}, "
+    )
+
+    # Enable cache-dit using BlockAdapter for transformer
+    cache_dit.enable_cache(
+        (
+            BlockAdapter(
+                transformer=pipeline.transformer,
+                blocks=[
+                    pipeline.transformer.transformer_blocks,
+                    pipeline.transformer.single_transformer_blocks,
+                ],
+                forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_1],
+                params_modifiers=[modifier],
+            )
+        ),
+        cache_config=db_cache_config,
+    )
+
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        """Refresh cache context for the transformer with new num_inference_steps.
+
+        Args:
+            pipeline: The Flux pipeline instance.
+            num_inference_steps: New number of inference steps.
+        """
+        if cache_config.scm_steps_mask_policy is None:
+            cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+        else:
+            cache_dit.refresh_context(
+                pipeline.transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_inference_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy,
+                        total_steps=num_inference_steps,
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
+            )
+
+    return refresh_cache_context
 
 
 def enable_cache_for_sd3(pipeline: Any, cache_config: Any) -> Callable[[int], None]:

vllm_omni/diffusion/distributed/parallel_state.py

@@ -503,10 +503,50 @@ def init_vae_group(
 
 
 # adapted from https://github.com/feifeibear/long-context-attention/blob/main/yunchang/globals.py
-def set_seq_parallel_pg(sp_ulysses_degree, sp_ring_degree, rank, world_size, use_ulysses_low=True):
+def set_seq_parallel_pg(
+    sp_ulysses_degree: int,
+    sp_ring_degree: int,
+    rank: int,
+    world_size: int,
+    use_ulysses_low: bool = True,
+    sp_group_ranks: list[list[int]] | None = None,
+) -> tuple[torch.distributed.ProcessGroup, torch.distributed.ProcessGroup]:
     """
-    sp_ulysses_degree x sp_ring_degree = seq_parallel_size
-    (ulysses_degree, dp_size)
+    Initialize sequence-parallel Ulysses and Ring process groups.
+
+    This builds sequence-parallel (SP) subgroups inside each data-parallel (DP)
+    slice. The SP group size is sp_ulysses_degree * sp_ring_degree, and
+    world_size must be divisible by that size.
+
+    Args:
+        sp_ulysses_degree: Size of each Ulysses subgroup.
+        sp_ring_degree: Size of each Ring subgroup.
+        rank: Global rank of the current process.
+        world_size: Total number of processes.
+        use_ulysses_low: If True, Ulysses groups are contiguous chunks and Ring
+            groups are strided within each SP group. If False, the opposite.
+        sp_group_ranks: Optional explicit SP groups. Each entry must be a list
+            of length sp_ulysses_degree * sp_ring_degree. When provided, groups
+            are built from these ranks instead of auto-generated contiguous
+            ranges.
+
+    Returns:
+        ulyssess_pg (torch.distributed.ProcessGroup): The Ulysses process group
+            for this rank.
+        ring_pg (torch.distributed.ProcessGroup): The Ring process group for
+            this rank.
+
+    Raises:
+        ValueError: If sp_group_ranks length does not match world_size or any
+            entry has the wrong size.
+        AssertionError: If world_size is not divisible by sp_size.
+
+    Behavior:
+        - If sp_group_ranks is provided, groups are built per entry and each
+          entry is further split into Ulysses/Ring groups according to
+          use_ulysses_low.
+        - If sp_group_ranks is None, groups are auto-generated within each DP
+          slice using offsets of size sp_size.
     """
     sp_size = sp_ring_degree * sp_ulysses_degree
     dp_size = world_size // sp_size
@@ -516,40 +556,96 @@ def set_seq_parallel_pg(sp_ulysses_degree, sp_ring_degree, rank, world_size, use
     num_ulysses_pgs = sp_ring_degree  # world_size // sp_ulysses_degree
     num_ring_pgs = sp_ulysses_degree  # world_size // sp_ring_degree
 
-    if use_ulysses_low:
-        for dp_rank in range(dp_size):
-            offset = dp_rank * sp_size
-            for i in range(num_ulysses_pgs):
-                ulysses_ranks = list(
-                    range(
-                        i * sp_ulysses_degree + offset,
-                        (i + 1) * sp_ulysses_degree + offset,
-                    )
-                )
-                group = torch.distributed.new_group(ulysses_ranks)
-                if rank in ulysses_ranks:
-                    ulyssess_pg = group
-
-            for i in range(num_ring_pgs):
-                ring_ranks = list(range(i + offset, sp_size + offset, num_ring_pgs))
-                group = torch.distributed.new_group(ring_ranks)
-                if rank in ring_ranks:
-                    ring_pg = group
-
+    if sp_group_ranks is not None:
+        if len(sp_group_ranks) * sp_size != world_size:
+            raise ValueError(
+                f"Invalid sp_group_ranks: expected {world_size // sp_size} groups of size {sp_size}, "
+                f"but got {len(sp_group_ranks)} groups."
+            )
+        logger.info(
+            "Building SP subgroups from explicit sp_group_ranks "
+            f"(sp_size={sp_size}, ulysses={sp_ulysses_degree}, ring={sp_ring_degree}, "
+            f"use_ulysses_low={use_ulysses_low})."
+        )
+        local_sp_group = None
+        local_ulysses = None
+        local_ring = None
+        for group_ranks in sp_group_ranks:
+            if len(group_ranks) != sp_size:
+                raise ValueError(f"Invalid sp_group_ranks entry: expected size {sp_size}, got {len(group_ranks)}.")
+            if rank in group_ranks:
+                local_sp_group = list(group_ranks)
+            if use_ulysses_low:
+                # Ulysses groups are contiguous chunks; Ring groups are strided.
+                for i in range(num_ulysses_pgs):
+                    ulysses_ranks = group_ranks[i * sp_ulysses_degree : (i + 1) * sp_ulysses_degree]
+                    group = torch.distributed.new_group(ulysses_ranks)
+                    if rank in ulysses_ranks:
+                        ulyssess_pg = group
+                        local_ulysses = list(ulysses_ranks)
+                for i in range(num_ring_pgs):
+                    ring_ranks = group_ranks[i::num_ring_pgs]
+                    group = torch.distributed.new_group(ring_ranks)
+                    if rank in ring_ranks:
+                        ring_pg = group
+                        local_ring = list(ring_ranks)
+            else:
+                # Ring groups are contiguous chunks; Ulysses groups are strided.
+                for i in range(num_ring_pgs):
+                    ring_ranks = group_ranks[i * sp_ring_degree : (i + 1) * sp_ring_degree]
+                    group = torch.distributed.new_group(ring_ranks)
+                    if rank in ring_ranks:
+                        ring_pg = group
+                        local_ring = list(ring_ranks)
+                for i in range(num_ulysses_pgs):
+                    ulysses_ranks = group_ranks[i::num_ulysses_pgs]
+                    group = torch.distributed.new_group(ulysses_ranks)
+                    if rank in ulysses_ranks:
+                        ulyssess_pg = group
+                        local_ulysses = list(ulysses_ranks)
+        if local_sp_group is not None:
+            logger.info(
+                "SP group details for rank %d: sp_group=%s, ulysses_group=%s, ring_group=%s",
+                rank,
+                local_sp_group,
+                local_ulysses,
+                local_ring,
+            )
     else:
-        for dp_rank in range(dp_size):
-            offset = dp_rank * sp_size
-            for i in range(num_ring_pgs):
-                ring_ranks = list(range(i * sp_ring_degree + offset, (i + 1) * sp_ring_degree + offset))
-                group = torch.distributed.new_group(ring_ranks)
-                if rank in ring_ranks:
-                    ring_pg = group
+        if use_ulysses_low:
+            for dp_rank in range(dp_size):
+                offset = dp_rank * sp_size
+                for i in range(num_ulysses_pgs):
+                    ulysses_ranks = list(
+                        range(
+                            i * sp_ulysses_degree + offset,
+                            (i + 1) * sp_ulysses_degree + offset,
+                        )
+                    )
+                    group = torch.distributed.new_group(ulysses_ranks)
+                    if rank in ulysses_ranks:
+                        ulyssess_pg = group
 
-            for i in range(num_ulysses_pgs):
-                ulysses_ranks = list(range(i + offset, sp_size + offset, num_ulysses_pgs))
-                group = torch.distributed.new_group(ulysses_ranks)
-                if rank in ulysses_ranks:
-                    ulyssess_pg = group
+                for i in range(num_ring_pgs):
+                    ring_ranks = list(range(i + offset, sp_size + offset, num_ring_pgs))
+                    group = torch.distributed.new_group(ring_ranks)
+                    if rank in ring_ranks:
+                        ring_pg = group
+
+        else:
+            for dp_rank in range(dp_size):
+                offset = dp_rank * sp_size
+                for i in range(num_ring_pgs):
+                    ring_ranks = list(range(i * sp_ring_degree + offset, (i + 1) * sp_ring_degree + offset))
+                    group = torch.distributed.new_group(ring_ranks)
+                    if rank in ring_ranks:
+                        ring_pg = group
+
+                for i in range(num_ulysses_pgs):
+                    ulysses_ranks = list(range(i + offset, sp_size + offset, num_ulysses_pgs))
+                    group = torch.distributed.new_group(ulysses_ranks)
+                    if rank in ulysses_ranks:
+                        ulyssess_pg = group
 
     return ulyssess_pg, ring_pg
 
@@ -652,6 +748,7 @@ def initialize_model_parallel(
         data_parallel_size,
         "tp-sp-pp-cfg-dp",
     )
+    sp_group_ranks = rank_generator.get_ranks("sp")
     global _DP
     assert _DP is None, "data parallel group is already initialized"
     _DP = init_model_parallel_group(
@@ -685,9 +782,10 @@ def initialize_model_parallel(
         sp_ring_degree=ring_degree,
         rank=get_world_group().rank_in_group,
         world_size=dit_parallel_size,
+        sp_group_ranks=sp_group_ranks,
     )
     _SP = init_model_parallel_group(
-        group_ranks=rank_generator.get_ranks("sp"),
+        group_ranks=sp_group_ranks,
         local_rank=get_world_group().local_rank,
         backend=backend,
         parallel_mode="sequence",

vllm_omni/diffusion/distributed/sp_sharding.py

@@ -63,6 +63,11 @@ def sp_shard(
             f"world_size ({world_size}) for sequence parallel sharding."
         )
 
+    if size < world_size:
+        raise ValueError(
+            f"Tensor size along dim {dim} ({size}) must be >= world_size ({world_size}). Tensor shape: {tensor.shape}"
+        )
+
     return tensor.chunk(world_size, dim=dim)[rank]
 
 

vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py

@@ -72,11 +72,15 @@ def get_longcat_image_edit_pre_process_func(
             if "additional_information" not in prompt:
                 prompt["additional_information"] = {}
 
-            if raw_image is None or isinstance(raw_image, list):
-                raise ValueError(
-                    """Received no image or a list of image. Only a single image is supported by this model."""
-                    """Please correctly set `"multi_modal_data": {"image": <an image object or file path>, …}`"""
-                )
+            if not raw_image:  # None or empty list
+                raise ValueError("""Received no input image. This model requires one input image to run.""")
+            elif isinstance(raw_image, list):
+                if len(raw_image) > 1:
+                    raise ValueError(
+                        """Received multiple input images. Only a single image is supported by this model."""
+                    )
+                else:
+                    raw_image = raw_image[0]
 
             if isinstance(raw_image, str):
                 image = PIL.Image.open(raw_image)

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py

@@ -75,11 +75,15 @@ def get_qwen_image_edit_pre_process_func(
                 prompt["additional_information"] = {}
 
             # Only handles single image
-            if raw_image is None or isinstance(raw_image, list):
-                raise ValueError(
-                    """Received no image or a list of image. Only a single image is supported by this model."""
-                    """Please correctly set `"multi_modal_data": {"image": <an image object or file path>, …}`"""
-                )
+            if not raw_image:  # None or empty list
+                raise ValueError("""Received no input image. This model requires one input image to run.""")
+            elif isinstance(raw_image, list):
+                if len(raw_image) > 1:
+                    raise ValueError(
+                        """Received multiple input images. Only a single image is supported by this model."""
+                    )
+                else:
+                    raw_image = raw_image[0]
 
             if isinstance(raw_image, str):
                 image = PIL.Image.open(raw_image)
@@ -628,6 +632,13 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi
         first_prompt = req.prompts[0]
         prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
         negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt")
+        if negative_prompt is None:
+            logger.warning(
+                "negative_prompt is not set. The official Qwen-Image-Edit model "
+                "may produce lower-quality results without a negative_prompt. "
+                "Qwen official repository recommends to use whitespace string as negative_prompt. "
+                "Note: some distilled variants may not be affected by this."
+            )
 
         # Get preprocessed image from request (pre-processing is done in DiffusionEngine)
         if not isinstance(first_prompt, str) and "preprocessed_image" in (

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py

@@ -563,6 +563,13 @@ class QwenImageEditPlusPipeline(nn.Module, SupportImageInput, QwenImageCFGParall
         first_prompt = req.prompts[0]
         prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
         negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt")
+        if negative_prompt is None:
+            logger.warning(
+                "negative_prompt is not set. The official Qwen-Image-Edit model "
+                "may produce lower-quality results without a negative_prompt. "
+                "Qwen official repository recommends to use whitespace string as negative_prompt. "
+                "Note: some distilled variants may not be affected by this."
+            )
 
         # Get preprocessed images from request (pre-processing is done in DiffusionEngine)
         if (

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py

@@ -79,11 +79,15 @@ def get_qwen_image_layered_pre_process_func(
             if "additional_information" not in prompt:
                 prompt["additional_information"] = {}
 
-            if raw_image is None or isinstance(raw_image, list):
-                raise ValueError(
-                    """Received no image or a list of image. Only a single image is supported by this model."""
-                    """Please correctly set `"multi_modal_data": {"image": <an image object or file path>, …}`"""
-                )
+            if not raw_image:  # None or empty list
+                raise ValueError("""Received no input image. This model requires one input image to run.""")
+            elif isinstance(raw_image, list):
+                if len(raw_image) > 1:
+                    raise ValueError(
+                        """Received multiple input images. Only a single image is supported by this model."""
+                    )
+                else:
+                    raw_image = raw_image[0]
 
             if isinstance(raw_image, str):
                 image = PIL.Image.open(raw_image)

vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py

@@ -1,8 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import functools
 from collections.abc import Iterable
+from functools import lru_cache
 from math import prod
 from typing import Any
 
@@ -249,7 +249,7 @@ class QwenEmbedLayer3DRope(nn.Module):
 
         return vid_freqs, txt_freqs
 
-    @functools.cache
+    @lru_cache(maxsize=16)
     def _compute_video_freqs(self, frame, height, width, idx=0):
         seq_lens = frame * height * width
         freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1)
@@ -268,7 +268,7 @@ class QwenEmbedLayer3DRope(nn.Module):
         freqs = torch.cat([freqs_frame, freqs_height, freqs_width], dim=-1).reshape(seq_lens, -1)
         return freqs.clone().contiguous()
 
-    @functools.cache
+    @lru_cache(maxsize=16)
     def _compute_condition_freqs(self, frame, height, width):
         seq_lens = frame * height * width
         freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1)
@@ -311,7 +311,6 @@ class QwenEmbedRope(nn.Module):
             ],
             dim=1,
         )
-        self.rope_cache = {}
 
         # DO NOT USING REGISTER BUFFER HERE, IT WILL CAUSE COMPLEX NUMBERS LOSE ITS IMAGINARY PART
         self.scale_rope = scale_rope
@@ -349,14 +348,7 @@ class QwenEmbedRope(nn.Module):
         max_vid_index = 0
         for idx, fhw in enumerate(video_fhw):
             frame, height, width = fhw
-            rope_key = f"{idx}_{height}_{width}"
-
-            if not torch.compiler.is_compiling():
-                if rope_key not in self.rope_cache:
-                    self.rope_cache[rope_key] = self._compute_video_freqs(frame, height, width, idx)
-                video_freq = self.rope_cache[rope_key]
-            else:
-                video_freq = self._compute_video_freqs(frame, height, width, idx)
+            video_freq = self._compute_video_freqs(frame, height, width, idx)
             video_freq = video_freq.to(device)
             vid_freqs.append(video_freq)
 
@@ -371,7 +363,7 @@ class QwenEmbedRope(nn.Module):
 
         return vid_freqs, txt_freqs
 
-    @functools.cache
+    @lru_cache(maxsize=16)
     def _compute_video_freqs(self, frame, height, width, idx=0):
         seq_lens = frame * height * width
         freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1)

vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py

@@ -7,14 +7,19 @@ from typing import Any
 
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 from diffusers.models.attention import FeedForward
 from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbedding, Timesteps
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.normalization import FP32LayerNorm
+from vllm.distributed import (
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+    tensor_model_parallel_all_reduce,
+)
 from vllm.logger import init_logger
 from vllm.model_executor.layers.conv import Conv3dLayer
-from vllm.model_executor.layers.layernorm import RMSNorm
-from vllm.model_executor.layers.linear import QKVParallelLinear, ReplicatedLinear
+from vllm.model_executor.layers.linear import ColumnParallelLinear, QKVParallelLinear, RowParallelLinear
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.layer import Attention
@@ -51,6 +56,94 @@ def apply_rotary_emb_wan(
     return out.type_as(hidden_states)
 
 
+class DistributedRMSNorm(nn.Module):
+    """
+    RMSNorm that computes global RMS across tensor parallel ranks.
+    This ensures mathematically equivalent results to non-TP execution.
+    """
+
+    def __init__(self, hidden_size: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        tp_size = get_tensor_model_parallel_world_size()
+
+        input_dtype = x.dtype
+        x_float = x.float()
+        local_sum_sq = (x_float**2).sum(dim=-1, keepdim=True)
+        local_count = x.shape[-1]
+
+        if tp_size > 1:
+            global_sum_sq = local_sum_sq.clone()
+            tensor_model_parallel_all_reduce(global_sum_sq)
+            global_count = local_count * tp_size
+        else:
+            global_sum_sq = local_sum_sq
+            global_count = local_count
+
+        rms = torch.sqrt(global_sum_sq / global_count + self.eps)
+
+        output = (x_float / rms) * self.weight.float()
+        return output.to(input_dtype)
+
+
+class ColumnParallelGELU(nn.Module):
+    """Column parallel linear with GELU activation."""
+
+    def __init__(self, dim_in: int, dim_out: int, *, approximate: str = "tanh", bias: bool = True):
+        super().__init__()
+        self.proj = ColumnParallelLinear(
+            dim_in,
+            dim_out,
+            bias=bias,
+            gather_output=False,
+            return_bias=False,
+        )
+        self.approximate = approximate
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.proj(x)
+        return F.gelu(x, approximate=self.approximate)
+
+
+class WanFeedForward(nn.Module):
+    """
+    TP-enabled FeedForward network for WAN2.2.
+    Replaces diffusers FeedForward with ColumnParallel + RowParallel.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        inner_dim: int,
+        dim_out: int | None = None,
+        bias: bool = True,
+    ) -> None:
+        super().__init__()
+        dim_out = dim_out or dim
+
+        # ColumnParallel: scatter to each tp_rank
+        self.net_0 = ColumnParallelGELU(dim, inner_dim, approximate="tanh", bias=bias)
+        # Placeholder for weight loading compatibility
+        self.net_1 = nn.Identity()
+        # RowParallel: gather from each tp_rank
+        self.net_2 = RowParallelLinear(
+            inner_dim,
+            dim_out,
+            bias=bias,
+            input_is_parallel=True,
+            return_bias=False,
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.net_0(hidden_states)
+        hidden_states = self.net_1(hidden_states)
+        hidden_states = self.net_2(hidden_states)
+        return hidden_states
+
+
 class WanRotaryPosEmbed(nn.Module):
     """
     Rotary position embeddings for 3D video data (temporal + spatial dimensions).
@@ -97,8 +190,8 @@ class WanRotaryPosEmbed(nn.Module):
         t = torch.arange(max_seq_len, dtype=freqs_dtype)
         freqs = torch.outer(t, freqs)
         # Repeat interleave for real representation
-        freqs_cos = freqs.cos().repeat_interleave(2, dim=-1)
-        freqs_sin = freqs.sin().repeat_interleave(2, dim=-1)
+        freqs_cos = freqs.cos().float().repeat_interleave(2, dim=-1)
+        freqs_sin = freqs.sin().float().repeat_interleave(2, dim=-1)
         return freqs_cos.float(), freqs_sin.float()
 
     def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
@@ -203,6 +296,47 @@ class WanTimeTextImageEmbedding(nn.Module):
         return temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image
 
 
+class TimestepProjPrepare(nn.Module):
+    """Prepares timestep_proj for sequence parallel in TI2V models.
+
+    Encapsulates the unflatten operation for timestep_proj to enable _sp_plan sharding.
+    """
+
+    def forward(
+        self,
+        timestep_proj: torch.Tensor,
+        ts_seq_len: int | None,
+    ) -> torch.Tensor:
+        if ts_seq_len is not None:
+            # TI2V mode: [batch, seq_len, 6, inner_dim]
+            timestep_proj = timestep_proj.unflatten(2, (6, -1))
+        else:
+            # T2V mode: [batch, 6, inner_dim]
+            timestep_proj = timestep_proj.unflatten(1, (6, -1))
+        return timestep_proj
+
+
+class OutputScaleShiftPrepare(nn.Module):
+    """Prepares output scale/shift for SP sharding in TI2V models."""
+
+    def __init__(self, inner_dim: int):
+        super().__init__()
+        self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5)
+
+    def forward(self, temb: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        if temb.ndim == 3:
+            # TI2V: [B, seq, D] -> 3D outputs for SP sharding
+            shift, scale = (self.scale_shift_table.unsqueeze(0).to(temb.device) + temb.unsqueeze(2)).chunk(2, dim=2)
+            shift = shift.squeeze(2)
+            scale = scale.squeeze(2)
+        else:
+            # T2V: [B, D] -> 2D outputs (skip SP sharding via expected_dims=3)
+            shift, scale = (self.scale_shift_table.to(temb.device) + temb.unsqueeze(1)).chunk(2, dim=1)
+            shift = shift.squeeze(1)
+            scale = scale.squeeze(1)
+        return shift, scale
+
+
 class WanSelfAttention(nn.Module):
     """
     Optimized self-attention module using vLLM layers.
@@ -229,25 +363,30 @@ class WanSelfAttention(nn.Module):
             head_size=head_dim,
             total_num_heads=num_heads,
             bias=True,
-            disable_tp=True,
         )
 
+        self.num_heads = self.to_qkv.num_heads
+        self.num_kv_heads = self.to_qkv.num_kv_heads
+        self.tp_inner_dim = self.num_heads * head_dim
+
         # QK normalization using vLLM's RMSNorm
-        self.norm_q = RMSNorm(self.inner_dim, eps=eps)
-        self.norm_k = RMSNorm(self.inner_dim, eps=eps)
+        self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
+        self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
 
-        # Output projection
-        self.to_out = nn.ModuleList(
-            [
-                ReplicatedLinear(self.inner_dim, dim, bias=True),
-                nn.Dropout(dropout),
-            ]
+        self.to_out = RowParallelLinear(
+            self.inner_dim,
+            dim,
+            bias=True,
+            input_is_parallel=True,
+            return_bias=False,
         )
+        self.dropout = nn.Dropout(dropout)
 
         # Unified attention layer
         self.attn = Attention(
-            num_heads=num_heads,
+            num_heads=self.num_heads,
             head_size=head_dim,
+            num_kv_heads=self.num_kv_heads,
             softmax_scale=1.0 / (head_dim**0.5),
             causal=False,
         )
@@ -259,16 +398,19 @@ class WanSelfAttention(nn.Module):
     ) -> torch.Tensor:
         # Fused QKV projection
         qkv, _ = self.to_qkv(hidden_states)
-        query, key, value = qkv.chunk(3, dim=-1)
+
+        q_size = self.num_heads * self.head_dim
+        kv_size = self.num_kv_heads * self.head_dim
+        query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1)
 
         # Apply QK normalization
         query = self.norm_q(query)
         key = self.norm_k(key)
 
         # Reshape for multi-head attention
-        query = query.unflatten(2, (self.num_heads, -1))
-        key = key.unflatten(2, (self.num_heads, -1))
-        value = value.unflatten(2, (self.num_heads, -1))
+        query = query.unflatten(2, (self.num_heads, self.head_dim))
+        key = key.unflatten(2, (self.num_kv_heads, self.head_dim))
+        value = value.unflatten(2, (self.num_kv_heads, self.head_dim))
 
         # Apply rotary embeddings
         if rotary_emb is not None:
@@ -282,8 +424,8 @@ class WanSelfAttention(nn.Module):
         hidden_states = hidden_states.type_as(query)
 
         # Output projection
-        hidden_states, _ = self.to_out[0](hidden_states)
-        hidden_states = self.to_out[1](hidden_states)
+        hidden_states = self.to_out(hidden_states)
+        hidden_states = self.dropout(hidden_states)
 
         return hidden_states
 
@@ -312,39 +454,77 @@ class WanCrossAttention(nn.Module):
         self.kv_inner_dim = head_dim * num_heads  # For cross-attention, K/V come from encoder
 
         # Query projection
-        self.to_q = ReplicatedLinear(dim, self.inner_dim, bias=True)
+        self.to_q = ColumnParallelLinear(
+            dim,
+            self.inner_dim,
+            bias=True,
+            gather_output=False,
+            return_bias=False,
+        )
 
         # Separate K and V projections for cross-attention
-        self.to_k = ReplicatedLinear(dim, self.kv_inner_dim, bias=True)
-        self.to_v = ReplicatedLinear(dim, self.kv_inner_dim, bias=True)
+        self.to_k = ColumnParallelLinear(
+            dim,
+            self.kv_inner_dim,
+            bias=True,
+            gather_output=False,
+            return_bias=False,
+        )
+
+        self.to_v = ColumnParallelLinear(
+            dim,
+            self.kv_inner_dim,
+            bias=True,
+            gather_output=False,
+            return_bias=False,
+        )
+
+        tp_size = get_tensor_model_parallel_world_size()
+        self.num_heads = num_heads // tp_size
+        self.tp_inner_dim = self.num_heads * head_dim
 
         # QK normalization
-        self.norm_q = RMSNorm(self.inner_dim, eps=eps)
-        self.norm_k = RMSNorm(self.kv_inner_dim, eps=eps)
+        self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
+        self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
 
         # Optional added KV projections for I2V (image embeddings)
         self.added_kv_proj_dim = added_kv_proj_dim
         if added_kv_proj_dim is not None:
-            self.add_k_proj = ReplicatedLinear(added_kv_proj_dim, self.inner_dim, bias=True)
-            self.add_v_proj = ReplicatedLinear(added_kv_proj_dim, self.inner_dim, bias=True)
-            self.norm_added_k = RMSNorm(self.inner_dim, eps=eps)
+            self.add_k_proj = ColumnParallelLinear(
+                added_kv_proj_dim,
+                self.inner_dim,
+                bias=True,
+                gather_output=False,
+                return_bias=False,
+            )
+            self.add_v_proj = ColumnParallelLinear(
+                added_kv_proj_dim,
+                self.inner_dim,
+                bias=True,
+                gather_output=False,
+                return_bias=False,
+            )
+            self.norm_added_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
         else:
             self.add_k_proj = None
             self.add_v_proj = None
             self.norm_added_k = None
 
         # Output projection
-        self.to_out = nn.ModuleList(
-            [
-                ReplicatedLinear(self.inner_dim, dim, bias=True),
-                nn.Dropout(dropout),
-            ]
+        self.to_out = RowParallelLinear(
+            self.inner_dim,
+            dim,
+            bias=True,
+            input_is_parallel=True,
+            return_bias=False,
         )
+        self.dropout = nn.Dropout(dropout)
 
         # Unified attention layer
         self.attn = Attention(
-            num_heads=num_heads,
+            num_heads=self.num_heads,
             head_size=head_dim,
+            num_kv_heads=self.num_heads,
             softmax_scale=1.0 / (head_dim**0.5),
             causal=False,
         )
@@ -363,28 +543,28 @@ class WanCrossAttention(nn.Module):
             encoder_hidden_states = encoder_hidden_states[:, image_context_length:]
 
         # Query projection
-        query, _ = self.to_q(hidden_states)
+        query = self.to_q(hidden_states)
         query = self.norm_q(query)
 
         # KV projection from encoder
-        key, _ = self.to_k(encoder_hidden_states)
-        value, _ = self.to_v(encoder_hidden_states)
+        key = self.to_k(encoder_hidden_states)
+        value = self.to_v(encoder_hidden_states)
         key = self.norm_k(key)
 
         # Reshape for multi-head attention
-        query = query.unflatten(2, (self.num_heads, -1))
-        key = key.unflatten(2, (self.num_heads, -1))
-        value = value.unflatten(2, (self.num_heads, -1))
+        query = query.unflatten(2, (self.num_heads, self.head_dim))
+        key = key.unflatten(2, (self.num_heads, self.head_dim))
+        value = value.unflatten(2, (self.num_heads, self.head_dim))
 
         # I2V: Additional attention with image embeddings
         hidden_states_img = None
         if encoder_hidden_states_img is not None:
-            key_img, _ = self.add_k_proj(encoder_hidden_states_img)
-            value_img, _ = self.add_v_proj(encoder_hidden_states_img)
+            key_img = self.add_k_proj(encoder_hidden_states_img)
+            value_img = self.add_v_proj(encoder_hidden_states_img)
             key_img = self.norm_added_k(key_img)
 
-            key_img = key_img.unflatten(2, (self.num_heads, -1))
-            value_img = value_img.unflatten(2, (self.num_heads, -1))
+            key_img = key_img.unflatten(2, (self.num_heads, self.head_dim))
+            value_img = value_img.unflatten(2, (self.num_heads, self.head_dim))
 
             hidden_states_img = self.attn(query, key_img, value_img)
             hidden_states_img = hidden_states_img.flatten(2, 3)
@@ -400,8 +580,8 @@ class WanCrossAttention(nn.Module):
             hidden_states = hidden_states + hidden_states_img
 
         # Output projection
-        hidden_states, _ = self.to_out[0](hidden_states)
-        hidden_states = self.to_out[1](hidden_states)
+        hidden_states = self.to_out(hidden_states)
+        hidden_states = self.dropout(hidden_states)
 
         return hidden_states
 
@@ -445,7 +625,7 @@ class WanTransformerBlock(nn.Module):
         self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()
 
         # 3. Feed-forward
-        self.ffn = FeedForward(dim, inner_dim=ffn_dim, activation_fn="gelu-approximate")
+        self.ffn = WanFeedForward(dim=dim, inner_dim=ffn_dim, dim_out=dim)
         self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False)
 
         # Scale-shift table for modulation
@@ -540,6 +720,7 @@ class WanTransformer3DModel(nn.Module):
     #
     # The _sp_plan specifies sharding/gathering at module boundaries:
     # - rope: Split both RoPE outputs (freqs_cos, freqs_sin) via split_output=True
+    # - timestep_proj_prepare: Split timestep_proj for TI2V models (4D tensor)
     # - blocks.0: Split hidden_states input at the first transformer block
     # - proj_out: Gather outputs after the final projection layer
     #
@@ -550,11 +731,22 @@ class WanTransformer3DModel(nn.Module):
             0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),  # freqs_cos [1, seq, 1, dim]
             1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),  # freqs_sin [1, seq, 1, dim]
         },
+        # Shard timestep_proj for TI2V models (4D tensor: [batch, seq_len, 6, inner_dim])
+        # This is only active when ts_seq_len is not None (TI2V mode)
+        # Output is a single tensor, shard along dim=1 (sequence dimension)
+        "timestep_proj_prepare": {
+            0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),  # [B, seq, 6, dim]
+        },
         # Shard hidden_states at first transformer block input
         # (after patch_embedding + flatten + transpose)
         "blocks.0": {
             "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),  # [B, seq, dim]
         },
+        # Shard output scale/shift for TI2V (3D); T2V outputs 2D and skips sharding
+        "output_scale_shift_prepare": {
+            0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True),
+            1: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True),
+        },
         # Gather at proj_out (final linear projection before unpatchify)
         "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
     }
@@ -635,7 +827,10 @@ class WanTransformer3DModel(nn.Module):
         # 4. Output norm & projection
         self.norm_out = FP32LayerNorm(inner_dim, eps, elementwise_affine=False)
         self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size))
-        self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5)
+
+        # SP helper modules
+        self.timestep_proj_prepare = TimestepProjPrepare()
+        self.output_scale_shift_prepare = OutputScaleShiftPrepare(inner_dim)
 
     @property
     def dtype(self) -> torch.dtype:
@@ -675,10 +870,10 @@ class WanTransformer3DModel(nn.Module):
         temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
             timestep, encoder_hidden_states, encoder_hidden_states_image, timestep_seq_len=ts_seq_len
         )
-        if ts_seq_len is not None:
-            timestep_proj = timestep_proj.unflatten(2, (6, -1))
-        else:
-            timestep_proj = timestep_proj.unflatten(1, (6, -1))
+        # Prepare timestep_proj via TimestepProjPrepare module
+        # _sp_plan will shard timestep_proj via split_output=True (when ts_seq_len is not None)
+        # This ensures timestep_proj sequence dimension matches sharded hidden_states
+        timestep_proj = self.timestep_proj_prepare(timestep_proj, ts_seq_len)
 
         if encoder_hidden_states_image is not None:
             encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1)
@@ -688,15 +883,12 @@ class WanTransformer3DModel(nn.Module):
             hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb)
 
         # Output norm, projection & unpatchify
-        if temb.ndim == 3:
-            shift, scale = (self.scale_shift_table.unsqueeze(0).to(temb.device) + temb.unsqueeze(2)).chunk(2, dim=2)
-            shift = shift.squeeze(2)
-            scale = scale.squeeze(2)
-        else:
-            shift, scale = (self.scale_shift_table.to(temb.device) + temb.unsqueeze(1)).chunk(2, dim=1)
-
+        shift, scale = self.output_scale_shift_prepare(temb)
         shift = shift.to(hidden_states.device)
         scale = scale.to(hidden_states.device)
+        if shift.ndim == 2:  # T2V mode: unsqueeze for broadcasting
+            shift = shift.unsqueeze(1)
+            scale = scale.unsqueeze(1)
 
         hidden_states = (self.norm_out(hidden_states.float()) * (1 + scale) + shift).type_as(hidden_states)
         hidden_states = self.proj_out(hidden_states)
@@ -725,32 +917,74 @@ class WanTransformer3DModel(nn.Module):
         Returns:
             Set of parameter names that were successfully loaded.
         """
+        tp_rank = get_tensor_model_parallel_rank()
+        tp_size = get_tensor_model_parallel_world_size()
         # Stacked params mapping for self-attention QKV fusion
-        # Format: (param_name, shard_name, shard_id)
-        # Note: Only fuse attn1 (self-attention), NOT attn2 (cross-attention)
         stacked_params_mapping = [
-            # self-attention QKV fusion (attn1 only)
+            # self-attention QKV fusion
             (".attn1.to_qkv", ".attn1.to_q", "q"),
             (".attn1.to_qkv", ".attn1.to_k", "k"),
             (".attn1.to_qkv", ".attn1.to_v", "v"),
         ]
 
+        # Remap scale_shift_table to new module location
+        weight_name_remapping = {
+            "scale_shift_table": "output_scale_shift_prepare.scale_shift_table",
+        }
+
         params_dict = dict(self.named_parameters())
         loaded_params: set[str] = set()
 
         for name, loaded_weight in weights:
+            name = weight_name_remapping.get(name, name)
+            original_name = name
+            lookup_name = name
+
+            # Handle QKV fusion
             for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
+                if weight_name not in original_name:
                     continue
-                name = name.replace(weight_name, param_name)
-                param = params_dict[name]
+                lookup_name = original_name.replace(weight_name, param_name)
+                param = params_dict[lookup_name]
                 weight_loader = param.weight_loader
                 weight_loader(param, loaded_weight, shard_id)
                 break
             else:
-                param = params_dict[name]
+                # diffusers: ffn.net.0.proj.weight -> our: ffn.net_0.proj.weight
+                if ".ffn.net.0." in lookup_name:
+                    lookup_name = lookup_name.replace(".ffn.net.0.", ".ffn.net_0.")
+                elif ".ffn.net.2." in lookup_name:
+                    lookup_name = lookup_name.replace(".ffn.net.2.", ".ffn.net_2.")
+
+                if ".to_out.0." in lookup_name:
+                    lookup_name = lookup_name.replace(".to_out.0.", ".to_out.")
+
+                if lookup_name not in params_dict:
+                    logger.warning(f"Skipping weight {original_name} -> {lookup_name}")
+                    continue
+
+                param = params_dict[lookup_name]
+
+                # Handle RMSNorm weights that need to be sharded for TP
+                # These norms are applied after ColumnParallelLinear outputs,
+                # so their weights must be sharded to match the sharded hidden dim
+                if tp_size > 1 and any(
+                    norm_name in lookup_name
+                    for norm_name in [
+                        ".attn1.norm_q.",
+                        ".attn1.norm_k.",
+                        ".attn2.norm_q.",
+                        ".attn2.norm_k.",
+                        ".attn2.norm_added_k.",
+                    ]
+                ):
+                    shard_size = loaded_weight.shape[0] // tp_size
+                    loaded_weight = loaded_weight[tp_rank * shard_size : (tp_rank + 1) * shard_size]
+
                 weight_loader = getattr(param, "weight_loader", default_weight_loader)
                 weight_loader(param, loaded_weight)
-            loaded_params.add(name)
+
+            loaded_params.add(original_name)
+            loaded_params.add(lookup_name)
 
         return loaded_params

vllm_omni/diffusion/models/z_image/pipeline_z_image.py

@@ -527,6 +527,16 @@ class ZImagePipeline(nn.Module):
         num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
         self._num_timesteps = len(timesteps)
 
+        # Precompute normalized timesteps once to avoid per-step GPU->CPU sync (.item() causes cudaStreamSynchronize)
+        if isinstance(timesteps, torch.Tensor):
+            timesteps_tensor = timesteps.to(device=device, dtype=torch.float32)
+        else:
+            timesteps_tensor = torch.as_tensor(timesteps, device=device, dtype=torch.float32)
+        norm_timesteps = (1000 - timesteps_tensor) / 1000
+        t_norm_list = norm_timesteps.cpu().tolist()
+        if not isinstance(t_norm_list, list):
+            t_norm_list = [t_norm_list]
+
         # 6. Denoising loop
         for i, t in enumerate(timesteps):
             if self.interrupt:
@@ -535,8 +545,9 @@ class ZImagePipeline(nn.Module):
             # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
             timestep = t.expand(latents.shape[0])
             timestep = (1000 - timestep) / 1000
-            # Normalized time for time-aware config (0 at start, 1 at end)
-            t_norm = timestep[0].item()
+            # Normalized time for time-aware config (0 at start, 1 at end);
+            # use precomputed to avoid .item() sync per step
+            t_norm = t_norm_list[i]
 
             # Handle cfg truncation
             current_guidance_scale = self.guidance_scale
@@ -582,13 +593,17 @@ class ZImagePipeline(nn.Module):
 
                     pred = pos + current_guidance_scale * (pos - neg)
 
-                    # Renormalization
+                    # Renormalization (torch.where avoids GPU->CPU sync from Python if/scalar comparison)
                     if self._cfg_normalization and float(self._cfg_normalization) > 0.0:
                         ori_pos_norm = torch.linalg.vector_norm(pos)
                         new_pos_norm = torch.linalg.vector_norm(pred)
                         max_new_norm = ori_pos_norm * float(self._cfg_normalization)
-                        if new_pos_norm > max_new_norm:
-                            pred = pred * (max_new_norm / new_pos_norm)
+                        scale = torch.where(
+                            new_pos_norm > max_new_norm,
+                            (max_new_norm / new_pos_norm.clamp(min=1e-12)).to(pred.dtype),
+                            pred.new_tensor(1.0),
+                        )
+                        pred = pred * scale
 
                     noise_pred.append(pred)
 

vllm_omni/diffusion/models/z_image/z_image_transformer.py

@@ -852,7 +852,13 @@ class ZImageTransformer2DModel(CachedTransformer):
 
         # Match t_embedder output dtype to x for layerwise casting compatibility
         adaln_input = t.type_as(x)
-        x[torch.cat(x_inner_pad_mask)] = self.x_pad_token
+        # Use torch.where instead of x[mask]= to avoid aten::index_put_/nonzero and cudaStreamSynchronize
+        x_pad_mask = torch.cat(x_inner_pad_mask)
+        x = torch.where(
+            x_pad_mask.unsqueeze(1).expand_as(x),
+            self.x_pad_token.expand(x.shape[0], -1),
+            x,
+        )
         x = list(x.split(x_item_seqlens, dim=0))
         x_cos, x_sin = self.rope_embedder(torch.cat(x_pos_ids, dim=0))
         x_cos = list(x_cos.split(x_item_seqlens, dim=0))
@@ -875,7 +881,13 @@ class ZImageTransformer2DModel(CachedTransformer):
 
         cap_feats = torch.cat(cap_feats, dim=0)
         cap_feats = self.cap_embedder(cap_feats)
-        cap_feats[torch.cat(cap_inner_pad_mask)] = self.cap_pad_token
+        # Use torch.where instead of cap_feats[mask]= to avoid aten::index_put_/nonzero and cudaStreamSynchronize
+        cap_pad_mask = torch.cat(cap_inner_pad_mask)
+        cap_feats = torch.where(
+            cap_pad_mask.unsqueeze(1).expand_as(cap_feats),
+            self.cap_pad_token.expand(cap_feats.shape[0], -1),
+            cap_feats,
+        )
         cap_feats = list(cap_feats.split(cap_item_seqlens, dim=0))
         cap_cos, cap_sin = self.rope_embedder(torch.cat(cap_pos_ids, dim=0))
         cap_cos = list(cap_cos.split(cap_item_seqlens, dim=0))

vllm_omni/diffusion/scheduler.py

@@ -13,14 +13,13 @@ logger = init_logger(__name__)
 
 class Scheduler:
     def initialize(self, od_config: OmniDiffusionConfig):
-        existing_context = getattr(self, "context", None)
-        if existing_context is not None and not existing_context.closed:
+        existing_mq = getattr(self, "mq", None)
+        if existing_mq is not None and not existing_mq.closed:
             logger.warning("SyncSchedulerClient is already initialized. Re-initializing.")
             self.close()
 
         self.num_workers = od_config.num_gpus
         self.od_config = od_config
-        self.context = zmq.Context()  # Standard synchronous context
 
         # Initialize single MessageQueue for all message types (generation & RPC)
         # Assuming all readers are local for now as per current launch_engine implementation
@@ -72,8 +71,5 @@ class Scheduler:
 
     def close(self):
         """Closes the socket and terminates the context."""
-        if hasattr(self, "context"):
-            self.context.term()
-        self.context = None
         self.mq = None
         self.result_mq = None

vllm_omni/distributed/omni_connectors/adapter.py

@@ -179,6 +179,27 @@ def try_recv_via_connector(
             return None, None
 
 
+def update_request_payload(connector: "OmniConnectorBase", req_id: str, payload_data: dict[str, Any]) -> dict[str, Any]:
+    """Update the payload data for a request in the connector.
+
+    Args:
+        connector: OmniConnectorBase instance
+        req_id: Request ID to update
+        payload_data: New payload data to store
+    """
+    origin_payload = connector.request_payload[req_id]
+    for key, value in payload_data.items():
+        if key == "finished":
+            continue
+        elif isinstance(value, torch.Tensor) and key in origin_payload:
+            payload_data[key] = torch.cat([origin_payload[key], value], dim=0)
+        elif isinstance(value, list) and key in origin_payload:
+            payload_data[key] = origin_payload[key] + value
+
+    connector.request_payload[req_id] = payload_data
+    return payload_data
+
+
 def get_chunk(
     connector: "OmniConnectorBase",
     scheduler_output: SchedulerOutput,
@@ -206,6 +227,7 @@ def get_chunk(
         payload_data = get_through_connector(connector, target_stage_id, stage_id, req_id, connector_get_key)
         if payload_data:
             new_req_data.additional_information = payload_data
+            connector.request_payload[req_id] = payload_data
             if payload_data.get("finished"):
                 connector.finished_requests.add(req_id)
 
@@ -222,15 +244,13 @@ def get_chunk(
         connector_get_key = f"{req_id}_{target_stage_id}_{chunk_id}"
         payload_data = get_through_connector(connector, target_stage_id, stage_id, req_id, connector_get_key)
         if payload_data:
+            payload_data = update_request_payload(connector, req_id, payload_data)
             cached_reqs.additional_information[cached_req_id] = payload_data
             if payload_data.get("finished"):
                 connector.finished_requests.add(req_id)
 
 
 def get_through_connector(connector, target_stage_id, stage_id, req_id, connector_get_key):
-    # Wait for data from previous stage
-    import time
-
     # TODO: add correct check mechanism for the payload_data
     max_wait = 300
     for _ in range(max_wait):
@@ -280,12 +300,8 @@ def get_chunk_for_generation(
     if payload_data.get("finished"):
         connector.finished_requests.add(request_id)
         request.status = RequestStatus.FINISHED_STOPPED
-
-    # TODO: remove special handling for prompt token ids ?
-    if chunk_id == 0:
-        request.prompt_token_ids = payload_data.get("code_predictor_codes", [])
-    else:
-        request.prompt_token_ids += payload_data.get("code_predictor_codes", [])
+    request.prompt_token_ids = payload_data.get("code_predictor_codes", [])
+    request.num_computed_tokens = 0
 
 
 def put_chunk(
@@ -308,8 +324,6 @@ def put_chunk(
     stage_id = connector.stage_id
     next_stage_id = stage_id + 1
     request_id = request.external_req_id
-    prompt_token_ids = request.prompt_token_ids
-    connector.request_prompt_token_ids[request_id] = prompt_token_ids
     chunk_id = connector.put_requests[request_id]
     connector_put_key = f"{request_id}_{stage_id}_{chunk_id}"
     payload_data = None
@@ -318,6 +332,7 @@ def put_chunk(
     if custom_process_input_func:
         try:
             payload_data = custom_process_input_func(
+                connector=connector,
                 pooling_output=pooling_output,
                 request=request,
             )
@@ -325,42 +340,8 @@ def put_chunk(
             logger.error(f"Failed to use custom_process_input_func for payload extraction: {e}")
 
         if not payload_data:
-            logger.warning("[Stage-%d] No payload data to send for request %s", stage_id, request_id)
             return
 
-        if stage_id == 0 and chunk_id == 0:
-            if connector.request_payload.get(request_id) is None:
-                if not payload_data.get("finished"):
-                    connector.request_payload[request_id] = payload_data
-                    return
-            else:
-                save_payload = connector.request_payload.pop(request_id)
-                payload_data["thinker_embeddings"] = torch.cat(
-                    (save_payload.get("thinker_embeddings"), payload_data.get("thinker_embeddings")), dim=0
-                )
-                payload_data["thinker_hidden_states"] = torch.cat(
-                    (save_payload.get("thinker_hidden_states"), payload_data.get("thinker_hidden_states")), dim=0
-                )
-                logger.debug("[Stage-%d] Merged embeddings and hidden states for request %s", stage_id, request_id)
-
-        if stage_id == 1:
-            # TODO: Make parameters configurable and optimize algorithms
-            chunk_size = left_context_size = 25
-            connector.code_prompt_token_ids[request_id].append(payload_data.get("code_predictor_codes", []))
-            length = len(connector.code_prompt_token_ids[request_id])
-            chunk_length = length % chunk_size
-            if chunk_length != 0 and not payload_data.get("finished"):
-                return
-
-            context_length = chunk_length if chunk_length != 0 else chunk_size
-            end_index = min(length, left_context_size + context_length)
-            payload_data["code_predictor_codes"] = (
-                torch.tensor(connector.code_prompt_token_ids[request_id][-end_index:])
-                .transpose(0, 1)
-                .reshape(-1)
-                .tolist()
-            )
-
         success, size, metadata = connector.put(
             from_stage=str(stage_id), to_stage=str(next_stage_id), put_key=connector_put_key, data=payload_data
         )

vllm_omni/entrypoints/async_omni.py

@@ -345,14 +345,22 @@ class AsyncOmni(OmniBase):
                 ):
                     yield output
 
-            logger.debug(f"[{self._name}] All requests completed")
-
-            # Summarize and print stats
+            logger.debug(f"[{self._name}] Request {request_id} finalized at stage-{final_stage_id_for_e2e}")
             try:
+                # Finalize E2E metrics if not already done
+                if str(request_id) not in metrics.e2e_done:
+                    metrics.on_finalize_request(
+                        final_stage_id_for_e2e,
+                        request_id,
+                        _req_start_ts.get(request_id, _wall_start_ts),
+                    )
+
+                logger.debug(f"[{self._name}] All requests completed")
+                # Summarize and print stats
                 summary = metrics.build_and_log_summary(final_stage_id_for_e2e)
                 logger.info("[Summary] %s", pformat(summary, sort_dicts=False))
             except Exception as e:
-                logger.exception(f"[{self._name}] Failed to build/log summary: {e}")
+                logger.exception(f"[{self._name}] Request {request_id} Failed to finalized/build/log summary: {e}")
             finally:
                 self.request_states.pop(request_id, None)
         except (asyncio.CancelledError, GeneratorExit):
@@ -382,9 +390,11 @@ class AsyncOmni(OmniBase):
                 except asyncio.QueueEmpty:
                     await asyncio.sleep(0.001)
                     continue
-
                 engine_outputs, finished, output_to_yield = self._process_single_result(
-                    result, stage, stage_id, metrics, req_start_ts, wall_start_ts, final_stage_id_for_e2e
+                    result,
+                    stage,
+                    stage_id,
+                    metrics,
                 )
                 if submit_flag and stage_id == 0:
                     submit_flag = False
@@ -421,9 +431,11 @@ class AsyncOmni(OmniBase):
             while not finished:
                 result = await req_state.queue.get()
                 assert stage_id == req_state.stage_id
-                req_id = result.get("request_id")
                 engine_outputs, finished, output_to_yield = self._process_single_result(
-                    result, stage, stage_id, metrics, req_start_ts, wall_start_ts, final_stage_id_for_e2e
+                    result,
+                    stage,
+                    stage_id,
+                    metrics,
                 )
                 if output_to_yield:
                     yield output_to_yield
@@ -432,7 +444,7 @@ class AsyncOmni(OmniBase):
             stage.set_engine_outputs(engine_outputs)
             # Forward to next stage if there is one
             next_stage_id = stage_id + 1
-            if next_stage_id <= final_stage_id_for_e2e and finished:
+            if next_stage_id <= final_stage_id_for_e2e:
                 next_stage: OmniStage = self.stage_list[next_stage_id]
                 next_inputs = next_stage.process_engine_inputs(self.stage_list, prompt)
                 sp_next: SamplingParams = sampling_params_list[next_stage_id]
@@ -447,7 +459,7 @@ class AsyncOmni(OmniBase):
                         connector=connector,
                         stage_id=stage_id,
                         next_stage_id=next_stage_id,
-                        req_id=req_id,
+                        req_id=request_id,
                         next_inputs=next_inputs,
                         sampling_params=sp_next,
                         original_prompt=prompt,
@@ -461,14 +473,14 @@ class AsyncOmni(OmniBase):
                     # because continuing would cause the request to be silently dropped
                     # and the orchestrator to hang waiting for completion.
                     error_msg = (
-                        f"[{self._name}] Failed to send request {req_id} to stage-{next_stage_id} via connector. "
+                        f"[{self._name}] Failed to send request {request_id} to stage-{next_stage_id} via connector. "
                         "Configure a connector for this edge or inspect connector logs for details."
                     )
                     logger.error(error_msg)
                     raise RuntimeError(error_msg)
-                logger.debug(f"[{self._name}] Forwarded request {req_id} to stage-{next_stage_id}")
+                logger.debug(f"[{self._name}] Forwarded request {request_id} to stage-{next_stage_id}")
             else:
-                logger.debug(f"[{self._name}] Request {req_id} fully completed")
+                logger.debug(f"[{self._name}] Request {request_id} fully completed")
 
     def _process_single_result(
         self,
@@ -476,9 +488,6 @@ class AsyncOmni(OmniBase):
         stage: OmniStage,
         stage_id: int,
         metrics: OrchestratorMetrics,
-        req_start_ts: dict[int, float],
-        wall_start_ts: float,
-        final_stage_id_for_e2e: int,
     ) -> tuple[Any, bool, OmniRequestOutput | None]:
         """
         Process a single result dictionary from a stage.
@@ -517,23 +526,8 @@ class AsyncOmni(OmniBase):
         )
 
         output_to_yield = None
+
         if getattr(stage, "final_output", False):
-            logger.debug(f"[{self._name}] Request {req_id} finalized at stage-{stage_id}")
-
-            # Finalize request metrics if this is the E2E final stage and it's finished
-            try:
-                rid_key = str(req_id)
-                if stage_id == final_stage_id_for_e2e and rid_key not in metrics.e2e_done and finished:
-                    metrics.on_finalize_request(
-                        stage_id,
-                        req_id,
-                        req_start_ts.get(req_id, wall_start_ts),
-                    )
-            except Exception as e:
-                logger.exception(
-                    f"[{self._name}] Finalize request handling error for req {req_id} at stage {stage_id}: {e}",
-                )
-
             # Construct output to yield
             images = []
             if stage.final_output_type == "image":

vllm_omni/entrypoints/cli/serve.py

@@ -194,17 +194,17 @@ class OmniServeCommand(CLISubcommand):
         )
 
         # diffusion model offload parameters
-        serve_parser.add_argument(
+        omni_config_group.add_argument(
             "--enable-cpu-offload",
             action="store_true",
             help="Enable CPU offloading for diffusion models.",
         )
-        serve_parser.add_argument(
+        omni_config_group.add_argument(
             "--enable-layerwise-offload",
             action="store_true",
             help="Enable layerwise (blockwise) offloading on DiT modules.",
         )
-        serve_parser.add_argument(
+        omni_config_group.add_argument(
             "--layerwise-num-gpu-layers",
             type=int,
             default=1,
@@ -245,7 +245,7 @@ class OmniServeCommand(CLISubcommand):
         # Diffusion model mixed precision
         omni_config_group.add_argument(
             "--max-generated-image-size",
-            type=float,
+            type=int,
             help="The max size of generate image (height * width).",
         )
         return serve_parser

vllm_omni/entrypoints/omni_stage.py

@@ -16,6 +16,7 @@ import sys
 import time
 import traceback
 from collections.abc import Sequence
+from contextlib import contextmanager
 from dataclasses import fields
 from typing import Any, Literal, cast
 
@@ -55,6 +56,144 @@ from vllm_omni.outputs import OmniRequestOutput
 logger = init_logger(__name__)
 
 
+@contextmanager
+def _sequential_init_lock(engine_args: dict[str, Any], stage_init_timeout: int = 300):
+    """Acquire device locks for sequential init if NVML is unavailable.
+
+    If process-scoped memory tracking is available (NVML works), stages can
+    safely initialize concurrently — each measures only its own GPU memory.
+    Otherwise, fall back to file-based locks to serialize initialization.
+    """
+    from vllm_omni.worker.gpu_memory_utils import is_process_scoped_memory_available
+
+    if is_process_scoped_memory_available():
+        logger.debug("NVML process-scoped memory available — concurrent init is safe, skipping locks")
+        yield
+        return
+
+    logger.debug("NVML unavailable — using sequential init locks")
+
+    from vllm_omni.platforms import current_omni_platform
+
+    # Get all parallel sizes from engine_args or parallel_config (defaults to 1)
+    if "parallel_config" in engine_args:
+        parallel_config = engine_args["parallel_config"]
+        tensor_parallel_size = parallel_config.get("tensor_parallel_size", 1)
+        pipeline_parallel_size = parallel_config.get("pipeline_parallel_size", 1)
+        data_parallel_size = parallel_config.get("data_parallel_size", 1)
+        prefill_context_parallel_size = parallel_config.get("prefill_context_parallel_size", 1)
+        sequence_parallel_size = parallel_config.get("sequence_parallel_size", 1)
+        cfg_parallel_size = parallel_config.get("cfg_parallel_size", 1)
+    else:
+        tensor_parallel_size = engine_args.get("tensor_parallel_size", 1)
+        pipeline_parallel_size = engine_args.get("pipeline_parallel_size", 1)
+        data_parallel_size = engine_args.get("data_parallel_size", 1)
+        prefill_context_parallel_size = engine_args.get("prefill_context_parallel_size", 1)
+        sequence_parallel_size = 1
+        cfg_parallel_size = 1
+
+    num_devices_per_stage = (
+        tensor_parallel_size
+        * pipeline_parallel_size
+        * data_parallel_size
+        * prefill_context_parallel_size
+        * sequence_parallel_size
+        * cfg_parallel_size
+    )
+
+    # Get physical device IDs from device control env var
+    device_control_env = current_omni_platform.device_control_env_var
+    visible_devices_str = os.environ.get(device_control_env)
+    physical_devices = []
+
+    if visible_devices_str:
+        try:
+            physical_devices = [int(x.strip()) for x in visible_devices_str.split(",") if x.strip()]
+        except (ValueError, IndexError):
+            pass
+
+    if not physical_devices:
+        num_devices = current_omni_platform.get_device_count()
+        physical_devices = list(range(num_devices))
+
+    num_devices_to_lock = min(num_devices_per_stage, len(physical_devices))
+    devices_to_lock = sorted(physical_devices[:num_devices_to_lock])
+
+    logger.debug(
+        "Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d, SP=%d, CFG=%d; will lock %d devices: %s",
+        tensor_parallel_size,
+        pipeline_parallel_size,
+        data_parallel_size,
+        prefill_context_parallel_size,
+        sequence_parallel_size,
+        cfg_parallel_size,
+        num_devices_to_lock,
+        devices_to_lock,
+    )
+
+    # Acquire exclusive locks for all devices using fcntl.flock
+    wait_start = time.time()
+    acquired_lock_fds = []
+
+    for device_id in devices_to_lock:
+        lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock"
+        lock_acquired = False
+
+        while not lock_acquired:
+            try:
+                lock_fd = os.open(lock_file, os.O_CREAT | os.O_RDWR, 0o644)
+
+                try:
+                    fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+                    os.ftruncate(lock_fd, 0)
+                    os.write(lock_fd, f"{os.getpid()}\n".encode())
+                    os.fsync(lock_fd)
+                    lock_acquired = True
+                    acquired_lock_fds.append(lock_fd)
+                    logger.debug("Acquired exclusive lock for device %s", device_id)
+                except BlockingIOError:
+                    os.close(lock_fd)
+
+                    if time.time() - wait_start > stage_init_timeout:
+                        logger.warning(
+                            "Timeout waiting for device %s initialization lock, proceeding anyway",
+                            device_id,
+                        )
+                        break
+
+                    time.sleep(0.1)
+            except OSError as e:
+                logger.debug(
+                    "Failed to acquire lock for device %s: %s, continuing anyway",
+                    device_id,
+                    e,
+                )
+                try:
+                    os.close(lock_fd)
+                except (OSError, NameError):
+                    pass
+                break
+
+    # Set FD_CLOEXEC to prevent child processes from inheriting locks
+    for lock_fd in acquired_lock_fds:
+        try:
+            flags = fcntl.fcntl(lock_fd, fcntl.F_GETFD)
+            fcntl.fcntl(lock_fd, fcntl.F_SETFD, flags | fcntl.FD_CLOEXEC)
+        except (OSError, ValueError):
+            pass
+
+    try:
+        yield
+    finally:
+        for lock_fd in acquired_lock_fds:
+            try:
+                fcntl.flock(lock_fd, fcntl.LOCK_UN)
+                os.close(lock_fd)
+                logger.debug("Released initialization lock (fd=%s)", lock_fd)
+            except (OSError, ValueError):
+                pass
+
+
 def _resolve_worker_cls(engine_args: dict[str, Any]) -> None:
     worker_type = engine_args.pop("worker_type", None)
     if not worker_type:
@@ -541,183 +680,32 @@ def _stage_worker(
     except Exception as e:
         logger.warning("Device setup failed: %s", e)
 
-    # Sequential initialization on the same device to avoid memory calculation errors
-    # when multiple instances start simultaneously
-    # For TP/PP/DP/SP, we need to lock ALL devices that will be used by this stage
-    lock_files = []
-    try:
-        # Get all parallel sizes from engine_args or parallel_config (defaults to 1)
-        if "parallel_config" in engine_args:
-            parallel_config = engine_args["parallel_config"]
-            tensor_parallel_size = parallel_config.get("tensor_parallel_size", 1)
-            pipeline_parallel_size = parallel_config.get("pipeline_parallel_size", 1)
-            data_parallel_size = parallel_config.get("data_parallel_size", 1)
-            prefill_context_parallel_size = parallel_config.get("prefill_context_parallel_size", 1)
-            sequence_parallel_size = parallel_config.get("sequence_parallel_size", 1)
-            cfg_parallel_size = parallel_config.get("cfg_parallel_size", 1)
+    # Use sequential init locks only when NVML is unavailable
+    with _sequential_init_lock(engine_args, stage_init_timeout):
+        # Init engine based on stage_type
+        logger.debug(
+            "[Stage-%s] Initializing %s engine with args keys=%s", stage_id, stage_type, list(engine_args.keys())
+        )
+        if engine_args.get("async_chunk", False):
+            logger.debug("[Stage-%s] Async chunk enabled, injecting connectors config", stage_id)
+            stage_connector_spec = {}
+            for v in connectors_config.values():
+                stage_connector_spec = dict(v.get("spec", {}))
+                break
+            engine_args["stage_connector_spec"] = stage_connector_spec
+            engine_args["stage_id"] = stage_id
+        if stage_type == "diffusion":
+            engine_args.pop("model_stage", None)
+            engine_args.pop("model", None)
+            stage_engine = OmniDiffusion(
+                model=model,
+                stage_id=stage_id,
+                engine_input_source=stage_payload.get("engine_input_source", []),
+                **engine_args,
+            )
         else:
-            tensor_parallel_size = engine_args.get("tensor_parallel_size", 1)
-            pipeline_parallel_size = engine_args.get("pipeline_parallel_size", 1)
-            data_parallel_size = engine_args.get("data_parallel_size", 1)
-            prefill_context_parallel_size = engine_args.get("prefill_context_parallel_size", 1)
-            sequence_parallel_size = 1  # not use in omni model
-            cfg_parallel_size = 1  # not used in omni model
-
-        # Calculate total number of devices needed for this stage
-        # For a single stage worker:
-        # - TP: splits model across devices (always needed)
-        # - PP: splits layers across pipeline stages, but each stage uses TP devices
-        # - DP: replicates model, but each replica uses TP devices
-        # - PCP: context parallelism, typically uses TP devices
-        # - SP: sequence parallelism, typically uses TP devices
-        # - CFG: Classifier-Free Guidance parallelism for diffusion models
-        # The number of devices per stage is determined by TP * PP * DP * PCP * SP * CFG size
-        # (PP/DP/PCP are higher-level parallelism that don't add devices per stage)
-        num_devices_per_stage = (
-            tensor_parallel_size
-            * pipeline_parallel_size
-            * data_parallel_size
-            * prefill_context_parallel_size
-            * sequence_parallel_size
-            * cfg_parallel_size
-        )
-
-        # Get physical device IDs from device control env var (e.g., CUDA_VISIBLE_DEVICES)
-        # After set_stage_devices, this env var is set to physical device(s)
-        device_control_env = current_omni_platform.device_control_env_var
-        visible_devices_str = _os.environ.get(device_control_env)
-        physical_devices = []
-
-        if visible_devices_str:
-            try:
-                physical_devices = [int(x.strip()) for x in visible_devices_str.split(",") if x.strip()]
-            except (ValueError, IndexError):
-                pass
-
-        if not physical_devices:
-            # Fallback: use logical device count if device control env var not set
-            num_devices = current_omni_platform.get_device_count()
-            physical_devices = list(range(num_devices))
-
-        # Determine which devices will be used (min of devices per stage and available devices)
-        num_devices_to_lock = min(num_devices_per_stage, len(physical_devices))
-        devices_to_lock = physical_devices[:num_devices_to_lock]
-
-        # Sort devices_to_lock to prevent deadlock (all processes acquire locks in same order)
-        devices_to_lock = sorted(devices_to_lock)
-
-        logger.debug(
-            "Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d, SP=%d, CFG=%d; will lock %d devices: %s",
-            tensor_parallel_size,
-            pipeline_parallel_size,
-            data_parallel_size,
-            prefill_context_parallel_size,
-            sequence_parallel_size,
-            cfg_parallel_size,
-            num_devices_to_lock,
-            devices_to_lock,
-        )
-
-        # Acquire exclusive locks for all devices using fcntl.flock
-        # Locks are automatically released when process dies
-        wait_start = _time.time()
-        acquired_lock_fds = []  # Store file descriptors to keep locks alive
-
-        for device_id in devices_to_lock:
-            lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock"
-            lock_acquired = False
-
-            while not lock_acquired:
-                try:
-                    # Open or create the lock file
-                    lock_fd = _os.open(lock_file, _os.O_CREAT | _os.O_RDWR, 0o644)
-
-                    # Try to acquire exclusive lock (non-blocking first)
-                    try:
-                        fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
-                        # Successfully acquired lock - write PID
-                        _os.ftruncate(lock_fd, 0)  # Clear file
-                        _os.write(lock_fd, f"{_os.getpid()}\n".encode())
-                        _os.fsync(lock_fd)  # Ensure written to disk
-                        lock_acquired = True
-                        acquired_lock_fds.append(lock_fd)
-                        logger.debug("Acquired exclusive lock for device %s", device_id)
-                    except BlockingIOError:
-                        # Lock is held by another process
-                        _os.close(lock_fd)
-
-                        # Check if we've been waiting too long
-                        if _time.time() - wait_start > stage_init_timeout:
-                            logger.warning(
-                                "Timeout waiting for device %s initialization lock, proceeding anyway",
-                                device_id,
-                            )
-                            break
-
-                        # Wait a bit before retrying
-                        _time.sleep(0.1)
-                except OSError as e:
-                    # Other error - log and continue without lock
-                    logger.debug(
-                        "Failed to acquire lock for device %s: %s, continuing anyway",
-                        device_id,
-                        e,
-                    )
-                    try:
-                        _os.close(lock_fd)
-                    except (OSError, NameError):
-                        pass
-                    break
-
-        lock_files = acquired_lock_fds
-
-        # Set FD_CLOEXEC on all lock file descriptors to prevent child processes
-        # (e.g., EngineCore) from inheriting them, which would cause deadlock
-        for lock_fd in acquired_lock_fds:
-            try:
-                flags = fcntl.fcntl(lock_fd, fcntl.F_GETFD)
-                fcntl.fcntl(lock_fd, fcntl.F_SETFD, flags | fcntl.FD_CLOEXEC)
-            except (OSError, ValueError):
-                pass
-    except Exception as e:
-        logger.debug(
-            "[Stage-%s] Failed to set up sequential initialization lock: %s",
-            stage_id,
-            e,
-        )
-
-    # Init engine based on stage_type
-    logger.debug("[Stage-%s] Initializing %s engine with args keys=%s", stage_id, stage_type, list(engine_args.keys()))
-    if engine_args.get("async_chunk", False):
-        logger.debug("[Stage-%s] Async chunk enabled, injecting connectors config", stage_id)
-        stage_connector_spec = {}
-        for v in connectors_config.values():
-            stage_connector_spec = dict(v.get("spec", {}))
-            break
-        engine_args["stage_connector_spec"] = stage_connector_spec
-        engine_args["stage_id"] = stage_id
-    if stage_type == "diffusion":
-        engine_args.pop("model_stage", None)
-        engine_args.pop("model", None)
-        stage_engine = OmniDiffusion(
-            model=model,
-            stage_id=stage_id,
-            engine_input_source=stage_payload.get("engine_input_source", []),
-            **engine_args,
-        )
-    else:
-        # Default to LLM engine
-        stage_engine = OmniLLM(model=model, **engine_args)
-
-    # Release all locks AFTER engine initialization completes
-    for lock_fd in lock_files:
-        try:
-            fcntl.flock(lock_fd, fcntl.LOCK_UN)
-            _os.close(lock_fd)
-            logger.debug("Released initialization lock (fd=%s)", lock_fd)
-        except (OSError, ValueError):
-            pass
-    lock_files = []  # Clear after release
+            # Default to LLM engine
+            stage_engine = OmniLLM(model=model, **engine_args)
 
     logger.debug("Engine initialized")
     # Initialize OmniConnectors if configured
@@ -1086,155 +1074,23 @@ async def _stage_worker_async(
             return
         connectors = built_connectors
 
-    # Sequential initialization on the same device to avoid memory calculation errors
-    # when multiple instances start simultaneously
-    # For TP/PP/DP/PCP, we need to lock ALL devices that will be used by this stage
-    lock_files = []
-    try:
-        # Get all parallel sizes from engine_args or parallel_config (defaults to 1)
-        if "parallel_config" in engine_args:
-            parallel_config = engine_args["parallel_config"]
-            tensor_parallel_size = parallel_config.get("tensor_parallel_size", 1)
-            pipeline_parallel_size = parallel_config.get("pipeline_parallel_size", 1)
-            data_parallel_size = parallel_config.get("data_parallel_size", 1)
-            prefill_context_parallel_size = parallel_config.get("prefill_context_parallel_size", 1)
-            sequence_parallel_size = parallel_config.get("sequence_parallel_size", 1)
-            cfg_parallel_size = parallel_config.get("cfg_parallel_size", 1)
-        else:
-            tensor_parallel_size = engine_args.get("tensor_parallel_size", 1)
-            pipeline_parallel_size = engine_args.get("pipeline_parallel_size", 1)
-            data_parallel_size = engine_args.get("data_parallel_size", 1)
-            prefill_context_parallel_size = engine_args.get("prefill_context_parallel_size", 1)
-            sequence_parallel_size = 1  # not use in omni model
-            cfg_parallel_size = 1  # not used in omni model
-
-        # Calculate total number of devices needed for this stage
-        # For a single stage worker:
-        # - TP: splits model across devices (always needed)
-        # - PP: splits layers across pipeline stages, but each stage uses TP devices
-        # - DP: replicates model, but each replica uses TP devices
-        # - PCP: context parallelism, typically uses TP devices
-        # - SP: sequence parallelism, typically uses TP devices
-        # - CFG: Classifier-Free Guidance parallelism for diffusion models
-        # The number of devices per stage is determined by TP * PP * DP * PCP * SP * CFG size
-        # (PP/DP/PCP are higher-level parallelism that don't add devices per stage)
-        num_devices_per_stage = (
-            tensor_parallel_size
-            * pipeline_parallel_size
-            * data_parallel_size
-            * prefill_context_parallel_size
-            * sequence_parallel_size
-            * cfg_parallel_size
-        )
-
-        # Get physical device IDs from device control env var (e.g., CUDA_VISIBLE_DEVICES)
-        # After set_stage_devices, this env var is set to physical device(s)
-        device_control_env = current_omni_platform.device_control_env_var
-        visible_devices_str = _os.environ.get(device_control_env)
-        physical_devices = []
-
-        if visible_devices_str:
-            try:
-                physical_devices = [int(x.strip()) for x in visible_devices_str.split(",") if x.strip()]
-            except (ValueError, IndexError):
-                pass
-
-        if not physical_devices:
-            # Fallback: use logical device count if device control env var not set
-            num_devices = current_omni_platform.get_device_count()
-            physical_devices = list(range(num_devices))
-
-        # Determine which devices will be used (min of devices per stage and available devices)
-        num_devices_to_lock = min(num_devices_per_stage, len(physical_devices))
-        devices_to_lock = physical_devices[:num_devices_to_lock]
-
-        # Sort devices_to_lock to prevent deadlock (all processes acquire locks in same order)
-        devices_to_lock = sorted(devices_to_lock)
-
+    # Use sequential init locks only when NVML is unavailable
+    with _sequential_init_lock(engine_args, stage_init_timeout):
+        # Init engine based on stage_type
         logger.debug(
-            "Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d, SP=%d, CFG=%d; will lock %d devices: %s",
-            tensor_parallel_size,
-            pipeline_parallel_size,
-            data_parallel_size,
-            prefill_context_parallel_size,
-            sequence_parallel_size,
-            cfg_parallel_size,
-            num_devices_to_lock,
-            devices_to_lock,
+            "[Stage-%s] Initializing %s engine with args keys=%s",
+            stage_id,
+            stage_type,
+            list(engine_args.keys()),
         )
-
-        # Acquire exclusive locks for all devices using fcntl.flock
-        # Locks are automatically released when process dies
-        wait_start = _time.time()
-        acquired_lock_fds = []  # Store file descriptors to keep locks alive
-
-        for device_id in devices_to_lock:
-            lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock"
-            lock_acquired = False
-
-            while not lock_acquired:
-                try:
-                    # Open or create the lock file
-                    lock_fd = _os.open(lock_file, _os.O_CREAT | _os.O_RDWR, 0o644)
-
-                    # Try to acquire exclusive lock (non-blocking first)
-                    try:
-                        fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
-                        # Successfully acquired lock - write PID
-                        _os.ftruncate(lock_fd, 0)  # Clear file
-                        _os.write(lock_fd, f"{_os.getpid()}\n".encode())
-                        _os.fsync(lock_fd)  # Ensure written to disk
-                        lock_acquired = True
-                        acquired_lock_fds.append(lock_fd)
-                        logger.debug("Acquired exclusive lock for device %s", device_id)
-                    except BlockingIOError:
-                        # Lock is held by another process
-                        _os.close(lock_fd)
-
-                        # Check if we've been waiting too long
-                        if _time.time() - wait_start > stage_init_timeout:
-                            logger.warning(
-                                "Timeout waiting for device %s initialization lock, proceeding anyway with timeout %s",
-                                device_id,
-                                stage_init_timeout,
-                            )
-                            break
-
-                        # Wait a bit before retrying
-                        _time.sleep(0.1)
-                except OSError as e:
-                    # Other error - log and continue without lock
-                    logger.debug(
-                        "Failed to acquire lock for device %s: %s, continuing anyway",
-                        device_id,
-                        e,
-                    )
-                    try:
-                        _os.close(lock_fd)
-                    except (OSError, NameError):
-                        pass
-                    break
-
-        lock_files = acquired_lock_fds
-    except Exception as e:
-        logger.debug("Failed to set up sequential initialization lock: %s", e)
-
-    # Init engine based on stage_type
-    logger.debug(
-        "[Stage-%s] Initializing %s engine with args keys=%s",
-        stage_id,
-        stage_type,
-        list(engine_args.keys()),
-    )
-    if engine_args.get("async_chunk", False):
-        logger.debug("[Stage-%s] Async chunk enabled, injecting connectors config", stage_id)
-        stage_connector_spec = {}
-        for v in connectors_config.values():
-            stage_connector_spec = dict(v.get("spec", {}))
-            break
-        engine_args["stage_connector_spec"] = stage_connector_spec
-        engine_args["stage_id"] = stage_id
-    try:
+        if engine_args.get("async_chunk", False):
+            logger.debug("[Stage-%s] Async chunk enabled, injecting connectors config", stage_id)
+            stage_connector_spec = {}
+            for v in connectors_config.values():
+                stage_connector_spec = dict(v.get("spec", {}))
+                break
+            engine_args["stage_connector_spec"] = stage_connector_spec
+            engine_args["stage_id"] = stage_id
         if stage_type == "diffusion":
             # For diffusion, we need to extract diffusion-specific config
             od_config = _build_od_config(engine_args, model)
@@ -1261,17 +1117,6 @@ async def _stage_worker_async(
                 usage_context=usage_context,
                 engine_args=omni_engine_args,
             )
-    finally:
-        # Release all locks by closing file descriptors
-        # Locks are automatically released when file descriptors are closed
-        # or when process dies
-        for lock_fd in lock_files:
-            try:
-                fcntl.flock(lock_fd, fcntl.LOCK_UN)
-                _os.close(lock_fd)
-                logger.debug("Released initialization lock (fd=%s)", lock_fd)
-            except (OSError, ValueError):
-                pass
     omni_stage.set_async_engine(stage_engine)
     if hasattr(omni_stage.async_engine, "log_stats") and omni_stage.async_engine.log_stats:
 

vllm_omni/entrypoints/openai/api_server.py

@@ -968,7 +968,9 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
         # a proper generator is initialized in the backend.
         # This fixes issues where using the default global generator
         # might produce blurry images in some environments.
-        _update_if_not_none(gen_params, "seed", random.randint(0, 2**32 - 1) if request.seed is None else request.seed)
+        _update_if_not_none(
+            gen_params, "seed", request.seed if request.seed is not None else random.randint(0, 2**32 - 1)
+        )
 
         request_id = f"img_gen_{uuid.uuid4().hex}"
 
@@ -1126,7 +1128,7 @@ async def edit_images(
         # a proper generator is initialized in the backend.
         # This fixes issues where using the default global generator
         # might produce blurry images in some environments.
-        _update_if_not_none(gen_params, "seed", seed or random.randint(0, 2**32 - 1))
+        _update_if_not_none(gen_params, "seed", seed if seed is not None else random.randint(0, 2**32 - 1))
 
         # 4. Generate images using AsyncOmni (multi-stage mode)
         request_id = f"img_edit_{int(time.time())}"