Merge branch 'main' into dev/rebase-v0.15.0

Update Dockerfile.ci to improve installation commands and ensure Python environment consistency
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-03 05:09:10 +00:00 · 2026-02-03 05:06:51 +00:00 · 2026-02-02 23:42:55 +08:00 · 2026-02-02 21:51:01 +08:00 · 2026-02-02 21:35:45 +08:00 · 2026-02-02 21:34:42 +08:00
117 changed files with 1268 additions and 4780 deletions
@@ -20,19 +20,12 @@ steps:
  - label: "Simple Unit Test"
    depends_on: image-build
    commands:
-    - |
-      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
+      - 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
    agents:
      queue: "gpu_1_queue"
    plugins:
@@ -1,45 +0,0 @@
-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
@@ -7,7 +7,6 @@ on:
    paths-ignore:
      - 'docs/**'
      - 'examples/**'
-      - 'apps/**'
      - '.github/**'
      - '.buildkite/**'
      - '**.md'
@@ -21,7 +20,6 @@ on:
    paths-ignore:
      - 'docs/**'
      - 'examples/**'
-      - 'apps/**'
      - '.github/**'
      - '.buildkite/**'
      - '*.yml'
@@ -1,27 +0,0 @@
-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 }}
@@ -1,40 +0,0 @@
-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
@@ -153,14 +153,6 @@ dmypy.json
 # Cython debug symbols
 cython_debug/

-# Claude
-CLAUDE.md
-.claude/
-
-# Codex
-AGENTS.md
-.codex/
-
 # cursor
 .cursor/

@@ -222,7 +214,6 @@ data/
 datasets/
 *.csv
 *.json
-!apps/ComfyUI-vLLM-Omni/example_workflows/*.json
 *.jsonl
 *.parquet

@@ -9,16 +9,14 @@ Easy, fast, and cheap omni-modality model serving for everyone
 </h3>

 <p align="center">
-| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true"><b>Slides</b></a> |
+| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> |
 </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).
+- [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.
+- [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).
 - [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.

 ---
@@ -69,19 +67,6 @@ Visit our [documentation](https://vllm-omni.readthedocs.io/en/latest/) to learn
 We welcome and value any contributions and collaborations.
 Please check out [Contributing to vLLM-Omni](https://vllm-omni.readthedocs.io/en/latest/contributing/) for how to get involved.

-## Citation
-
-If you use vLLM-Omni for your research, please cite our [paper](https://arxiv.org/abs/2602.02204):
-
-```bibtex
-@article{yin2026vllmomni,
-  title={vLLM-Omni: Fully Disaggregated Serving for Any-to-Any Multimodal Models},
-  author={Peiqi Yin, Jiangyun Zhu, Han Gao, Chenguang Zheng, Yongxiang Huang, Taichang Zhou, Ruirui Yang, Weizhi Liu, Weiqing Chen, Canlin Guo, Didan Deng, Zifeng Mo, Cong Wang, James Cheng, Roger Wang, Hongsheng Liu},
-  journal={arXiv preprint arXiv:2602.02204},
-  year={2026}
-}
-```
-
 ## Join the Community
 Feel free to ask questions, provide feedbacks and discuss with fellow users of vLLM-Omni in `#sig-omni` slack channel at [slack.vllm.ai](https://slack.vllm.ai) or vLLM user forum at [discuss.vllm.ai](https://discuss.vllm.ai).

@@ -1,115 +0,0 @@
-# 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/
@@ -1,15 +0,0 @@
-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.
@@ -1,170 +0,0 @@
-# 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.
@@ -1,52 +0,0 @@
-"""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"
@@ -1 +0,0 @@
-# noqa: N999   # This is not a python library intended to be imported
@@ -1,569 +0,0 @@
-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},)
@@ -1,473 +0,0 @@
-"""
-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
@@ -1,230 +0,0 @@
-"""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)
@@ -1,125 +0,0 @@
-"""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)
@@ -1,119 +0,0 @@
-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")
@@ -1,39 +0,0 @@
-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]
@@ -1,84 +0,0 @@
-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
@@ -1 +0,0 @@
-{"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
@@ -1 +0,0 @@
-{"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}
@@ -1 +0,0 @@
-{"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}
@@ -1,21 +0,0 @@
-/**
- * @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!")
-    },
-})
@@ -3,18 +3,19 @@ ARG VLLM_BASE_TAG=v0.15.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}
-
 COPY . .
-
 # Install system dependencies
 RUN apt-get update && \
-    apt-get install -y ffmpeg && \
-    apt-get clean && \
-    rm -rf /var/lib/apt/lists/*
+  apt-get install -y ffmpeg && \
+  apt-get clean && \
+  rm -rf /var/lib/apt/lists/*
+
+# Install vLLM wheel for the specific commit
+# URL format: https://wheels.vllm.ai/<commit_hash>/
+# Current vLLM commit: f176443446f659dbab5315e056e605d8984fd976 (v0.15.0)
+RUN uv pip install --system --upgrade --force-reinstall vllm --extra-index-url https://wheels.vllm.ai/f176443446f659dbab5315e056e605d8984fd976

 # Install vllm-omni into the same uv-managed Python environment used by the base image.
-RUN uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]"
-
+RUN uv pip install --system --no-cache-dir ".[dev]"
 RUN ln -sf /usr/bin/python3 /usr/bin/python
-
 ENTRYPOINT []
@@ -1,5 +1,5 @@
-ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.14.0rc1
+ARG VLLM_ASCEND_IMAGE=quay.nju.edu.cn/ascend/vllm-ascend
+ARG VLLM_ASCEND_TAG=v0.11.0rc2
 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 . --no-build-isolation
+RUN pip install --no-cache-dir -e ".[dev]"

 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn

@@ -1,15 +0,0 @@
-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 []
@@ -15,7 +15,11 @@ 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]" --no-build-isolation
+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 ln -sf /usr/bin/python3 /usr/bin/python

@@ -11,23 +11,17 @@ 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/*
@@ -44,7 +38,6 @@ 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
@@ -61,7 +54,6 @@ 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
@@ -10,9 +10,6 @@ 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][]
@@ -29,9 +26,7 @@ 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][]
@@ -63,7 +58,6 @@ 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][]
@@ -94,23 +88,20 @@ Configuration classes.

 Worker classes and model runners for distributed inference.

- [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.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.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.mixins.OmniWorkerMixin][]
+- [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][]
@@ -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", "audio_duration". '
+                    'Allowed metric names are "ttft", "tpot", "itl", "e2el", "audio_ttfp", "audio_rtf". '

 - `--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):                  3.89
-Request throughput (req/s):              0.51
+Benchmark duration (s):                  24.35
+Request throughput (req/s):              0.08
 Peak concurrent requests:                2.00
 ----------------End-to-end Latency----------------
-Mean E2EL (ms):                          3824.76
-Median E2EL (ms):                        3824.76
-P99 E2EL (ms):                           3888.54
+Mean E2EL (ms):                          22576.23
+Median E2EL (ms):                        22576.23
+P99 E2EL (ms):                           24205.72
 ================== Text Result ===================
 Total input tokens:                      30
-Total generated tokens:                  10101
-Output token throughput (tok/s):         2595.57
-Peak output token throughput (tok/s):    111.00
+Total generated tokens:                  8973
+Output token throughput (tok/s):         368.52
+Peak output token throughput (tok/s):    81.00
 Peak concurrent requests:                2.00
-Total Token throughput (tok/s):          2603.28
+Total Token throughput (tok/s):          369.76
 ---------------Time to First Token----------------
-Mean TTFT (ms):                          117.15
-Median TTFT (ms):                        117.15
-P99 TTFT (ms):                           142.69
+Mean TTFT (ms):                          125.16
+Median TTFT (ms):                        125.16
+P99 TTFT (ms):                           155.88
 -----Time per Output Token (excl. 1st token)------
-Mean TPOT (ms):                          0.73
-Median TPOT (ms):                        0.73
-P99 TPOT (ms):                           0.74
+Mean TPOT (ms):                          5.01
+Median TPOT (ms):                        5.01
+P99 TPOT (ms):                           5.42
 ---------------Inter-token Latency----------------
-Mean ITL (ms):                           16.47
-Median ITL (ms):                         16.19
-P99 ITL (ms):                            52.55
+Mean ITL (ms):                           34.15
+Median ITL (ms):                         0.01
+P99 ITL (ms):                            376.19
 ================== Audio Result ==================
-Total audio duration generated(s):       15.79
-Total audio frames generated:            379050
-Audio throughput(audio duration/s):      4.06
+Total audio duration generated(s):       3.95
+Total audio frames generated:            94890
+Audio throughput(audio duration/s):      0.16
 ---------------Time to First Packet---------------
-Mean AUDIO_TTFP (ms):                    3701.37
-Median AUDIO_TTFP (ms):                  3701.37
-P99 AUDIO_TTFP (ms):                     3762.25
+Mean AUDIO_TTFP (ms):                    11756.89
+Median AUDIO_TTFP (ms):                  11756.89
+P99 AUDIO_TTFP (ms):                     20854.25
 -----------------Real Time Factor-----------------
-Mean AUDIO_RTF:                          0.47
-Median AUDIO_RTF:                        0.47
-P99 AUDIO_RTF:                           0.48
+Mean AUDIO_RTF:                          3.75
+Median AUDIO_RTF:                        3.75
+P99 AUDIO_RTF:                           7.39
 ==================================================
 ```
 Notes:
-We use audio generation time / audio duration to calculate RTF.
+We use (audio generation time - first packet latency) / audio duration to calculate RTF.

 </details>

@@ -1,70 +0,0 @@
-# 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.
@@ -26,8 +26,6 @@ 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"
@@ -17,10 +17,12 @@ 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.15.0 --torch-backend=auto
+uv pip install vllm==0.14.0 --torch-backend=auto
 ```

 #### Installation of vLLM-Omni
@@ -34,10 +36,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.15.0 release version of the library
+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.15.0 --torch-backend=auto
+uv pip install vllm==0.14.0 --torch-backend=auto
 ```

 The release 0.14.0 of vLLM is based on PyTorch 2.9.0 which requires CUDA 12.9 environment.
@@ -56,11 +58,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.15.0
+git checkout v0.14.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.15.0/vllm-0.15.0-cp38-abi3-manylinux_2_31_x86_64.whl
+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
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -9,60 +9,10 @@ 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]
@@ -16,13 +16,7 @@ vLLM-Omni supports NPU through the vLLM Ascend Plugin (vllm-ascend). This is a c

 ## Installation

-### Set up using Docker
-
-=== "NPU"
-
-    --8<-- "docs/getting_started/installation/npu/npu.inc.md:pre-built-images"
-
-### Build wheel from source
+### Recommended

 === "NPU"

@@ -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.0rc1
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0
 # Atlas A3:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1-a3
-export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0-a3
+export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0
 docker run --rm \
    --name vllm-omni-npu \
    --shm-size=1g \
@@ -34,79 +34,26 @@ 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
-VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
-# OR pip install -v -e . --no-build-isolation
+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
 ```

 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]
@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
 source .venv/bin/activate

 # On CUDA
-uv pip install vllm==0.15.0 --torch-backend=auto
+uv pip install vllm==0.14.0 --torch-backend=auto

 # On ROCm
-uv pip install vllm==0.15.0 --extra-index-url https://wheels.vllm.ai/rocm/0.15.0/rocm700
+uv pip install vllm==0.14.0 --extra-index-url https://wheels.vllm.ai/rocm/0.14.0/rocm700

 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
@@ -25,7 +25,6 @@ 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` |
@@ -52,7 +51,6 @@ 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` |
@@ -61,8 +59,3 @@ 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` |
@@ -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: 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).
+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!

 > Q: Does vLLM-Omni support multimodal streaming input and output?

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

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

@@ -2,7 +2,6 @@

 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.
@@ -100,7 +99,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`](https://github.com/vllm-project/vllm-omni/tree/main/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`](../../../vllm_omni/model_executor/stage_configs/bagel.yaml)

 #### 📌 Command Line Arguments (end2end.py)

@@ -178,10 +177,3 @@ 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"
-    ``````
@@ -47,15 +47,10 @@ 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

@@ -7,19 +7,12 @@ 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 cherry_blossom.jpg \
-  --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
+  --image input.png \
+  --prompt "A cat playing with yarn, smooth motion" \
  --negative_prompt "<optional quality filter>" \
  --height 480 \
  --width 832 \
@@ -34,12 +27,11 @@ 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 cherry_blossom.jpg \
-  --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
+  --image input.png \
+  --prompt "A cat playing with yarn, smooth motion" \
  --negative_prompt "<optional quality filter>" \
  --height 480 \
  --width 832 \
@@ -60,17 +52,12 @@ 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

@@ -1,9 +1,8 @@
-# LoRA Inference(Diffusion)
+# LoRA-Inference

 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/lora_inference>.

-
-This directory contains examples for using LoRA (Low-Rank Adaptation) adapters with vLLM-omni diffusion models for offline inference.
+This 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
@@ -1,4 +1,4 @@
-# Qwen3-TTS
+# Qwen3-TTS Offline Inference

 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen3_tts>.

@@ -16,15 +16,6 @@ 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:
@@ -1,47 +0,0 @@
-# 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"
-    ``````
@@ -51,8 +51,7 @@ 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.
@@ -96,16 +95,11 @@ 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.

@@ -12,11 +12,10 @@ 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 832 \
-  --num_frames 33 \
+  --width 640 \
+  --num_frames 32 \
  --guidance_scale 4.0 \
  --guidance_scale_high 3.0 \
-  --flow_shift 12.0 \
  --num_inference_steps 40 \
  --fps 16 \
  --output t2v_out.mp4
@@ -25,19 +24,14 @@ python text_to_video.py \
 Key arguments:

 - `--prompt`: text description (string).
- `--height/--width`: output resolution (defaults 480x832, i.e. 480P). Dimensions should align with Wan VAE downsampling (multiples of 8).
+- `--height/--width`: output resolution (defaults 720x1280). 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).
- `--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).
+- `--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.
 - `--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

@@ -2,10 +2,9 @@

 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/bagel>.

-
 ## 🛠️ Installation

-Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
+Please refer to [README.md](../../../README.md)

 ## Run examples (BAGEL-7B-MoT)

@@ -75,7 +74,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`](https://github.com/vllm-project/vllm-omni/tree/main/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`](../../../vllm_omni/model_executor/stage_configs/bagel.yaml)

 | Modality    | Input        | Output | Description                            |
 | ----------- | ------------ | ------ | -------------------------------------- |
@@ -231,14 +230,3 @@ 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"
-    ``````
@@ -1,8 +1,7 @@
-# LoRA Inference(Diffusion)
+# LoRA-Inference

 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).
@@ -56,17 +55,15 @@ 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"
-    ``````sh
+    ``````py
    --8<-- "examples/online_serving/lora_inference/run_curl_lora_inference.sh"
    ``````
 ??? abstract "run_server.sh"
-    ``````sh
+    ``````py
    --8<-- "examples/online_serving/lora_inference/run_server.sh"
    ``````
@@ -1,189 +0,0 @@
-# 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"
-    ``````
@@ -4,19 +4,12 @@ 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 cherry_blossom.jpg \
-  --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
+  --image input.png \
+  --prompt "A cat playing with yarn, smooth motion" \
  --negative_prompt "<optional quality filter>" \
  --height 480 \
  --width 832 \
@@ -31,12 +24,11 @@ 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 cherry_blossom.jpg \
-  --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
+  --image input.png \
+  --prompt "A cat playing with yarn, smooth motion" \
  --negative_prompt "<optional quality filter>" \
  --height 480 \
  --width 832 \
@@ -1,4 +1,4 @@
-# LoRA Inference(Diffusion)
+# LoRA Inference Examples

 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.
@@ -1,4 +1,4 @@
-# Qwen3-TTS
+# Qwen3-TTS Offline Inference

 This directory contains an offline demo for running Qwen3 TTS models with vLLM Omni. It builds task-specific inputs and generates WAV files locally.

@@ -109,12 +109,7 @@ 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()


@@ -146,7 +141,6 @@ 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
@@ -179,7 +173,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}, tensor_parallel_size={args.tensor_parallel_size}"
+        f"  Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}"
    )
    print(f"  Video size: {args.width}x{args.height}")
    print(f"{'=' * 60}\n")
@@ -12,9 +12,6 @@ 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
@@ -1,4 +1,4 @@
-# LoRA Inference(Diffusion)
+# Online LoRA Inference (Diffusion)

 This example shows how to use **per-request LoRA** with vLLM-Omni diffusion models via the OpenAI-compatible Chat Completions API.

@@ -308,24 +308,19 @@ 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)
-    # Note: audio is extracted from video automatically when use_audio_in_video=True
-    # Do not include a separate audio_url here
+    audio_url = get_audio_url_from_path(audio_path)
    return {
        "role": "user",
        "content": [
            {"type": "video_url", "video_url": {"url": video_url}},
+            {"type": "audio_url", "audio_url": {"url": audio_url}},
            {"type": "text", "text": question},
        ],
    }
@@ -402,6 +397,7 @@ 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:
@@ -488,7 +484,7 @@ def parse_args():
        "--query-type",
        "-q",
        type=str,
-        default="use_audio_in_video",
+        default="use_mixed_modalities",
        choices=query_map.keys(),
        help="Query type.",
    )
@@ -1,4 +1,4 @@
-# Qwen3-TTS
+# Qwen3-TTS Online Serving

 This directory contains examples for running Qwen3-TTS models with vLLM-Omni's online serving API.

@@ -111,14 +111,9 @@ def run_tts_generation(args) -> None:
        print(response.text)
        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
+    if response.content.decode("utf-8").startswith('{"error"'):
+        print(f"Error: {response.content.decode('utf-8')}")
+        return

    # Save audio response
    output_path = args.output or "tts_output.wav"
@@ -1,9 +1,5 @@
 [build-system]
-requires = [
-    "setuptools>=77.0.3,<81.0.0",
-    "wheel",
-    "setuptools-scm>=8.0",
-]
+requires = ["setuptools>=61.0", "wheel"]
 build-backend = "setuptools.build_meta"

 [project]
@@ -12,16 +8,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 = "Apache-2.0"
+license = {text = "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",
@@ -30,10 +26,26 @@ classifiers = [
    "Topic :: Software Development :: Libraries :: Python Modules",
 ]

-# 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.
+
+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
+]

 [project.optional-dependencies]

@@ -143,7 +155,11 @@ python_classes = ["Test*"]
 python_functions = ["test_*"]
 addopts = [
    "--strict-markers",
-    "--strict-config"
+    "--strict-config",
+    "--cov=vllm_omni",
+    "--cov-report=term-missing",
+    "--cov-report=html",
+    "--cov-report=xml",
 ]
 markers = [
    # ci/cd required
@@ -188,5 +204,4 @@ extend-ignore-identifiers-re = [
    ".*NOTHINK.*",
    ".*nin.*",
    "Ono_Anna",
-    ".*cann.*",
 ]
@@ -1,14 +0,0 @@
-# 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
@@ -1,2 +0,0 @@
-r common.txt
-onnxruntime>=1.23.2
@@ -1,3 +0,0 @@
-r common.txt
-onnxruntime>=1.23.2
-fa3-fwd==0.0.1
@@ -1,3 +0,0 @@
-r common.txt
-onnxruntime-cann>=1.23.2
-torchaudio==2.9.0
@@ -1,2 +0,0 @@
-r common.txt
-onnxruntime-rocm>=1.22.2
@@ -1,2 +0,0 @@
-r common.txt
-onnxruntime>=1.23.2
@@ -1,171 +0,0 @@
-"""
-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,
-    )
@@ -60,12 +60,6 @@ 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")
@@ -1,71 +0,0 @@
-# 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]],
-        )
@@ -24,7 +24,6 @@ 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"

@@ -63,7 +62,7 @@ def test_cache_dit(model_name: str):
                width=width,
                num_inference_steps=num_inference_steps,
                guidance_scale=0.0,
-                generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
+                generator=torch.Generator("cuda").manual_seed(42),
                num_outputs_per_prompt=1,  # Single output for speed
            ),
        )
@@ -47,7 +47,7 @@ def test_diffusion_model(model_name: str):
                width=width,
                num_inference_steps=2,
                guidance_scale=0.0,
-                generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
+                generator=torch.Generator("cuda").manual_seed(42),
                num_outputs_per_prompt=2,
            ),
        )
@@ -14,7 +14,6 @@ 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"
@@ -43,7 +42,7 @@ def test_video_diffusion_model(model_name: str):
            num_frames=num_frames,
            num_inference_steps=2,
            guidance_scale=1.0,
-            generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
+            generator=torch.Generator("cuda").manual_seed(42),
        ),
    )
    first_output = outputs[0]
@@ -8,7 +8,6 @@ import concurrent.futures
 import os

 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-import threading
 import time
 from pathlib import Path

@@ -22,6 +21,7 @@ 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,29 +33,6 @@ 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(
@@ -109,180 +86,227 @@ def get_deploy_config(deploy_type="TP1"):
    return result.get(deploy_type, result["TP1"])


-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_text_to_text_001(client: openai.OpenAI, omni_server) -> None:
+@pytest.mark.parametrize("test_config", test_params)
+def test_text_to_text_001(test_config: tuple[str, str]) -> None:
    """Test processing text, generating text output via OpenAI API."""
-    messages = dummy_messages_from_mix_data(system_prompt=get_system_prompt(), content_text=get_prompt())
+    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())

-    # 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 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("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."""
-    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 "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.
-
-
-@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."""
-
-    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()
-    # 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)
-        ]
+        # Test single completion
+        api_client = client(server)
        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)
+        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.

-    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 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_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"
+        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("omni_server", test_params, indirect=True)
-def test_image_to_text_001(client: openai.OpenAI, omni_server) -> None:
+@pytest.mark.parametrize("test_config", test_params)
+def test_audio_to_text_001(test_config: tuple[str, str]) -> 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.
+    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."

-    # 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]
-    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."
+        # Verify E2E
+        print(f"the request e2e is: {time.perf_counter() - start_time}")
+        # TODO: Verify the E2E latency after confirmation baseline.


-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_image_to_text_audio_001(client: openai.OpenAI, omni_server) -> None:
+@pytest.mark.parametrize("test_config", test_params)
+def test_audio_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 = 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']}")

-    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(1280, 720)['base64']}"
+        messages = dummy_messages_from_mix_data(audio_data_url=audio_data_url)

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

-    # 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)
-        ]
+        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)
        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)
+        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.

-    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 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."

-        # 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_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"
@@ -338,25 +338,6 @@ 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(
@@ -833,55 +814,3 @@ 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."
-    )
@@ -1,154 +0,0 @@
-"""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
@@ -14,17 +14,13 @@ 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_s: float = 0.0
+    total_audio_duration_ms: 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(
@@ -77,7 +73,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_s))
+    print("{:<40} {:<10.2f}".format("Total audio duration generated(s):", metrics.total_audio_duration_ms))
    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:
@@ -96,23 +92,16 @@ 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 = "_ms"
-    unit_suffix = " (ms)"
-    if is_audio_duration:
-        suffix = "_s"
-        unit_suffix = " (s)"
-    elif is_audio_rtf:
-        suffix = ""
-        unit_suffix = ""
+    suffix = "" if is_audio_rtf else "_ms"
+    unit_suffix = "" if is_audio_rtf else " (ms)"
+
    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}")
@@ -223,9 +212,6 @@ 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,
@@ -310,11 +296,7 @@ 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],
-        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_duration_ms=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),
@@ -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,6 +133,7 @@ async def async_request_openai_chat_omni_completions(
                            continue

                        chunk = message.removeprefix("data: ")
+
                        if chunk != "[DONE]":
                            timestamp = time.perf_counter()
                            data = json.loads(chunk)
@@ -147,11 +148,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 - st
+                                    audio_generate_time = timestamp - audio_first_timestamp
                                    if content != "":
                                        audio_bytes = base64.b64decode(content)
                                        seg = AudioSegment.from_file(io.BytesIO(audio_bytes))
@@ -163,8 +164,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
@@ -180,7 +181,9 @@ 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
@@ -470,9 +473,6 @@ 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,20 +505,8 @@ 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 = "_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
+        suffix = "" if is_audio_rtf else "_ms"
        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
@@ -67,11 +67,9 @@ class OmniGenerationScheduler(VLLMScheduler):
            if self.omni_connector is not None:
                get_chunk_for_generation(self.omni_connector, request)

-            # 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
-            ):
+            # 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:
                already_finished_reqs.add(request)
                req_index += 1
                continue
@@ -110,10 +108,9 @@ class OmniGenerationScheduler(VLLMScheduler):
            if self.omni_connector is not None:
                get_chunk_for_generation(self.omni_connector, request)

-            # 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
-            ):
+            # 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:
                # Pop the finished request from waiting queue and don't schedule it
                self.waiting.pop_request()
                continue
@@ -273,77 +273,19 @@ 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.1-dev pipeline.
+    """Enable cache-dit for Flux dual-transformer architecture.
+
+    Flux uses two transformers (transformer and transformer_2) that need
+    to be enabled together using BlockAdapter.

    Args:
        pipeline: The Flux pipeline instance.
        cache_config: DiffusionCacheConfig instance with cache configuration.
+
    Returns:
-        A refresh function that can be called with a new ``num_inference_steps``
-        to update the cache context for the pipeline.
+        A refresh function that can be called to update cache context with new num_inference_steps.
    """
-    # 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
+    raise NotImplementedError("cache-dit is not implemented for Flux pipeline.")


 def enable_cache_for_sd3(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
@@ -253,7 +253,7 @@ class DiffusionEngine:

        try:
            # Give worker enough time — export + compression + table can be slow
-            results = self.collective_rpc(method="stop_profile", timeout=600)
+            results = self.collective_rpc(method="stop_profile", timeout=60000)
        except Exception:
            logger.error("Failed to stop profiling on workers", exc_info=True)
            return {"traces": [], "tables": []}
@@ -503,50 +503,10 @@ 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: 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]:
+def set_seq_parallel_pg(sp_ulysses_degree, sp_ring_degree, rank, world_size, use_ulysses_low=True):
    """
-    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_ulysses_degree x sp_ring_degree = seq_parallel_size
+    (ulysses_degree, dp_size)
    """
    sp_size = sp_ring_degree * sp_ulysses_degree
    dp_size = world_size // sp_size
@@ -556,96 +516,40 @@ def set_seq_parallel_pg(
    num_ulysses_pgs = sp_ring_degree  # world_size // sp_ulysses_degree
    num_ring_pgs = sp_ulysses_degree  # world_size // sp_ring_degree

-    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:
-        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,
-                        )
+    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
+                )
+                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
+            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
+    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
+            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

@@ -748,7 +652,6 @@ 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(
@@ -782,10 +685,9 @@ 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=sp_group_ranks,
+        group_ranks=rank_generator.get_ranks("sp"),
        local_rank=get_world_group().local_rank,
        backend=backend,
        parallel_mode="sequence",
@@ -63,11 +63,6 @@ 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]


@@ -72,15 +72,11 @@ def get_longcat_image_edit_pre_process_func(
            if "additional_information" not in prompt:
                prompt["additional_information"] = {}

-            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 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 isinstance(raw_image, str):
                image = PIL.Image.open(raw_image)
@@ -75,15 +75,11 @@ def get_qwen_image_edit_pre_process_func(
                prompt["additional_information"] = {}

            # Only handles single image
-            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 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 isinstance(raw_image, str):
                image = PIL.Image.open(raw_image)
@@ -632,13 +628,6 @@ 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 (
@@ -563,13 +563,6 @@ 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 (
@@ -79,15 +79,11 @@ def get_qwen_image_layered_pre_process_func(
            if "additional_information" not in prompt:
                prompt["additional_information"] = {}

-            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 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 isinstance(raw_image, str):
                image = PIL.Image.open(raw_image)
@@ -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

-    @lru_cache(maxsize=16)
+    @functools.cache
    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()

-    @lru_cache(maxsize=16)
+    @functools.cache
    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,6 +311,7 @@ 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
@@ -348,7 +349,14 @@ class QwenEmbedRope(nn.Module):
        max_vid_index = 0
        for idx, fhw in enumerate(video_fhw):
            frame, height, width = fhw
-            video_freq = self._compute_video_freqs(frame, height, width, idx)
+            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 = video_freq.to(device)
            vid_freqs.append(video_freq)

@@ -363,7 +371,7 @@ class QwenEmbedRope(nn.Module):

        return vid_freqs, txt_freqs

-    @lru_cache(maxsize=16)
+    @functools.cache
    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)
@@ -7,19 +7,14 @@ 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.linear import ColumnParallelLinear, QKVParallelLinear, RowParallelLinear
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import QKVParallelLinear, ReplicatedLinear
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

 from vllm_omni.diffusion.attention.layer import Attention
@@ -56,94 +51,6 @@ 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).
@@ -190,8 +97,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().float().repeat_interleave(2, dim=-1)
-        freqs_sin = freqs.sin().float().repeat_interleave(2, dim=-1)
+        freqs_cos = freqs.cos().repeat_interleave(2, dim=-1)
+        freqs_sin = freqs.sin().repeat_interleave(2, dim=-1)
        return freqs_cos.float(), freqs_sin.float()

    def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
@@ -296,47 +203,6 @@ 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.
@@ -363,30 +229,25 @@ 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 = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
-        self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
+        self.norm_q = RMSNorm(self.inner_dim, eps=eps)
+        self.norm_k = RMSNorm(self.inner_dim, eps=eps)

-        self.to_out = RowParallelLinear(
-            self.inner_dim,
-            dim,
-            bias=True,
-            input_is_parallel=True,
-            return_bias=False,
+        # Output projection
+        self.to_out = nn.ModuleList(
+            [
+                ReplicatedLinear(self.inner_dim, dim, bias=True),
+                nn.Dropout(dropout),
+            ]
        )
-        self.dropout = nn.Dropout(dropout)

        # Unified attention layer
        self.attn = Attention(
-            num_heads=self.num_heads,
+            num_heads=num_heads,
            head_size=head_dim,
-            num_kv_heads=self.num_kv_heads,
            softmax_scale=1.0 / (head_dim**0.5),
            causal=False,
        )
@@ -398,19 +259,16 @@ class WanSelfAttention(nn.Module):
    ) -> torch.Tensor:
        # Fused QKV projection
        qkv, _ = self.to_qkv(hidden_states)
-
-        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)
+        query, key, value = qkv.chunk(3, 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, 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))
+        query = query.unflatten(2, (self.num_heads, -1))
+        key = key.unflatten(2, (self.num_heads, -1))
+        value = value.unflatten(2, (self.num_heads, -1))

        # Apply rotary embeddings
        if rotary_emb is not None:
@@ -424,8 +282,8 @@ class WanSelfAttention(nn.Module):
        hidden_states = hidden_states.type_as(query)

        # Output projection
-        hidden_states = self.to_out(hidden_states)
-        hidden_states = self.dropout(hidden_states)
+        hidden_states, _ = self.to_out[0](hidden_states)
+        hidden_states = self.to_out[1](hidden_states)

        return hidden_states

@@ -454,77 +312,39 @@ 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 = ColumnParallelLinear(
-            dim,
-            self.inner_dim,
-            bias=True,
-            gather_output=False,
-            return_bias=False,
-        )
+        self.to_q = ReplicatedLinear(dim, self.inner_dim, bias=True)

        # Separate K and V projections for cross-attention
-        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
+        self.to_k = ReplicatedLinear(dim, self.kv_inner_dim, bias=True)
+        self.to_v = ReplicatedLinear(dim, self.kv_inner_dim, bias=True)

        # QK normalization
-        self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
-        self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
+        self.norm_q = RMSNorm(self.inner_dim, eps=eps)
+        self.norm_k = RMSNorm(self.kv_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 = 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)
+            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)
        else:
            self.add_k_proj = None
            self.add_v_proj = None
            self.norm_added_k = None

        # Output projection
-        self.to_out = RowParallelLinear(
-            self.inner_dim,
-            dim,
-            bias=True,
-            input_is_parallel=True,
-            return_bias=False,
+        self.to_out = nn.ModuleList(
+            [
+                ReplicatedLinear(self.inner_dim, dim, bias=True),
+                nn.Dropout(dropout),
+            ]
        )
-        self.dropout = nn.Dropout(dropout)

        # Unified attention layer
        self.attn = Attention(
-            num_heads=self.num_heads,
+            num_heads=num_heads,
            head_size=head_dim,
-            num_kv_heads=self.num_heads,
            softmax_scale=1.0 / (head_dim**0.5),
            causal=False,
        )
@@ -543,28 +363,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, self.head_dim))
-        key = key.unflatten(2, (self.num_heads, self.head_dim))
-        value = value.unflatten(2, (self.num_heads, self.head_dim))
+        query = query.unflatten(2, (self.num_heads, -1))
+        key = key.unflatten(2, (self.num_heads, -1))
+        value = value.unflatten(2, (self.num_heads, -1))

        # 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, self.head_dim))
-            value_img = value_img.unflatten(2, (self.num_heads, self.head_dim))
+            key_img = key_img.unflatten(2, (self.num_heads, -1))
+            value_img = value_img.unflatten(2, (self.num_heads, -1))

            hidden_states_img = self.attn(query, key_img, value_img)
            hidden_states_img = hidden_states_img.flatten(2, 3)
@@ -580,8 +400,8 @@ class WanCrossAttention(nn.Module):
            hidden_states = hidden_states + hidden_states_img

        # Output projection
-        hidden_states = self.to_out(hidden_states)
-        hidden_states = self.dropout(hidden_states)
+        hidden_states, _ = self.to_out[0](hidden_states)
+        hidden_states = self.to_out[1](hidden_states)

        return hidden_states

@@ -625,7 +445,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 = WanFeedForward(dim=dim, inner_dim=ffn_dim, dim_out=dim)
+        self.ffn = FeedForward(dim, inner_dim=ffn_dim, activation_fn="gelu-approximate")
        self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False)

        # Scale-shift table for modulation
@@ -720,7 +540,6 @@ 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
    #
@@ -731,22 +550,11 @@ 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),
    }
@@ -827,10 +635,7 @@ 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))
-
-        # SP helper modules
-        self.timestep_proj_prepare = TimestepProjPrepare()
-        self.output_scale_shift_prepare = OutputScaleShiftPrepare(inner_dim)
+        self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5)

    @property
    def dtype(self) -> torch.dtype:
@@ -870,10 +675,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
        )
-        # 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 ts_seq_len is not None:
+            timestep_proj = timestep_proj.unflatten(2, (6, -1))
+        else:
+            timestep_proj = timestep_proj.unflatten(1, (6, -1))

        if encoder_hidden_states_image is not None:
            encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1)
@@ -883,12 +688,15 @@ class WanTransformer3DModel(nn.Module):
            hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb)

        # Output norm, projection & unpatchify
-        shift, scale = self.output_scale_shift_prepare(temb)
+        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 = 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)
@@ -917,74 +725,32 @@ 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
+            # self-attention QKV fusion (attn1 only)
            (".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 original_name:
+                if weight_name not in name:
                    continue
-                lookup_name = original_name.replace(weight_name, param_name)
-                param = params_dict[lookup_name]
+                name = name.replace(weight_name, param_name)
+                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
-                # 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]
-
+                param = params_dict[name]
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
                weight_loader(param, loaded_weight)
-
-            loaded_params.add(original_name)
-            loaded_params.add(lookup_name)
+            loaded_params.add(name)

        return loaded_params
@@ -527,16 +527,6 @@ 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:
@@ -545,9 +535,8 @@ 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);
-            # use precomputed to avoid .item() sync per step
-            t_norm = t_norm_list[i]
+            # Normalized time for time-aware config (0 at start, 1 at end)
+            t_norm = timestep[0].item()

            # Handle cfg truncation
            current_guidance_scale = self.guidance_scale
@@ -593,17 +582,13 @@ class ZImagePipeline(nn.Module):

                    pred = pos + current_guidance_scale * (pos - neg)

-                    # Renormalization (torch.where avoids GPU->CPU sync from Python if/scalar comparison)
+                    # Renormalization
                    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)
-                        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
+                        if new_pos_norm > max_new_norm:
+                            pred = pred * (max_new_norm / new_pos_norm)

                    noise_pred.append(pred)

@@ -852,13 +852,7 @@ class ZImageTransformer2DModel(CachedTransformer):

        # Match t_embedder output dtype to x for layerwise casting compatibility
        adaln_input = t.type_as(x)
-        # 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[torch.cat(x_inner_pad_mask)] = self.x_pad_token
        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))
@@ -881,13 +875,7 @@ class ZImageTransformer2DModel(CachedTransformer):

        cap_feats = torch.cat(cap_feats, dim=0)
        cap_feats = self.cap_embedder(cap_feats)
-        # 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[torch.cat(cap_inner_pad_mask)] = self.cap_pad_token
        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))
@@ -13,13 +13,14 @@ logger = init_logger(__name__)

 class Scheduler:
    def initialize(self, od_config: OmniDiffusionConfig):
-        existing_mq = getattr(self, "mq", None)
-        if existing_mq is not None and not existing_mq.closed:
+        existing_context = getattr(self, "context", None)
+        if existing_context is not None and not existing_context.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
@@ -71,5 +72,8 @@ 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
@@ -179,27 +179,6 @@ 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,
@@ -227,7 +206,6 @@ 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)

@@ -244,13 +222,15 @@ 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):
@@ -300,8 +280,12 @@ def get_chunk_for_generation(
    if payload_data.get("finished"):
        connector.finished_requests.add(request_id)
        request.status = RequestStatus.FINISHED_STOPPED
-    request.prompt_token_ids = payload_data.get("code_predictor_codes", [])
-    request.num_computed_tokens = 0
+
+    # 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", [])


 def put_chunk(
@@ -324,6 +308,8 @@ 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
@@ -332,7 +318,6 @@ def put_chunk(
    if custom_process_input_func:
        try:
            payload_data = custom_process_input_func(
-                connector=connector,
                pooling_output=pooling_output,
                request=request,
            )
@@ -340,8 +325,42 @@ 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
        )
@@ -345,22 +345,14 @@ class AsyncOmni(OmniBase):
                ):
                    yield output

-            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")

-                logger.debug(f"[{self._name}] All requests completed")
-                # Summarize and print stats
+            # Summarize and print stats
+            try:
                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}] Request {request_id} Failed to finalized/build/log summary: {e}")
+                logger.exception(f"[{self._name}] Failed to build/log summary: {e}")
            finally:
                self.request_states.pop(request_id, None)
        except (asyncio.CancelledError, GeneratorExit):
@@ -390,11 +382,9 @@ 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,
+                    result, stage, stage_id, metrics, req_start_ts, wall_start_ts, final_stage_id_for_e2e
                )
                if submit_flag and stage_id == 0:
                    submit_flag = False
@@ -431,11 +421,9 @@ 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,
+                    result, stage, stage_id, metrics, req_start_ts, wall_start_ts, final_stage_id_for_e2e
                )
                if output_to_yield:
                    yield output_to_yield
@@ -444,7 +432,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:
+            if next_stage_id <= final_stage_id_for_e2e and finished:
                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]
@@ -459,7 +447,7 @@ class AsyncOmni(OmniBase):
                        connector=connector,
                        stage_id=stage_id,
                        next_stage_id=next_stage_id,
-                        req_id=request_id,
+                        req_id=req_id,
                        next_inputs=next_inputs,
                        sampling_params=sp_next,
                        original_prompt=prompt,
@@ -473,14 +461,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 {request_id} to stage-{next_stage_id} via connector. "
+                        f"[{self._name}] Failed to send request {req_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 {request_id} to stage-{next_stage_id}")
+                logger.debug(f"[{self._name}] Forwarded request {req_id} to stage-{next_stage_id}")
            else:
-                logger.debug(f"[{self._name}] Request {request_id} fully completed")
+                logger.debug(f"[{self._name}] Request {req_id} fully completed")

    def _process_single_result(
        self,
@@ -488,6 +476,9 @@ 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.
@@ -526,8 +517,23 @@ 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":
@@ -97,7 +97,7 @@ class OmniServeCommand(CLISubcommand):
        omni_config_group.add_argument(
            "--init-timeout",
            type=int,
-            default=600,
+            default=60000,
            help="The timeout for initializing the stages.",
        )
        omni_config_group.add_argument(
@@ -194,17 +194,17 @@ class OmniServeCommand(CLISubcommand):
        )

        # diffusion model offload parameters
-        omni_config_group.add_argument(
+        serve_parser.add_argument(
            "--enable-cpu-offload",
            action="store_true",
            help="Enable CPU offloading for diffusion models.",
        )
-        omni_config_group.add_argument(
+        serve_parser.add_argument(
            "--enable-layerwise-offload",
            action="store_true",
            help="Enable layerwise (blockwise) offloading on DiT modules.",
        )
-        omni_config_group.add_argument(
+        serve_parser.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=int,
+            type=float,
            help="The max size of generate image (height * width).",
        )
        return serve_parser
Author	SHA1	Message	Date
tzhouam	7f2035634b	Merge branch 'main' into dev/rebase-v0.15.0	2026-02-03 05:09:10 +00:00
tzhouam	4cbde7a460	Update Dockerfile.ci to improve installation commands and ensure Python environment consistency Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-02-03 05:06:51 +00:00
Taichang Zhou	e1026a6bc0	Update according to Cyrus comments Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>	2026-02-02 23:42:55 +08:00
Taichang Zhou	2eda72b5dc	Refactor comments and formatting in Qwen3 TTS configuration and modeling files for improved readability. Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>	2026-02-02 21:51:01 +08:00
Taichang Zhou	22ca74432b	Merge branch 'main' into dev/vllm-align	2026-02-02 21:35:45 +08:00
Taichang Zhou	354777d5e7	Merge branch 'main' of https://github.com/vllm-project/vllm-omni	2026-02-02 21:34:42 +08:00
汪志鹏	b6c846cc95	Add stage_init_timeout=300 (#1153 ) Signed-off-by: princepride <wangzhipeng628@gmail.com>	2026-02-02 17:58:48 +08:00
tzhouam	9194c1dd84	Update base image in Dockerfile.rocm from v0.14.0 to v0.15.0 Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-02-02 08:41:35 +00:00
tzhouam	ce9105267f	Enhance error handling in TTS generation and improve audio output compatibility by supporting both "audio" and "model_outputs" keys in response. Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-02-02 08:40:56 +00:00
tzhouam	39f7cade27	Fix audio sample rate access in end2end.py for correct output retrieval Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-02-02 08:38:24 +00:00
tzhouam	08fa85a1ad	Update audio tensor access in offline inference tests and increase image fetch timeout for CI Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-02-02 06:10:54 +00:00
tzhouam	6568a78a26	Remove unused tokenizer argument from OmniInputProcessor class docstring Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-02-02 03:12:43 +00:00
tzhouam	ff8b3dbbc8	update qwen 2.5 Omni Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-02-02 03:01:11 +00:00
tzhouam	1297a46040	Merge remote-tracking branch 'origin/main' into dev/vllm-align Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-02-02 01:54:12 +00:00
Taichang Zhou	e033498600	Merge branch 'main' of https://github.com/vllm-project/vllm-omni	2026-01-31 10:59:36 +08:00
tzhouam	7aea8e8d91	Add OmniMRotaryEmbedding for multimodal position computation - Introduced OmniMRotaryEmbedding to extend MRotaryEmbedding with methods for computing input positions across various multimodal scenarios, including image, video, and audio inputs. - Updated patch.py to replace the original MRotaryEmbedding with OmniMRotaryEmbedding in the module patching logic. - Created new files for the OmniMRotaryEmbedding implementation and its associated methods. Signed-off-by: [Your Name] <[Your Email]> Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-01-30 10:03:02 +00:00
tzhouam	3208afd82b	Refactor OmniRequestOutput to encapsulate multimodal output - Changed the multimodal_output field to a private attribute _multimodal_output for better encapsulation. - Added a property method to access multimodal_output, allowing for a clearer interface and improved handling of request outputs. Signed-off-by: [Your Name] <[Your Email]> Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-01-30 10:02:22 +00:00
tzhouam	2b009b3cef	Fix audio tensor extraction in end2end.py files across multiple examples - Updated the audio tensor extraction logic to access the correct output structure in Qwen2.5 Omni, Qwen3 Omni, and Qwen3 TTS implementations. - Ensured consistency in handling audio outputs for improved functionality. Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-01-30 10:00:14 +00:00
tzhouam	4faf190e28	Merge remote-tracking branch 'origin/main' into dev/vllm-align	2026-01-30 03:49:04 +00:00
tzhouam	ca57f510c7	Refactor OmniInputPreprocessor to remove unused embed processing logic - Eliminated the _process_embeds method and related imports from preprocess.py, streamlining the input preprocessing workflow. - Updated the main processing logic to focus solely on token inputs, enhancing clarity and maintainability. Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-01-30 03:02:38 +00:00
tzhouam	4b895459e4	Merge branch 'main' into dev/vllm-align Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-01-30 02:53:03 +00:00
tzhouam	62dbf07fb6	Refactor OpenAI serving modules for improved functionality and organization - Reintroduced the load_chat_template import in api_server.py for better template handling. - Updated the logic in serving_chat.py to streamline the resolution of chat template content formats, enhancing clarity and functionality. - Adjusted the application of chat templates to ensure proper handling of tools and chat templates. Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-01-30 02:44:13 +00:00
tzhouam	104e105c93	Refactor multimodal processing imports and update configuration handling in Qwen3OmniMoeThinker - Changed the import path for multimodal processing to improve module organization. - Removed the unused multimodal configuration parameter from the Qwen3OmniMoeThinker initialization, streamlining the class definition. Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-01-30 02:43:41 +00:00
tzhouam	ef21879c3a	Refactor GPU model runner to streamline multimodal configuration handling - Removed the unused import of supports_mm_encoder_only from interfaces. - Updated the dummy run check to utilize the multimodal configuration directly, enhancing clarity and future extensibility. Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-01-30 02:43:26 +00:00
tzhouam	45880a30f0	Refactor chat_utils and OpenAI modules for improved functionality and organization - Updated import statements in chat_utils.py to replace the internal content format type with the public one for better clarity. - Refactored api_server.py to streamline the chat template loading process by replacing the previous method with a direct function call. - Enhanced serving_chat.py by reintroducing safe_apply_chat_template functions for better template handling. Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-01-29 09:00:43 +00:00
tzhouam	2ac33489b5	Update Dockerfile.ci to install vllm package using the correct URL format - Modified the pip install command for vllm to use the correct URL format for direct installation from the specified commit hash, ensuring a more reliable installation process. - Removed the fallback option for pre-release installation to streamline the installation command. Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-01-29 08:25:51 +00:00
tzhouam	63fd163efb	Refactor import statement in OpenAI API server for improved module organization - Updated the import path for AnthropicServingMessages in api_server.py to enhance clarity and maintain consistency with the new module structure. Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-01-29 08:23:17 +00:00
tzhouam	fea328364a	Update Dockerfile.ci to install vllm package directly from wheel file - Changed the pip install command for vllm to download the wheel file directly instead of using an extra index URL, ensuring a more reliable installation process. - Added a fallback option to install the pre-release version of vllm if the direct installation fails. Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-01-29 08:05:57 +00:00
tzhouam	478c036a01	[Refactor] Update OpenAI API server and related modules for improved structure and functionality - Refactored import statements in api_server.py to enhance module organization, including changes to the routing and server setup. - Updated the method for removing routes to ensure compatibility with the new app structure. - Adjusted the handling of chat completion requests to improve error responses and streamline processing. - Modified serving modules to align with the new import paths and enhance code clarity. Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-01-29 15:28:51 +08:00
tzhouam	8407818368	Update simple_test.sh to force reinstall vllm package with extra index URL - Modified the pip install command for vllm to include --upgrade, --force-reinstall flags, and an extra index URL, ensuring the latest version is installed from the specified source during testing. Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-01-28 17:34:22 +08:00
tzhouam	6e7cb01afe	Update Dockerfile.ci to force reinstall vllm package during CI build - Modified the pip install command for vllm to include --upgrade and --force-reinstall flags, ensuring the latest version is installed during the CI process. Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-01-28 17:20:29 +08:00
tzhouam	acb28f3d1d	Remove commented-out import and related code for MRotaryEmbedding in patch.py to clean up the codebase. Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-01-28 16:52:10 +08:00
tzhouam	bf51318176	debug Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-01-28 16:40:06 +08:00
tzhouam	bd8c0443cf	build image from commit id Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-01-28 16:21:29 +08:00
tzhouam	554d14fd28	debug docker file Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-01-28 15:57:13 +08:00
tzhouam	49b9d0d6b1	debug docker build Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-01-28 15:47:30 +08:00
tzhouam	2c2d06d920	[Enhancement] Update GPUARWorker and GPUGenerationWorker to enforce v1 model runner - Added a warning in GPUARWorker and GPUGenerationWorker to indicate that the v2 model runner does not include omni hooks, forcing the use of the v1 model runner. - Updated import statements in gpu_ar_model_runner.py and gpu_generation_model_runner.py to include necessary functions from parallel_state. - Removed unused imports from both model runner files to streamline the code. Signed-off-by: [Your Name] <your.email@example.com> Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-01-28 15:43:27 +08:00
tzhouam	31d127edcd	[Refactor] Remove unused MRotaryEmbedding and vision position ID functions - Deleted the MRotaryEmbedding class and its associated methods from mrope.py to streamline the codebase. - Removed the get_llm_pos_ids_for_vision function from vision.py as it is no longer needed. - Updated import statements in qwen2_5_omni.py to reflect the removal of the vision module. Signed-off-by: [Your Name] <your.email@example.com> Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-01-28 15:42:35 +08:00
tzhouam	52d0122fb5	[Enhancement] Refactor OmniInputProcessor and related classes for improved prompt handling - Updated OmniInputProcessor to streamline input processing, including enhanced handling of prompt embeddings and additional information. - Introduced OmniTextPrompt and OmniEmbedsInputs to support new input structures with optional metadata. - Modified input preprocessing methods to accommodate new prompt types and ensure compatibility with multimodal inputs. - Improved error handling and type checking for data parallel configurations. Signed-off-by: [Your Name] <your.email@example.com> Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-01-28 15:42:05 +08:00
tzhouam	b26a9735e6	[Enhancement] Update type hint for stage_connector_spec in OmniEngineArgs and AsyncOmniEngineArgs - Changed type hint for stage_connector_spec from 'any' to 'Any' for improved type checking and consistency. - Ensured compatibility with typing standards in the arg_utils module. Signed-off-by: [Your Name] <your.email@example.com> Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-01-28 15:41:20 +08:00
tzhouam	d7cfe195ea	[Enhancement] Refactor OmniARScheduler and OmniGenerationScheduler to include OmniSchedulerOutput - Updated OmniARScheduler to wrap output in OmniSchedulerOutput, including transfer metadata for finished requests needing KV cache transfer. - Enhanced OmniGenerationScheduler to handle routed experts and finish reasons more effectively, ensuring proper output structure. - Introduced performance statistics handling in both schedulers for improved monitoring. Signed-off-by: [Your Name] <your.email@example.com> Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-01-28 15:40:59 +08:00
tzhouam	e1c6639f18	[Enhancement] Add mm_encoder_only parameter and deprecation warning for convert option - Introduced mm_encoder_only parameter to OmniModelConfig. - Added a deprecation warning for the `--convert mm_encoder_only` option, advising users to switch to `--mm-encoder-only`. Signed-off-by: [Your Name] <your.email@example.com> Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-01-28 15:40:29 +08:00
tzhouam	ce7430b562	[Enhancement] Update OmniRequest to support torch.Tensor for prompt_embeds and add decoding method - Modified the constructor to accept torch.Tensor as an option for prompt_embeds. - Introduced a new static method _maybe_decode_prompt_embeds to handle the conversion from PromptEmbedsPayload to torch.Tensor. - Preserved the original prompt_embeds payload for reference. Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-01-28 15:40:17 +08:00
tzhouam	ed8b03224b	update the docker setup to vllm nightly Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-01-28 11:15:54 +08:00
Taichang Zhou	af8501dcc7	debug Qwen TTS Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>	2026-01-22 22:46:06 +08:00
				`@@ -1 +0,0 @@`
				`# noqa: N999 # This is not a python library intended to be imported`
				`@@ -1 +0,0 @@`
				{"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
				`@@ -1 +0,0 @@`
				{"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}
				`@@ -1 +0,0 @@`
				{"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}