typo

Signed-off-by: Roger Wang <hey@rogerw.io>

.github/ISSUE_TEMPLATE/config.yml

@@ -1,5 +1,5 @@
 blank_issues_enabled: false
 contact_links:
   - name: Questions
-    url: https://github.com/vllm-project/vllm-omni/discussions
-    about: Ask questions and discuss with other vLLM-Omni community members
+    url: https://slack.vllm.ai
+    about: Ask questions and discuss with other vLLM-Omni community members in our #sig-omni slack channel.

.gitignore

@@ -232,4 +232,3 @@ configs/development.yaml
 # Docker
 .dockerignore
 Dockerfile.dev
-discussion

CONTRIBUTING.md

@@ -1,3 +1,3 @@
 # Contributing to vLLM-Omni
 
-You may find information about contributing to vLLM-Omni on [Contributing](docs/contributing/)
+You may find information about contributing to vLLM-Omni on [Contributing](https://vllm-omni.readthedocs.io/en/latest/contributing/)

README.md

@@ -1,7 +1,7 @@
 <p align="center">
   <picture>
-    <source media="(prefers-color-scheme: dark)" src="./docs/source/logos/vllm-omni-logo-text-dark.png">
-    <img alt="vllm-omni" src="./docs/source/logos/vllm-omni-logo-text-dark.png" width=55%>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png">
+    <img alt="vllm-omni" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png" width=55%>
   </picture>
 </p>
 <h3 align="center">
@@ -9,7 +9,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 </h3>
 
 <p align="center">
-| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://tinyurl.com/vllm-omni-meeting"><b>Weekly Meeting</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="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> |
 </p>
 
 ---
@@ -30,7 +30,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 
 <p align="center">
   <picture>
-    <img alt="vllm-omni" src="./docs/source/architecture/omni-modality-model-architecture.png" width=55%>
+    <img alt="vllm-omni" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/omni-modality-model-architecture.png" width=55%>
   </picture>
 </p>
 
@@ -66,10 +66,8 @@ 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.
 
-## Weekly Meeting
-
-- vLLM Omni Weekly Meeting: https://tinyurl.com/vllm-omni-meeting
-- Wednesday, 11:30 - 12:30 (UTC+8, [Convert to your timezone](https://dateful.com/convert/gmt8?t=15))
+## 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).
 
 ## License
 

docs/.nav.yml

@@ -7,11 +7,11 @@ nav:
   - Examples:
     - examples/README.md
     - Offline Inference:
-      - Offline Example of vLLM-Omni for Qwen2.5-omni: user_guide/examples/offline_inference/qwen2_5_omni.md
-      - Offline Example of vLLM-Omni for Qwen3-omni: user_guide/examples/offline_inference/qwen3_omni.md
+      - Offline Example of vLLM-Omni for Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
+      - Offline Example of vLLM-Omni for Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
       - Qwen-Image Offline Inference: user_guide/examples/offline_inference/qwen_image.md
     - Online Serving:
-      - Online serving Example of vLLM-Omni for Qwen2.5-omni: user_guide/examples/online_serving/qwen2_5_omni.md
+      - Online serving Example of vLLM-Omni for Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
       - Online serving Example of vLLM-Omni for Qwen3-omni: user_guide/examples/online_serving/qwen3_omni.md
   - General:
     - usage/*
@@ -33,7 +33,6 @@ nav:
     - design/architecture_overview.md
     - design/vllm_omni_design.md
     - design/mrs_design.md
-    - design/api_design_doc.md
   - Docs Guide: contributing/DOCS_GUIDE.md
 - API Reference:
   - api/README.md

docs/README.md

@@ -8,8 +8,8 @@ hide:
 
 <p align="center">
   <picture>
-    <source media="(prefers-color-scheme: dark)" src="./source/logos/vllm-omni-logo-text-dark.png">
-    <img alt="vllm-omni" src="./source/logos/vllm-omni-logo-text-dark.png" width=55%>
+    <source media="(prefers-color-scheme: dark)" src="./source/logos/vllm-omni-logo.png">
+    <img alt="vllm-omni" src="./source/logos/vllm-omni-logo.png" width=55%>
   </picture>
 </p>
 <h3 align="center">
@@ -34,7 +34,8 @@ Easy, fast, and cheap omni-modality model serving for everyone
 
 <p align="center">
   <picture>
-    <img alt="vllm-omni" src="./source/architecture/omni-modality model architecture.png" width=55%>
+    <source media="(prefers-color-scheme: dark)" src="./source/architecture/omni-modality-model-architecture.png">
+    <img alt="vllm-omni-arch" src="./source/architecture/omni-modality-model-architecture.png" width=55%>
   </picture>
 </p>
 
@@ -58,4 +59,4 @@ vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, in
 - Multi-modality generation models (e.g. Qwen-Image)
 
 For more information, checkout the following:
-- [Placeholder for vllm-omni blogpost]()
+- [vllm-omni announcement blogpost](https://blog.vllm.ai/2025/11/30/vllm-omni.html)

docs/api/README.md

@@ -102,6 +102,10 @@ Configuration classes.
 
 Worker classes and model runners for distributed inference.
 
+- [vllm_omni.diffusion.worker.gpu_worker.GPUWorker][]
+- [vllm_omni.diffusion.worker.gpu_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.GPUARModelRunner][]
 - [vllm_omni.worker.gpu_ar_worker.GPUARWorker][]
 - [vllm_omni.worker.gpu_diffusion_model_runner.GPUDiffusionModelRunner][]

docs/community/contact_us.md

@@ -2,4 +2,4 @@
 
 - For technical questions and feature requests, please use GitHub [Issues](https://github.com/vllm-project/vllm-omni/issues)
 - For coordinating contributions and development and discussing with other users and developers, please join `sig-omni` channel in our [Slack](https://slack.vllm.ai/) or use the [vLLM Forum](https://discuss.vllm.ai/)
-- For security disclosures, please use GitHub's [Security Advisories](https://github.com/vllm-project/vllm/security/advisories) feature
+- For security disclosures, please use GitHub's [Security Advisories](https://github.com/vllm-project/vllm-omni/security/advisories) feature

docs/configuration/stage_configs.md

@@ -2,6 +2,9 @@
 
 In vLLM-Omni, the target model is separated into multiple stages, which are processed by different LLMEngines, DiffusionEngines or other types of engines. Depending on different types of stages, such as Autoregressive (AR) stage or Diffusion transformer (DiT) stage, each can choose corresponding schedulers, model workers to load with the Engines in a plug-in fashion.
 
+!!! note
+    Default stage config YAMLs (for example, `vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml` and `vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml`) are bundled and loaded automatically when `stage_configs_path` is not provided. They have been verified to work on 1xH100 for Qwen2.5-Omni and 2xH100 for Qwen3-Omni.
+
 Therefore, as a core part of vLLM-Omni, the stage configs for a model have several main functions:
 
 - Claim partition of stages and their corresponding class implementation in `model_executor/models`.

docs/contributing/README.md

@@ -2,6 +2,9 @@
 
 Thank you for your interest in contributing to vLLM-Omni! This document provides guidelines and instructions for contributing.
 
+!!! note
+    We host weekly developer-facing online meetings to discuss milestones and updates **every Tuesday at 19:30 PDT**. Meeting link as well as the past meeting notes can be found [here](https://tinyurl.com/vllm-omni-meeting).
+
 ## Getting Started
 
 vLLM-Omni uses `uv` as the environment manager, to create and manage Python environments. Please follow the documentation to install `uv`. After installing `uv`, you can create a new Python environment using the following commands:
@@ -20,7 +23,7 @@ Install vLLM-Omni from source with development dependencies:
 
 ```bash
 git clone https://github.com/vllm-project/vllm-omni.git
-cd vllm_omni
+cd vllm-omni
 uv pip install -e ".[dev]"
 ```
 

docs/design/index.md

@@ -7,7 +7,3 @@ This section contains design documents and architecture specifications for vLLM-
 - [Architecture Overview](architecture_overview.md)
 - [vLLM-Omni Software Design Document](vllm_omni_design.md)
 - [Multi-Request Streaming](mrs_design.md)
-
-## API Design Documentation
-
-- [vLLM-Omni API Design Documentation](api_design_doc.md)

docs/getting_started/installation/gpu.md

@@ -37,4 +37,9 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
 ## Set up using Docker
 
 ### Pre-built images
-To be released... Stay tuned!
+vLLM-Omni offers an official docker image for deployment. These images are built on top of vLLM docker images
+and available on Docker Hub as [vllm/vllm-omni](https://hub.docker.com/r/vllm/vllm-omni/tags). The version of vLLM-Omni
+indicates which release of vLLM it is based on.
+
+!!! tip
+    You can use this docker image to serve models the same way you would do 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.

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

@@ -44,7 +44,7 @@ uv pip install vllm==0.11.0 --torch-backend=auto
 Install additional requirements for vLLM-Omni
 ```bash
 git clone https://github.com/vllm-project/vllm-omni.git
-cd vllm_omni
+cd vllm-omni
 uv pip install -e .
 ```
 

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

@@ -19,6 +19,7 @@ export DEVICE1=/dev/davinci1
 export IMAGE=quay.io/ascend/vllm-ascend:v0.11.0rc2
 docker run --rm \
     --name vllm-omni-npu \
+    --shm-size=1g \
     --device $DEVICE0 \
     --device $DEVICE1 \
     --device /dev/davinci_manager \

docs/mkdocs/hooks/generate_examples.py

@@ -173,7 +173,7 @@ class Example:
             gh_file = gh_file.relative_to(ROOT_DIR)
 
             # Make GitHub URL
-            url = "https://github.com/vllm-project/vllm/"
+            url = "https://github.com/vllm-project/vllm-omni/"
             url += "tree/main" if self.path.is_dir() else "blob/main"
             gh_url = f"{url}/{gh_file}"
 
@@ -183,7 +183,7 @@ class Example:
 
     def generate(self) -> str:
         content = f"# {self.title}\n\n"
-        url = "https://github.com/vllm-project/vllm/"
+        url = "https://github.com/vllm-project/vllm-omni/"
         url += "tree/main" if self.path.is_dir() else "blob/main"
         content += f"Source <{url}/{self.path.relative_to(ROOT_DIR)}>.\n\n"
 
@@ -192,7 +192,8 @@ class Example:
         code_fence = "``````"
 
         if self.is_code:
-            content += f'{code_fence}{self.main_file.suffix[1:]}\n--8<-- "{self.main_file}"\n{code_fence}\n'
+            main_file_rel = self.main_file.relative_to(ROOT_DIR)
+            content += f'{code_fence}{self.main_file.suffix[1:]}\n--8<-- "{main_file_rel}"\n{code_fence}\n'
         else:
             with open(self.main_file) as f:
                 # Skip the title from md snippets as it's been included above
@@ -208,7 +209,7 @@ class Example:
             content += f'??? abstract "{file.relative_to(self.path)}"\n'
             if file.suffix != ".md":
                 content += f"    {code_fence}{file.suffix[1:]}\n"
-            content += f'    --8<-- "{file}"\n'
+            content += f'    --8<-- "{file.relative_to(ROOT_DIR)}"\n'
             if file.suffix != ".md":
                 content += f"    {code_fence}\n"
 

docs/models/supported_models.md

@@ -20,6 +20,7 @@ th {
 | `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` |
+|`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |
 
 
 ## List of Supported Models for Ascend NPU

docs/user_guide/examples/offline_inference/qwen2_5_omni.md

@@ -1,13 +1,13 @@
-# Offline Example of vLLM-Omni for Qwen2.5-omni
+# Offline Example of vLLM-Omni for Qwen2.5-Omni
 
-Source <https://github.com/vllm-project/vllm/tree/main/examples/offline_inference/qwen2_5_omni>.
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen2_5_omni>.
 
 
 ## 🛠️ Installation
 
-Please refer to [README.md](https://github.com/vllm-project/vllm/tree/main/README.md)
+Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
 
-## Run examples (Qwen2.5-omni)
+## Run examples (Qwen2.5-Omni)
 ### Multiple Prompts
 Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). To get the prompt, you can:
 ```bash
@@ -48,17 +48,17 @@ sudo apt install ffmpeg
 
 ??? abstract "end2end.py"
     ``````py
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen2_5_omni/end2end.py"
+    --8<-- "examples/offline_inference/qwen2_5_omni/end2end.py"
     ``````
 ??? abstract "extract_prompts.py"
     ``````py
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen2_5_omni/extract_prompts.py"
+    --8<-- "examples/offline_inference/qwen2_5_omni/extract_prompts.py"
     ``````
 ??? abstract "run_multiple_prompts.sh"
     ``````sh
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen2_5_omni/run_multiple_prompts.sh"
+    --8<-- "examples/offline_inference/qwen2_5_omni/run_multiple_prompts.sh"
     ``````
 ??? abstract "run_single_prompt.sh"
     ``````sh
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen2_5_omni/run_single_prompt.sh"
+    --8<-- "examples/offline_inference/qwen2_5_omni/run_single_prompt.sh"
     ``````

docs/user_guide/examples/offline_inference/qwen3_omni.md

@@ -1,15 +1,15 @@
-# Offline Example of vLLM-Omni for Qwen3-omni
+# Offline Example of vLLM-Omni for Qwen3-Omni
 
-Source <https://github.com/vllm-project/vllm/tree/main/examples/offline_inference/qwen3_omni>.
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen3_omni>.
 
 
 ## 🛠️ Installation
 
-Please refer to [README.md](https://github.com/vllm-project/vllm/tree/main/README.md)
+Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
 
-## Run examples (Qwen3-omni)
+## Run examples (Qwen3-Omni)
 ### Multiple Prompts
-Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). For processing dataset please refer to [Qwen2.5-omni README.md](https://github.com/vllm-project/vllm/tree/main/examples/offline_inference/qwen2_5_omni/README.md)
+Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). For processing dataset please refer to [Qwen2.5-Omni README.md](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen2_5_omni/README.md)
 Get into the example folder
 ```bash
 cd examples/offline_inference/qwen3_omni
@@ -44,21 +44,17 @@ sudo apt install ffmpeg
 
 ??? abstract "end2end.py"
     ``````py
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen3_omni/end2end.py"
-    ``````
-??? abstract "qwen3_omni_moe_tp.yaml"
-    ``````yaml
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen3_omni/qwen3_omni_moe_tp.yaml"
+    --8<-- "examples/offline_inference/qwen3_omni/end2end.py"
     ``````
 ??? abstract "run_multiple_prompts.sh"
     ``````sh
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen3_omni/run_multiple_prompts.sh"
+    --8<-- "examples/offline_inference/qwen3_omni/run_multiple_prompts.sh"
     ``````
 ??? abstract "run_single_prompt.sh"
     ``````sh
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen3_omni/run_single_prompt.sh"
+    --8<-- "examples/offline_inference/qwen3_omni/run_single_prompt.sh"
     ``````
 ??? abstract "run_single_prompt_tp.sh"
     ``````sh
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen3_omni/run_single_prompt_tp.sh"
+    --8<-- "examples/offline_inference/qwen3_omni/run_single_prompt_tp.sh"
     ``````

docs/user_guide/examples/offline_inference/qwen_image.md

@@ -1,6 +1,6 @@
 # Qwen-Image Offline Inference
 
-Source <https://github.com/vllm-project/vllm/tree/main/examples/offline_inference/qwen_image>.
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen_image>.
 
 
 This folder provides two simple entrypoints for experimenting with `Qwen/Qwen-Image` using vLLM-Omni:
@@ -49,9 +49,9 @@ Then open `http://localhost:7862/` on your local browser to interact with the we
 
 ??? abstract "gradio_demo.py"
     ``````py
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen_image/gradio_demo.py"
+    --8<-- "examples/offline_inference/qwen_image/gradio_demo.py"
     ``````
 ??? abstract "text_to_image.py"
     ``````py
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen_image/text_to_image.py"
+    --8<-- "examples/offline_inference/qwen_image/text_to_image.py"
     ``````

docs/user_guide/examples/online_serving/qwen2_5_omni.md

@@ -1,15 +1,16 @@
-# Online serving Example of vLLM-Omni for Qwen2.5-omni
+# Online serving Example of vLLM-Omni for Qwen2.5-Omni
 
-Source <https://github.com/vllm-project/vllm/tree/main/examples/online_serving/qwen2_5_omni>.
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/qwen2_5_omni>.
 
 
 ## 🛠️ Installation
 
-Please refer to [README.md](https://github.com/vllm-project/vllm/tree/main/examples/README.md)
+Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
 
-## Run examples (Qwen2.5-omni)
+## Run examples (Qwen2.5-Omni)
+
+### Launch the Server
 
-Launch the server
 ```bash
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
 ```
@@ -19,29 +20,59 @@ If you have custom stage configs file, launch the server with command below
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
 ```
 
+### Send Multi-modal Request
+
 Get into the example folder
 ```bash
-cd examples/online_serving
+cd examples/online_serving/qwen2_5_omni
 ```
 
-Send request via python
+####  Send request via python
+
 ```bash
 python openai_chat_completion_client_for_multimodal_generation.py --query-type mixed_modalities
 ```
 
-Send request via curl
+The Python client supports the following command-line arguments:
+
+- `--query-type` (or `-q`): Query type (default: `mixed_modalities`)
+  - Options: `mixed_modalities`, `use_audio_in_video`, `multi_audios`, `text`
+- `--video-path` (or `-v`): Path to local video file or URL
+  - If not provided and query-type uses video, uses default video URL
+  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
+  - Example: `--video-path /path/to/video.mp4` or `--video-path https://example.com/video.mp4`
+- `--image-path` (or `-i`): Path to local image file or URL
+  - If not provided and query-type uses image, uses default image URL
+  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
+  - Supports common image formats: JPEG, PNG, GIF, WebP
+  - Example: `--image-path /path/to/image.jpg` or `--image-path https://example.com/image.png`
+- `--audio-path` (or `-a`): Path to local audio file or URL
+  - If not provided and query-type uses audio, uses default audio URL
+  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
+  - Supports common audio formats: MP3, WAV, OGG, FLAC, M4A
+  - Example: `--audio-path /path/to/audio.wav` or `--audio-path https://example.com/audio.mp3`
+- `--prompt` (or `-p`): Custom text prompt/question
+  - If not provided, uses default prompt for the selected query type
+  - Example: `--prompt "What are the main activities shown in this video?"`
+
+
+For example, to use mixed modalities with all local files:
+
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py \
+    --query-type mixed_modalities \
+    --video-path /path/to/your/video.mp4 \
+    --image-path /path/to/your/image.jpg \
+    --audio-path /path/to/your/audio.wav \
+    --prompt "Analyze all the media content and provide a comprehensive summary."
+```
+
+####  Send request via curl
+
 ```bash
 bash run_curl_multimodal_generation.sh mixed_modalities
 ```
 
-### FAQ
-
-If you encounter error about backend of librosa, try to install ffmpeg with command below.
-```
-sudo apt update
-sudo apt install ffmpeg
-```
-
 ## Run Local Web UI Demo
 
 This Web UI demo allows users to interact with the model through a web browser.
@@ -56,36 +87,33 @@ python gradio_demo.py  --model Qwen/Qwen2.5-Omni-7B --port 7861
 
 Then open `http://localhost:7861/` on your local browser to interact with the web UI.
 
+The gradio script supports the following arguments:
 
-### Options
+- `--model`: Model name
+- `--ip`: Host/IP for Gradio server (default: 127.0.0.1)
+- `--port`: Port for Gradio server (default: 7861)
+- `--stage-configs-path`: Path to custom stage configs YAML file (optional)
+- `--share`: Share the Gradio demo publicly (creates a public link)
 
-You can customize its basic launch parameters:
+### FAQ
 
-```bash
-python gradio_demo.py \
-    --model Qwen/Qwen2.5-Omni-7B \
-    --ip 127.0.0.1 \
-    --port 7861 \
-    --stage-configs-path /path/to/stage_configs.yaml
+If you encounter error about backend of librosa, try to install ffmpeg with command below.
+```
+sudo apt update
+sudo apt install ffmpeg
 ```
-
-- `--model`: Local model checkpoint to load (default `Qwen/Qwen2.5-Omni-7B`).
-- `--ip`: Host/IP for the Gradio server (default `127.0.0.1`).
-- `--port`: Port for the Gradio server (default `7861`).
-- `--stage-configs-path`: Optional path to custom stage configs YAML.
-- `--share`: Set to expose a temporary public link via Gradio.
 
 ## Example materials
 
 ??? abstract "gradio_demo.py"
     ``````py
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/online_serving/qwen2_5_omni/gradio_demo.py"
+    --8<-- "examples/online_serving/qwen2_5_omni/gradio_demo.py"
     ``````
 ??? abstract "openai_chat_completion_client_for_multimodal_generation.py"
     ``````py
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/online_serving/qwen2_5_omni/openai_chat_completion_client_for_multimodal_generation.py"
+    --8<-- "examples/online_serving/qwen2_5_omni/openai_chat_completion_client_for_multimodal_generation.py"
     ``````
 ??? abstract "run_curl_multimodal_generation.sh"
     ``````sh
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh"
+    --8<-- "examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh"
     ``````

docs/user_guide/examples/online_serving/qwen3_omni.md

@@ -1,15 +1,16 @@
 # Online serving Example of vLLM-Omni for Qwen3-omni
 
-Source <https://github.com/vllm-project/vllm/tree/main/examples/online_serving/qwen3_omni>.
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/qwen3_omni>.
 
 
 ## 🛠️ Installation
 
-Please refer to [README.md](https://github.com/vllm-project/vllm/tree/main/examples/README.md)
+Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
 
 ## Run examples (Qwen3-Omni)
 
-Launch the server
+### Launch the Server
+
 ```bash
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
@@ -19,21 +20,59 @@ If you have custom stage configs file, launch the server with command below
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
 ```
 
+### Send Multi-modal Request
+
 Get into the example folder
 ```bash
-cd examples/online_serving
+cd examples/online_serving/qwen3_omni
 ```
 
-Send request via python
+####  Send request via python
+
 ```bash
 python openai_chat_completion_client_for_multimodal_generation.py --query-type use_image
 ```
 
-Send request via curl
+The Python client supports the following command-line arguments:
+
+- `--query-type` (or `-q`): Query type (default: `use_video`)
+  - Options: `text`, `use_audio`, `use_image`, `use_video`
+- `--model` (or `-m`): Model name/path (default: `Qwen/Qwen3-Omni-30B-A3B-Instruct`)
+- `--video-path` (or `-v`): Path to local video file or URL
+  - If not provided and query-type is `use_video`, uses default video URL
+  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
+  - Example: `--video-path /path/to/video.mp4` or `--video-path https://example.com/video.mp4`
+- `--image-path` (or `-i`): Path to local image file or URL
+  - If not provided and query-type is `use_image`, uses default image URL
+  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
+  - Supports common image formats: JPEG, PNG, GIF, WebP
+  - Example: `--image-path /path/to/image.jpg` or `--image-path https://example.com/image.png`
+- `--audio-path` (or `-a`): Path to local audio file or URL
+  - If not provided and query-type is `use_audio`, uses default audio URL
+  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
+  - Supports common audio formats: MP3, WAV, OGG, FLAC, M4A
+  - Example: `--audio-path /path/to/audio.wav` or `--audio-path https://example.com/audio.mp3`
+- `--prompt` (or `-p`): Custom text prompt/question
+  - If not provided, uses default prompt for the selected query type
+  - Example: `--prompt "What are the main activities shown in this video?"`
+
+
+For example, to use a local video file with custom prompt:
+
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py \
+    --query-type use_video \
+    --video-path /path/to/your/video.mp4 \
+    --prompt "What are the main activities shown in this video?"
+```
+
+####  Send request via curl
+
 ```bash
 bash run_curl_multimodal_generation.sh use_image
 ```
 
+
 ### FAQ
 
 If you encounter error about backend of librosa, try to install ffmpeg with command below.
@@ -57,23 +96,9 @@ python gradio_demo.py  --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 7861
 Then open `http://localhost:7861/` on your local browser to interact with the web UI.
 
 
-### Options
-
-The gradio demo also supports running with an existing API server and can be customized with the following arguments.
-
-
-```bash
-python gradio_demo.py \
-    --model Qwen/Qwen3-Omni-30B-A3B-Instruct \
-    --use-api-server \
-    --api-base http://localhost:8091/v1 \
-    --ip 127.0.0.1 \
-    --port 7861
-```
+The gradio script supports the following arguments:
 
 - `--model`: Model name
-- `--use-api-server`: If set, connect to an existing vLLM HTTP API server instead of running AsyncOmni locally.
-- `--api-base`: Base URL for vllm serve (only used when `use-api-server` is set, default: http://localhost:8091/v1)
 - `--ip`: Host/IP for Gradio server (default: 127.0.0.1)
 - `--port`: Port for Gradio server (default: 7861)
 - `--stage-configs-path`: Path to custom stage configs YAML file (optional)
@@ -83,13 +108,13 @@ python gradio_demo.py \
 
 ??? abstract "gradio_demo.py"
     ``````py
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/online_serving/qwen3_omni/gradio_demo.py"
+    --8<-- "examples/online_serving/qwen3_omni/gradio_demo.py"
     ``````
 ??? abstract "openai_chat_completion_client_for_multimodal_generation.py"
     ``````py
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py"
+    --8<-- "examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py"
     ``````
 ??? abstract "run_curl_multimodal_generation.sh"
     ``````sh
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh"
+    --8<-- "examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh"
     ``````

examples/offline_inference/qwen2_5_omni/README.md

@@ -5,6 +5,7 @@
 Please refer to [README.md](../../../README.md)
 
 ## Run examples (Qwen2.5-Omni)
+
 ### Multiple Prompts
 Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). To get the prompt, you can:
 ```bash
@@ -23,6 +24,7 @@ Then run the command below.
 ```bash
 bash run_multiple_prompts.sh
 ```
+
 ### Single Prompt
 Get into the example folder
 ```bash
@@ -33,6 +35,34 @@ Then run the command below.
 bash run_single_prompt.sh
 ```
 
+#### Using Local Media Files
+The `end2end.py` script supports local media files (audio, video, image) via CLI arguments:
+
+```bash
+# Use single local media files
+python end2end.py --query-type use_image --image-path /path/to/image.jpg
+python end2end.py --query-type use_video --video-path /path/to/video.mp4
+python end2end.py --query-type use_audio --audio-path /path/to/audio.wav
+
+# Combine multiple local media files
+python end2end.py --query-type mixed_modalities \
+    --video-path /path/to/video.mp4 \
+    --image-path /path/to/image.jpg \
+    --audio-path /path/to/audio.wav
+
+# Use audio from video file
+python end2end.py --query-type use_audio_in_video --video-path /path/to/video.mp4
+
+```
+
+If media file paths are not provided, the script will use default assets. Supported query types:
+- `use_image`: Image input only
+- `use_video`: Video input only
+- `use_audio`: Audio input only
+- `mixed_modalities`: Audio + image + video
+- `use_audio_in_video`: Extract audio from video
+- `text`: Text-only query
+
 ### FAQ
 
 If you encounter error about backend of librosa, try to install ffmpeg with command below.

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -6,12 +6,15 @@ with the correct prompt format on Qwen2.5-Omni
 """
 
 import os
-from typing import NamedTuple
+from typing import NamedTuple, Optional
 
+import librosa
+import numpy as np
 import soundfile as sf
+from PIL import Image
 from vllm.assets.audio import AudioAsset
 from vllm.assets.image import ImageAsset
-from vllm.assets.video import VideoAsset
+from vllm.assets.video import VideoAsset, video_to_ndarrays
 from vllm.multimodal.image import convert_image_mode
 from vllm.sampling_params import SamplingParams
 from vllm.utils import FlexibleArgumentParser
@@ -54,7 +57,13 @@ def get_text_query(question: str = None) -> QueryResult:
     )
 
 
-def get_mixed_modalities_query() -> QueryResult:
+def get_mixed_modalities_query(
+    video_path: Optional[str] = None,
+    image_path: Optional[str] = None,
+    audio_path: Optional[str] = None,
+    num_frames: int = 16,
+    sampling_rate: int = 16000,
+) -> QueryResult:
     question = "What is recited in the audio? What is the content of this image? Why is this video funny?"
     prompt = (
         f"<|im_start|>system\n{default_system}<|im_end|>\n"
@@ -64,20 +73,49 @@ def get_mixed_modalities_query() -> QueryResult:
         f"{question}<|im_end|>\n"
         f"<|im_start|>assistant\n"
     )
+
+    # Load video
+    if video_path:
+        if not os.path.exists(video_path):
+            raise FileNotFoundError(f"Video file not found: {video_path}")
+        video_frames = video_to_ndarrays(video_path, num_frames=num_frames)
+    else:
+        video_frames = VideoAsset(name="baby_reading", num_frames=num_frames).np_ndarrays
+
+    # Load image
+    if image_path:
+        if not os.path.exists(image_path):
+            raise FileNotFoundError(f"Image file not found: {image_path}")
+        pil_image = Image.open(image_path)
+        image_data = convert_image_mode(pil_image, "RGB")
+    else:
+        image_data = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB")
+
+    # Load audio
+    if audio_path:
+        if not os.path.exists(audio_path):
+            raise FileNotFoundError(f"Audio file not found: {audio_path}")
+        audio_signal, sr = librosa.load(audio_path, sr=sampling_rate)
+        audio_data = (audio_signal.astype(np.float32), sr)
+    else:
+        audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
+
     return QueryResult(
         inputs={
             "prompt": prompt,
             "multi_modal_data": {
-                "audio": AudioAsset("mary_had_lamb").audio_and_sample_rate,
-                "image": convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB"),
-                "video": VideoAsset(name="baby_reading", num_frames=16).np_ndarrays,
+                "audio": audio_data,
+                "image": image_data,
+                "video": video_frames,
             },
         },
         limit_mm_per_prompt={"audio": 1, "image": 1, "video": 1},
     )
 
 
-def get_use_audio_in_video_query() -> QueryResult:
+def get_use_audio_in_video_query(
+    video_path: Optional[str] = None, num_frames: int = 16, sampling_rate: int = 16000
+) -> QueryResult:
     question = "Describe the content of the video, then convert what the baby say into text."
     prompt = (
         f"<|im_start|>system\n{default_system}<|im_end|>\n"
@@ -85,14 +123,24 @@ def get_use_audio_in_video_query() -> QueryResult:
         f"{question}<|im_end|>\n"
         f"<|im_start|>assistant\n"
     )
-    asset = VideoAsset(name="baby_reading", num_frames=16)
-    audio = asset.get_audio(sampling_rate=16000)
+
+    if video_path:
+        if not os.path.exists(video_path):
+            raise FileNotFoundError(f"Video file not found: {video_path}")
+        video_frames = video_to_ndarrays(video_path, num_frames=num_frames)
+        # Extract audio from video file
+        audio_signal, sr = librosa.load(video_path, sr=sampling_rate)
+        audio = (audio_signal.astype(np.float32), sr)
+    else:
+        asset = VideoAsset(name="baby_reading", num_frames=num_frames)
+        video_frames = asset.np_ndarrays
+        audio = asset.get_audio(sampling_rate=sampling_rate)
 
     return QueryResult(
         inputs={
             "prompt": prompt,
             "multi_modal_data": {
-                "video": asset.np_ndarrays,
+                "video": video_frames,
                 "audio": audio,
             },
             "mm_processor_kwargs": {
@@ -103,7 +151,7 @@ def get_use_audio_in_video_query() -> QueryResult:
     )
 
 
-def get_multi_audios_query() -> QueryResult:
+def get_multi_audios_query(audio_path: Optional[str] = None, sampling_rate: int = 16000) -> QueryResult:
     question = "Are these two audio clips the same?"
     prompt = (
         f"<|im_start|>system\n{default_system}<|im_end|>\n"
@@ -112,14 +160,28 @@ def get_multi_audios_query() -> QueryResult:
         f"{question}<|im_end|>\n"
         f"<|im_start|>assistant\n"
     )
+
+    if audio_path:
+        if not os.path.exists(audio_path):
+            raise FileNotFoundError(f"Audio file not found: {audio_path}")
+        audio_signal, sr = librosa.load(audio_path, sr=sampling_rate)
+        audio_data = (audio_signal.astype(np.float32), sr)
+        # Use the provided audio as the first audio, default as second
+        audio_list = [
+            audio_data,
+            AudioAsset("mary_had_lamb").audio_and_sample_rate,
+        ]
+    else:
+        audio_list = [
+            AudioAsset("winning_call").audio_and_sample_rate,
+            AudioAsset("mary_had_lamb").audio_and_sample_rate,
+        ]
+
     return QueryResult(
         inputs={
             "prompt": prompt,
             "multi_modal_data": {
-                "audio": [
-                    AudioAsset("winning_call").audio_and_sample_rate,
-                    AudioAsset("mary_had_lamb").audio_and_sample_rate,
-                ],
+                "audio": audio_list,
             },
         },
         limit_mm_per_prompt={
@@ -128,17 +190,135 @@ def get_multi_audios_query() -> QueryResult:
     )
 
 
+def get_image_query(question: str = None, image_path: Optional[str] = None) -> QueryResult:
+    if question is None:
+        question = "What is the content of this image?"
+    prompt = (
+        f"<|im_start|>system\n{default_system}<|im_end|>\n"
+        "<|im_start|>user\n<|vision_bos|><|IMAGE|><|vision_eos|>"
+        f"{question}<|im_end|>\n"
+        f"<|im_start|>assistant\n"
+    )
+
+    if image_path:
+        if not os.path.exists(image_path):
+            raise FileNotFoundError(f"Image file not found: {image_path}")
+        pil_image = Image.open(image_path)
+        image_data = convert_image_mode(pil_image, "RGB")
+    else:
+        image_data = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB")
+
+    return QueryResult(
+        inputs={
+            "prompt": prompt,
+            "multi_modal_data": {
+                "image": image_data,
+            },
+        },
+        limit_mm_per_prompt={"image": 1},
+    )
+
+
+def get_video_query(question: str = None, video_path: Optional[str] = None, num_frames: int = 16) -> QueryResult:
+    if question is None:
+        question = "Why is this video funny?"
+    prompt = (
+        f"<|im_start|>system\n{default_system}<|im_end|>\n"
+        "<|im_start|>user\n<|vision_bos|><|VIDEO|><|vision_eos|>"
+        f"{question}<|im_end|>\n"
+        f"<|im_start|>assistant\n"
+    )
+
+    if video_path:
+        if not os.path.exists(video_path):
+            raise FileNotFoundError(f"Video file not found: {video_path}")
+        video_frames = video_to_ndarrays(video_path, num_frames=num_frames)
+    else:
+        video_frames = VideoAsset(name="baby_reading", num_frames=num_frames).np_ndarrays
+
+    return QueryResult(
+        inputs={
+            "prompt": prompt,
+            "multi_modal_data": {
+                "video": video_frames,
+            },
+        },
+        limit_mm_per_prompt={"video": 1},
+    )
+
+
+def get_audio_query(question: str = None, audio_path: Optional[str] = None, sampling_rate: int = 16000) -> QueryResult:
+    if question is None:
+        question = "What is the content of this audio?"
+    prompt = (
+        f"<|im_start|>system\n{default_system}<|im_end|>\n"
+        "<|im_start|>user\n<|audio_bos|><|AUDIO|><|audio_eos|>"
+        f"{question}<|im_end|>\n"
+        f"<|im_start|>assistant\n"
+    )
+
+    if audio_path:
+        if not os.path.exists(audio_path):
+            raise FileNotFoundError(f"Audio file not found: {audio_path}")
+        audio_signal, sr = librosa.load(audio_path, sr=sampling_rate)
+        audio_data = (audio_signal.astype(np.float32), sr)
+    else:
+        audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
+
+    return QueryResult(
+        inputs={
+            "prompt": prompt,
+            "multi_modal_data": {
+                "audio": audio_data,
+            },
+        },
+        limit_mm_per_prompt={"audio": 1},
+    )
+
+
 query_map = {
     "mixed_modalities": get_mixed_modalities_query,
     "use_audio_in_video": get_use_audio_in_video_query,
     "multi_audios": get_multi_audios_query,
+    "use_image": get_image_query,
+    "use_video": get_video_query,
+    "use_audio": get_audio_query,
     "text": get_text_query,
 }
 
 
 def main(args):
     model_name = "Qwen/Qwen2.5-Omni-7B"
-    query_result = query_map[args.query_type]()
+
+    # Get paths from args
+    video_path = getattr(args, "video_path", None)
+    image_path = getattr(args, "image_path", None)
+    audio_path = getattr(args, "audio_path", None)
+    num_frames = getattr(args, "num_frames", 16)
+    sampling_rate = getattr(args, "sampling_rate", 16000)
+
+    # Get the query function and call it with appropriate parameters
+    query_func = query_map[args.query_type]
+    if args.query_type == "mixed_modalities":
+        query_result = query_func(
+            video_path=video_path,
+            image_path=image_path,
+            audio_path=audio_path,
+            num_frames=num_frames,
+            sampling_rate=sampling_rate,
+        )
+    elif args.query_type == "use_audio_in_video":
+        query_result = query_func(video_path=video_path, num_frames=num_frames, sampling_rate=sampling_rate)
+    elif args.query_type == "multi_audios":
+        query_result = query_func(audio_path=audio_path, sampling_rate=sampling_rate)
+    elif args.query_type == "use_image":
+        query_result = query_func(image_path=image_path)
+    elif args.query_type == "use_video":
+        query_result = query_func(video_path=video_path, num_frames=num_frames)
+    elif args.query_type == "use_audio":
+        query_result = query_func(audio_path=audio_path, sampling_rate=sampling_rate)
+    else:
+        query_result = query_func()
 
     omni_llm = Omni(
         model=model_name,
@@ -284,18 +464,38 @@ def parse_args():
         help="Path to a .txt file with one prompt per line (preferred).",
     )
     parser.add_argument(
-        "--worker-backend",
-        type=str,
-        default="process",
-        choices=["process","ray"],
-        help="backend"
-    )
-    parser.add_argument(
-        "--ray-address",
+        "--video-path",
+        "-v",
         type=str,
         default=None,
-        help="Path to a .txt file with one prompt per line (preferred).",
-
+        help="Path to local video file. If not provided, uses default video asset.",
+    )
+    parser.add_argument(
+        "--image-path",
+        "-i",
+        type=str,
+        default=None,
+        help="Path to local image file. If not provided, uses default image asset.",
+    )
+    parser.add_argument(
+        "--audio-path",
+        "-a",
+        type=str,
+        default=None,
+        help="Path to local audio file. If not provided, uses default audio asset.",
+    )
+    parser.add_argument(
+        "--num-frames",
+        type=int,
+        default=16,
+        help="Number of frames to extract from video (default: 16).",
+    )
+    parser.add_argument(
+        "--sampling-rate",
+        type=int,
+        default=16000,
+        help="Sampling rate for audio loading (default: 16000).",
+    )
 
     return parser.parse_args()
 

examples/offline_inference/qwen3_omni/README.md

@@ -5,6 +5,7 @@
 Please refer to [README.md](../../../README.md)
 
 ## Run examples (Qwen3-Omni)
+
 ### Multiple Prompts
 Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). For processing dataset please refer to [Qwen2.5-Omni README.md](../qwen2_5_omni/README.md)
 Get into the example folder
@@ -29,6 +30,26 @@ If you have not enough memory, you can set thinker with tensor parallel. Just ru
 bash run_single_prompt_tp.sh
 ```
 
+#### Using Local Media Files
+The `end2end.py` script supports local media files (audio, video, image) via command-line arguments:
+
+```bash
+# Use local video file
+python end2end.py --query-type use_video --video-path /path/to/video.mp4
+
+# Use local image file
+python end2end.py --query-type use_image --image-path /path/to/image.jpg
+
+# Use local audio file
+python end2end.py --query-type use_audio --audio-path /path/to/audio.wav
+```
+
+If media file paths are not provided, the script will use default assets. Supported query types:
+- `use_video`: Video input
+- `use_image`: Image input
+- `use_audio`: Audio input
+- `text`: Text-only query
+
 ### FAQ
 
 If you encounter error about backend of librosa, try to install ffmpeg with command below.

examples/offline_inference/qwen3_omni/end2end.py

@@ -6,13 +6,16 @@ with the correct prompt format on Qwen3-Omni (thinker only).
 """
 
 import os
-from typing import NamedTuple
+from typing import NamedTuple, Optional
 
+import librosa
+import numpy as np
 import soundfile as sf
+from PIL import Image
 from vllm import SamplingParams
 from vllm.assets.audio import AudioAsset
 from vllm.assets.image import ImageAsset
-from vllm.assets.video import VideoAsset
+from vllm.assets.video import VideoAsset, video_to_ndarrays
 from vllm.multimodal.image import convert_image_mode
 from vllm.utils import FlexibleArgumentParser
 
@@ -54,7 +57,7 @@ def get_text_query(question: str = None) -> QueryResult:
     )
 
 
-def get_video_query(question: str = None) -> QueryResult:
+def get_video_query(question: str = None, video_path: Optional[str] = None, num_frames: int = 16) -> QueryResult:
     if question is None:
         question = "Why is this video funny?"
     prompt = (
@@ -64,18 +67,25 @@ def get_video_query(question: str = None) -> QueryResult:
         f"<|im_start|>assistant\n"
     )
 
+    if video_path:
+        if not os.path.exists(video_path):
+            raise FileNotFoundError(f"Video file not found: {video_path}")
+        video_frames = video_to_ndarrays(video_path, num_frames=num_frames)
+    else:
+        video_frames = VideoAsset(name="baby_reading", num_frames=num_frames).np_ndarrays
+
     return QueryResult(
         inputs={
             "prompt": prompt,
             "multi_modal_data": {
-                "video": VideoAsset(name="baby_reading", num_frames=16).np_ndarrays,
+                "video": video_frames,
             },
         },
         limit_mm_per_prompt={"video": 1},
     )
 
 
-def get_image_query(question: str = None) -> QueryResult:
+def get_image_query(question: str = None, image_path: Optional[str] = None) -> QueryResult:
     if question is None:
         question = "What is the content of this image?"
     prompt = (
@@ -84,18 +94,27 @@ def get_image_query(question: str = None) -> QueryResult:
         f"{question}<|im_end|>\n"
         f"<|im_start|>assistant\n"
     )
+
+    if image_path:
+        if not os.path.exists(image_path):
+            raise FileNotFoundError(f"Image file not found: {image_path}")
+        pil_image = Image.open(image_path)
+        image_data = convert_image_mode(pil_image, "RGB")
+    else:
+        image_data = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB")
+
     return QueryResult(
         inputs={
             "prompt": prompt,
             "multi_modal_data": {
-                "image": convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB"),
+                "image": image_data,
             },
         },
         limit_mm_per_prompt={"image": 1},
     )
 
 
-def get_audio_query(question: str = None) -> QueryResult:
+def get_audio_query(question: str = None, audio_path: Optional[str] = None, sampling_rate: int = 16000) -> QueryResult:
     if question is None:
         question = "What is the content of this audio?"
     prompt = (
@@ -104,11 +123,20 @@ def get_audio_query(question: str = None) -> QueryResult:
         f"{question}<|im_end|>\n"
         f"<|im_start|>assistant\n"
     )
+
+    if audio_path:
+        if not os.path.exists(audio_path):
+            raise FileNotFoundError(f"Audio file not found: {audio_path}")
+        audio_signal, sr = librosa.load(audio_path, sr=sampling_rate)
+        audio_data = (audio_signal.astype(np.float32), sr)
+    else:
+        audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
+
     return QueryResult(
         inputs={
             "prompt": prompt,
             "multi_modal_data": {
-                "audio": AudioAsset("mary_had_lamb").audio_and_sample_rate,
+                "audio": audio_data,
             },
         },
         limit_mm_per_prompt={"audio": 1},
@@ -125,7 +153,22 @@ query_map = {
 
 def main(args):
     model_name = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
-    query_result = query_map[args.query_type]()
+
+    # Get paths from args
+    video_path = getattr(args, "video_path", None)
+    image_path = getattr(args, "image_path", None)
+    audio_path = getattr(args, "audio_path", None)
+
+    # Get the query function and call it with appropriate parameters
+    query_func = query_map[args.query_type]
+    if args.query_type == "use_video":
+        query_result = query_func(video_path=video_path, num_frames=getattr(args, "num_frames", 16))
+    elif args.query_type == "use_image":
+        query_result = query_func(image_path=image_path)
+    elif args.query_type == "use_audio":
+        query_result = query_func(audio_path=audio_path, sampling_rate=getattr(args, "sampling_rate", 16000))
+    else:
+        query_result = query_func()
 
     omni_llm = Omni(
         model=model_name,
@@ -283,6 +326,39 @@ def parse_args():
         default=None,
         help="Path to a stage configs file.",
     )
+    parser.add_argument(
+        "--video-path",
+        "-v",
+        type=str,
+        default=None,
+        help="Path to local video file. If not provided, uses default video asset.",
+    )
+    parser.add_argument(
+        "--image-path",
+        "-i",
+        type=str,
+        default=None,
+        help="Path to local image file. If not provided, uses default image asset.",
+    )
+    parser.add_argument(
+        "--audio-path",
+        "-a",
+        type=str,
+        default=None,
+        help="Path to local audio file. If not provided, uses default audio asset.",
+    )
+    parser.add_argument(
+        "--num-frames",
+        type=int,
+        default=16,
+        help="Number of frames to extract from video (default: 16).",
+    )
+    parser.add_argument(
+        "--sampling-rate",
+        type=int,
+        default=16000,
+        help="Sampling rate for audio loading (default: 16000).",
+    )
 
     return parser.parse_args()
 

examples/offline_inference/qwen3_omni/qwen3_omni_moe_tp.yaml

@@ -1,94 +0,0 @@
-# Stage config for running Qwen3-Omni-MoE with 3-stage architecture
-# Stage 0: Thinker (multimodal understanding + text generation)
-# Stage 1: Talker (text embeddings → 8-layer RVQ codec codes)
-# Stage 2: Code2Wav (8-layer RVQ codes → audio waveform)
-
-stage_args:
-  - stage_id: 0
-    runtime:
-      devices: "0,1"
-      max_batch_size: 1
-    engine_args:
-      model_stage: thinker
-      model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
-      scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
-      gpu_memory_utilization: 0.8
-      enforce_eager: true
-      trust_remote_code: true
-      engine_output_type: latent  # Output hidden states for talker
-      distributed_executor_backend: "mp"
-      enable_prefix_caching: false
-      hf_config_name: thinker_config
-      tensor_parallel_size: 2
-    final_output: true
-    final_output_type: text
-    is_comprehension: true
-    default_sampling_params:
-      temperature: 0.4
-      top_p: 0.9
-      top_k: 1
-      max_tokens: 2048
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.05
-
-  - stage_id: 1
-    runtime:
-       devices: "1"
-       max_batch_size: 1
-    engine_args:
-       model_stage: talker
-       model_arch: Qwen3OmniMoeForConditionalGeneration
-       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
-       scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
-       gpu_memory_utilization: 0.5
-       enforce_eager: true
-       trust_remote_code: true
-       engine_output_type: latent  # Output codec codes for code2wav
-      #  tensor_parallel_size: 2
-       enable_prefix_caching: false
-       distributed_executor_backend: "mp"
-       hf_config_name: talker_config
-    engine_input_source: [0]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
-    # final_output: true
-    # final_output_type: text
-    default_sampling_params:
-      temperature: 0.9
-      top_k: 50
-      max_tokens: 4096
-      seed: 42
-      detokenize: False
-      repetition_penalty: 1.05
-      stop_token_ids: [2150]
-
-  - stage_id: 2
-    runtime:
-      devices: "1"
-      max_batch_size: 1
-    engine_args:
-      model_stage: code2wav
-      model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
-      scheduler_cls: vllm_omni.core.sched.generation_scheduler.GenerationScheduler
-      enforce_eager: true
-      trust_remote_code: true
-      enable_prefix_caching: false
-      engine_output_type: audio  # Final output: audio waveform
-      gpu_memory_utilization: 0.1
-      distributed_executor_backend: "mp"
-      max_num_batched_tokens: 1000000
-      hf_config_name: thinker_config
-    engine_input_source: [1]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
-    final_output: true
-    final_output_type: audio
-    default_sampling_params:
-      temperature: 0.0
-      top_p: 1.0
-      top_k: -1
-      max_tokens: 65536
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.1

examples/offline_inference/qwen3_omni/run_single_prompt_tp.sh

@@ -1,6 +1,5 @@
 python end2end.py --output-wav output_audio \
                   --query-type use_audio \
-                  --init-sleep-seconds 90 \
-                  --stage-configs-path qwen3_omni_moe_tp.yaml
+                  --init-sleep-seconds 90
 
 # init-sleep-seconds works to avoid two vLLM stages initialized at the same time within a card.

examples/offline_inference/qwen_image/gradio_demo.py

@@ -5,7 +5,7 @@ import gradio as gr
 import torch
 
 from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.utils.platform_utils import detect_device_type
+from vllm_omni.utils.platform_utils import detect_device_type, is_npu
 
 ASPECT_RATIOS: dict[str, tuple[int, int]] = {
     "1:1": (1328, 1328),
@@ -59,7 +59,14 @@ def parse_args() -> argparse.Namespace:
 
 @lru_cache(maxsize=1)
 def get_omni(model_name: str) -> Omni:
-    return Omni(model=model_name)
+    # Enable VAE memory optimizations on NPU
+    vae_use_slicing = is_npu()
+    vae_use_tiling = is_npu()
+    return Omni(
+        model=model_name,
+        vae_use_slicing=vae_use_slicing,
+        vae_use_tiling=vae_use_tiling,
+    )
 
 
 def build_demo(args: argparse.Namespace) -> gr.Blocks:

examples/offline_inference/qwen_image/text_to_image.py

@@ -7,7 +7,7 @@ from pathlib import Path
 import torch
 
 from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.utils.platform_utils import detect_device_type
+from vllm_omni.utils.platform_utils import detect_device_type, is_npu
 
 
 def parse_args() -> argparse.Namespace:
@@ -49,7 +49,15 @@ def main():
     device = detect_device_type()
     generator = torch.Generator(device=device).manual_seed(args.seed)
 
-    omni = Omni(model=args.model)
+    # Enable VAE memory optimizations on NPU
+    vae_use_slicing = is_npu()
+    vae_use_tiling = is_npu()
+
+    omni = Omni(
+        model=args.model,
+        vae_use_slicing=vae_use_slicing,
+        vae_use_tiling=vae_use_tiling,
+    )
     images = omni.generate(
         args.prompt,
         height=args.height,

examples/offline_inference/tmp_test/README.md

@@ -1,34 +0,0 @@
-# Offline Example of vLLM-omni for Qwen2.5-omni
-
-## 🛠️ Installation
-
-Please refer to [README.md](../../../README.md)
-
-## Run examples (Qwen2.5-omni)
-### Multiple Prompts
-Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). To get the prompt, you can:
-```bash
-tar -xf <Your Download Path>/seedtts_testset.tar
-cp seedtts_testset/en/meta.lst examples/offline_inference/qwen2_5_omni/meta.lst
-python3 examples/offline_inference/qwen2_5_omni/extract_prompts.py \
-  --input examples/offline_inference/qwen2_5_omni/meta.lst \
-  --output examples/offline_inference/qwen2_5_omni/top100.txt \
-  --topk 100
-```
-Get into the example folder
-```bash
-cd examples/offline_inference/qwen2_5_omni
-```
-Then run the command below.
-```bash
-bash run_multiple_prompts.sh
-```
-### Single Prompts
-Get into the example folder
-```bash
-cd examples/offline_inference/qwen2_5_omni
-```
-Then run the command below.
-```bash
-bash run_single_prompt.sh
-```

examples/offline_inference/tmp_test/end2end.py

@@ -1,264 +0,0 @@
-import argparse
-import os
-import os as _os_env_toggle
-import random
-
-import numpy as np
-import soundfile as sf
-import torch
-from utils import make_omni_prompt
-from vllm.sampling_params import SamplingParams
-
-from vllm_omni.entrypoints.omni_llm import OmniLLM
-
-_os_env_toggle.environ["VLLM_USE_V1"] = "1"
-
-SEED = 42
-# Set all random seeds
-random.seed(SEED)
-np.random.seed(SEED)
-torch.manual_seed(SEED)
-torch.cuda.manual_seed(SEED)
-torch.cuda.manual_seed_all(SEED)
-
-# Make PyTorch deterministic
-torch.backends.cudnn.deterministic = True
-torch.backends.cudnn.benchmark = False
-
-# Set environment variables for deterministic behavior
-os.environ["PYTHONHASHSEED"] = str(SEED)
-os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
-
-
-def parse_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--model",
-        required=True,
-        help="Path to merged model directory (will be created if downloading).",
-    )
-    parser.add_argument("--thinker-model", type=str, default=None)
-    parser.add_argument("--talker-model", type=str, default=None)
-    parser.add_argument("--code2wav-model", type=str, default=None)
-    parser.add_argument(
-        "--hf-hub-id",
-        default="Qwen/Qwen2.5-Omni-7B",
-        help="Hugging Face repo id to download if needed.",
-    )
-    parser.add_argument("--hf-revision", default=None, help="Optional HF revision (branch/tag/commit).")
-    parser.add_argument("--prompts", nargs="+", default=None, help="Input text prompts.")
-    parser.add_argument("--voice-type", default="default", help="Voice type, e.g., m02, f030, default.")
-    parser.add_argument(
-        "--code2wav-dir",
-        default=None,
-        help="Path to code2wav folder (contains spk_dict.pt).",
-    )
-    parser.add_argument("--dit-ckpt", default=None, help="Path to DiT checkpoint file (e.g., dit.pt).")
-    parser.add_argument("--bigvgan-ckpt", default=None, help="Path to BigVGAN checkpoint file.")
-    parser.add_argument("--dtype", default="bfloat16", choices=["float16", "bfloat16", "float32"])
-    parser.add_argument("--max-model-len", type=int, default=32768)
-    parser.add_argument(
-        "--init-sleep-seconds",
-        type=int,
-        default=20,
-        help="Sleep seconds after starting each stage process to allow initialization (default: 20)",
-    )
-
-    parser.add_argument("--thinker-only", action="store_true")
-    parser.add_argument("--text-only", action="store_true")
-    parser.add_argument("--do-wave", action="store_true")
-    parser.add_argument(
-        "--prompt_type",
-        choices=[
-            "text",
-            "audio",
-            "audio-long",
-            "audio-long-chunks",
-            "audio-long-expand-chunks",
-            "image",
-            "video",
-            "video-frames",
-            "audio-in-video",
-            "audio-in-video-v2",
-            "audio-multi-round",
-            "badcase-vl",
-            "badcase-text",
-            "badcase-image-early-stop",
-            "badcase-two-audios",
-            "badcase-two-videos",
-            "badcase-multi-round",
-            "badcase-voice-type",
-            "badcase-voice-type-v2",
-            "badcase-audio-tower-1",
-            "badcase-audio-only",
-        ],
-        default="text",
-    )
-    parser.add_argument("--use-torchvision", action="store_true")
-    parser.add_argument("--tokenize", action="store_true")
-    parser.add_argument(
-        "--output-wav",
-        default="output.wav",
-        help="[Deprecated] Output wav directory (use --output-dir).",
-    )
-    parser.add_argument(
-        "--output-dir",
-        default="outputs",
-        help="Output directory to save text and wav files together.",
-    )
-    parser.add_argument(
-        "--thinker-hidden-states-dir",
-        default="thinker_hidden_states",
-        help="Path to thinker hidden states directory.",
-    )
-    parser.add_argument(
-        "--batch-timeout",
-        type=int,
-        default=5,
-        help="Timeout for batching in seconds (default: 5)",
-    )
-    parser.add_argument(
-        "--init-timeout",
-        type=int,
-        default=300,
-        help="Timeout for initializing stages in seconds (default: 300)",
-    )
-    parser.add_argument(
-        "--shm-threshold-bytes",
-        type=int,
-        default=65536,
-        help="Threshold for using shared memory in bytes (default: 65536)",
-    )
-    parser.add_argument(
-        "--enable-stats",
-        action="store_true",
-        default=False,
-        help="Enable writing detailed statistics (default: disabled)",
-    )
-    parser.add_argument(
-        "--txt-prompts",
-        type=str,
-        default=None,
-        help="Path to a .txt file with one prompt per line (preferred).",
-    )
-    parser.add_argument(
-        "--worker-backend",
-        type=str,
-        default="process",
-        choices=["process","ray"],
-        help="backend"
-    )
-    parser.add_argument(
-        "--ray-address",
-        type=str,
-        default=None,
-        help="Path to a .txt file with one prompt per line (preferred)."
-        )
-
-    args = parser.parse_args()
-    return args
-
-
-def main():
-    args = parse_args()
-    model_name = args.model
-    try:
-        # Preferred: load from txt file (one prompt per line)
-        if getattr(args, "txt_prompts", None) and args.prompt_type == "text":
-            with open(args.txt_prompts, encoding="utf-8") as f:
-                lines = [ln.strip() for ln in f.readlines()]
-            args.prompts = [ln for ln in lines if ln != ""]
-            print(f"[Info] Loaded {len(args.prompts)} prompts from {args.txt_prompts}")
-    except Exception as e:
-        print(f"[Error] Failed to load prompts: {e}")
-        raise
-
-    if args.prompts is None:
-        raise ValueError("No prompts provided. Use --prompts ... or --txt-prompts <file.txt> (with --prompt_type text)")
-    omni_llm = OmniLLM(
-        model=model_name,
-        log_stats=args.enable_stats,
-        log_file=("omni_llm_pipeline.log" if args.enable_stats else None),
-        init_sleep_seconds=args.init_sleep_seconds,
-        batch_timeout=args.batch_timeout,
-        init_timeout=args.init_timeout,
-        shm_threshold_bytes=args.shm_threshold_bytes,
-        worker_backend=args.worker_backend,
-        ray_address=args.ray_address
-    )
-    thinker_sampling_params = SamplingParams(
-        temperature=0.0,  # Deterministic - no randomness
-        top_p=1.0,  # Disable nucleus sampling
-        top_k=-1,  # Disable top-k sampling
-        max_tokens=2048,
-        seed=SEED,  # Fixed seed for sampling
-        detokenize=True,
-        repetition_penalty=1.1,
-    )
-    talker_sampling_params = SamplingParams(
-        temperature=0.9,
-        top_p=0.8,
-        top_k=40,
-        max_tokens=2048,
-        seed=SEED,  # Fixed seed for sampling
-        detokenize=True,
-        repetition_penalty=1.05,
-        stop_token_ids=[8294],
-    )
-    code2wav_sampling_params = SamplingParams(
-        temperature=0.0,  # Deterministic - no randomness
-        top_p=1.0,  # Disable nucleus sampling
-        top_k=-1,  # Disable top-k sampling
-        max_tokens=2048,
-        seed=SEED,  # Fixed seed for sampling
-        detokenize=True,
-        repetition_penalty=1.1,
-    )
-
-    sampling_params_list = [
-        thinker_sampling_params,
-        talker_sampling_params,
-        code2wav_sampling_params,
-    ]
-    import time
-    for i in range(1):
-        t1 = time.time()
-        prompt = [make_omni_prompt(args, prompt) for prompt in args.prompts]
-        print(f"prompt:{prompt}")
-        omni_outputs = omni_llm.generate(prompt, sampling_params_list)
-        t2 = time.time()
-        print(f"==========> time:{t2-t1}")
-
-    # Determine output directory: prefer --output-dir; fallback to --output-wav
-    output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
-    os.makedirs(output_dir, exist_ok=True)
-    for stage_outputs in omni_outputs:
-        if stage_outputs.final_output_type == "text":
-            for output in stage_outputs.request_output:
-                request_id = int(output.request_id)
-                text_output = output.outputs[0].text
-                # Save aligned text file per request
-                prompt_text = args.prompts[request_id]
-                out_txt = os.path.join(output_dir, f"{request_id:05d}.txt")
-                lines = []
-                lines.append("Prompt:\n")
-                lines.append(str(prompt_text) + "\n")
-                lines.append("vllm_text_output:\n")
-                lines.append(str(text_output).strip() + "\n")
-                try:
-                    with open(out_txt, "w", encoding="utf-8") as f:
-                        f.writelines(lines)
-                except Exception as e:
-                    print(f"[Warn] Failed writing text file {out_txt}: {e}")
-                print(f"Request ID: {request_id}, Text saved to {out_txt}")
-        elif stage_outputs.final_output_type == "audio":
-            for output in stage_outputs.request_output:
-                request_id = int(output.request_id)
-                audio_tensor = output.multimodal_output["audio"]
-                output_wav = os.path.join(output_dir, f"output_{output.request_id}.wav")
-                sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000)
-                print(f"Request ID: {request_id}, Saved audio to {output_wav}")
-
-
-if __name__ == "__main__":
-    main()

examples/offline_inference/tmp_test/extract_prompts.py

@@ -1,45 +0,0 @@
-#!/usr/bin/env python3
-import argparse
-from typing import Optional
-
-
-def extract_prompt(line: str) -> Optional[str]:
-    # 提取第一个 '|' 与第二个 '|' 之间的内容
-    i = line.find("|")
-    if i == -1:
-        return None
-    j = line.find("|", i + 1)
-    if j == -1:
-        return None
-    return line[i + 1 : j].strip()
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--input", "-i", required=True, help="Input .lst file path")
-    parser.add_argument("--output", "-o", required=True, help="Output file path")
-    parser.add_argument(
-        "--topk",
-        "-k",
-        type=int,
-        default=100,
-        help="Extract the top K prompts (default: 100)",
-    )
-    args = parser.parse_args()
-
-    prompts = []
-    with open(args.input, encoding="utf-8", errors="ignore") as f:
-        for line in f:
-            if len(prompts) >= args.topk:
-                break
-            p = extract_prompt(line.rstrip("\n"))
-            if p:
-                prompts.append(p)
-
-    with open(args.output, "w", encoding="utf-8") as f:
-        for p in prompts:
-            f.write(p + "\n")
-
-
-if __name__ == "__main__":
-    main()

examples/offline_inference/tmp_test/processing_omni.py

@@ -1,367 +0,0 @@
-from __future__ import annotations
-
-import base64
-import logging
-import math
-import os
-import time
-import warnings
-from functools import lru_cache
-from io import BytesIO
-
-import requests
-import torch
-import torchvision
-from packaging import version
-from PIL import Image
-from torchvision import io, transforms
-from torchvision.transforms import InterpolationMode
-
-logger = logging.getLogger(__name__)
-
-IMAGE_FACTOR = 28
-MIN_PIXELS = 4 * 28 * 28
-MAX_PIXELS = 16384 * 28 * 28
-MAX_RATIO = 200
-
-VIDEO_MIN_PIXELS = 128 * 28 * 28
-VIDEO_MAX_PIXELS = 768 * 28 * 28
-VIDEO_TOTAL_PIXELS = 24576 * 28 * 28
-FRAME_FACTOR = 2
-FPS = 2.0
-FPS_MIN_FRAMES = 4
-FPS_MAX_FRAMES = 768
-
-temporal_patch_size = 2
-spatial_patch_size = 14
-spatial_merge_size = 2
-
-
-def round_by_factor(number: int, factor: int) -> int:
-    """Returns the closest integer to 'number' that is divisible by 'factor'."""
-    return round(number / factor) * factor
-
-
-def ceil_by_factor(number: int, factor: int) -> int:
-    """Returns the smallest integer greater than or equal to 'number' that is
-    divisible by 'factor'."""
-    return math.ceil(number / factor) * factor
-
-
-def floor_by_factor(number: int, factor: int) -> int:
-    """Returns the largest integer less than or equal to 'number' that is
-    divisible by 'factor'."""
-    return math.floor(number / factor) * factor
-
-
-def smart_resize(
-    height: int,
-    width: int,
-    factor: int = IMAGE_FACTOR,
-    min_pixels: int = MIN_PIXELS,
-    max_pixels: int = MAX_PIXELS,
-) -> tuple[int, int]:
-    """
-    Rescales the image so that the following conditions are met:
-
-    1. Both dimensions (height and width) are divisible by 'factor'.
-
-    2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
-
-    3. The aspect ratio of the image is maintained as closely as possible.
-    """
-    if max(height, width) / min(height, width) > MAX_RATIO:
-        raise ValueError(
-            f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}"
-        )
-    h_bar = max(factor, round_by_factor(height, factor))
-    w_bar = max(factor, round_by_factor(width, factor))
-    if h_bar * w_bar > max_pixels:
-        beta = math.sqrt((height * width) / max_pixels)
-        h_bar = floor_by_factor(height / beta, factor)
-        w_bar = floor_by_factor(width / beta, factor)
-    elif h_bar * w_bar < min_pixels:
-        beta = math.sqrt(min_pixels / (height * width))
-        h_bar = ceil_by_factor(height * beta, factor)
-        w_bar = ceil_by_factor(width * beta, factor)
-    return h_bar, w_bar
-
-
-def fetch_image(ele: dict[str, str | Image.Image], size_factor: int = IMAGE_FACTOR) -> Image.Image:
-    if "image" in ele:
-        image = ele["image"]
-    else:
-        image = ele["image_url"]
-    image_obj = None
-    if isinstance(image, Image.Image):
-        image_obj = image
-    elif image.startswith("http://") or image.startswith("https://"):
-        image_obj = Image.open(requests.get(image, stream=True).raw)
-    elif image.startswith("file://"):
-        image_obj = Image.open(image[7:])
-    elif image.startswith("data:image"):
-        if "base64," in image:
-            _, base64_data = image.split("base64,", 1)
-            data = base64.b64decode(base64_data)
-            image_obj = Image.open(BytesIO(data))
-    else:
-        image_obj = Image.open(image)
-    if image_obj is None:
-        raise ValueError(f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}")
-    image = image_obj.convert("RGB")
-    # resize
-    if "resized_height" in ele and "resized_width" in ele:
-        resized_height, resized_width = smart_resize(
-            ele["resized_height"],
-            ele["resized_width"],
-            factor=size_factor,
-        )
-    else:
-        width, height = image.size
-        min_pixels = ele.get("min_pixels", MIN_PIXELS)
-        max_pixels = ele.get("max_pixels", MAX_PIXELS)
-        resized_height, resized_width = smart_resize(
-            height,
-            width,
-            factor=size_factor,
-            min_pixels=min_pixels,
-            max_pixels=max_pixels,
-        )
-    image = image.resize((resized_width, resized_height))
-
-    return image
-
-
-def smart_nframes(
-    ele: dict,
-    total_frames: int,
-    video_fps: int | float,
-) -> int:
-    """calculate the number of frames for video used for model inputs.
-
-    Args:
-        ele (dict): a dict contains the configuration of video.
-            support either `fps` or `nframes`:
-                - nframes: the number of frames to extract for model inputs.
-                - fps: the fps to extract frames for model inputs.
-                    - min_frames: the minimum number of frames of the video,
-                      only used when fps is provided.
-                    - max_frames: the maximum number of frames of the video,
-                      only used when fps is provided.
-        total_frames (int): the original total number of frames of the video.
-        video_fps (int | float): the original fps of the video.
-
-    Raises:
-        ValueError: nframes should in interval [FRAME_FACTOR, total_frames].
-
-    Returns:
-        int: the number of frames for video used for model inputs.
-    """
-    assert not ("fps" in ele and "nframes" in ele), "Only accept either `fps` or `nframes`"
-    if "nframes" in ele:
-        nframes = round_by_factor(ele["nframes"], FRAME_FACTOR)
-    else:
-        fps = ele.get("fps", FPS)
-        min_frames = ceil_by_factor(ele.get("min_frames", FPS_MIN_FRAMES), FRAME_FACTOR)
-        max_frames = floor_by_factor(
-            ele.get("max_frames", min(FPS_MAX_FRAMES, total_frames)),
-            FRAME_FACTOR,
-        )
-        nframes = total_frames / video_fps * fps
-        nframes = min(max(nframes, min_frames), max_frames)
-        nframes = round_by_factor(nframes, FRAME_FACTOR)
-    if not (FRAME_FACTOR <= nframes and nframes <= total_frames):
-        raise ValueError(f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}.")
-    return nframes
-
-
-def _read_video_torchvision(
-    ele: dict,
-) -> torch.Tensor:
-    """read video using torchvision.io.read_video
-
-    Args:
-        ele (dict): a dict contains the configuration of video.
-        support keys:
-            - video: the path of video. support "file://", "http://",
-              "https://" and local path.
-            - video_start: the start time of video.
-            - video_end: the end time of video.
-    Returns:
-        torch.Tensor: the video tensor with shape (T, C, H, W).
-    """
-    video_path = ele["video"]
-    if version.parse(torchvision.__version__) < version.parse("0.19.0"):
-        if "http://" in video_path or "https://" in video_path:
-            warnings.warn("torchvision < 0.19.0 does not support http/https video path, please upgrade to 0.19.0.")
-        if "file://" in video_path:
-            video_path = video_path[7:]
-    st = time.time()
-    video, audio, info = io.read_video(
-        video_path,
-        start_pts=ele.get("video_start", 0.0),
-        end_pts=ele.get("video_end", None),
-        pts_unit="sec",
-        output_format="TCHW",
-    )
-    total_frames, video_fps = video.size(0), info["video_fps"]
-    total_duration = round(total_frames / video_fps, 3)
-    logger.info(
-        f"torchvision:  {video_path=}, {total_frames=}, {video_fps=}, duration={total_duration}s, time={time.time() - st:.3f}s"
-    )
-    nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
-    idx = torch.linspace(0, total_frames - 1, nframes).round().long()
-    video = video[idx]
-    return video, total_duration, nframes
-
-
-def is_decord_available() -> bool:
-    import importlib.util
-
-    return importlib.util.find_spec("decord") is not None
-
-
-def _read_video_decord(
-    ele: dict,
-) -> torch.Tensor:
-    """read video using decord.VideoReader
-
-    Args:
-        ele (dict): a dict contains the configuration of video.
-        support keys:
-            - video: the path of video. support "file://", "http://",
-              "https://" and local path.
-            - video_start: the start time of video.
-            - video_end: the end time of video.
-    Returns:
-        torch.Tensor: the video tensor with shape (T, C, H, W).
-    """
-    import decord
-
-    video_path = ele["video"]
-    st = time.time()
-    vr = decord.VideoReader(video_path)
-    # TODO: support start_pts and end_pts
-    if "video_start" in ele or "video_end" in ele:
-        raise NotImplementedError("not support start_pts and end_pts in decord for now.")
-    total_frames, video_fps = len(vr), vr.get_avg_fps()
-    total_duration = round(total_frames / video_fps, 3)
-    logger.info(f"decord:  {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s")
-    nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
-    idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist()
-    video = vr.get_batch(idx).asnumpy()
-    video = torch.tensor(video).permute(0, 3, 1, 2)  # Convert to TCHW format
-    return video, total_duration, nframes
-
-
-VIDEO_READER_BACKENDS = {
-    "decord": _read_video_decord,
-    "torchvision": _read_video_torchvision,
-}
-
-FORCE_QWENVL_VIDEO_READER = os.getenv("FORCE_QWENVL_VIDEO_READER", None)
-
-
-@lru_cache(maxsize=1)
-def get_video_reader_backend() -> str:
-    if FORCE_QWENVL_VIDEO_READER is not None:
-        video_reader_backend = FORCE_QWENVL_VIDEO_READER
-    elif is_decord_available():
-        video_reader_backend = "decord"
-    else:
-        video_reader_backend = "torchvision"
-    # print(f"qwen-vl-utils using {video_reader_backend} to read video.",
-    #       file=sys.stderr)
-    return video_reader_backend
-
-
-def fetch_video(ele: dict, image_factor: int = IMAGE_FACTOR) -> torch.Tensor | list[Image.Image]:
-    if isinstance(ele["video"], str):
-        video_reader_backend = get_video_reader_backend()
-        video, total_dur, nframes = VIDEO_READER_BACKENDS[video_reader_backend](ele)
-        frame_timestamps = total_dur * torch.arange(1, nframes + 1) / nframes
-        grid_timestamps = frame_timestamps[::FRAME_FACTOR]
-        second_per_grid = grid_timestamps[1] - grid_timestamps[0]
-        nframes, _, height, width = video.shape
-        min_pixels = ele.get("min_pixels", VIDEO_MIN_PIXELS)
-        total_pixels = ele.get("total_pixels", VIDEO_TOTAL_PIXELS)
-        max_pixels = max(
-            min(VIDEO_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR),
-            int(min_pixels * 1.05),
-        )
-        max_pixels = ele.get("max_pixels", max_pixels)
-        # min_pixels = (factor ** 2) * 52
-        # max_pixels = (factor ** 2) * min(768, (16384 / nframes * temporal_patch_size))
-        if "resized_height" in ele and "resized_width" in ele:
-            resized_height, resized_width = smart_resize(
-                ele["resized_height"],
-                ele["resized_width"],
-                factor=image_factor,
-            )
-        else:
-            resized_height, resized_width = smart_resize(
-                height,
-                width,
-                factor=image_factor,
-                min_pixels=min_pixels,
-                max_pixels=max_pixels,
-            )
-        video = transforms.functional.resize(
-            video,
-            [resized_height, resized_width],
-            interpolation=InterpolationMode.BICUBIC,
-            antialias=True,
-        ).float()
-        return video, total_dur, nframes, second_per_grid
-    else:
-        assert isinstance(ele["video"], (list, tuple))
-        process_info = ele.copy()
-        process_info.pop("type", None)
-        process_info.pop("video", None)
-        images = [
-            fetch_image({"image": video_element, **process_info}, size_factor=image_factor)
-            for video_element in ele["video"]
-        ]
-        nframes = ceil_by_factor(len(images), FRAME_FACTOR)
-        if len(images) < nframes:
-            images.extend([images[-1]] * (nframes - len(images)))
-        return images, None, None, None
-
-
-def extract_vision_info(conversations: list[dict] | list[list[dict]]) -> list[dict]:
-    vision_infos = []
-    if isinstance(conversations[0], dict):
-        conversations = [conversations]
-    for conversation in conversations:
-        for message in conversation:
-            if isinstance(message["content"], list):
-                for ele in message["content"]:
-                    if (
-                        "image" in ele
-                        or "image_url" in ele
-                        or "video" in ele
-                        or ele["type"] in ("image", "image_url", "video")
-                    ):
-                        vision_infos.append(ele)
-    return vision_infos
-
-
-def process_vision_info(
-    conversations: list[dict] | list[list[dict]],
-) -> tuple[list[Image.Image] | None, list[torch.Tensor | list[Image.Image]] | None]:
-    vision_infos = extract_vision_info(conversations)
-    # Read images or videos
-    image_inputs = []
-    video_inputs = []
-    for vision_info in vision_infos:
-        if "image" in vision_info or "image_url" in vision_info:
-            image_inputs.append(fetch_image(vision_info))
-        elif "video" in vision_info:
-            video_inputs.append(fetch_video(vision_info))
-        else:
-            raise ValueError("image, image_url or video should in content.")
-    if len(image_inputs) == 0:
-        image_inputs = None
-    if len(video_inputs) == 0:
-        video_inputs = None
-    return image_inputs, video_inputs

examples/offline_inference/tmp_test/ray_head.sh

@@ -1 +0,0 @@
-ray start --head --port=6399 --num-gpus=1

examples/offline_inference/tmp_test/ray_worker.sh

@@ -1 +0,0 @@
-ray start --address=10.90.67.86:6399

examples/offline_inference/tmp_test/run_multiple_prompts.sh

@@ -1,8 +0,0 @@
-python end2end.py --model Qwen/Qwen2.5-Omni-7B \
-                                 --voice-type "m02" \
-                                 --dit-ckpt none \
-                                 --bigvgan-ckpt none \
-                                 --output-wav output_audio \
-                                 --prompt_type text \
-                                 --init-sleep-seconds 0 \
-                                 --txt-prompts top100.txt

examples/offline_inference/tmp_test/run_single_prompt.sh

@@ -1,11 +0,0 @@
-export PYTHONPATH=/workspace/omni/vllm-omni/:$PYTHONPATH
-CUDA_VISIBLE_DEVICES=4 python end2end.py --model /workspace/Qwen2.5-Omni-7B/ \
-                                 --voice-type "m02" \
-                                 --dit-ckpt none \
-                                 --bigvgan-ckpt none \
-                                 --output-wav output_audio \
-                                 --prompt_type text \
-                                 --init-sleep-seconds 0 \
-				 --worker-backend ray \
-				 --ray-address auto \
-                                 --prompts "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."

examples/offline_inference/tmp_test/utils.py

@@ -1,312 +0,0 @@
-import tempfile
-from typing import Optional, Union
-from urllib.request import urlopen
-
-import librosa
-import requests
-import resampy
-import soundfile as sf
-import torch
-import torchvision.io
-from processing_omni import fetch_image, fetch_video
-from transformers import AutoConfig, AutoProcessor
-from vllm.inputs import TextPrompt
-
-from vllm_omni.inputs.data import OmniTokensPrompt
-
-# Simple caches to avoid repeated heavy HF loads per prompt
-_PROCESSOR_CACHE: dict[str, "AutoProcessor"] = {}
-_CONFIG_CACHE: dict[str, "AutoConfig"] = {}
-
-
-def get_system_prompt():
-    return {
-        "role": "system",
-        "content": [
-            {
-                "type": "text",
-                "text": (
-                    "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."
-                ),
-            }
-        ],
-    }
-
-
-def resample_wav_to_16khz(input_filepath):
-    data, original_sample_rate = sf.read(input_filepath)
-    # Only use the first channel
-    if len(data.shape) > 1:
-        data = data[:, 0]
-    # resample to 16kHz
-    data_resampled = resampy.resample(data, sr_orig=original_sample_rate, sr_new=16000)
-    return data_resampled
-
-
-def fetch_and_read_video(args, video_url: str, fps=2):
-    def read_video_with_torchvision(video_file_name: str):
-        video, audio, info = torchvision.io.read_video(
-            video_file_name,
-            start_pts=0.0,
-            end_pts=None,
-            pts_unit="sec",
-            output_format="TCHW",
-        )
-
-        total_frames, video_fps = video.size(0), info["video_fps"]
-        total_duration = round(total_frames / video_fps, 3)
-        nframes = int(total_frames / video_fps * fps)
-
-        frame_timestamps = total_duration * torch.arange(1, nframes + 1) / nframes
-        grid_timestamps = frame_timestamps[::2]
-        second_per_grid = grid_timestamps[1] - grid_timestamps[0]
-
-        idx = torch.linspace(0, video.size(0) - 1, nframes).round().long()
-        video = video[idx]
-
-        if args.legacy_omni_video:
-            return [video, total_duration, nframes, second_per_grid.item()]
-        else:
-            return video
-
-    def read_video_with_transformers(video_file_name: Union[str, list[str]]):
-        video, total_duration, nframes, second_per_grid = fetch_video({"video": video_file_name})
-        if total_duration is None and nframes is None:
-            nframes = len(video)
-            total_duration = 0.5 * nframes
-            second_per_grid = 1.0
-        if args.legacy_omni_video:
-            return [video, total_duration, nframes, second_per_grid]
-        else:
-            return video
-
-    def read_video(video_file_name: str):
-        if args.use_torchvision:
-            return read_video_with_torchvision(video_file_name)
-        else:
-            return read_video_with_transformers(video_file_name)
-
-    if isinstance(video_url, str) and video_url.startswith("http"):
-        with tempfile.NamedTemporaryFile(delete=True) as temp_video_file:
-            resp = requests.get(video_url)
-            assert resp.status_code == requests.codes.ok, (
-                f"Failed to fetch video from {video_url}, status_code:{resp.status_code}, resp:{resp}"
-            )
-
-            temp_video_file.write(urlopen(video_url).read())
-            temp_video_file_path = temp_video_file.name
-            video_file_name = temp_video_file_path
-            return read_video(video_file_name)
-    else:
-        video_file_name = video_url
-        return read_video(video_file_name)
-
-
-def make_inputs_qwen2_omni(
-    args,
-    messages: list[dict[str, Union[str, list[dict[str, str]]]]],
-    use_audio_in_video: Optional[bool] = False,
-    tokenize: bool = False,
-) -> Union[OmniTokensPrompt, TextPrompt]:
-    from transformers import AutoConfig, AutoProcessor
-
-    # Cached processor/config to prevent per-prompt reloading and repeated warnings
-    if args.model not in _PROCESSOR_CACHE:
-        _PROCESSOR_CACHE[args.model] = AutoProcessor.from_pretrained(args.model)
-    processor = _PROCESSOR_CACHE[args.model]
-
-    config = _CONFIG_CACHE.get(args.model)
-    if config is None:
-        try:
-            config = AutoConfig.from_pretrained(args.model)
-        except Exception:
-            config = None
-        _CONFIG_CACHE[args.model] = config  # cache even if None to avoid retry storms
-
-    # Decide legacy flag only once based on config (default True if unknown)
-    if getattr(args, "legacy_omni_video", None) is None:
-        if config is not None and hasattr(config, "architectures"):
-            args.legacy_omni_video = "Qwen2_5OmniModel" not in config.architectures
-        else:
-            args.legacy_omni_video = True
-
-    audios, images, videos = [], [], []
-    for message in messages:
-        if not isinstance(message["content"], list):
-            message["content"] = [
-                {
-                    "type": "text",
-                    "text": message["content"],
-                }
-            ]
-        index, num_contents = 0, len(message["content"])
-        while index < num_contents:
-            ele = message["content"][index]
-            if "type" not in ele:
-                if "text" in ele:
-                    ele["type"] = "text"
-                elif "audio" in ele:
-                    ele["type"] = "audio"
-                elif "audio_url" in ele:
-                    ele["type"] = "audio_url"
-                elif "image" in ele:
-                    ele["type"] = "image"
-                elif "image_url" in ele:
-                    ele["type"] = "image_url"
-                elif "video" in ele:
-                    ele["type"] = "video"
-                elif "video_url" in ele:
-                    ele["type"] = "video_url"
-                else:
-                    raise ValueError(f"Unknown ele: {ele}")
-
-            if ele["type"] == "audio" or ele["type"] == "audio_url":
-                if "audio_url" in ele:
-                    audio_key = "audio_url"
-                    with tempfile.NamedTemporaryFile(delete=True) as temp_audio_file:
-                        temp_audio_file.write(urlopen(ele[audio_key]).read())
-                        temp_audio_file_path = temp_audio_file.name
-                        audios.append(resample_wav_to_16khz(temp_audio_file_path))
-                        ele["audio"] = temp_audio_file_path
-                elif "audio" in ele:
-                    audio_key = "audio"
-                    audios.append(resample_wav_to_16khz(ele[audio_key]))
-                else:
-                    raise ValueError(f"Unknown ele {ele}")
-            elif use_audio_in_video and (ele["type"] == "video" or ele["type"] == "video_url"):
-                # use video as audio as well
-                if "video_url" in ele:
-                    audio_key = "video_url"
-                    with tempfile.NamedTemporaryFile(delete=True) as temp_video_file:
-                        temp_video_file.write(urlopen(ele[audio_key]).read())
-                        temp_video_file_path = temp_video_file.name
-                        ele[audio_key] = temp_video_file_path
-                        audios.append(librosa.load(temp_video_file_path, sr=16000)[0])
-                        videos.append(fetch_and_read_video(args, temp_video_file_path))
-                        ele["video"] = temp_video_file_path
-                elif "video" in ele:
-                    audio_key = "video"
-                    audios.append(librosa.load(ele[audio_key], sr=16000)[0])
-                    videos.append(fetch_and_read_video(args, audio_key))
-                else:
-                    raise ValueError(f"Unknown ele {ele}")
-                # insert a audio after the video
-                message["content"].insert(
-                    index + 1,
-                    {
-                        "type": "audio",
-                        "audio": ele[audio_key],
-                    },
-                )
-                # no need to load the added audio again
-                index += 1
-            elif ele["type"] == "video" or ele["type"] == "video_url":
-                if "video_url" in ele:
-                    video_key = "video_url"
-                    with tempfile.NamedTemporaryFile(delete=True) as temp_video_file:
-                        temp_video_file.write(urlopen(ele["video_url"]).read())
-                        temp_video_file_path = temp_video_file.name
-                        videos.append(fetch_and_read_video(args, temp_video_file))
-                        ele["video"] = temp_video_file_path
-                else:
-                    video_key = "video"
-                    videos.append(fetch_and_read_video(args, ele[video_key]))
-            elif ele["type"] == "image" or ele["type"] == "image_url":
-                images.append(fetch_image(ele))
-
-            # move to the next content
-            index += 1
-
-    prompt = processor.apply_chat_template(
-        messages,
-        tokenize=tokenize,
-        add_generation_prompt=True,
-        add_vision_id=True,
-    )
-
-    audios = audios if len(audios) > 0 else None
-    images = images if len(images) > 0 else None
-    videos = videos if len(videos) > 0 else None
-
-    multi_modal_data = {}
-    if audios:
-        multi_modal_data["audio"] = audios
-    if images:
-        multi_modal_data["image"] = images
-    if videos:
-        multi_modal_data["video"] = videos
-
-    if isinstance(prompt, list) and isinstance(prompt[0], (list, str)):
-        prompt = prompt[0]
-
-    if tokenize:
-        return OmniTokensPrompt(
-            prompt_token_ids=prompt,
-            multi_modal_data=multi_modal_data,
-        )
-    else:
-        return TextPrompt(
-            prompt=prompt,
-            multi_modal_data=multi_modal_data,
-        )
-
-
-def make_text_prompt(args, prompt):
-    messages = [
-        get_system_prompt(),
-        {
-            "role": "user",
-            "content": [
-                {"type": "text", "text": prompt},
-            ],
-        },
-    ]
-
-    prompt = make_inputs_qwen2_omni(args, messages, tokenize=args.tokenize)
-    return prompt
-
-
-def make_audio_in_video_v2_prompt(args):
-    messages = [
-        {
-            "role": "system",
-            "content": [
-                {
-                    "type": "text",
-                    "text": (
-                        "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": "video_url",
-                    "video_url": ("https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-Omni/draw_small.mp4"),
-                },
-            ],
-        },
-    ]
-    prompt = make_inputs_qwen2_omni(
-        args,
-        messages,
-        use_audio_in_video=True,
-        tokenize=args.tokenize,
-    )
-    return prompt
-
-
-def make_omni_prompt(args, prompt=None) -> Union[OmniTokensPrompt, list[OmniTokensPrompt]]:
-    if args.prompt_type == "text":
-        prompt = make_text_prompt(args, prompt)
-    elif args.prompt_type == "audio-in-video-v2":
-        prompt = make_audio_in_video_v2_prompt(args)
-    else:
-        raise ValueError(f"Unsupported prompt type: {args.prompt_type}")
-    return prompt

examples/online_serving/qwen2_5_omni/README.md

@@ -6,7 +6,8 @@ Please refer to [README.md](../../../README.md)
 
 ## Run examples (Qwen2.5-Omni)
 
-Launch the server
+### Launch the Server
+
 ```bash
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
 ```
@@ -16,29 +17,59 @@ If you have custom stage configs file, launch the server with command below
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
 ```
 
+### Send Multi-modal Request
+
 Get into the example folder
 ```bash
 cd examples/online_serving/qwen2_5_omni
 ```
 
-Send request via python
+####  Send request via python
+
 ```bash
 python openai_chat_completion_client_for_multimodal_generation.py --query-type mixed_modalities
 ```
 
-Send request via curl
+The Python client supports the following command-line arguments:
+
+- `--query-type` (or `-q`): Query type (default: `mixed_modalities`)
+  - Options: `mixed_modalities`, `use_audio_in_video`, `multi_audios`, `text`
+- `--video-path` (or `-v`): Path to local video file or URL
+  - If not provided and query-type uses video, uses default video URL
+  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
+  - Example: `--video-path /path/to/video.mp4` or `--video-path https://example.com/video.mp4`
+- `--image-path` (or `-i`): Path to local image file or URL
+  - If not provided and query-type uses image, uses default image URL
+  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
+  - Supports common image formats: JPEG, PNG, GIF, WebP
+  - Example: `--image-path /path/to/image.jpg` or `--image-path https://example.com/image.png`
+- `--audio-path` (or `-a`): Path to local audio file or URL
+  - If not provided and query-type uses audio, uses default audio URL
+  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
+  - Supports common audio formats: MP3, WAV, OGG, FLAC, M4A
+  - Example: `--audio-path /path/to/audio.wav` or `--audio-path https://example.com/audio.mp3`
+- `--prompt` (or `-p`): Custom text prompt/question
+  - If not provided, uses default prompt for the selected query type
+  - Example: `--prompt "What are the main activities shown in this video?"`
+
+
+For example, to use mixed modalities with all local files:
+
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py \
+    --query-type mixed_modalities \
+    --video-path /path/to/your/video.mp4 \
+    --image-path /path/to/your/image.jpg \
+    --audio-path /path/to/your/audio.wav \
+    --prompt "Analyze all the media content and provide a comprehensive summary."
+```
+
+####  Send request via curl
+
 ```bash
 bash run_curl_multimodal_generation.sh mixed_modalities
 ```
 
-### FAQ
-
-If you encounter error about backend of librosa, try to install ffmpeg with command below.
-```
-sudo apt update
-sudo apt install ffmpeg
-```
-
 ## Run Local Web UI Demo
 
 This Web UI demo allows users to interact with the model through a web browser.
@@ -53,21 +84,18 @@ python gradio_demo.py  --model Qwen/Qwen2.5-Omni-7B --port 7861
 
 Then open `http://localhost:7861/` on your local browser to interact with the web UI.
 
+The gradio script supports the following arguments:
 
-### Options
+- `--model`: Model name
+- `--ip`: Host/IP for Gradio server (default: 127.0.0.1)
+- `--port`: Port for Gradio server (default: 7861)
+- `--stage-configs-path`: Path to custom stage configs YAML file (optional)
+- `--share`: Share the Gradio demo publicly (creates a public link)
 
-You can customize its basic launch parameters:
+### FAQ
 
-```bash
-python gradio_demo.py \
-    --model Qwen/Qwen2.5-Omni-7B \
-    --ip 127.0.0.1 \
-    --port 7861 \
-    --stage-configs-path /path/to/stage_configs.yaml
+If you encounter error about backend of librosa, try to install ffmpeg with command below.
+```
+sudo apt update
+sudo apt install ffmpeg
 ```
-
-- `--model`: Local model checkpoint to load (default `Qwen/Qwen2.5-Omni-7B`).
-- `--ip`: Host/IP for the Gradio server (default `127.0.0.1`).
-- `--port`: Port for the Gradio server (default `7861`).
-- `--stage-configs-path`: Optional path to custom stage configs YAML.
-- `--share`: Set to expose a temporary public link via Gradio.

examples/online_serving/qwen2_5_omni/openai_chat_completion_client_for_multimodal_generation.py

@@ -1,4 +1,6 @@
 import base64
+import os
+from typing import Optional
 
 import requests
 from openai import OpenAI
@@ -28,6 +30,129 @@ def encode_base64_content_from_url(content_url: str) -> str:
     return result
 
 
+def encode_base64_content_from_file(file_path: str) -> str:
+    """Encode a local file to base64 format."""
+    with open(file_path, "rb") as f:
+        content = f.read()
+        result = base64.b64encode(content).decode("utf-8")
+    return result
+
+
+def get_video_url_from_path(video_path: Optional[str]) -> str:
+    """Convert a video path (local file or URL) to a video URL format for the API.
+
+    If video_path is None or empty, returns the default URL.
+    If video_path is a local file path, encodes it to base64 data URL.
+    If video_path is a URL, returns it as-is.
+    """
+    if not video_path:
+        # Default video URL
+        return "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
+
+    # Check if it's a URL (starts with http:// or https://)
+    if video_path.startswith(("http://", "https://")):
+        return video_path
+
+    # Otherwise, treat it as a local file path
+    if not os.path.exists(video_path):
+        raise FileNotFoundError(f"Video file not found: {video_path}")
+
+    # Detect video MIME type from file extension
+    video_path_lower = video_path.lower()
+    if video_path_lower.endswith(".mp4"):
+        mime_type = "video/mp4"
+    elif video_path_lower.endswith(".webm"):
+        mime_type = "video/webm"
+    elif video_path_lower.endswith(".mov"):
+        mime_type = "video/quicktime"
+    elif video_path_lower.endswith(".avi"):
+        mime_type = "video/x-msvideo"
+    elif video_path_lower.endswith(".mkv"):
+        mime_type = "video/x-matroska"
+    else:
+        # Default to mp4 if extension is unknown
+        mime_type = "video/mp4"
+
+    video_base64 = encode_base64_content_from_file(video_path)
+    return f"data:{mime_type};base64,{video_base64}"
+
+
+def get_image_url_from_path(image_path: Optional[str]) -> str:
+    """Convert an image path (local file or URL) to an image URL format for the API.
+
+    If image_path is None or empty, returns the default URL.
+    If image_path is a local file path, encodes it to base64 data URL.
+    If image_path is a URL, returns it as-is.
+    """
+    if not image_path:
+        # Default image URL
+        return "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
+
+    # Check if it's a URL (starts with http:// or https://)
+    if image_path.startswith(("http://", "https://")):
+        return image_path
+
+    # Otherwise, treat it as a local file path
+    if not os.path.exists(image_path):
+        raise FileNotFoundError(f"Image file not found: {image_path}")
+
+    # Detect image MIME type from file extension
+    image_path_lower = image_path.lower()
+    if image_path_lower.endswith((".jpg", ".jpeg")):
+        mime_type = "image/jpeg"
+    elif image_path_lower.endswith(".png"):
+        mime_type = "image/png"
+    elif image_path_lower.endswith(".gif"):
+        mime_type = "image/gif"
+    elif image_path_lower.endswith(".webp"):
+        mime_type = "image/webp"
+    else:
+        # Default to jpeg if extension is unknown
+        mime_type = "image/jpeg"
+
+    image_base64 = encode_base64_content_from_file(image_path)
+    return f"data:{mime_type};base64,{image_base64}"
+
+
+def get_audio_url_from_path(audio_path: Optional[str]) -> str:
+    """Convert an audio path (local file or URL) to an audio URL format for the API.
+
+    If audio_path is None or empty, returns the default URL.
+    If audio_path is a local file path, encodes it to base64 data URL.
+    If audio_path is a URL, returns it as-is.
+    """
+    if not audio_path:
+        # Default audio URL
+        return AudioAsset("mary_had_lamb").url
+
+    # Check if it's a URL (starts with http:// or https://)
+    if audio_path.startswith(("http://", "https://")):
+        return audio_path
+
+    # Otherwise, treat it as a local file path
+    if not os.path.exists(audio_path):
+        raise FileNotFoundError(f"Audio file not found: {audio_path}")
+
+    # Detect audio MIME type from file extension
+    audio_path_lower = audio_path.lower()
+    if audio_path_lower.endswith((".mp3", ".mpeg")):
+        mime_type = "audio/mpeg"
+    elif audio_path_lower.endswith(".wav"):
+        mime_type = "audio/wav"
+    elif audio_path_lower.endswith(".ogg"):
+        mime_type = "audio/ogg"
+    elif audio_path_lower.endswith(".flac"):
+        mime_type = "audio/flac"
+    elif audio_path_lower.endswith(".m4a"):
+        mime_type = "audio/mp4"
+    else:
+        # Default to wav if extension is unknown
+        mime_type = "audio/wav"
+
+    audio_base64 = encode_base64_content_from_file(audio_path)
+    return f"data:{mime_type};base64,{audio_base64}"
+
+
 def get_system_prompt():
     return {
         "role": "system",
@@ -44,8 +169,10 @@ def get_system_prompt():
     }
 
 
-def get_text_query():
-    question = "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
+def get_text_query(custom_prompt: Optional[str] = None):
+    question = (
+        custom_prompt or "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
+    )
     prompt = {
         "role": "user",
         "content": [
@@ -58,26 +185,32 @@ def get_text_query():
     return prompt
 
 
-def get_mixed_modalities_query():
-    question = "What is recited in the audio? What is the content of this image? Why is this video funny?"
+def get_mixed_modalities_query(
+    video_path: Optional[str] = None,
+    image_path: Optional[str] = None,
+    audio_path: Optional[str] = None,
+    custom_prompt: Optional[str] = None,
+):
+    question = (
+        custom_prompt or "What is recited in the audio? What is the content of this image? Why is this video funny?"
+    )
+    video_url = get_video_url_from_path(video_path)
+    image_url = get_image_url_from_path(image_path)
+    audio_url = get_audio_url_from_path(audio_path)
     prompt = {
         "role": "user",
         "content": [
             {
                 "type": "audio_url",
-                "audio_url": {"url": AudioAsset("mary_had_lamb").url},
+                "audio_url": {"url": audio_url},
             },
             {
                 "type": "image_url",
-                "image_url": {
-                    "url": "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
-                },
+                "image_url": {"url": image_url},
             },
             {
                 "type": "video_url",
-                "video_url": {
-                    "url": "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
-                },
+                "video_url": {"url": video_url},
             },
             {
                 "type": "text",
@@ -89,8 +222,9 @@ def get_mixed_modalities_query():
     return prompt
 
 
-def get_use_audio_in_video_query():
-    question = "Describe the content of the video, then convert what the baby say into text."
+def get_use_audio_in_video_query(video_path: Optional[str] = None, custom_prompt: Optional[str] = None):
+    question = custom_prompt or "Describe the content of the video, then convert what the baby say into text."
+    video_url = get_video_url_from_path(video_path)
 
     prompt = {
         "role": "user",
@@ -98,7 +232,7 @@ def get_use_audio_in_video_query():
             {
                 "type": "video_url",
                 "video_url": {
-                    "url": "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4",
+                    "url": video_url,
                     "num_frames": 16,
                 },
             },
@@ -112,14 +246,15 @@ def get_use_audio_in_video_query():
     return prompt
 
 
-def get_multi_audios_query():
-    question = "Are these two audio clips the same?"
+def get_multi_audios_query(audio_path: Optional[str] = None, custom_prompt: Optional[str] = None):
+    question = custom_prompt or "Are these two audio clips the same?"
+    audio_url = get_audio_url_from_path(audio_path)
     prompt = {
         "role": "user",
         "content": [
             {
                 "type": "audio_url",
-                "audio_url": {"url": AudioAsset("mary_had_lamb").url},
+                "audio_url": {"url": audio_url},
             },
             {
                 "type": "audio_url",
@@ -179,7 +314,27 @@ def run_multimodal_generation(args) -> None:
         code2wav_sampling_params,
     ]
 
-    prompt = query_map[args.query_type]()
+    # Get paths and custom prompt from args
+    video_path = getattr(args, "video_path", None)
+    image_path = getattr(args, "image_path", None)
+    audio_path = getattr(args, "audio_path", None)
+    custom_prompt = getattr(args, "prompt", None)
+
+    # Get the query function and call it with appropriate parameters
+    query_func = query_map[args.query_type]
+    if args.query_type == "mixed_modalities":
+        prompt = query_func(
+            video_path=video_path, image_path=image_path, audio_path=audio_path, custom_prompt=custom_prompt
+        )
+    elif args.query_type == "use_audio_in_video":
+        prompt = query_func(video_path=video_path, custom_prompt=custom_prompt)
+    elif args.query_type == "multi_audios":
+        prompt = query_func(audio_path=audio_path, custom_prompt=custom_prompt)
+    elif args.query_type == "text":
+        prompt = query_func(custom_prompt=custom_prompt)
+    else:
+        prompt = query_func()
+
     extra_body = {
         "sampling_params_list": sampling_params_list  # Optional, it has a default setting in stage_configs of the corresponding model.
     }
@@ -219,6 +374,34 @@ def parse_args():
         choices=query_map.keys(),
         help="Query type.",
     )
+    parser.add_argument(
+        "--video-path",
+        "-v",
+        type=str,
+        default=None,
+        help="Path to local video file or URL. If not provided and query-type uses video, uses default video URL.",
+    )
+    parser.add_argument(
+        "--image-path",
+        "-i",
+        type=str,
+        default=None,
+        help="Path to local image file or URL. If not provided and query-type uses image, uses default image URL.",
+    )
+    parser.add_argument(
+        "--audio-path",
+        "-a",
+        type=str,
+        default=None,
+        help="Path to local audio file or URL. If not provided and query-type uses audio, uses default audio URL.",
+    )
+    parser.add_argument(
+        "--prompt",
+        "-p",
+        type=str,
+        default=None,
+        help="Custom text prompt/question to use instead of the default prompt for the selected query type.",
+    )
 
     return parser.parse_args()
 

examples/online_serving/qwen3_omni/README.md

@@ -6,7 +6,8 @@ Please refer to [README.md](../../../README.md)
 
 ## Run examples (Qwen3-Omni)
 
-Launch the server
+### Launch the Server
+
 ```bash
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
@@ -16,21 +17,59 @@ If you have custom stage configs file, launch the server with command below
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
 ```
 
+### Send Multi-modal Request
+
 Get into the example folder
 ```bash
 cd examples/online_serving/qwen3_omni
 ```
 
-Send request via python
+####  Send request via python
+
 ```bash
 python openai_chat_completion_client_for_multimodal_generation.py --query-type use_image
 ```
 
-Send request via curl
+The Python client supports the following command-line arguments:
+
+- `--query-type` (or `-q`): Query type (default: `use_video`)
+  - Options: `text`, `use_audio`, `use_image`, `use_video`
+- `--model` (or `-m`): Model name/path (default: `Qwen/Qwen3-Omni-30B-A3B-Instruct`)
+- `--video-path` (or `-v`): Path to local video file or URL
+  - If not provided and query-type is `use_video`, uses default video URL
+  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
+  - Example: `--video-path /path/to/video.mp4` or `--video-path https://example.com/video.mp4`
+- `--image-path` (or `-i`): Path to local image file or URL
+  - If not provided and query-type is `use_image`, uses default image URL
+  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
+  - Supports common image formats: JPEG, PNG, GIF, WebP
+  - Example: `--image-path /path/to/image.jpg` or `--image-path https://example.com/image.png`
+- `--audio-path` (or `-a`): Path to local audio file or URL
+  - If not provided and query-type is `use_audio`, uses default audio URL
+  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
+  - Supports common audio formats: MP3, WAV, OGG, FLAC, M4A
+  - Example: `--audio-path /path/to/audio.wav` or `--audio-path https://example.com/audio.mp3`
+- `--prompt` (or `-p`): Custom text prompt/question
+  - If not provided, uses default prompt for the selected query type
+  - Example: `--prompt "What are the main activities shown in this video?"`
+
+
+For example, to use a local video file with custom prompt:
+
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py \
+    --query-type use_video \
+    --video-path /path/to/your/video.mp4 \
+    --prompt "What are the main activities shown in this video?"
+```
+
+####  Send request via curl
+
 ```bash
 bash run_curl_multimodal_generation.sh use_image
 ```
 
+
 ### FAQ
 
 If you encounter error about backend of librosa, try to install ffmpeg with command below.
@@ -54,23 +93,9 @@ python gradio_demo.py  --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 7861
 Then open `http://localhost:7861/` on your local browser to interact with the web UI.
 
 
-### Options
-
-The gradio demo also supports running with an existing API server and can be customized with the following arguments.
-
-
-```bash
-python gradio_demo.py \
-    --model Qwen/Qwen3-Omni-30B-A3B-Instruct \
-    --use-api-server \
-    --api-base http://localhost:8091/v1 \
-    --ip 127.0.0.1 \
-    --port 7861
-```
+The gradio script supports the following arguments:
 
 - `--model`: Model name
-- `--use-api-server`: If set, connect to an existing vLLM HTTP API server instead of running AsyncOmni locally.
-- `--api-base`: Base URL for vllm serve (only used when `use-api-server` is set, default: http://localhost:8091/v1)
 - `--ip`: Host/IP for Gradio server (default: 127.0.0.1)
 - `--port`: Port for Gradio server (default: 7861)
 - `--stage-configs-path`: Path to custom stage configs YAML file (optional)

examples/online_serving/qwen3_omni/gradio_demo.py

@@ -264,7 +264,7 @@ async def run_inference_async_omni(
         # Process audio
         audio_data = process_audio_file(audio_file)
         if audio_data is not None:
-            multi_modal_data["audio"] = audio_data
+            multi_modal_data["audio"] = [audio_data] if not isinstance(audio_data, list) else audio_data
 
         # Process image
         image_data = process_image_file(image_file)
@@ -282,7 +282,9 @@ async def run_inference_async_omni(
                 video_entry = video_frames
             multi_modal_data["video"] = video_entry
             if use_audio_in_video and extracted_audio is not None and "audio" not in multi_modal_data:
-                multi_modal_data["audio"] = extracted_audio
+                multi_modal_data["audio"] = (
+                    [extracted_audio] if not isinstance(extracted_audio, list) else extracted_audio
+                )
                 mm_processor_kwargs["use_audio_in_video"] = True
 
         # Build the prompt input

examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py

@@ -1,5 +1,6 @@
 import base64
-from typing import NamedTuple
+import os
+from typing import NamedTuple, Optional
 
 import requests
 from openai import OpenAI
@@ -34,6 +35,129 @@ def encode_base64_content_from_url(content_url: str) -> str:
     return result
 
 
+def encode_base64_content_from_file(file_path: str) -> str:
+    """Encode a local file to base64 format."""
+    with open(file_path, "rb") as f:
+        content = f.read()
+        result = base64.b64encode(content).decode("utf-8")
+    return result
+
+
+def get_video_url_from_path(video_path: Optional[str]) -> str:
+    """Convert a video path (local file or URL) to a video URL format for the API.
+
+    If video_path is None or empty, returns the default URL.
+    If video_path is a local file path, encodes it to base64 data URL.
+    If video_path is a URL, returns it as-is.
+    """
+    if not video_path:
+        # Default video URL
+        return "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
+
+    # Check if it's a URL (starts with http:// or https://)
+    if video_path.startswith(("http://", "https://")):
+        return video_path
+
+    # Otherwise, treat it as a local file path
+    if not os.path.exists(video_path):
+        raise FileNotFoundError(f"Video file not found: {video_path}")
+
+    # Detect video MIME type from file extension
+    video_path_lower = video_path.lower()
+    if video_path_lower.endswith(".mp4"):
+        mime_type = "video/mp4"
+    elif video_path_lower.endswith(".webm"):
+        mime_type = "video/webm"
+    elif video_path_lower.endswith(".mov"):
+        mime_type = "video/quicktime"
+    elif video_path_lower.endswith(".avi"):
+        mime_type = "video/x-msvideo"
+    elif video_path_lower.endswith(".mkv"):
+        mime_type = "video/x-matroska"
+    else:
+        # Default to mp4 if extension is unknown
+        mime_type = "video/mp4"
+
+    video_base64 = encode_base64_content_from_file(video_path)
+    return f"data:{mime_type};base64,{video_base64}"
+
+
+def get_image_url_from_path(image_path: Optional[str]) -> str:
+    """Convert an image path (local file or URL) to an image URL format for the API.
+
+    If image_path is None or empty, returns the default URL.
+    If image_path is a local file path, encodes it to base64 data URL.
+    If image_path is a URL, returns it as-is.
+    """
+    if not image_path:
+        # Default image URL
+        return "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
+
+    # Check if it's a URL (starts with http:// or https://)
+    if image_path.startswith(("http://", "https://")):
+        return image_path
+
+    # Otherwise, treat it as a local file path
+    if not os.path.exists(image_path):
+        raise FileNotFoundError(f"Image file not found: {image_path}")
+
+    # Detect image MIME type from file extension
+    image_path_lower = image_path.lower()
+    if image_path_lower.endswith((".jpg", ".jpeg")):
+        mime_type = "image/jpeg"
+    elif image_path_lower.endswith(".png"):
+        mime_type = "image/png"
+    elif image_path_lower.endswith(".gif"):
+        mime_type = "image/gif"
+    elif image_path_lower.endswith(".webp"):
+        mime_type = "image/webp"
+    else:
+        # Default to jpeg if extension is unknown
+        mime_type = "image/jpeg"
+
+    image_base64 = encode_base64_content_from_file(image_path)
+    return f"data:{mime_type};base64,{image_base64}"
+
+
+def get_audio_url_from_path(audio_path: Optional[str]) -> str:
+    """Convert an audio path (local file or URL) to an audio URL format for the API.
+
+    If audio_path is None or empty, returns the default URL.
+    If audio_path is a local file path, encodes it to base64 data URL.
+    If audio_path is a URL, returns it as-is.
+    """
+    if not audio_path:
+        # Default audio URL
+        return AudioAsset("mary_had_lamb").url
+
+    # Check if it's a URL (starts with http:// or https://)
+    if audio_path.startswith(("http://", "https://")):
+        return audio_path
+
+    # Otherwise, treat it as a local file path
+    if not os.path.exists(audio_path):
+        raise FileNotFoundError(f"Audio file not found: {audio_path}")
+
+    # Detect audio MIME type from file extension
+    audio_path_lower = audio_path.lower()
+    if audio_path_lower.endswith((".mp3", ".mpeg")):
+        mime_type = "audio/mpeg"
+    elif audio_path_lower.endswith(".wav"):
+        mime_type = "audio/wav"
+    elif audio_path_lower.endswith(".ogg"):
+        mime_type = "audio/ogg"
+    elif audio_path_lower.endswith(".flac"):
+        mime_type = "audio/flac"
+    elif audio_path_lower.endswith(".m4a"):
+        mime_type = "audio/mp4"
+    else:
+        # Default to wav if extension is unknown
+        mime_type = "audio/wav"
+
+    audio_base64 = encode_base64_content_from_file(audio_path)
+    return f"data:{mime_type};base64,{audio_base64}"
+
+
 def get_system_prompt():
     return {
         "role": "system",
@@ -50,8 +174,10 @@ def get_system_prompt():
     }
 
 
-def get_text_query():
-    question = "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
+def get_text_query(custom_prompt: Optional[str] = None):
+    question = (
+        custom_prompt or "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
+    )
     prompt = {
         "role": "user",
         "content": [
@@ -71,16 +197,15 @@ default_system = (
 )
 
 
-def get_video_query() -> QueryResult:
-    question = "Why is this video funny?"
+def get_video_query(video_path: Optional[str] = None, custom_prompt: Optional[str] = None):
+    question = custom_prompt or "Why is this video funny?"
+    video_url = get_video_url_from_path(video_path)
     prompt = {
         "role": "user",
         "content": [
             {
                 "type": "video_url",
-                "video_url": {
-                    "url": "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
-                },
+                "video_url": {"url": video_url},
             },
             {
                 "type": "text",
@@ -91,16 +216,15 @@ def get_video_query() -> QueryResult:
     return prompt
 
 
-def get_image_query() -> QueryResult:
-    question = "What is the content of this image?"
+def get_image_query(image_path: Optional[str] = None, custom_prompt: Optional[str] = None):
+    question = custom_prompt or "What is the content of this image?"
+    image_url = get_image_url_from_path(image_path)
     prompt = {
         "role": "user",
         "content": [
             {
                 "type": "image_url",
-                "image_url": {
-                    "url": "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
-                },
+                "image_url": {"url": image_url},
             },
             {
                 "type": "text",
@@ -111,14 +235,15 @@ def get_image_query() -> QueryResult:
     return prompt
 
 
-def get_audio_query() -> QueryResult:
-    question = "What is the content of this audio?"
+def get_audio_query(audio_path: Optional[str] = None, custom_prompt: Optional[str] = None):
+    question = custom_prompt or "What is the content of this audio?"
+    audio_url = get_audio_url_from_path(audio_path)
     prompt = {
         "role": "user",
         "content": [
             {
                 "type": "audio_url",
-                "audio_url": {"url": AudioAsset("mary_had_lamb").url},
+                "audio_url": {"url": audio_url},
             },
             {
                 "type": "text",
@@ -178,7 +303,25 @@ def run_multimodal_generation(args) -> None:
         code2wav_sampling_params,
     ]
 
-    prompt = query_map[args.query_type]()
+    # Get paths and custom prompt from args
+    video_path = getattr(args, "video_path", None)
+    image_path = getattr(args, "image_path", None)
+    audio_path = getattr(args, "audio_path", None)
+    custom_prompt = getattr(args, "prompt", None)
+
+    # Get the query function and call it with appropriate parameters
+    query_func = query_map[args.query_type]
+    if args.query_type == "use_video":
+        prompt = query_func(video_path=video_path, custom_prompt=custom_prompt)
+    elif args.query_type == "use_image":
+        prompt = query_func(image_path=image_path, custom_prompt=custom_prompt)
+    elif args.query_type == "use_audio":
+        prompt = query_func(audio_path=audio_path, custom_prompt=custom_prompt)
+    elif args.query_type == "text":
+        prompt = query_func(custom_prompt=custom_prompt)
+    else:
+        prompt = query_func()
+
     extra_body = {
         "sampling_params_list": sampling_params_list  # Optional, it has a default setting in stage_configs of the corresponding model.
     }
@@ -225,6 +368,34 @@ def parse_args():
         default="Qwen/Qwen3-Omni-30B-A3B-Instruct",
         help="Model Name / Path",
     )
+    parser.add_argument(
+        "--video-path",
+        "-v",
+        type=str,
+        default=None,
+        help="Path to local video file or URL. If not provided and query-type is 'use_video', uses default video URL.",
+    )
+    parser.add_argument(
+        "--image-path",
+        "-i",
+        type=str,
+        default=None,
+        help="Path to local image file or URL. If not provided and query-type is 'use_image', uses default image URL.",
+    )
+    parser.add_argument(
+        "--audio-path",
+        "-a",
+        type=str,
+        default=None,
+        help="Path to local audio file or URL. If not provided and query-type is 'use_audio', uses default audio URL.",
+    )
+    parser.add_argument(
+        "--prompt",
+        "-p",
+        type=str,
+        default=None,
+        help="Custom text prompt/question to use instead of the default prompt for the selected query type.",
+    )
 
     return parser.parse_args()
 

examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh

@@ -41,7 +41,7 @@ code2wav_sampling_params='{
   "temperature": 0.0,
   "top_p": 1.0,
   "top_k": -1,
-  "max_tokens": 4096 * 16,
+  "max_tokens": 65536,
   "seed": 42,
   "detokenize": true,
   "repetition_penalty": 1.1

mkdocs.yml

@@ -1,5 +1,5 @@
 site_name: vLLM-Omni
-site_description: Multi-modality models inference and serving with non-autoregressive structures
+site_description: Efficient omni-modality model serving for everyone
 site_author: vLLM-Omni Team
 site_url: https://vllm-project.github.io/vllm-omni/
 

pyproject.toml

@@ -5,12 +5,12 @@ build-backend = "setuptools.build_meta"
 [project]
 name = "vllm-omni"
 version = "0.11.0rc1"
-description = "vLLM-Omni: Multi-modality models inference and serving with non-autoregressive structures"
+description = "A framework for efficient model inference with omni-modality models"
 readme = "README.md"
 requires-python = ">=3.9,<3.14"  # Align with vLLM v0.11 supported Python versions
 license = {text = "Apache-2.0"}
 authors = [
-    {name = "vLLM-Omni Team", email = "vllm-project@gmail.com"}
+    {name = "vLLM-Omni Team"}
 ]
 keywords = ["vllm", "multimodal", "diffusion", "transformer", "inference", "serving"]
 classifiers = [
@@ -27,6 +27,7 @@ classifiers = [
     "Topic :: Software Development :: Libraries :: Python Modules",
 ]
 
+
 dependencies = [
     # Core runtime dependencies (required for actual usage)
     "omegaconf>=2.3.0",

vllm_omni/diffusion/data.py

@@ -21,6 +21,8 @@ class OmniDiffusionConfig:
 
     model_class_name: str | None = None
 
+    dtype: torch.dtype = torch.bfloat16
+
     # Attention
     # attention_backend: str = None
 
@@ -80,6 +82,10 @@ class OmniDiffusionConfig:
     vae_cpu_offload: bool = True
     pin_cpu_memory: bool = True
 
+    # VAE memory optimization parameters
+    vae_use_slicing: bool = False
+    vae_use_tiling: bool = False
+
     # STA (Sliding Tile Attention) parameters
     mask_strategy_file_path: str | None = None
     # STA_mode: STA_Mode = STA_Mode.STA_INFERENCE

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py

@@ -20,7 +20,6 @@ from diffusers.schedulers.scheduling_flow_match_euler_discrete import (
 from diffusers.utils.torch_utils import randn_tensor
 from torch import nn
 from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer
-from vllm.model_executor.model_loader.utils import set_default_torch_dtype
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
@@ -236,7 +235,7 @@ class QwenImagePipeline(
     def __init__(
         self,
         *,
-        od_config: OmniDiffusionConfig = None,
+        od_config: OmniDiffusionConfig,
         prefix: str = "",
     ):
         super().__init__()
@@ -250,23 +249,14 @@ class QwenImagePipeline(
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             model, subfolder="scheduler", local_files_only=local_files_only
         )
-        logger.info("Loaded Qwen-Image scheduler successfully")
-
-        with set_default_torch_dtype(torch.bfloat16):
-            self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-                model, subfolder="text_encoder", local_files_only=local_files_only
-            )
-            logger.info("Loaded Qwen-Image text encoder successfully")
-            self.vae = AutoencoderKLQwenImage.from_pretrained(
-                model, subfolder="vae", local_files_only=local_files_only
-            ).to(self.device)
-            logger.info("Loaded Qwen-Image VAE successfully")
-            self.transformer = QwenImageTransformer2DModel()
-            logger.info("Initialized Qwen-Image transformer successfully.")
-            self.tokenizer = Qwen2Tokenizer.from_pretrained(
-                model, subfolder="tokenizer", local_files_only=local_files_only
-            )
-            logger.info("Loaded Qwen-Image tokenizer successfully.")
+        self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+            model, subfolder="text_encoder", local_files_only=local_files_only
+        )
+        self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
+            self.device
+        )
+        self.transformer = QwenImageTransformer2DModel()
+        self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
 
         self.stage = None
 

vllm_omni/diffusion/models/z_image/pipeline_z_image.py

@@ -0,0 +1,648 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright 2025 Alibaba Z-Image Team and The HuggingFace Team. All rights reserved.
+#
+# 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.
+
+import inspect
+import json
+import os
+from typing import Any, Callable
+
+import torch
+import torch.nn as nn
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.models.autoencoders import AutoencoderKL
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import logging
+from diffusers.utils.torch_utils import randn_tensor
+from transformers import AutoModel, AutoTokenizer
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.models.z_image.z_image_transformer import (
+    ZImageTransformer2DModel,
+)
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.model_executor.model_loader.weight_utils import (
+    download_weights_from_hf_specific,
+)
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+def get_post_process_func(
+    od_config: OmniDiffusionConfig,
+):
+    model_name = od_config.model
+    if os.path.exists(model_name):
+        model_path = model_name
+    else:
+        model_path = download_weights_from_hf_specific(model_name, None, ["*"])
+    vae_config_path = os.path.join(model_path, "vae/config.json")
+    with open(vae_config_path) as f:
+        vae_config = json.load(f)
+        vae_scale_factor = 2 ** (len(vae_config["block_out_channels"]) - 1) if "block_out_channels" in vae_config else 8
+
+    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2)
+
+    def post_process_func(
+        images: torch.Tensor,
+    ):
+        return image_processor.postprocess(images)
+
+    return post_process_func
+
+
+# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.15,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: int | None = None,
+    device: str | torch.device | None = None,
+    timesteps: list[int] | None = None,
+    sigmas: list[float] | None = None,
+    **kwargs,
+) -> tuple[torch.Tensor, int]:
+    r"""
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`list[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`list[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+class ZImagePipeline(nn.Module):
+    def __init__(
+        self,
+        *,
+        od_config: OmniDiffusionConfig = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.od_config = od_config
+        self._execution_device = get_local_device()
+        model = od_config.model
+        local_files_only = os.path.exists(model)
+        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+            model, subfolder="scheduler", local_files_only=local_files_only
+        )
+
+        self.text_encoder = AutoModel.from_pretrained(
+            model, subfolder="text_encoder", local_files_only=local_files_only
+        )
+        self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
+            self._execution_device
+        )
+        self.transformer = ZImageTransformer2DModel()
+        self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
+
+        self.vae_scale_factor = (
+            2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self, "vae") and self.vae is not None else 8
+        )
+
+    def encode_prompt(
+        self,
+        prompt: str | list[str],
+        device: torch.device | None = None,
+        do_classifier_free_guidance: bool = True,
+        negative_prompt: str | list[str] | None = None,
+        prompt_embeds: list[torch.FloatTensor] | None = None,
+        negative_prompt_embeds: torch.FloatTensor | None = None,
+        max_sequence_length: int = 512,
+    ):
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        prompt_embeds = self._encode_prompt(
+            prompt=prompt,
+            device=device,
+            prompt_embeds=prompt_embeds,
+            max_sequence_length=max_sequence_length,
+        )
+
+        if do_classifier_free_guidance:
+            if negative_prompt is None:
+                negative_prompt = ["" for _ in prompt]
+            else:
+                negative_prompt = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+            assert len(prompt) == len(negative_prompt)
+            negative_prompt_embeds = self._encode_prompt(
+                prompt=negative_prompt,
+                device=device,
+                prompt_embeds=negative_prompt_embeds,
+                max_sequence_length=max_sequence_length,
+            )
+        else:
+            negative_prompt_embeds = []
+        return prompt_embeds, negative_prompt_embeds
+
+    def _encode_prompt(
+        self,
+        prompt: str | list[str],
+        device: torch.device | None = None,
+        prompt_embeds: list[torch.FloatTensor] | None = None,
+        max_sequence_length: int = 512,
+    ) -> list[torch.FloatTensor]:
+        device = device or self._execution_device
+
+        if prompt_embeds is not None:
+            return prompt_embeds
+
+        if isinstance(prompt, str):
+            prompt = [prompt]
+
+        for i, prompt_item in enumerate(prompt):
+            messages = [
+                {"role": "user", "content": prompt_item},
+            ]
+            prompt_item = self.tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True,
+                enable_thinking=True,
+            )
+            prompt[i] = prompt_item
+
+        text_inputs = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=max_sequence_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+
+        text_input_ids = text_inputs.input_ids.to(device)
+        prompt_masks = text_inputs.attention_mask.to(device).bool()
+
+        prompt_embeds = self.text_encoder(
+            input_ids=text_input_ids,
+            attention_mask=prompt_masks,
+            output_hidden_states=True,
+        ).hidden_states[-2]
+
+        embeddings_list = []
+
+        for i in range(len(prompt_embeds)):
+            embeddings_list.append(prompt_embeds[i][prompt_masks[i]])
+
+        return embeddings_list
+
+    def prepare_latents(
+        self,
+        batch_size,
+        num_channels_latents,
+        height,
+        width,
+        dtype,
+        device,
+        generator,
+        latents=None,
+    ):
+        height = 2 * (int(height) // (self.vae_scale_factor * 2))
+        width = 2 * (int(width) // (self.vae_scale_factor * 2))
+
+        shape = (batch_size, num_channels_latents, height, width)
+
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        else:
+            if latents.shape != shape:
+                raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")
+            latents = latents.to(device)
+        return latents
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def do_classifier_free_guidance(self):
+        return self._guidance_scale > 1
+
+    @property
+    def joint_attention_kwargs(self):
+        return self._joint_attention_kwargs
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    @torch.no_grad()
+    def forward(
+        self,
+        req: OmniDiffusionRequest,
+        prompt: str | list[str] | None = None,
+        height: int | None = None,
+        width: int | None = None,
+        num_inference_steps: int = 50,
+        sigmas: list[float] | None = None,
+        guidance_scale: float = 5.0,
+        cfg_normalization: bool = False,
+        cfg_truncation: float = 1.0,
+        negative_prompt: str | list[str] | None = None,
+        num_images_per_prompt: int | None = 1,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.FloatTensor | None = None,
+        prompt_embeds: list[torch.FloatTensor] | None = None,
+        negative_prompt_embeds: list[torch.FloatTensor] | None = None,
+        output_type: str | None = "pil",
+        return_dict: bool = True,
+        joint_attention_kwargs: dict[str, Any] | None = None,
+        callback_on_step_end: Callable[[int, int, dict], None] | None = None,
+        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
+        max_sequence_length: int = 512,
+    ) -> DiffusionOutput:
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `list[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            height (`int`, *optional*, defaults to 1024):
+                The height in pixels of the generated image.
+            width (`int`, *optional*, defaults to 1024):
+                The width in pixels of the generated image.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            sigmas (`list[float]`, *optional*):
+                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
+                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
+                will be used.
+            guidance_scale (`float`, *optional*, defaults to 5.0):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            cfg_normalization (`bool`, *optional*, defaults to False):
+                Whether to apply configuration normalization.
+            cfg_truncation (`float`, *optional*, defaults to 1.0):
+                The truncation value for configuration.
+            negative_prompt (`str` or `list[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will be generated by sampling using the supplied random `generator`.
+            prompt_embeds (`list[torch.FloatTensor]`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`list[torch.FloatTensor]`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.stable_diffusion.ZImagePipelineOutput`] instead of a plain
+                tuple.
+            joint_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            callback_on_step_end (`Callable`, *optional*):
+                A function that calls at the end of each denoising steps during the inference. The function is called
+                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
+                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`list`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+            max_sequence_length (`int`, *optional*, defaults to 512):
+                Maximum sequence length to use with the `prompt`.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.z_image.ZImagePipelineOutput`] or `tuple`: [`~pipelines.z_image.ZImagePipelineOutput`] if
+            `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the
+            generated images.
+        """
+        prompt = req.prompt
+        negative_prompt = req.negative_prompt
+        height: int = req.height or 1024
+        width: int = req.width or 1024
+        num_inference_steps = req.num_inference_steps or 50
+        generator = req.generator
+        guidance_scale = req.guidance_scale if req.guidance_rescale is not None else guidance_scale
+        req_num_outputs = getattr(req, "num_outputs_per_prompt", None)
+        if req_num_outputs and req_num_outputs > 0:
+            num_images_per_prompt = req_num_outputs
+
+        vae_scale = self.vae_scale_factor * 2
+        if height % vae_scale != 0:
+            raise ValueError(
+                f"Height must be divisible by {vae_scale} (got {height}). "
+                f"Please adjust the height to a multiple of {vae_scale}."
+            )
+        if width % vae_scale != 0:
+            raise ValueError(
+                f"Width must be divisible by {vae_scale} (got {width}). "
+                f"Please adjust the width to a multiple of {vae_scale}."
+            )
+
+        device = self._execution_device
+
+        self._guidance_scale = guidance_scale
+        self._joint_attention_kwargs = joint_attention_kwargs
+        self._interrupt = False
+        self._cfg_normalization = cfg_normalization
+        self._cfg_truncation = cfg_truncation
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = len(prompt_embeds)
+
+        # If prompt_embeds is provided and prompt is None, skip encoding
+        if prompt_embeds is not None and prompt is None:
+            if self.do_classifier_free_guidance and negative_prompt_embeds is None:
+                raise ValueError(
+                    "When `prompt_embeds` is provided without `prompt`, "
+                    "`negative_prompt_embeds` must also be provided for classifier-free guidance."
+                )
+        else:
+            (
+                prompt_embeds,
+                negative_prompt_embeds,
+            ) = self.encode_prompt(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                do_classifier_free_guidance=self.do_classifier_free_guidance,
+                prompt_embeds=prompt_embeds,
+                negative_prompt_embeds=negative_prompt_embeds,
+                device=device,
+                max_sequence_length=max_sequence_length,
+            )
+
+        # 4. Prepare latent variables
+        num_channels_latents = self.transformer.in_channels
+
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            torch.float32,
+            device,
+            generator,
+            latents,
+        )
+
+        # Repeat prompt_embeds for num_images_per_prompt
+        if num_images_per_prompt > 1:
+            prompt_embeds = [pe for pe in prompt_embeds for _ in range(num_images_per_prompt)]
+            if self.do_classifier_free_guidance and negative_prompt_embeds:
+                negative_prompt_embeds = [npe for npe in negative_prompt_embeds for _ in range(num_images_per_prompt)]
+
+        actual_batch_size = batch_size * num_images_per_prompt
+        image_seq_len = (latents.shape[2] // 2) * (latents.shape[3] // 2)
+
+        # 5. Prepare timesteps
+        mu = calculate_shift(
+            image_seq_len,
+            self.scheduler.config.get("base_image_seq_len", 256),
+            self.scheduler.config.get("max_image_seq_len", 4096),
+            self.scheduler.config.get("base_shift", 0.5),
+            self.scheduler.config.get("max_shift", 1.15),
+        )
+        self.scheduler.sigma_min = 0.0
+        scheduler_kwargs = {"mu": mu}
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            sigmas=sigmas,
+            **scheduler_kwargs,
+        )
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+        self._num_timesteps = len(timesteps)
+
+        # 6. Denoising loop
+        for i, t in enumerate(timesteps):
+            if self.interrupt:
+                continue
+
+            # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+            timestep = t.expand(latents.shape[0])
+            timestep = (1000 - timestep) / 1000
+            # Normalized time for time-aware config (0 at start, 1 at end)
+            t_norm = timestep[0].item()
+
+            # Handle cfg truncation
+            current_guidance_scale = self.guidance_scale
+            if (
+                self.do_classifier_free_guidance
+                and self._cfg_truncation is not None
+                and float(self._cfg_truncation) <= 1
+            ):
+                if t_norm > self._cfg_truncation:
+                    current_guidance_scale = 0.0
+
+            # Run CFG only if configured AND scale is non-zero
+            apply_cfg = self.do_classifier_free_guidance and current_guidance_scale > 0
+
+            if apply_cfg:
+                latents_typed = latents.to(self.transformer.dtype)
+                latent_model_input = latents_typed.repeat(2, 1, 1, 1)
+                prompt_embeds_model_input = prompt_embeds + negative_prompt_embeds
+                timestep_model_input = timestep.repeat(2)
+            else:
+                latent_model_input = latents.to(self.transformer.dtype)
+                prompt_embeds_model_input = prompt_embeds
+                timestep_model_input = timestep
+
+            latent_model_input = latent_model_input.unsqueeze(2)
+            latent_model_input_list = list(latent_model_input.unbind(dim=0))
+
+            model_out_list = self.transformer(
+                latent_model_input_list,
+                timestep_model_input,
+                prompt_embeds_model_input,
+            )[0]
+
+            if apply_cfg:
+                # Perform CFG
+                pos_out = model_out_list[:actual_batch_size]
+                neg_out = model_out_list[actual_batch_size:]
+
+                noise_pred = []
+                for j in range(actual_batch_size):
+                    pos = pos_out[j].float()
+                    neg = neg_out[j].float()
+
+                    pred = pos + current_guidance_scale * (pos - neg)
+
+                    # 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)
+                        if new_pos_norm > max_new_norm:
+                            pred = pred * (max_new_norm / new_pos_norm)
+
+                    noise_pred.append(pred)
+
+                noise_pred = torch.stack(noise_pred, dim=0)
+            else:
+                noise_pred = torch.stack([t.float() for t in model_out_list], dim=0)
+
+            noise_pred = noise_pred.squeeze(2)
+            noise_pred = -noise_pred
+
+            # compute the previous noisy sample x_t -> x_t-1
+            latents = self.scheduler.step(noise_pred.to(torch.float32), t, latents, return_dict=False)[0]
+            assert latents.dtype == torch.float32
+
+            if callback_on_step_end is not None:
+                callback_kwargs = {}
+                for k in callback_on_step_end_tensor_inputs:
+                    callback_kwargs[k] = locals()[k]
+                callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                latents = callback_outputs.pop("latents", latents)
+                prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+                negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
+
+        if output_type == "latent":
+            image = latents
+        else:
+            latents = latents.to(self.vae.dtype)
+            latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+
+            image = self.vae.decode(latents, return_dict=False)[0]
+            # image = self.image_processor.postprocess(image, output_type=output_type)
+
+        return DiffusionOutput(output=image)
+
+    def load_weights(self):
+        self.load_transformer()
+
+    def load_transformer(self):
+        import glob
+        import os
+
+        # Define the weight iterator
+        def weight_iterator(transformer_path):
+            if not os.path.exists(transformer_path):
+                logger.warning(f"Path {transformer_path} does not exist.")
+                return
+
+            # Look for safetensors first
+            safetensors_files = glob.glob(os.path.join(transformer_path, "*.safetensors"))
+            if safetensors_files:
+                try:
+                    from safetensors.torch import load_file
+                except ImportError:
+                    logger.warning("safetensors not installed, cannot load .safetensors files.")
+                    return
+
+                for file_path in safetensors_files:
+                    state_dict = load_file(file_path)
+                    for name, tensor in state_dict.items():
+                        yield name, tensor
+            else:
+                # Fallback to bin
+                bin_files = glob.glob(os.path.join(transformer_path, "*.bin"))
+                for file_path in bin_files:
+                    state_dict = torch.load(file_path)
+                    for name, tensor in state_dict.items():
+                        yield name, tensor
+
+        try:
+            # Get model path from config or download from HF
+            model_name = self.od_config.model
+            if os.path.exists(model_name):
+                model_path = model_name
+            else:
+                model_path = download_weights_from_hf_specific(model_name, None, ["*"])
+
+            transformer_path = os.path.join(model_path, "transformer")
+            self.transformer.load_weights(weight_iterator(transformer_path))
+
+        except Exception as e:
+            print(f"An error occurred: {e}")
+            raise e

vllm_omni/diffusion/models/z_image/z_image_transformer.py

@@ -0,0 +1,668 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright 2025 Alibaba Z-Image Team and The HuggingFace Team. All rights reserved.
+#
+# 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.
+
+import math
+from collections.abc import Iterable
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.models.attention_dispatch import dispatch_attention_fn
+from torch.nn.utils.rnn import pad_sequence
+from vllm.model_executor.layers.layernorm import RMSNorm
+
+ADALN_EMBED_DIM = 256
+SEQ_MULTI_OF = 32
+
+
+class TimestepEmbedder(nn.Module):
+    def __init__(self, out_size, mid_size=None, frequency_embedding_size=256):
+        super().__init__()
+        if mid_size is None:
+            mid_size = out_size
+        self.mlp = nn.Sequential(
+            nn.Linear(
+                frequency_embedding_size,
+                mid_size,
+                bias=True,
+            ),
+            nn.SiLU(),
+            nn.Linear(
+                mid_size,
+                out_size,
+                bias=True,
+            ),
+        )
+
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half
+        )
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        weight_dtype = self.mlp[0].weight.dtype
+        if weight_dtype.is_floating_point:
+            t_freq = t_freq.to(weight_dtype)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+class ZImageAttention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        num_kv_heads: int,
+        qk_norm: bool = True,
+        eps: float = 1e-6,
+    ) -> None:
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_dim = dim // num_heads
+        self.qk_norm = qk_norm
+
+        self.to_q = nn.Linear(dim, dim, bias=False)
+        self.to_k = nn.Linear(dim, self.head_dim * num_kv_heads, bias=False)
+        self.to_v = nn.Linear(dim, self.head_dim * num_kv_heads, bias=False)
+        # TODO enable this when we refactor weight loader
+        # self.to_qkv = QKVParallelLinear(
+        #     hidden_size=dim,
+        #     head_size=self.head_dim,
+        #     total_num_heads=num_heads,
+        #     disable_tp=True,
+        # )
+
+        assert qk_norm is True
+        self.norm_q = RMSNorm(self.head_dim, eps=eps)
+        self.norm_k = RMSNorm(self.head_dim, eps=eps)
+
+        self.to_out = nn.ModuleList([nn.Linear(dim, dim, bias=False)])
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ):
+        query = self.to_q(hidden_states)
+        key = self.to_k(hidden_states)
+        value = self.to_v(hidden_states)
+        # qkv, _ = self.to_qkv(hidden_states)
+        # query, key, value = qkv.chunk(3, dim=-1)
+
+        query = query.unflatten(-1, (self.num_heads, -1))
+        key = key.unflatten(-1, (self.num_heads, -1))
+        value = value.unflatten(-1, (self.num_heads, -1))
+
+        query = self.norm_q(query)
+        key = self.norm_k(key)
+
+        def apply_rotary_emb(x_in: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
+            x = torch.view_as_complex(x_in.float().reshape(*x_in.shape[:-1], -1, 2))
+            freqs_cis = freqs_cis.unsqueeze(2)
+            x_out = torch.view_as_real(x * freqs_cis).flatten(3)
+            return x_out.type_as(x_in)
+
+        if freqs_cis is not None:
+            query = apply_rotary_emb(query, freqs_cis)
+            key = apply_rotary_emb(key, freqs_cis)
+
+        # Cast to correct dtype
+        dtype = query.dtype
+        query, key = query.to(dtype), key.to(dtype)
+
+        # From [batch, seq_len] to [batch, 1, 1, seq_len] -> broadcast to [batch, heads, seq_len, seq_len]
+        if attention_mask is not None and attention_mask.ndim == 2:
+            attention_mask = attention_mask[:, None, None, :]
+
+        # Compute joint attention
+        hidden_states = dispatch_attention_fn(
+            query,
+            key,
+            value,
+            attn_mask=attention_mask,
+            dropout_p=0.0,
+            is_causal=False,
+        )
+
+        # Reshape back
+        hidden_states = hidden_states.flatten(2, 3)
+        hidden_states = hidden_states.to(dtype)
+
+        hidden_states = self.to_out[0](hidden_states)
+
+        return hidden_states
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim: int, hidden_dim: int):
+        super().__init__()
+        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
+        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
+
+    def _forward_silu_gating(self, x1, x3):
+        return F.silu(x1) * x3
+
+    def forward(self, x):
+        return self.w2(self._forward_silu_gating(self.w1(x), self.w3(x)))
+
+
+class ZImageTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        layer_id: int,
+        dim: int,
+        n_heads: int,
+        n_kv_heads: int,
+        norm_eps: float,
+        qk_norm: bool,
+        modulation=True,
+    ):
+        super().__init__()
+        self.dim = dim
+
+        self.attention = ZImageAttention(
+            dim=dim,
+            num_heads=n_heads,
+            num_kv_heads=n_kv_heads,
+            qk_norm=qk_norm,
+            eps=1e-5,
+        )
+
+        self.feed_forward = FeedForward(dim=dim, hidden_dim=int(dim / 3 * 8))
+        self.layer_id = layer_id
+
+        self.attention_norm1 = RMSNorm(dim, eps=norm_eps)
+        self.ffn_norm1 = RMSNorm(dim, eps=norm_eps)
+
+        self.attention_norm2 = RMSNorm(dim, eps=norm_eps)
+        self.ffn_norm2 = RMSNorm(dim, eps=norm_eps)
+
+        self.modulation = modulation
+        if modulation:
+            self.adaLN_modulation = nn.Sequential(
+                nn.Linear(min(dim, ADALN_EMBED_DIM), 4 * dim, bias=True),
+            )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attn_mask: torch.Tensor,
+        freqs_cis: torch.Tensor,
+        adaln_input: Optional[torch.Tensor] = None,
+    ):
+        if self.modulation:
+            assert adaln_input is not None
+            scale_msa, gate_msa, scale_mlp, gate_mlp = self.adaLN_modulation(adaln_input).unsqueeze(1).chunk(4, dim=2)
+            gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh()
+            scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp
+
+            # Attention block
+            attn_out = self.attention(
+                self.attention_norm1(x) * scale_msa,
+                attention_mask=attn_mask,
+                freqs_cis=freqs_cis,
+            )
+            x = x + gate_msa * self.attention_norm2(attn_out)
+
+            # FFN block
+            x = x + gate_mlp * self.ffn_norm2(
+                self.feed_forward(
+                    self.ffn_norm1(x) * scale_mlp,
+                )
+            )
+        else:
+            # Attention block
+            attn_out = self.attention(
+                self.attention_norm1(x),
+                attention_mask=attn_mask,
+                freqs_cis=freqs_cis,
+            )
+            x = x + self.attention_norm2(attn_out)
+
+            # FFN block
+            x = x + self.ffn_norm2(
+                self.feed_forward(
+                    self.ffn_norm1(x),
+                )
+            )
+
+        return x
+
+
+class FinalLayer(nn.Module):
+    def __init__(self, hidden_size, out_channels):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, out_channels, bias=True)
+
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(min(hidden_size, ADALN_EMBED_DIM), hidden_size, bias=True),
+        )
+
+    def forward(self, x, c):
+        scale = 1.0 + self.adaLN_modulation(c)
+        x = self.norm_final(x) * scale.unsqueeze(1)
+        x = self.linear(x)
+        return x
+
+
+class RopeEmbedder:
+    def __init__(
+        self,
+        theta: float = 256.0,
+        axes_dims: list[int] = (16, 56, 56),
+        axes_lens: list[int] = (64, 128, 128),
+    ):
+        self.theta = theta
+        self.axes_dims = axes_dims
+        self.axes_lens = axes_lens
+        assert len(axes_dims) == len(axes_lens), "axes_dims and axes_lens must have the same length"
+        self.freqs_cis = None
+
+    @staticmethod
+    def precompute_freqs_cis(dim: list[int], end: list[int], theta: float = 256.0):
+        with torch.device("cpu"):
+            freqs_cis = []
+            for i, (d, e) in enumerate(zip(dim, end)):
+                freqs = 1.0 / (theta ** (torch.arange(0, d, 2, dtype=torch.float64, device="cpu") / d))
+                timestep = torch.arange(e, device=freqs.device, dtype=torch.float64)
+                freqs = torch.outer(timestep, freqs).float()
+                freqs_cis_i = torch.polar(torch.ones_like(freqs), freqs).to(torch.complex64)  # complex64
+                freqs_cis.append(freqs_cis_i)
+
+            return freqs_cis
+
+    def __call__(self, ids: torch.Tensor):
+        assert ids.ndim == 2
+        assert ids.shape[-1] == len(self.axes_dims)
+        device = ids.device
+
+        if self.freqs_cis is None:
+            self.freqs_cis = self.precompute_freqs_cis(self.axes_dims, self.axes_lens, theta=self.theta)
+            self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis]
+        else:
+            # Ensure freqs_cis are on the same device as ids
+            if self.freqs_cis[0].device != device:
+                self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis]
+
+        result = []
+        for i in range(len(self.axes_dims)):
+            index = ids[:, i]
+            result.append(self.freqs_cis[i][index])
+        return torch.cat(result, dim=-1)
+
+
+class ZImageTransformer2DModel(nn.Module):
+    def __init__(
+        self,
+        all_patch_size=(2,),
+        all_f_patch_size=(1,),
+        in_channels=16,
+        dim=3840,
+        n_layers=30,
+        n_refiner_layers=2,
+        n_heads=30,
+        n_kv_heads=30,
+        norm_eps=1e-5,
+        qk_norm=True,
+        cap_feat_dim=2560,
+        rope_theta=256.0,
+        t_scale=1000.0,
+        axes_dims=[32, 48, 48],
+        axes_lens=[1024, 512, 512],
+    ) -> None:
+        super().__init__()
+        self.dtype = torch.bfloat16
+        self.in_channels = in_channels
+        self.out_channels = in_channels
+        self.all_patch_size = all_patch_size
+        self.all_f_patch_size = all_f_patch_size
+        self.dim = dim
+        self.n_heads = n_heads
+
+        self.rope_theta = rope_theta
+        self.t_scale = t_scale
+        self.gradient_checkpointing = False
+
+        assert len(all_patch_size) == len(all_f_patch_size)
+
+        all_x_embedder = {}
+        all_final_layer = {}
+        for patch_idx, (patch_size, f_patch_size) in enumerate(zip(all_patch_size, all_f_patch_size)):
+            x_embedder = nn.Linear(f_patch_size * patch_size * patch_size * in_channels, dim, bias=True)
+            all_x_embedder[f"{patch_size}-{f_patch_size}"] = x_embedder
+
+            final_layer = FinalLayer(dim, patch_size * patch_size * f_patch_size * self.out_channels)
+            all_final_layer[f"{patch_size}-{f_patch_size}"] = final_layer
+
+        self.all_x_embedder = nn.ModuleDict(all_x_embedder)
+        self.all_final_layer = nn.ModuleDict(all_final_layer)
+        self.noise_refiner = nn.ModuleList(
+            [
+                ZImageTransformerBlock(
+                    1000 + layer_id,
+                    dim,
+                    n_heads,
+                    n_kv_heads,
+                    norm_eps,
+                    qk_norm,
+                    modulation=True,
+                )
+                for layer_id in range(n_refiner_layers)
+            ]
+        )
+        self.context_refiner = nn.ModuleList(
+            [
+                ZImageTransformerBlock(
+                    layer_id,
+                    dim,
+                    n_heads,
+                    n_kv_heads,
+                    norm_eps,
+                    qk_norm,
+                    modulation=False,
+                )
+                for layer_id in range(n_refiner_layers)
+            ]
+        )
+        self.t_embedder = TimestepEmbedder(min(dim, ADALN_EMBED_DIM), mid_size=1024)
+        self.cap_embedder = nn.Sequential(
+            RMSNorm(cap_feat_dim, eps=norm_eps),
+            nn.Linear(cap_feat_dim, dim, bias=True),
+        )
+
+        self.x_pad_token = nn.Parameter(torch.empty((1, dim)))
+        self.cap_pad_token = nn.Parameter(torch.empty((1, dim)))
+
+        self.layers = nn.ModuleList(
+            [
+                ZImageTransformerBlock(layer_id, dim, n_heads, n_kv_heads, norm_eps, qk_norm)
+                for layer_id in range(n_layers)
+            ]
+        )
+        self.axes_dims = axes_dims
+        self.axes_lens = axes_lens
+
+        self.rope_embedder = RopeEmbedder(theta=rope_theta, axes_dims=axes_dims, axes_lens=axes_lens)
+
+    def unpatchify(self, x: list[torch.Tensor], size: list[tuple], patch_size, f_patch_size) -> list[torch.Tensor]:
+        pH = pW = patch_size
+        pF = f_patch_size
+        bsz = len(x)
+        assert len(size) == bsz
+        for i in range(bsz):
+            F, H, W = size[i]
+            ori_len = (F // pF) * (H // pH) * (W // pW)
+            # "f h w pf ph pw c -> c (f pf) (h ph) (w pw)"
+            x[i] = (
+                x[i][:ori_len]
+                .view(F // pF, H // pH, W // pW, pF, pH, pW, self.out_channels)
+                .permute(6, 0, 3, 1, 4, 2, 5)
+                .reshape(self.out_channels, F, H, W)
+            )
+        return x
+
+    @staticmethod
+    def create_coordinate_grid(size, start=None, device=None):
+        if start is None:
+            start = (0 for _ in size)
+
+        axes = [torch.arange(x0, x0 + span, dtype=torch.int32, device=device) for x0, span in zip(start, size)]
+        grids = torch.meshgrid(axes, indexing="ij")
+        return torch.stack(grids, dim=-1)
+
+    def patchify_and_embed(
+        self,
+        all_image: list[torch.Tensor],
+        all_cap_feats: list[torch.Tensor],
+        patch_size: int,
+        f_patch_size: int,
+    ):
+        pH = pW = patch_size
+        pF = f_patch_size
+        device = all_image[0].device
+
+        all_image_out = []
+        all_image_size = []
+        all_image_pos_ids = []
+        all_image_pad_mask = []
+        all_cap_pos_ids = []
+        all_cap_pad_mask = []
+        all_cap_feats_out = []
+
+        for i, (image, cap_feat) in enumerate(zip(all_image, all_cap_feats)):
+            ### Process Caption
+            cap_ori_len = len(cap_feat)
+            cap_padding_len = (-cap_ori_len) % SEQ_MULTI_OF
+            # padded position ids
+            cap_padded_pos_ids = self.create_coordinate_grid(
+                size=(cap_ori_len + cap_padding_len, 1, 1),
+                start=(1, 0, 0),
+                device=device,
+            ).flatten(0, 2)
+            all_cap_pos_ids.append(cap_padded_pos_ids)
+            # pad mask
+            all_cap_pad_mask.append(
+                torch.cat(
+                    [
+                        torch.zeros((cap_ori_len,), dtype=torch.bool, device=device),
+                        torch.ones((cap_padding_len,), dtype=torch.bool, device=device),
+                    ],
+                    dim=0,
+                )
+            )
+            # padded feature
+            cap_padded_feat = torch.cat(
+                [cap_feat, cap_feat[-1:].repeat(cap_padding_len, 1)],
+                dim=0,
+            )
+            all_cap_feats_out.append(cap_padded_feat)
+
+            ### Process Image
+            C, F, H, W = image.size()
+            all_image_size.append((F, H, W))
+            F_tokens, H_tokens, W_tokens = F // pF, H // pH, W // pW
+
+            image = image.view(C, F_tokens, pF, H_tokens, pH, W_tokens, pW)
+            # "c f pf h ph w pw -> (f h w) (pf ph pw c)"
+            image = image.permute(1, 3, 5, 2, 4, 6, 0).reshape(F_tokens * H_tokens * W_tokens, pF * pH * pW * C)
+
+            image_ori_len = len(image)
+            image_padding_len = (-image_ori_len) % SEQ_MULTI_OF
+
+            image_ori_pos_ids = self.create_coordinate_grid(
+                size=(F_tokens, H_tokens, W_tokens),
+                start=(cap_ori_len + cap_padding_len + 1, 0, 0),
+                device=device,
+            ).flatten(0, 2)
+            image_padding_pos_ids = (
+                self.create_coordinate_grid(
+                    size=(1, 1, 1),
+                    start=(0, 0, 0),
+                    device=device,
+                )
+                .flatten(0, 2)
+                .repeat(image_padding_len, 1)
+            )
+            image_padded_pos_ids = torch.cat([image_ori_pos_ids, image_padding_pos_ids], dim=0)
+            all_image_pos_ids.append(image_padded_pos_ids)
+            # pad mask
+            all_image_pad_mask.append(
+                torch.cat(
+                    [
+                        torch.zeros((image_ori_len,), dtype=torch.bool, device=device),
+                        torch.ones((image_padding_len,), dtype=torch.bool, device=device),
+                    ],
+                    dim=0,
+                )
+            )
+            # padded feature
+            image_padded_feat = torch.cat([image, image[-1:].repeat(image_padding_len, 1)], dim=0)
+            all_image_out.append(image_padded_feat)
+
+        return (
+            all_image_out,
+            all_cap_feats_out,
+            all_image_size,
+            all_image_pos_ids,
+            all_cap_pos_ids,
+            all_image_pad_mask,
+            all_cap_pad_mask,
+        )
+
+    def forward(
+        self,
+        x: list[torch.Tensor],
+        t,
+        cap_feats: list[torch.Tensor],
+        patch_size=2,
+        f_patch_size=1,
+    ):
+        assert patch_size in self.all_patch_size
+        assert f_patch_size in self.all_f_patch_size
+
+        bsz = len(x)
+        device = x[0].device
+        t = t * self.t_scale
+        t = self.t_embedder(t)
+
+        (
+            x,
+            cap_feats,
+            x_size,
+            x_pos_ids,
+            cap_pos_ids,
+            x_inner_pad_mask,
+            cap_inner_pad_mask,
+        ) = self.patchify_and_embed(x, cap_feats, patch_size, f_patch_size)
+
+        # x embed & refine
+        x_item_seqlens = [len(_) for _ in x]
+        assert all(_ % SEQ_MULTI_OF == 0 for _ in x_item_seqlens)
+        x_max_item_seqlen = max(x_item_seqlens)
+
+        x = torch.cat(x, dim=0)
+        x = self.all_x_embedder[f"{patch_size}-{f_patch_size}"](x)
+
+        # Match t_embedder output dtype to x for layerwise casting compatibility
+        adaln_input = t.type_as(x)
+        x[torch.cat(x_inner_pad_mask)] = self.x_pad_token
+        x = list(x.split(x_item_seqlens, dim=0))
+        x_freqs_cis = list(self.rope_embedder(torch.cat(x_pos_ids, dim=0)).split(x_item_seqlens, dim=0))
+
+        x = pad_sequence(x, batch_first=True, padding_value=0.0)
+        x_freqs_cis = pad_sequence(x_freqs_cis, batch_first=True, padding_value=0.0)
+        x_attn_mask = torch.zeros((bsz, x_max_item_seqlen), dtype=torch.bool, device=device)
+        for i, seq_len in enumerate(x_item_seqlens):
+            x_attn_mask[i, :seq_len] = 1
+
+        for layer in self.noise_refiner:
+            x = layer(x, x_attn_mask, x_freqs_cis, adaln_input)
+
+        # cap embed & refine
+        cap_item_seqlens = [len(_) for _ in cap_feats]
+        assert all(_ % SEQ_MULTI_OF == 0 for _ in cap_item_seqlens)
+        cap_max_item_seqlen = max(cap_item_seqlens)
+
+        cap_feats = torch.cat(cap_feats, dim=0)
+        cap_feats = self.cap_embedder(cap_feats)
+        cap_feats[torch.cat(cap_inner_pad_mask)] = self.cap_pad_token
+        cap_feats = list(cap_feats.split(cap_item_seqlens, dim=0))
+        cap_freqs_cis = list(self.rope_embedder(torch.cat(cap_pos_ids, dim=0)).split(cap_item_seqlens, dim=0))
+
+        cap_feats = pad_sequence(cap_feats, batch_first=True, padding_value=0.0)
+        cap_freqs_cis = pad_sequence(cap_freqs_cis, batch_first=True, padding_value=0.0)
+        cap_attn_mask = torch.zeros((bsz, cap_max_item_seqlen), dtype=torch.bool, device=device)
+        for i, seq_len in enumerate(cap_item_seqlens):
+            cap_attn_mask[i, :seq_len] = 1
+
+        for layer in self.context_refiner:
+            cap_feats = layer(cap_feats, cap_attn_mask, cap_freqs_cis)
+
+        # unified
+        unified = []
+        unified_freqs_cis = []
+        for i in range(bsz):
+            x_len = x_item_seqlens[i]
+            cap_len = cap_item_seqlens[i]
+            unified.append(torch.cat([x[i][:x_len], cap_feats[i][:cap_len]]))
+            unified_freqs_cis.append(torch.cat([x_freqs_cis[i][:x_len], cap_freqs_cis[i][:cap_len]]))
+        unified_item_seqlens = [a + b for a, b in zip(cap_item_seqlens, x_item_seqlens)]
+        assert unified_item_seqlens == [len(_) for _ in unified]
+        unified_max_item_seqlen = max(unified_item_seqlens)
+
+        unified = pad_sequence(unified, batch_first=True, padding_value=0.0)
+        unified_freqs_cis = pad_sequence(unified_freqs_cis, batch_first=True, padding_value=0.0)
+        unified_attn_mask = torch.zeros((bsz, unified_max_item_seqlen), dtype=torch.bool, device=device)
+        for i, seq_len in enumerate(unified_item_seqlens):
+            unified_attn_mask[i, :seq_len] = 1
+
+        for layer in self.layers:
+            unified = layer(unified, unified_attn_mask, unified_freqs_cis, adaln_input)
+
+        unified = self.all_final_layer[f"{patch_size}-{f_patch_size}"](unified, adaln_input)
+        unified = list(unified.unbind(dim=0))
+        x = self.unpatchify(unified, x_size, patch_size, f_patch_size)
+
+        return x, {}
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
+        weights_map: dict[str, torch.Tensor] = {name: tensor for name, tensor in weights}
+        params_dict = dict(self.named_parameters())
+
+        for name, loaded_weight in weights_map.items():
+            # In ZImageAttention, to_q, to_k, to_v are separate.
+            # The weight names are like '...attention.to_q.weight'.
+            # The logic from qwen_image_transformer for stacking is not needed here.
+            # We can do a direct mapping.
+
+            # The names in the state dict seem to match the model structure.
+            # Let's try a direct load.
+            if name in params_dict:
+                param = params_dict[name]
+                with torch.no_grad():
+                    param.copy_(loaded_weight)
+            else:
+                # The provided weight list might use a different naming convention.
+                # For example, `layers.0.attention.to_out.0.weight` vs `...`
+                # Let's try to handle simple cases.
+
+                # Replace dots with hyphens for ModuleDict keys
+                # e.g. all_x_embedder.2-1.weight -> all_x_embedder.2-1.weight
+                # The names seem to match already.
+
+                # If there are any specific mismatches, they can be handled here.
+                # For now, we assume a direct mapping and warn if a weight is not found.
+                import logging
+
+                logging.warning(f"Weight {name} not found in model parameters.")

vllm_omni/diffusion/registry.py

@@ -12,6 +12,11 @@ _DIFFUSION_MODELS = {
         "pipeline_qwen_image",
         "QwenImagePipeline",
     ),
+    "ZImagePipeline": (
+        "z_image",
+        "pipeline_z_image",
+        "ZImagePipeline",
+    ),
 }
 
 
@@ -24,6 +29,13 @@ def initialize_model(
         module = importlib.import_module(module_name)
         model_class = getattr(module, cls_name)
         model = model_class(od_config=od_config, prefix=mod_relname)
+
+        # Configure VAE memory optimization settings from config
+        if hasattr(model.vae, "use_slicing"):
+            model.vae.use_slicing = od_config.vae_use_slicing
+        if hasattr(model.vae, "use_tiling"):
+            model.vae.use_tiling = od_config.vae_use_tiling
+
         return model
     else:
         raise ValueError(f"Model class {od_config.model_class_name} not found in diffusion model registry.")
@@ -34,6 +46,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     # `post_process_func` function must be placed in {mod_folder}/{mod_relname}.py,
     # where mod_folder and mod_relname are  defined and mapped using `_DIFFUSION_MODELS` via the `arch` key
     "QwenImagePipeline": "get_qwen_image_post_process_func",
+    "ZImagePipeline": "get_post_process_func",
 }
 
 

vllm_omni/diffusion/worker/gpu_worker.py

@@ -12,6 +12,7 @@ from vllm.distributed.parallel_state import (
     initialize_model_parallel,
 )
 from vllm.logger import init_logger
+from vllm.model_executor.model_loader.utils import set_default_torch_dtype
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.registry import initialize_model
@@ -61,9 +62,10 @@ class GPUWorker:
         initialize_model_parallel(tensor_model_parallel_size=world_size)
 
         with device:
-            self.pipeline = initialize_model(self.od_config)
-            self.pipeline.load_weights()
-            self.pipeline.eval()
+            with set_default_torch_dtype(self.od_config.dtype):
+                self.pipeline = initialize_model(self.od_config)
+                self.pipeline.load_weights()
+                self.pipeline.eval()
         logger.info(f"Worker {self.rank}: Initialized device, model, and distributed environment.")
         logger.info(f"Worker {self.rank}: Model loaded successfully.")
 
@@ -129,7 +131,7 @@ class WorkerProc:
         """
         Receive requests from broadcast queue
         """
-        return self.mq.dequeue()
+        return self.mq.dequeue(indefinite=True)
 
     # TODO: queueing, cancellation
     def worker_busy_loop(self) -> None:

vllm_omni/diffusion/worker/npu/npu_worker.py

@@ -5,13 +5,15 @@ import os
 
 import torch
 import zmq
-from vllm.config import VllmConfig, set_current_vllm_config
+from transformers import PretrainedConfig
+from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
 from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
 from vllm.distributed.parallel_state import (
     init_distributed_environment,
     initialize_model_parallel,
 )
 from vllm.logger import init_logger
+from vllm.model_executor.model_loader.utils import set_default_torch_dtype
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.registry import initialize_model
@@ -53,7 +55,12 @@ class NPUWorker:
         torch.npu.set_device(device)
 
         # hack
-        vllm_config = VllmConfig()
+        # set hf_config to a fake one to avolid get attr error
+        class _FakePretrainedConfig(PretrainedConfig):
+            def __getattr__(self, name):
+                return "fake"
+
+        vllm_config = VllmConfig(model_config=ModelConfig(hf_config=_FakePretrainedConfig()))
         vllm_config.parallel_config.tensor_parallel_size = self.od_config.num_gpus
         set_current_vllm_config(vllm_config)
 
@@ -61,9 +68,10 @@ class NPUWorker:
         initialize_model_parallel(tensor_model_parallel_size=world_size)
 
         with device:
-            self.pipeline = initialize_model(self.od_config)
-            self.pipeline.load_weights()
-            self.pipeline.eval()
+            with set_default_torch_dtype(self.od_config.dtype):
+                self.pipeline = initialize_model(self.od_config)
+                self.pipeline.load_weights()
+                self.pipeline.eval()
         logger.info(f"Worker {self.rank}: Initialized device, model, and distributed environment.")
         logger.info(f"Worker {self.rank}: Model loaded successfully.")
 
@@ -129,7 +137,7 @@ class NPUWorkerProc:
         """
         Receive requests from broadcast queue
         """
-        return self.mq.dequeue()
+        return self.mq.dequeue(indefinite=True)
 
     # TODO: queueing, cancellation
     def worker_busy_loop(self) -> None:

vllm_omni/distributed/__init__.py

@@ -1,25 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-from .connectors import (
-    ConnectorSpec, OmniTransferConfig,
-    OmniConnectorBase, InMemoryOmniConnector,
-    OmniConnectorFactory, MooncakeConnector,
-    load_omni_transfer_config, create_simple_config, create_mooncake_config
-)
-
-__all__ = [
-    # Config
-    "ConnectorSpec",
-    "OmniTransferConfig",
-
-    # Connectors
-    "OmniConnectorBase",
-    "InMemoryOmniConnector",
-    "OmniConnectorFactory",
-    "MooncakeConnector",
-
-    # Utilities
-    "load_omni_transfer_config",
-    "create_simple_config",
-    "create_mooncake_config",
-]

vllm_omni/distributed/connectors/__init__.py

@@ -1,44 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-from .config import ConnectorSpec, OmniTransferConfig
-from .base import OmniConnectorBase, InMemoryOmniConnector
-from .factory import OmniConnectorFactory
-from .mooncake_connector import MooncakeConnector
-from .utils import (
-    load_omni_transfer_config,
-    initialize_connectors_from_config,
-    create_simple_config,
-    create_mooncake_config,
-    get_connectors_config_for_stage,
-    initialize_orchestrator_connectors,
-    get_stage_connector_config,
-    build_stage_connectors,
-)
-
-__all__ = [
-    # Config
-    "ConnectorSpec",
-    "OmniTransferConfig",
-
-    # Base classes and implementations
-    "OmniConnectorBase",
-    "InMemoryOmniConnector",
-
-    # Factory
-    "OmniConnectorFactory",
-
-    # Specific implementations
-    "MooncakeConnector",
-
-    # Utilities
-    "load_omni_transfer_config",
-    "initialize_connectors_from_config",
-    "create_simple_config",
-    "create_mooncake_config",
-    "get_connectors_config_for_stage",
-
-    # Manager helpers
-    "initialize_orchestrator_connectors",
-    "get_stage_connector_config",
-    "build_stage_connectors",
-]

vllm_omni/distributed/connectors/adapter.py

@@ -1,145 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# temporary for compatibility with vllm_omni.entrypoints.omni_stage.py 
-# and vllm_omni.entrypoints.omni_llm.py
-
-import time
-import logging
-from typing import Any, Optional, Dict, Callable
-
-logger = logging.getLogger(__name__)
-
-
-def try_send_via_connector(
-    connector: Any,
-    stage_id: int,
-    next_stage_id: int,
-    req_id: str,
-    next_inputs: Any,
-    sampling_params: Any,
-    original_prompt: Any,
-    next_stage_queue_submit_fn: Callable[[Dict[str, Any]], None],
-    metrics: Any
-) -> bool:
-    """
-    Attempts to send data via OmniConnector.
-    Returns True if successful, False otherwise.
-    Encapsulates the logic of preparing payload, sending via connector, 
-    sending notification, and recording metrics.
-    """
-    try:
-        t0 = time.time()
-
-        # Prepare data for connector
-        payload_data = {
-            "engine_inputs": next_inputs,
-            "sampling_params": sampling_params,
-            "metadata": {
-                "original_prompt": original_prompt,
-                "stage_transition": f"{stage_id}->{next_stage_id}",
-                "timestamp": time.time()
-            }
-        }
-
-        # Send data via connector
-        success, serialized_size = connector.put(str(stage_id), str(next_stage_id), str(req_id), payload_data)
-        
-        if success:
-            # Send lightweight notification via queue
-            notify_payload = {
-                "request_id": req_id,
-                "sampling_params": sampling_params,
-                "from_connector": True,
-                "from_stage": str(stage_id),
-                "to_stage": str(next_stage_id),
-                "sent_ts": time.time(),
-            }
-            next_stage_queue_submit_fn(notify_payload)
-
-            t1 = time.time()
-            tx_ms = (t1 - t0) * 1000.0
-            
-            metrics.on_forward(
-                stage_id,
-                next_stage_id,
-                req_id,
-                serialized_size,  # Use size from connector
-                float(tx_ms),
-                True,  # Mark as using connector
-            )
-            return True
-        else:
-            # If put returned False, we let the caller handle fallback
-            return False
-
-    except Exception as e:
-        logger.warning(
-            "[Orchestrator] OmniConnector failed for req %s: %s; falling back to queue",
-            req_id,
-            e,
-        )
-        return False
-
-
-def try_recv_via_connector(
-    task: Dict[str, Any], 
-    connectors: Dict[Any, Any], 
-    stage_id: int,
-) -> tuple[Any, Optional[Dict[str, Any]]]:
-    """
-    Attempts to resolve input data from either connector or IPC.
-    Returns (engine_inputs, rx_metrics) or (None, None) if failed/skipped.
-    """
-    rid = task["request_id"]
-    
-    if task.get("from_connector"):
-        from_stage = task.get("from_stage")
-        to_stage = str(stage_id)
-
-        if not from_stage:
-            logger.error(
-                "[Stage-%s] 'from_connector' is true but 'from_stage' is missing for request %s", stage_id, rid
-            )
-            return None, None
-
-        # Get connector for this edge
-        connector_key = (from_stage, to_stage)
-        connector = connectors.get(connector_key)
-
-        if connector:
-            try:
-                # Get data from connector with timeout
-                _t_start = time.time()
-                payload_data = connector.get(from_stage, to_stage, str(rid))
-                _t_end = time.time()
-                
-                if payload_data and isinstance(payload_data, dict):
-                    ein = payload_data.get("engine_inputs")
-                    # Use connector's serialization size for consistency
-                    serialized_size = len(connector.serialize_obj(payload_data))
-                    decode_ms = (_t_end - _t_start) * 1000.0
-                    
-                    rx_metrics = {
-                        "rx_decode_time_ms": decode_ms,
-                        "rx_transfer_bytes": serialized_size
-                    }
-                    return ein, rx_metrics
-                else:
-                    logger.error(
-                        "[Stage-%s] Failed to get data from connector for request %s or payload is empty", stage_id, rid
-                    )
-                    return None, None
-            except Exception as e:
-                logger.error(
-                    "[Stage-%s] Error retrieving data from connector for request %s: %s", stage_id, rid, e
-                )
-                return None, None
-        else:
-            logger.error(
-                "[Stage-%s] No connector found for edge %s -> %s for request %s", stage_id, from_stage, to_stage, rid
-            )
-            return None, None
-    else:
-        # Data comes from queue as usual
-        #return maybe_load_from_ipc_with_metrics_fn(task, obj_key="engine_inputs", shm_key="engine_inputs_shm")
-        return None, None
-

vllm_omni/distributed/connectors/base.py

@@ -1,119 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-import logging
-from abc import ABC, abstractmethod
-from typing import Any, Optional
-
-logger = logging.getLogger(__name__)
-
-
-class OmniConnectorBase(ABC):
-    """Base class for all OmniConnectors."""
-
-    @abstractmethod
-    def put(self, from_stage: str, to_stage: str, request_id: str, data: Any) -> tuple[bool, int]:
-        """Store Python object, internal serialization handled by connector.
-
-        Args:
-            from_stage: Source stage identifier
-            to_stage: Destination stage identifier
-            request_id: Unique request identifier
-            data: Python object to store
-
-        Returns:
-            tuple: (success: bool, serialized_size: int)
-        """
-        pass
-
-    @abstractmethod
-    def get(self, from_stage: str, to_stage: str, request_id: str) -> Optional[Any]:
-        """Retrieve Python object, internal deserialization handled by connector.
-
-        Args:
-            from_stage: Source stage identifier
-            to_stage: Destination stage identifier
-            request_id: Unique request identifier
-
-        Returns:
-            Python object if found, None otherwise
-        """
-        pass
-
-    @abstractmethod
-    def cleanup(self, request_id: str) -> None:
-        """Clean up resources for a request."""
-        pass
-
-    @abstractmethod
-    def health(self) -> dict[str, Any]:
-        """Return health status and metrics."""
-        pass
-
-    @staticmethod
-    def serialize_obj(obj: Any) -> bytes:
-        """Serialize a Python object to bytes using torch.save (consistent with stage_utils)."""
-        import io
-        import torch
-        buf = io.BytesIO()
-        torch.save(obj, buf)
-        return buf.getvalue()
-
-    @staticmethod
-    def deserialize_obj(data: bytes) -> Any:
-        """Deserialize bytes to Python object using torch.load (consistent with stage_utils)."""
-        import io
-        import torch
-        return torch.load(io.BytesIO(data), map_location="cpu", weights_only=False)
-
-
-class InMemoryOmniConnector(OmniConnectorBase):
-    """Simple in-memory connector for testing and co-located stages."""
-
-    def __init__(self):
-        self._storage: dict[str, Any] = {}
-        self._metrics = {
-            "puts": 0,
-            "gets": 0,
-            "errors": 0,
-        }
-
-    def _make_key(self, from_stage: str, to_stage: str, request_id: str) -> str:
-        return f"{from_stage}:{to_stage}:{request_id}"
-
-    def put(self, from_stage: str, to_stage: str, request_id: str, data: Any) -> tuple[bool, int]:
-        try:
-            key = self._make_key(from_stage, to_stage, request_id)
-            self._storage[key] = data
-            serialized_size = len(self.serialize_obj(data))
-            self._metrics["puts"] += 1
-            logger.debug(f"InMemoryConnector: stored {key}")
-            return True, serialized_size
-        except Exception as e:
-            self._metrics["errors"] += 1
-            logger.error(f"InMemoryConnector put failed: {e}")
-            return False, 0
-
-    def get(self, from_stage: str, to_stage: str, request_id: str) -> Optional[Any]:
-        try:
-            key = self._make_key(from_stage, to_stage, request_id)
-            data = self._storage.get(key)
-            if data is not None:
-                self._metrics["gets"] += 1
-                return data
-            return None
-        except Exception as e:
-            logger.error(f"InMemoryConnector get failed: {e}")
-            return None
-
-    def cleanup(self, request_id: str) -> None:
-        keys_to_remove = [k for k in self._storage.keys() if k.endswith(f":{request_id}")]
-        for key in keys_to_remove:
-            del self._storage[key]
-        logger.debug(f"InMemoryConnector: cleaned up {len(keys_to_remove)} keys for {request_id}")
-
-    def health(self) -> dict[str, Any]:
-        return {
-            "status": "healthy",
-            "storage_size": len(self._storage),
-            **self._metrics,
-        }

vllm_omni/distributed/connectors/config.py

@@ -1,31 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-import logging
-from dataclasses import dataclass, field
-from typing import Any, Optional
-
-logger = logging.getLogger(__name__)
-
-
-@dataclass
-class ConnectorSpec:
-    """Specification for a connector instance."""
-    name: str  # e.g., "MooncakeConnector", "InMemoryConnector"
-    extra: dict[str, Any] = field(default_factory=dict)  # backend-specific config
-
-
-@dataclass
-class OmniTransferConfig:
-    """Top-level configuration for OmniConnector system."""
-    # Direct mapping: (from_stage, to_stage) -> connector
-    connectors: dict[tuple[str, str], ConnectorSpec] = field(default_factory=dict)
-    default_connector: Optional[ConnectorSpec] = None
-
-    def get_connector_for_edge(self, from_stage: str, to_stage: str) -> Optional[ConnectorSpec]:
-        """Get connector spec for a specific edge."""
-        edge_key = (from_stage, to_stage)
-        return self.connectors.get(edge_key, self.default_connector)
-
-    def has_connector_for_edge(self, from_stage: str, to_stage: str) -> bool:
-        """Check if there's a connector configured for the edge."""
-        return self.get_connector_for_edge(from_stage, to_stage) is not None

vllm_omni/distributed/connectors/connector_tests/run_single_prompt.sh

@@ -1,9 +0,0 @@
-export PYTHONPATH=/workspace/omni/vllm-omni/:$PYTHONPATH
-CUDA_VISIBLE_DEVICES=0,1,2,3 python end2end.py --model /workspace/Qwen2.5-Omni-7B/ \
-                                 --voice-type "m02" \
-                                 --dit-ckpt none \
-                                 --bigvgan-ckpt none \
-                                 --output-wav output_audio \
-                                 --prompt_type text \
-                                 --init-sleep-seconds 0 \
-                                 --prompts "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."

vllm_omni/distributed/connectors/connector_tests/test_connector_integration.py

@@ -1,78 +0,0 @@
-#!/usr/bin/env python3
-"""
-Test OmniConnector integration with OmniLLM.
-"""
-
-import sys
-import os
-sys.path.insert(0, os.path.dirname(__file__))
-
-from vllm_omni.distributed.connectors import create_mooncake_config
-
-def test_connector_config():
-    """Test connector configuration creation."""
-    print("Testing connector configuration...")
-
-    # Create Mooncake config
-    config = create_mooncake_config(
-        host="127.0.0.1",
-        metadata_server="http://127.0.0.1:8080/metadata",
-        master="127.0.0.1:50051"
-    )
-
-    print(f"Created config with {len(config.connectors)} connectors")
-    for edge, connector_spec in config.connectors.items():
-        print(f"  {edge}: {connector_spec.name}")
-
-    return config
-
-def test_omni_llm_initialization():
-    """Test OmniLLM with connector config."""
-    print("\nTesting OmniLLM initialization with connectors...")
-
-    try:
-        from vllm_omni.entrypoints.omni_llm import OmniLLM
-
-        # Create connector config
-        connector_config = {
-            "connectors": {
-                "0_to_1": {
-                    "name": "InMemoryConnector",  # Use InMemory for testing
-                    "extra": {}
-                }
-            }
-        }
-
-        # Test connector initialization logic
-        omni = OmniLLM.__new__(OmniLLM)  # Create without calling __init__
-        omni.omni_transfer_config = connector_config
-        omni.connectors = {}
-        omni._initialize_connectors()
-
-        # Check if connector was created
-        if ("0", "1") in omni.connectors:
-            print(f"✅ Connector created: {type(omni.connectors[('0', '1')]).__name__}")
-            return True
-        else:
-            print("❌ Connector not created")
-            return False
-
-    except Exception as e:
-        print(f"Error: {e}")
-        return False
-
-if __name__ == "__main__":
-    print("OmniConnector Integration Test")
-    print("=" * 40)
-
-    # Test config creation
-    config = test_connector_config()
-
-    # Test OmniLLM integration
-    success = test_omni_llm_initialization()
-
-    if success:
-        print("\n✅ Integration test passed!")
-        print("Connector system is ready for use with OmniLLM")
-    else:
-        print("\n❌ Integration test failed!")

vllm_omni/distributed/connectors/connector_tests/test_epdg_pipeline.py

@@ -1,332 +0,0 @@
-#!/usr/bin/env python3
-"""
-Test EPD (Encode-Prefill-Decode) 3-stage pipeline with connectors.
-Tests data flow between stages using configured connectors.
-"""
-
-import os
-import sys
-from typing import Any, Dict, Optional
-from dataclasses import dataclass
-import uuid
-
-try:
-    import yaml
-except ImportError:
-    print("PyYAML not available, using basic YAML parsing")
-    yaml = None
-
-# Add the connectors package to path
-current_dir = os.path.dirname(os.path.abspath(__file__))
-parent_dir = os.path.dirname(current_dir)
-sys.path.insert(0, parent_dir)
-
-# Import from parent package
-try:
-    from connectors.config import ConnectorSpec, OmniTransferConfig
-    from connectors.base import OmniConnectorBase, InMemoryOmniConnector
-    from connectors.factory import OmniConnectorFactory
-    from connectors.utils import load_omni_transfer_config
-except ImportError:
-    # Fallback for direct execution
-    sys.path.insert(0, current_dir)
-    from config import ConnectorSpec, OmniTransferConfig
-    from base import OmniConnectorBase, InMemoryOmniConnector
-    from factory import OmniConnectorFactory
-    from utils import load_omni_transfer_config
-
-
-@dataclass
-class StageResult:
-    """Result from a stage processing."""
-    stage_name: str
-    request_id: str
-    output_data: Any
-    metadata: Dict[str, Any]
-
-
-class EncodeStage:
-    """Simple Encode stage for testing."""
-
-    def __init__(self, stage_id: int, output_connector: Optional[OmniConnectorBase] = None):
-        self.stage_id = stage_id
-        self.stage_name = "encode"
-        self.output_connector = output_connector
-
-    def process(self, request_id: str, input_data: Any) -> StageResult:
-        """Process input data and optionally send to next stage."""
-        print(f"[{self.stage_name}] Processing request {request_id}")
-
-        # Simple encoding simulation
-        encoded_data = {
-            "original_text": input_data,
-            "encoded_tokens": [ord(c) for c in str(input_data)][:10],  # Simple tokenization
-            "embeddings": [0.1 * i for i in range(10)],  # Mock embeddings
-            "stage": self.stage_name
-        }
-
-        result = StageResult(
-            stage_name=self.stage_name,
-            request_id=request_id,
-            output_data=encoded_data,
-            metadata={"tokens_count": len(encoded_data["encoded_tokens"])}
-        )
-
-        # Send to next stage if connector is available
-        if self.output_connector:
-            success = self.output_connector.put(
-                str(self.stage_id), "1", request_id, encoded_data
-            )
-            print(f"[{self.stage_name}] Sent data to stage 1 via connector: {success}")
-
-        return result
-
-
-class PrefillStage:
-    """Simple Prefill stage for testing."""
-
-    def __init__(self, stage_id: int,
-                 input_connector: Optional[OmniConnectorBase] = None,
-                 output_connector: Optional[OmniConnectorBase] = None):
-        self.stage_id = stage_id
-        self.stage_name = "prefill"
-        self.input_connector = input_connector
-        self.output_connector = output_connector
-
-    def process(self, request_id: str) -> Optional[StageResult]:
-        """Process data from previous stage."""
-        print(f"[{self.stage_name}] Processing request {request_id}")
-
-        # Try to get data from input connector
-        input_data = None
-        if self.input_connector:
-            input_data = self.input_connector.get("0", str(self.stage_id), request_id)
-
-        if input_data is None:
-            print(f"[{self.stage_name}] No input data received for {request_id}")
-            return None
-
-        print(f"[{self.stage_name}] Received data from encode")
-
-        # Simple prefill simulation
-        prefills_data = {
-            **input_data,
-            "kv_cache": {"layer_0": [0.2 * i for i in range(20)]},  # Mock KV cache
-            "stage": self.stage_name
-        }
-
-        result = StageResult(
-            stage_name=self.stage_name,
-            request_id=request_id,
-            output_data=prefills_data,
-            metadata={"kv_size": len(prefills_data["kv_cache"])}
-        )
-
-        # Send to next stage if connector is available
-        if self.output_connector:
-            success = self.output_connector.put(
-                str(self.stage_id), "2", request_id, prefills_data
-            )
-            print(f"[{self.stage_name}] Sent data to stage 2 via connector: {success}")
-
-        return result
-
-
-class DecodeStage:
-    """Simple Decode stage for testing (final stage in 3-stage pipeline)."""
-
-    def __init__(self, stage_id: int,
-                 input_connector: Optional[OmniConnectorBase] = None):
-        self.stage_id = stage_id
-        self.stage_name = "decode"
-        self.input_connector = input_connector
-
-    def process(self, request_id: str) -> Optional[StageResult]:
-        """Process data from previous stage (final output)."""
-        print(f"[{self.stage_name}] Processing request {request_id}")
-
-        # Try to get data from input connector
-        input_data = None
-        if self.input_connector:
-            input_data = self.input_connector.get("1", str(self.stage_id), request_id)
-
-        if input_data is None:
-            print(f"[{self.stage_name}] No input data received for {request_id}")
-            return None
-
-        print(f"[{self.stage_name}] Received data from prefill")
-
-        # Simple decode simulation (final output)
-        decoded_data = {
-            **input_data,
-            "hidden_states": [0.3 * i for i in range(15)],  # Mock hidden states
-            "final_output": f"Decoded text from: {input_data.get('original_text', 'unknown')}",
-            "stage": self.stage_name
-        }
-
-        result = StageResult(
-            stage_name=self.stage_name,
-            request_id=request_id,
-            output_data=decoded_data,
-            metadata={"hidden_size": len(decoded_data["hidden_states"]), "final": True}
-        )
-
-        return result
-
-
-class GeneratorStage:
-    """Simple Generator stage for testing (removed for 3-stage pipeline)."""
-    pass  # Not used in current 3-stage setup
-
-
-def load_config_from_yaml(yaml_path: str) -> OmniTransferConfig:
-    """Load connector configuration from YAML file."""
-    if yaml is None:
-        # Fallback: create a simple test config
-        print("Using fallback InMemoryConnector config for testing")
-        return create_test_config()
-
-    with open(yaml_path, 'r', encoding='utf-8') as f:
-        config_dict = yaml.safe_load(f)
-
-    try:
-        config = load_omni_transfer_config(config_dict=config_dict)
-        print(f"Successfully loaded config from YAML with {len(config.connectors)} connectors")
-        return config
-    except Exception as e:
-        print(f"Failed to load YAML config: {e}, using fallback config")
-        return create_test_config()
-
-
-def create_test_config() -> OmniTransferConfig:
-    """Create a test configuration with InMemoryConnectors."""
-    from .config import ConnectorSpec
-
-    connectors = {}
-    # Create connectors for EPD (3-stage) pipeline: 0->1->2
-    pipeline_stages = [("0", "1"), ("1", "2")]
-
-    for from_stage, to_stage in pipeline_stages:
-        connectors[(from_stage, to_stage)] = ConnectorSpec(
-            name="InMemoryConnector",
-            extra={}
-        )
-
-    return OmniTransferConfig(connectors=connectors)
-
-
-def create_stages_with_connectors(config: OmniTransferConfig) -> Dict[str, Any]:
-    """Create EPD (3-stage) pipeline with appropriate connectors."""
-
-    # Create connector instances based on config
-    connectors = {}
-    for edge_key, connector_spec in config.connectors.items():
-        from_stage, to_stage = edge_key
-        connectors[edge_key] = OmniConnectorFactory.create_connector(connector_spec)
-
-    # Create stages with connectors (3-stage pipeline: 0->1->2)
-    stages = {}
-
-    # Encode stage (stage 0)
-    encode_output_connector = connectors.get(("0", "1"))  # to prefill
-    stages["encode"] = EncodeStage(0, encode_output_connector)
-
-    # Prefill stage (stage 1)
-    prefill_input_connector = connectors.get(("0", "1"))   # from encode
-    prefill_output_connector = connectors.get(("1", "2"))  # to decode
-    stages["prefill"] = PrefillStage(1, prefill_input_connector, prefill_output_connector)
-
-    # Decode stage (stage 2) - final stage
-    decode_input_connector = connectors.get(("1", "2"))    # from prefill
-    stages["decode"] = DecodeStage(2, decode_input_connector)
-
-    return stages
-
-
-def test_pipeline_with_config(yaml_path: str):
-    """Test the EPDG pipeline with connector configuration."""
-
-    print("🔧 Loading configuration from YAML...")
-    try:
-        config = load_config_from_yaml(yaml_path)
-        print(f"✓ Loaded config with {len(config.connectors)} connector configurations")
-    except Exception as e:
-        print(f"✗ Failed to load config: {e}")
-        return False
-
-    print("\n🏗️  Creating EPDG stages with connectors...")
-    try:
-        stages = create_stages_with_connectors(config)
-        print("✓ Created all stages with connectors")
-    except Exception as e:
-        print(f"✗ Failed to create stages: {e}")
-        return False
-
-    print("\n🚀 Testing data flow through pipeline...")
-
-    # Test request
-    request_id = f"test_req_{uuid.uuid4()}"
-    input_text = "Hello, EPDG pipeline!"
-
-    try:
-        # Stage 1: Encode
-        print(f"\n--- Stage 1: Encode ---")
-        encode_result = stages["encode"].process(request_id, input_text)
-        print(f"✓ Encode completed: {encode_result.metadata}")
-
-        # Stage 2: Prefill
-        print(f"\n--- Stage 2: Prefill ---")
-        prefill_result = stages["prefill"].process(request_id)
-        if prefill_result:
-            print(f"✓ Prefill completed: {prefill_result.metadata}")
-        else:
-            print("✗ Prefill failed - no input data")
-            return False
-
-        # Stage 3: Decode
-        print(f"\n--- Stage 3: Decode ---")
-        decode_result = stages["decode"].process(request_id)
-        if decode_result:
-            print(f"✓ Decode completed: {decode_result.metadata}")
-        else:
-            print("✗ Decode failed - no input data")
-            return False
-
-        # Decode stage produces final output in 3-stage pipeline
-
-    except Exception as e:
-        print(f"✗ Pipeline test failed: {e}")
-        import traceback
-        traceback.print_exc()
-        return False
-
-    print("\n🎉 Pipeline test completed successfully!")
-    return True
-
-
-def main():
-    """Main test function."""
-    # Path to the YAML config file
-    yaml_path = "/workspace/omni/vllm-omni/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml"
-
-    if not os.path.exists(yaml_path):
-        print(f"✗ YAML config file not found: {yaml_path}")
-        return False
-
-    print("Testing EPDG Pipeline with Connectors")
-    print("=" * 50)
-    print(f"Config file: {yaml_path}")
-
-    success = test_pipeline_with_config(yaml_path)
-
-    if success:
-        print("\n✅ All tests passed!")
-    else:
-        print("\n❌ Tests failed!")
-
-    return success
-
-
-if __name__ == "__main__":
-    success = main()
-    sys.exit(0 if success else 1)

vllm_omni/distributed/connectors/factory.py

@@ -1,86 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-import logging
-from collections.abc import Callable
-from typing import Any
-
-try:
-    from .base import OmniConnectorBase
-    from .config import ConnectorSpec
-except ImportError:
-    # Fallback for direct execution
-    import sys
-    import os
-    sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
-    from connectors.base import OmniConnectorBase
-    from connectors.config import ConnectorSpec
-
-logger = logging.getLogger(__name__)
-
-
-class OmniConnectorFactory:
-    """Factory for creating OmniConnectors."""
-
-    _registry: dict[str, Callable[[dict[str, Any]], OmniConnectorBase]] = {}
-
-    @classmethod
-    def register_connector(
-        cls,
-        name: str,
-        constructor: Callable[[dict[str, Any]], OmniConnectorBase]
-    ) -> None:
-        """Register a connector constructor."""
-        if name in cls._registry:
-            raise ValueError(f"Connector '{name}' is already registered.")
-        cls._registry[name] = constructor
-        logger.debug(f"Registered connector: {name}")
-
-    @classmethod
-    def create_connector(cls, spec: ConnectorSpec) -> OmniConnectorBase:
-        """Create a connector from specification."""
-        if spec.name not in cls._registry:
-            raise ValueError(f"Unknown connector: {spec.name}. Available: {list(cls._registry.keys())}")
-
-        constructor = cls._registry[spec.name]
-        try:
-            connector = constructor(spec.extra)
-            logger.info(f"Created connector: {spec.name}")
-            return connector
-        except Exception as e:
-            logger.error(f"Failed to create connector {spec.name}: {e}")
-            raise
-
-    @classmethod
-    def list_registered_connectors(cls) -> list[str]:
-        """List all registered connector names."""
-        return list(cls._registry.keys())
-
-
-# Register built-in connectors with lazy imports
-def _create_mooncake_connector(config: dict[str, Any]) -> OmniConnectorBase:
-    try:
-        from .mooncake_connector import MooncakeConnector
-    except ImportError:
-        # Fallback import
-        import sys
-        import os
-        sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
-        from connectors.mooncake_connector import MooncakeConnector
-    return MooncakeConnector(config)
-
-
-def _create_in_memory_connector(config: dict[str, Any]) -> OmniConnectorBase:
-    try:
-        from .base import InMemoryOmniConnector
-    except ImportError:
-        # Fallback import
-        import sys
-        import os
-        sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
-        from connectors.base import InMemoryOmniConnector
-    return InMemoryOmniConnector()
-
-
-# Register connectors
-OmniConnectorFactory.register_connector("MooncakeConnector", _create_mooncake_connector)
-OmniConnectorFactory.register_connector("InMemoryConnector", _create_in_memory_connector)

vllm_omni/distributed/connectors/mooncake_connector.py

@@ -1,157 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-import io
-import logging
-import time
-from typing import Any, Optional
-
-import torch
-
-from .base import OmniConnectorBase
-
-logger = logging.getLogger(__name__)
-
-try:
-    from mooncake.store import MooncakeDistributedStore, ReplicateConfig
-except ImportError:
-    try:
-        from mooncake import MooncakeDistributedStore, ReplicateConfig
-    except ImportError:
-        logger.warning("Mooncake not available, MooncakeOmniConnector will not work")
-        MooncakeDistributedStore = None
-        ReplicateConfig = None
-
-
-class MooncakeConnector(OmniConnectorBase):
-    """Mooncake-based distributed connector for OmniConnector."""
-
-    def __init__(self, config: dict[str, Any]):
-        if MooncakeDistributedStore is None:
-            raise ImportError("Mooncake not available")
-
-        self.config = config
-        self.host = config.get("host", "127.0.0.1")
-        self.metadata = config.get("metadata_server", "http://127.0.0.1:8080/metadata")
-        self.master = config.get("master", "127.0.0.1:50051")
-        self.segment = config.get("segment", 512 * 1024 * 1024)  # 512MB
-        self.localbuf = config.get("localbuf", 64 * 1024 * 1024)  # 64MB
-        self.proto = config.get("proto", "tcp")
-        self.rdma = config.get("rdma", "")
-
-        self.store: Optional[MooncakeDistributedStore] = None
-        self.pin: Optional[ReplicateConfig] = None
-
-        self._metrics = {
-            "puts": 0,
-            "gets": 0,
-            "bytes_transferred": 0,
-            "errors": 0,
-            "timeouts": 0,
-        }
-
-        self._init_store()
-
-    def _make_key(self, rid: str, from_stage: str, to_stage: str) -> str:
-        """Generate store key for request between stages."""
-        return f"{rid}/{from_stage}_to_{to_stage}"
-
-    def _init_store(self):
-        """Initialize Mooncake store."""
-        try:
-            self.store = MooncakeDistributedStore()
-            rc = self.store.setup(
-                self.host, self.metadata, self.segment,
-                self.localbuf, self.proto, self.rdma, self.master
-            )
-            if rc != 0:
-                raise RuntimeError(f"Mooncake setup failed: {rc}")
-
-            self.pin = ReplicateConfig()
-            self.pin.with_soft_pin = True
-            logger.info("MooncakeConnector initialized successfully")
-        except Exception as e:
-            logger.error(f"Failed to initialize Mooncake store: {e}")
-            raise
-
-    # Use base class serialization methods for consistency
-
-    def put(self, from_stage: str, to_stage: str, request_id: str, data: Any) -> tuple[bool, int]:
-        if not self.store:
-            logger.error("Store not initialized")
-            return False, 0
-
-        try:
-            serialized_data = self.serialize_obj(data)
-            key = self._make_key(request_id, from_stage, to_stage)
-            self.store.put(key, serialized_data, self.pin)
-
-            self._metrics["puts"] += 1
-            self._metrics["bytes_transferred"] += len(serialized_data)
-
-            logger.debug(f"MooncakeConnector: stored {key}, {len(serialized_data)} bytes")
-            print(f"========= mooncake put {from_stage}-{to_stage}-{request_id}, with size {len(serialized_data)}")
-            return True, len(serialized_data)
-
-        except Exception as e:
-            self._metrics["errors"] += 1
-            logger.error(f"MooncakeConnector put failed: {e}")
-            return False, 0
-
-    def get(self, from_stage: str, to_stage: str, request_id: str) -> Optional[Any]:
-        if not self.store:
-            logger.error("Store not initialized")
-            return None
-
-        retries = 20
-        sleep_s = 0.05
-        key = self._make_key(request_id, from_stage, to_stage)
-
-        for attempt in range(retries):
-            try:
-                raw_data = self.store.get(key)
-
-                if raw_data:
-                    data = self.deserialize_obj(raw_data)
-                    self._metrics["gets"] += 1
-                    print(f"========= mooncake get {from_stage}-{to_stage}-{request_id}, with size {len(raw_data)}")
-                    return data
-
-            except Exception as e:
-                logger.debug(f"MooncakeConnector get attempt {attempt} failed: {e}")
-
-            if attempt < retries - 1:
-                time.sleep(sleep_s)
-
-        self._metrics["timeouts"] += 1
-        logger.warning(f"MooncakeConnector: timeout waiting for {key}")
-        return None
-
-    def cleanup(self, request_id: str) -> None:
-        if not self.store:
-            return
-
-        # Note: Mooncake doesn't have explicit delete, data will be garbage collected
-        # We could implement a cleanup mechanism by storing deletion markers
-        logger.debug(f"MooncakeConnector: cleanup requested for {request_id} (no-op)")
-
-    def health(self) -> dict[str, Any]:
-        if not self.store:
-            return {"status": "unhealthy", "error": "Store not initialized"}
-
-        return {
-            "status": "healthy",
-            "host": self.host,
-            "metadata_server": self.metadata,
-            "master": self.master,
-            **self._metrics,
-        }
-
-    def close(self):
-        """Clean shutdown."""
-        if self.store:
-            try:
-                self.store.close()
-                self.store = None
-                logger.info("MooncakeConnector closed")
-            except Exception as e:
-                logger.error(f"Error closing Mooncake store: {e}")

vllm_omni/distributed/connectors/utils.py

@@ -1,314 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-"""Utilities for OmniConnector configuration and validation."""
-
-import json
-import logging
-from pathlib import Path
-from typing import Any, Dict, Optional, Tuple, Union
-
-from vllm_omni.distributed.connectors.config import OmniTransferConfig, ConnectorSpec
-from vllm_omni.distributed.connectors.factory import OmniConnectorFactory
-from vllm_omni.distributed.connectors.base import OmniConnectorBase
-
-logger = logging.getLogger(__name__)
-
-
-def initialize_connectors_from_config(
-    config_path: Optional[Union[str, Path]] = None
-) -> tuple[Optional[OmniTransferConfig], dict[tuple[str, str], OmniConnectorBase]]:
-    """
-    Initialize connectors from configuration file.
-
-    Returns:
-        tuple: (OmniTransferConfig, dict of {(from, to): connector_instance})
-    """
-    try:
-        transfer_config = load_omni_transfer_config(config_path)
-    except Exception as e:
-        logger.warning(f"Failed to load OmniTransferConfig from {config_path}: {e}")
-        return None, {}
-
-    if not transfer_config:
-        logger.info("No OmniTransferConfig provided, using default IPC")
-        return None, {}
-
-    # 使用统一的连接器创建逻辑
-    connectors = create_connectors_from_config(transfer_config.connectors)
-    return transfer_config, connectors
-
-
-def create_connectors_from_config(
-    connectors_config: dict[tuple[str, str], ConnectorSpec]
-) -> dict[tuple[str, str], OmniConnectorBase]:
-    """通用连接器创建逻辑，从配置字典创建连接器实例。"""
-    connectors = {}
-    try:
-        for edge_key, connector_spec in connectors_config.items():
-            connector = OmniConnectorFactory.create_connector(connector_spec)
-            connectors[edge_key] = connector
-            logger.info(f"Created connector for {edge_key[0]} -> {edge_key[1]}: {type(connector).__name__}")
-    except Exception as e:
-        logger.error(f"Failed to initialize connectors: {e}")
-
-    return connectors
-
-
-def get_connectors_config_for_stage(
-    transfer_config: Optional[OmniTransferConfig],
-    stage_id: Union[str, int]
-) -> dict[str, Any]:
-    """
-    Extract connector configurations relevant for a specific stage worker.
-    
-    Returns a dict compatible with worker initialization:
-    {
-        "from_stage_X": {
-            "spec": {
-                "name": "ConnectorName",
-                "extra": {...}
-            }
-        },
-        ...
-    }
-    """
-    if not transfer_config:
-        return {}
-        
-    stage_connectors_config = {}
-    target_stage = str(stage_id)
-    
-    # Iterate through all configured edges
-    for (from_stage, to_stage), spec in transfer_config.connectors.items():
-        # We only care about incoming edges for the worker process
-        # (Worker needs to create connectors to receive data)
-        if to_stage == target_stage:
-            stage_connectors_config[f"from_stage_{from_stage}"] = {
-                "spec": {
-                    "name": spec.name,
-                    "extra": spec.extra
-                }
-            }
-            
-    return stage_connectors_config
-
-
-def load_omni_transfer_config(
-    config_path: Optional[Union[str, Path]] = None,
-    config_dict: Optional[dict[str, Any]] = None
-) -> Optional[OmniTransferConfig]:
-    """Load OmniTransferConfig from file or dict."""
-    if config_path is None and config_dict is None:
-        return None
-
-    if config_path is not None:
-        config_path = Path(config_path)
-        if not config_path.exists():
-            raise FileNotFoundError(f"Config file not found: {config_path}")
-
-        with open(config_path, 'r', encoding='utf-8') as f:
-            if config_path.suffix.lower() == '.json':
-                config_dict = json.load(f)
-            elif config_path.suffix.lower() in ['.yaml', '.yml']:
-                try:
-                    import yaml
-                    config_dict = yaml.safe_load(f)
-                except ImportError:
-                    raise ImportError("PyYAML required for YAML config files")
-            else:
-                raise ValueError(f"Unsupported config file format: {config_path.suffix}")
-
-    if config_dict is None:
-        return None
-
-    # Parse connectors
-    connectors = {}
-    runtime_config = config_dict.get('runtime', {})
-
-    # Parse global connectors (from runtime.connectors)
-    global_connectors = runtime_config.get('connectors', {})
-    for conn_name, conn_config in global_connectors.items():
-        connector = ConnectorSpec(
-            name=conn_config['name'],
-            extra=conn_config.get('extra', {})
-        )
-        # Store globally defined connectors for reference
-        # These will be resolved to specific edges later
-
-    # Parse stage-level connectors
-    for stage_config in config_dict.get('stage_args', []):
-        stage_id = str(stage_config['stage_id'])
-
-        # Input connectors
-        for input_key, conn_ref in stage_config.get('input_connectors', {}).items():
-            if isinstance(conn_ref, str):
-                # Reference to global connector
-                if conn_ref in global_connectors:
-                    conn_config = global_connectors[conn_ref]
-                    connector = ConnectorSpec(
-                        name=conn_config['name'],
-                        extra=conn_config.get('extra', {})
-                    )
-                else:
-                    raise ValueError(f"Undefined connector reference: {conn_ref}")
-            else:
-                # Inline connector definition
-                connector = ConnectorSpec(
-                    name=conn_ref['name'],
-                    extra=conn_ref.get('extra', {})
-                )
-
-            # Parse from_stage from key (e.g., "from_stage_0" -> "0")
-            from_stage = input_key.replace("from_stage_", "")
-            edge_key = (from_stage, stage_id)
-            connectors[edge_key] = connector
-
-        # Output connectors
-        for output_key, conn_ref in stage_config.get('output_connectors', {}).items():
-            if isinstance(conn_ref, str):
-                # Reference to global connector
-                if conn_ref in global_connectors:
-                    conn_config = global_connectors[conn_ref]
-                    connector = ConnectorSpec(
-                        name=conn_config['name'],
-                        extra=conn_config.get('extra', {})
-                    )
-                else:
-                    raise ValueError(f"Undefined connector reference: {conn_ref}")
-            else:
-                # Inline connector definition
-                connector = ConnectorSpec(
-                    name=conn_ref['name'],
-                    extra=conn_ref.get('extra', {})
-                )
-
-            # Parse to_stage from key (e.g., "to_stage_1" -> "1")
-            to_stage = output_key.replace("to_stage_", "")
-            edge_key = (stage_id, to_stage)
-            connectors[edge_key] = connector
-
-    config = OmniTransferConfig(connectors=connectors)
-
-    logger.info(f"Loaded OmniTransferConfig with {len(connectors)} connector configurations")
-    return config
-
-
-def create_simple_config(
-    connector_name: str = "InMemoryConnector",
-    connector_config: Optional[dict[str, Any]] = None
-) -> OmniTransferConfig:
-    """Create a simple OmniTransferConfig for testing."""
-    if connector_config is None:
-        connector_config = {}
-
-    # Create connectors for typical EPDG pipeline
-    connectors = {}
-    pipeline_edges = [
-        ("0", "1"),  # thinker -> talker
-        ("1", "2"),  # talker -> code2wav
-    ]
-
-    for from_stage, to_stage in pipeline_edges:
-        connectors[(from_stage, to_stage)] = ConnectorSpec(
-            name=connector_name,
-            extra=connector_config
-        )
-
-    return OmniTransferConfig(connectors=connectors)
-
-
-def create_mooncake_config(
-    host: str = "127.0.0.1",
-    metadata_server: str = "http://127.0.0.1:8080/metadata",
-    master: str = "127.0.0.1:50051",
-    segment: int = 512 * 1024 * 1024,
-    localbuf: int = 64 * 1024 * 1024,
-) -> OmniTransferConfig:
-    """Create OmniTransferConfig for Mooncake-based setup."""
-    connector_config = {
-        "host": host,
-        "metadata_server": metadata_server,
-        "master": master,
-        "segment": segment,
-        "localbuf": localbuf,
-    }
-
-    return create_simple_config("MooncakeConnector", connector_config)
-
-
-# High-level management functions (moved from manager.py)
-
-def initialize_orchestrator_connectors(
-    config_path: Optional[str],
-) -> Tuple[Optional[OmniTransferConfig], Dict[tuple[str, str], OmniConnectorBase]]:
-    """Initialize connectors shared at orchestrator level."""
-    try:
-        transfer_config, connectors = initialize_connectors_from_config(config_path)
-        return transfer_config, connectors
-    except Exception as exc:  # pragma: no cover - defensive logging
-        logger.error("Unexpected error initializing connectors: %s", exc)
-        return None, {}
-
-
-def get_stage_connector_config(
-    transfer_config: Optional[OmniTransferConfig],
-    stage_id: int,
-) -> dict[str, Any]:
-    """Return the serialized connector config payload for a specific stage."""
-    if transfer_config is None:
-        return {}
-
-    try:
-        return get_connectors_config_for_stage(transfer_config, stage_id)
-    except Exception as exc:  # pragma: no cover - defensive logging
-        logger.warning(
-            "Failed to build connector config for stage %s: %s. Using IPC fallback.",
-            stage_id,
-            exc,
-        )
-        return {}
-
-
-def build_stage_connectors(
-    stage_id: int,
-    connectors_config: dict[str, Any],
-) -> Optional[dict[tuple[str, str], Any]]:
-    """Instantiate OmniConnectors for a stage based on config."""
-    if not connectors_config:
-        return {}
-
-    logger.info(
-        "[Stage-%s] Initializing OmniConnectors with config keys: %s",
-        stage_id,
-        list(connectors_config.keys()),
-    )
-
-    from .factory import OmniConnectorFactory
-    from .config import ConnectorSpec
-
-    connectors: dict[tuple[str, str], Any] = {}
-    try:
-        # 将字典格式的配置转换为ConnectorSpec对象
-        stage_connector_specs = {}
-        for input_key, config in connectors_config.items():
-            if not input_key.startswith("from_stage_"):
-                continue
-
-            from_stage = input_key.replace("from_stage_", "")
-            spec_dict = config.get("spec", {})
-            if not spec_dict:
-                continue
-
-            connector_spec = ConnectorSpec(
-                name=spec_dict.get("name", "InMemoryConnector"),
-                extra=spec_dict.get("extra", {}),
-            )
-            stage_connector_specs[(str(from_stage), str(stage_id))] = connector_spec
-
-        # 使用统一的连接器创建逻辑
-        connectors = create_connectors_from_config(stage_connector_specs)
-    except Exception as exc:  # pragma: no cover - defensive logging
-        logger.exception("[Stage-%s] Failed to initialize connectors: %s", stage_id, exc)
-        return None
-
-    return connectors

vllm_omni/entrypoints/async_omni.py

@@ -90,7 +90,7 @@ class AsyncOmni(EngineClient):
         stage_configs_path: Optional[str] = None,
         log_stats: bool = False,
         log_file: Optional[str] = None,
-        init_sleep_seconds: int = 20,
+        init_sleep_seconds: int = 30,
         shm_threshold_bytes: int = 65536,
         batch_timeout: int = 10,
         init_timeout: int = 60000,

vllm_omni/entrypoints/omni.py

@@ -1,16 +1,38 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import os
 
 from vllm_omni.diffusion.omni_diffusion import OmniDiffusion
 from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model
 from vllm_omni.entrypoints.omni_llm import OmniLLM
 
 
+def _dummy_snapshot_download(model_id):
+    return model_id
+
+
+def omni_snapshot_download(model_id) -> str:
+    # TODO: this is just a workaround for quickly use modelscope, we should support
+    # modelscope in weight loading feature instead of using `snapshot_download`
+    if os.environ.get("VLLM_USE_MODELSCOPE", False):
+        from modelscope.hub.snapshot_download import snapshot_download
+
+        return snapshot_download(model_id)
+    else:
+        return _dummy_snapshot_download(model_id)
+
+
 class Omni:
     """Unified entrypoint for both LLM and Diffusion models for better usability."""
 
     def __init__(self, *args, **kwargs):
         model = args[0] if args else kwargs.get("model", "")
+        assert model != "", "Null model id detected, please specify a model id."
+        model = omni_snapshot_download(model)
+        if args:
+            args[0] = model
+        elif kwargs.get("model", "") != "":
+            kwargs["model"] = model
         if is_diffusion_model(model):
             self.instance: OmniLLM | OmniDiffusion = OmniDiffusion(*args, **kwargs)
         else:

vllm_omni/entrypoints/omni_llm.py

@@ -1,8 +1,6 @@
 import multiprocessing as mp
 import os
-import sys
 import time
-import uuid
 from collections.abc import Sequence
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import Any, Optional, Union
@@ -22,11 +20,6 @@ from vllm.utils import Counter
 from vllm.v1.engine.llm_engine import LLMEngine
 
 # Internal imports (our code)
-from vllm_omni.distributed.connectors.adapter import try_send_via_connector
-from vllm_omni.distributed.connectors import (
-    initialize_orchestrator_connectors,
-    get_stage_connector_config,
-)
 from vllm_omni.engine.arg_utils import OmniEngineArgs
 from vllm_omni.engine.output_processor import MultimodalOutputProcessor
 from vllm_omni.engine.processor import OmniProcessor
@@ -38,14 +31,9 @@ from vllm_omni.entrypoints.log_utils import (
 )
 from vllm_omni.entrypoints.omni_stage import OmniStage
 from vllm_omni.entrypoints.stage_utils import encode_for_ipc as _encode
-from vllm_omni.entrypoints.stage_utils import _to_dict
 from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
 from vllm_omni.entrypoints.stage_utils import serialize_obj as _set
-from vllm_omni.entrypoints.utils import (
-    load_stage_configs_from_model,
-    load_stage_configs_from_yaml,
-    resolve_model_config_path,
-)
+from vllm_omni.entrypoints.utils import load_stage_configs_from_model, load_stage_configs_from_yaml
 from vllm_omni.outputs import OmniRequestOutput
 
 logger = init_logger(__name__)
@@ -96,23 +84,14 @@ class OmniLLM:
         init_timeout: int = 300,
         **kwargs: Any,
     ):
-        self.worker_backend = kwargs.get("worker_backend", "process")
-        self.ray_address = kwargs.get("ray_address", None)
         self.batch_timeout = batch_timeout
         self._enable_stats: bool = bool(log_stats)
-        self._ray_pg = None
-
         # Do NOT call super().__init__ to avoid creating OmniStageLLM instances in parent.
         if stage_configs_path is None:
-            self.config_path = resolve_model_config_path(model)
-            self.stage_configs = load_stage_configs_from_yaml(self.config_path)
+            self.stage_configs = load_stage_configs_from_model(model)
         else:
-            self.config_path = stage_configs_path
             self.stage_configs = load_stage_configs_from_yaml(stage_configs_path)
 
-        # Initialize connectors
-        self.omni_transfer_config,self.connectors = initialize_orchestrator_connectors(self.config_path)
-
         # Optional file handler for orchestrator
         self._log_file = log_file
 
@@ -145,140 +124,36 @@ class OmniLLM:
         self.stage_list = [st for _, st in results]
         logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list))
 
-        if self.worker_backend == "ray":
-            try:
-                import ray
-                from ray.util.queue import Queue as RayQueue
-            except ImportError:
-                raise ImportError("ray is required for worker_backend='ray'")
-            self._queue_cls = lambda: RayQueue(maxsize=0)
-        else:
-            self._ctx = mp.get_context("spawn")
-            self._queue_cls = lambda: self._ctx.Queue(maxsize=0)
-
+        self._ctx = mp.get_context("spawn")
         self._stage_in_queues: list[mp.Queue] = []
         self._stage_out_queues: list[mp.Queue] = []
         self._init_sleep_seconds = max(0, int(init_sleep_seconds))
         self._shm_threshold_bytes = max(0, int(shm_threshold_bytes))
-        self._start_stages(model)
+        self._start_stage_processes(model)
         # Wait for all stages to report readiness before seeding
         self._stages_ready: set[int] = set()
         self._wait_for_stages_ready(timeout=init_timeout)
 
-    def _start_stages(self, model: str) -> None:
-        if self.worker_backend == "ray":
-            self._start_stage_actors_ray(model)
-        else:
-            self._start_stage_processes_mp(model)
-
-    def _start_stage_processes_mp(self, model: str) -> None:
+    def _start_stage_processes(self, model: str) -> None:
         for stage_id, stage in enumerate(self.stage_list):
             # Use unbounded queues to avoid deadlock when seeding many requests
-            in_q: mp.Queue = self._queue_cls()
-            out_q: mp.Queue = self._queue_cls()
+            in_q: mp.Queue = self._ctx.Queue(maxsize=0)
+            out_q: mp.Queue = self._ctx.Queue(maxsize=0)
             self._stage_in_queues.append(in_q)
             self._stage_out_queues.append(out_q)
 
             # Attach queues and start Stage-owned worker process
             stage.attach_queues(in_q, out_q)
-
-            # Build connectors config for this stage
-            stage_connectors_config = get_stage_connector_config(
-                self.omni_transfer_config,
-                stage_id,
-            )
-
             stage.init_stage_worker(
                 model,
                 log_file=self._log_file,
                 shm_threshold_bytes=self._shm_threshold_bytes,
                 ctx=self._ctx,
                 batch_timeout=self.batch_timeout,
-                connectors_config=stage_connectors_config,
             )
             logger.debug("[Orchestrator] Stage-%s process started", stage_id)
             time.sleep(self._init_sleep_seconds)
 
-    def _start_stage_actors_ray(self, model: str) -> None:
-        import ray
-        import os
-        from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
-
-        # Initialize Ray Cluster
-        if not ray.is_initialized():
-             # Create a dummy parallel config just to reuse vLLM's init logic if needed
-             # or just call ray.init()
-             # Pass current PYTHONPATH to workers to ensure they can find vllm_omni
-             runtime_env = {
-                 "env_vars": {"PYTHONPATH": os.environ.get("PYTHONPATH", "")}
-             }
-             ray.init(address=self.ray_address, ignore_reinit_error=True, runtime_env=runtime_env)
-
-        # Create Placement Group
-        # We allocate 1 bundle per stage. Assume 1 GPU per stage for now.
-        # Must include CPU=1 because default Actor request includes 1 CPU.
-        bundles = [{"GPU": 1.0, "CPU": 1.0} for _ in self.stage_list]
-        self._ray_pg = ray.util.placement_group(bundles, strategy="SPREAD")
-        ray.get(self._ray_pg.ready())
-        logger.info("[Orchestrator] Ray Placement Group created")
-
-        # Define Actor Wrapper
-        @ray.remote(num_gpus=1)
-        class OmniStageRayWorker:
-            def run(self, func, *args, **kwargs):
-                return func(*args, **kwargs)
-
-        from vllm_omni.entrypoints.omni_stage import _stage_worker
-
-        for stage_id, stage in enumerate(self.stage_list):
-            in_q = self._queue_cls()
-            out_q = self._queue_cls()
-            self._stage_in_queues.append(in_q)
-            self._stage_out_queues.append(out_q)
-            stage.attach_queues(in_q, out_q)
-
-            stage_connectors_config = get_stage_connector_config(
-                self.omni_transfer_config,
-                stage_id,
-            )
-
-            # Prepare payload
-            # Reuse logic from OmniStage.init_stage_worker but do not start process
-            engine_args = _to_dict(stage.engine_args)
-            runtime_cfg = _to_dict(getattr(stage.stage_config, "runtime", {}))
-            stage_payload = {
-                "stage_id": stage.stage_id,
-                "engine_args": engine_args,
-                "runtime": runtime_cfg,
-                "shm_threshold_bytes": sys.maxsize, # Disable SHM for Ray mode by setting huge threshold
-                "connectors_config": stage_connectors_config or {},
-            }
-
-            # Start Actor
-            # Inject PYTHONPATH again for the specific actor to be safe
-            worker_actor = OmniStageRayWorker.options(
-                scheduling_strategy=PlacementGroupSchedulingStrategy(
-                    placement_group=self._ray_pg,
-                    placement_group_bundle_index=stage_id
-                ),
-                runtime_env={"env_vars": {"PYTHONPATH": os.environ.get("PYTHONPATH", "")}}
-            ).remote()
-
-            # Start worker loop asynchronously
-            worker_actor.run.remote(
-                _stage_worker,
-                model=model,
-                stage_payload=stage_payload,
-                in_q=in_q,
-                out_q=out_q,
-                log_file=self._log_file,
-                batch_timeout=self.batch_timeout,
-            )
-            
-            stage._ray_actor = worker_actor
-            logger.debug("[Orchestrator] Stage-%s Ray Actor started", stage_id)
-            time.sleep(self._init_sleep_seconds)
-
     def close(self) -> None:
         """Close all stage processes and clean up resources.
 
@@ -296,21 +171,10 @@ class OmniLLM:
                 )
         for stage in self.stage_list:
             try:
-                if self.worker_backend == "ray" and hasattr(stage, "_ray_actor"):
-                    import ray
-                    # Ray kill
-                    ray.kill(stage._ray_actor)
-                else:
-                    stage.stop_stage_worker()
+                stage.stop_stage_worker()
             except Exception as e:
                 logger.warning("[Orchestrator] Failed to stop stage worker: %s", e)
 
-        if self.worker_backend == "ray" and self._ray_pg:
-            try:
-                ray.util.remove_placement_group(self._ray_pg)
-            except Exception as e:
-                logger.warning("[Orchestrator] Failed to remove placement group: %s", e)
-
     def __del__(self) -> None:  # best-effort
         try:
             self.close()
@@ -373,14 +237,11 @@ class OmniLLM:
         # Orchestrator keeps stage objects for input derivation
         num_stages = len(self.stage_list)
 
-        # Generate globally unique request IDs and map them to original prompts
-        request_ids: list[str] = [f"{i}_{uuid.uuid4()}" for i in range(len(request_prompts))]
-        request_id_to_prompt: dict[str, PromptType] = {
-            rid: p for rid, p in zip(request_ids, request_prompts)
-        }
+        # Map from request_id to original prompt
+        request_id_to_prompt: dict[int, PromptType] = {i: p for i, p in enumerate(request_prompts)}
 
         # Track per-request start time for end-to-end timing
-        _req_start_ts: dict[str, float] = {}
+        _req_start_ts: dict[int, float] = {}
         _wall_start_ts: float = time.time()
 
         # Determine the final stage for E2E stats (highest stage_id with final_output=True; fallback to last stage)
@@ -526,90 +387,51 @@ class OmniLLM:
                         )
                         continue
                     sp_next: SamplingParams = sampling_params_list[next_stage_id]  # type: ignore[index]
-
-                    # Check if we have a connector for this edge
-                    connector_key = (str(stage_id), str(next_stage_id))
-                    connector = self.connectors.get(connector_key)
-                    print(f"connectors: {self.connectors}")
-
-                    sent_via_connector = False
-                    if connector:
-                        sent_via_connector = try_send_via_connector(
-                            connector=connector,
-                            stage_id=stage_id,
-                            next_stage_id=next_stage_id,
-                            req_id=req_id,
-                            next_inputs=next_inputs,
-                            sampling_params=sp_next,
-                            original_prompt=request_id_to_prompt[req_id],
-                            next_stage_queue_submit_fn=self.stage_list[next_stage_id].submit,
-                            metrics=metrics
-                        )
-
-                    if not sent_via_connector:
-                        # Use original queue mechanism
+                    try:
+                        # Measure transfer size and time (encode + enqueue)
+                        size_bytes = 0
                         try:
-                            if self.worker_backend == "ray":
-                                # Ray mode: Skip SHM fallback, just send object via Queue (Ray handles serialization)
-                                # This prevents SHM access errors across nodes
-                                ipc_payload = {
-                                    "request_id": req_id,
-                                    "engine_inputs": next_inputs,
-                                    "sampling_params": sp_next,
-                                    "sent_ts": time.time(),
-                                }
-                                size_bytes = 0 # Approximation or calculate size if needed
-                                try:
-                                    # Approximate size for stats
-                                    size_bytes = len(_set(next_inputs))
-                                except:
-                                    pass
-                                t0 = time.time()
-                                # Direct submit to Ray Queue
-                            else:
-                                size_bytes = 0
-                                try:
-                                    size_bytes = len(_set(next_inputs))
-                                except Exception:
-                                    size_bytes = 0
-                                t0 = time.time()
-                                ipc_payload = _encode(
-                                    next_inputs,
-                                    getattr(self, "_shm_threshold_bytes", 65536),
-                                    obj_key="engine_inputs",
-                                    shm_key="engine_inputs_shm",
-                                )
-                                ipc_payload.update(
-                                    {
-                                        "request_id": req_id,
-                                        "sampling_params": sp_next,
-                                        "sent_ts": time.time(),
-                                    }
-                                )
-                            self.stage_list[next_stage_id].submit(ipc_payload)
-                            t1 = time.time()
-                            tx_ms = (t1 - t0) * 1000.0
-                            metrics.on_forward(
-                                stage_id,
-                                next_stage_id,
-                                req_id,
-                                int(size_bytes),
-                                float(tx_ms),
-                                bool("engine_inputs_shm" in ipc_payload),
-                            )
-                        except Exception as e:
-                            logger.warning(
-                                "[Orchestrator] IPC encode failed for req %s: %s; falling back to inline payload",
-                                req_id,
-                                e,
-                            )
-                            self.stage_list[next_stage_id].submit(
-                                {
-                                    "request_id": req_id,
-                                    "engine_inputs": next_inputs,
-                                    "sampling_params": sp_next,
-                                }
-                            )
+                            size_bytes = len(_set(next_inputs))
+                        except Exception:
+                            size_bytes = 0
+                        t0 = time.time()
+                        ipc_payload = _encode(
+                            next_inputs,
+                            getattr(self, "_shm_threshold_bytes", 65536),
+                            obj_key="engine_inputs",
+                            shm_key="engine_inputs_shm",
+                        )
+                        ipc_payload.update(
+                            {
+                                "request_id": req_id,
+                                "sampling_params": sp_next,
+                                "sent_ts": time.time(),
+                            }
+                        )
+                        self.stage_list[next_stage_id].submit(ipc_payload)
+                        t1 = time.time()
+                        tx_ms = (t1 - t0) * 1000.0
+                        metrics.on_forward(
+                            stage_id,
+                            next_stage_id,
+                            req_id,
+                            int(size_bytes),
+                            float(tx_ms),
+                            bool("engine_inputs_shm" in ipc_payload),
+                        )
+                    except Exception as e:
+                        logger.warning(
+                            "[Orchestrator] IPC encode failed for req %s: %s; falling back to inline payload",
+                            req_id,
+                            e,
+                        )
+                        self.stage_list[next_stage_id].submit(
+                            {
+                                "request_id": req_id,
+                                "engine_inputs": next_inputs,
+                                "sampling_params": sp_next,
+                            }
+                        )
                     logger.debug(
                         "[Orchestrator] Forwarded request %s to stage-%s",
                         req_id,

vllm_omni/entrypoints/omni_stage.py

@@ -163,7 +163,6 @@ class OmniStage:
         shm_threshold_bytes: int = 65536,
         ctx: Optional[mp.context.BaseContext] = None,
         batch_timeout: int = 10,
-        connectors_config: Optional[dict] = None,  # New parameter
     ) -> None:
         """Initialize and start the stage worker process.
 
@@ -193,7 +192,6 @@ class OmniStage:
             "engine_args": engine_args,
             "runtime": runtime_cfg,
             "shm_threshold_bytes": self._shm_threshold_bytes,
-            "connectors_config": connectors_config or {},
         }
         if is_async:
             self._proc = ctx.Process(
@@ -342,8 +340,6 @@ def _stage_worker(
         log_stage_batch_stats,
         log_stage_running_avg,
     )
-    from vllm_omni.distributed.connectors.adapter import try_recv_via_connector
-    from vllm_omni.distributed.connectors.utils import build_stage_connectors
     from vllm_omni.entrypoints.omni_llm import OmniStageLLM
     # no inline JSONL/serialization imports; logging handled by utilities
 
@@ -351,7 +347,6 @@ def _stage_worker(
     engine_args = stage_payload.get("engine_args", {})
     runtime_cfg = stage_payload.get("runtime", {})
     shm_threshold_bytes = int(stage_payload.get("shm_threshold_bytes", 65536))
-    connectors_config = stage_payload.get("connectors_config", {})
 
     # Per-stage logger: clear inherited handlers to avoid broken parent streams
     try:
@@ -445,18 +440,6 @@ def _stage_worker(
     )
     stage_engine = OmniStageLLM(model=model, **engine_args)
     _logging.getLogger(__name__).debug("[Stage-%s] Engine initialized", stage_id)
-
-    # Initialize OmniConnectors if configured
-    connectors = {}
-    if connectors_config:
-        built_connectors = build_stage_connectors(
-            stage_id=stage_id,
-            connectors_config=connectors_config,
-        )
-        if built_connectors is None:
-            return
-        connectors = built_connectors
-
     # Signal readiness to orchestrator
     try:
         out_q.put({"type": "stage_ready", "stage_id": stage_id})
@@ -513,21 +496,9 @@ def _stage_worker(
                     _in_flight_ms_by_rid[rid] = 0.0
             except Exception:
                 _in_flight_ms_by_rid[rid] = 0.0
-
-            # Resolve input data (via Connector or IPC)
-            ein, _rx_metrics = try_recv_via_connector(
-                task=t,
-                connectors=connectors,
-                stage_id=stage_id
-            )
-
-            if ein is None or _rx_metrics is None:
-                ein, _rx_metrics = maybe_load_from_ipc_with_metrics(task, obj_key="engine_inputs", shm_key="engine_inputs_shm")
-
-            if _rx_metrics:
-                _rx_decode_ms_by_rid[rid] = float(_rx_metrics.get("rx_decode_time_ms", 0.0))
-                _rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0))
-
+            ein, _rx_metrics = maybe_load_from_ipc_with_metrics(t, obj_key="engine_inputs", shm_key="engine_inputs_shm")
+            _rx_decode_ms_by_rid[rid] = float(_rx_metrics.get("rx_decode_time_ms", 0.0))
+            _rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0))
             batch_request_ids.append(rid)
             if isinstance(ein, list):
                 batch_engine_inputs.extend(ein)

vllm_omni/entrypoints/utils.py

@@ -10,41 +10,6 @@ from vllm_omni.utils import detect_device_type
 PROJECT_ROOT = Path(__file__).parent.parent.parent
 
 
-def resolve_model_config_path(model: str) -> str:
-    """Resolve the stage config file path from the model name.
-
-    Resolves stage configuration path based on the model type and device type.
-    First tries to find a device-specific YAML file from stage_configs/{device_type}/
-    directory. If not found, falls back to the default config file.
-
-    Args:
-        model: Model name or path (used to determine model_type)
-
-    Returns:
-        String path to the stage configuration file
-
-    Raises:
-        FileNotFoundError: If no stage config file exists for the model type
-    """
-    hf_config = get_config(model, trust_remote_code=True)
-    model_type = hf_config.model_type
-    device_type = detect_device_type()
-
-    # Try device-specific config first
-    if device_type != "cuda":
-        device_config_file = f"vllm_omni/model_executor/stage_configs/{device_type}/{model_type}.yaml"
-        device_config_path = PROJECT_ROOT / device_config_file
-        if os.path.exists(device_config_path):
-            return str(device_config_path)
-
-    # Fall back to default config
-    stage_config_file = f"vllm_omni/model_executor/stage_configs/{model_type}.yaml"
-    stage_config_path = PROJECT_ROOT / stage_config_file
-    if not os.path.exists(stage_config_path):
-        raise FileNotFoundError(f"Stage config file {stage_config_path} not found")
-    return str(stage_config_path)
-
-
 def load_stage_configs_from_model(model: str) -> list:
     """Load stage configurations from model's default config file.
 
@@ -61,8 +26,24 @@ def load_stage_configs_from_model(model: str) -> list:
     Raises:
         FileNotFoundError: If no stage config file exists for the model type
     """
-    stage_config_path = resolve_model_config_path(model)
-    stage_configs = load_stage_configs_from_yaml(config_path=stage_config_path)
+    hf_config = get_config(model, trust_remote_code=True)
+    model_type = hf_config.model_type
+    device_type = detect_device_type()
+
+    # Try device-specific config first
+    if device_type != "cuda":
+        device_config_file = f"vllm_omni/model_executor/stage_configs/{device_type}/{model_type}.yaml"
+        device_config_path = PROJECT_ROOT / device_config_file
+        if os.path.exists(device_config_path):
+            stage_configs = load_stage_configs_from_yaml(config_path=str(device_config_path))
+            return stage_configs
+
+    # Fall back to default config
+    stage_config_file = f"vllm_omni/model_executor/stage_configs/{model_type}.yaml"
+    stage_config_path = PROJECT_ROOT / stage_config_file
+    if not os.path.exists(stage_config_path):
+        raise FileNotFoundError(f"Stage config file {stage_config_path} not found")
+    stage_configs = load_stage_configs_from_yaml(config_path=str(stage_config_path))
     return stage_configs
 
 

vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml

@@ -1,4 +1,6 @@
 # stage config for running qwen2.5-omni with architecture of OmniLLM.
+
+# The following config has been verified on 1x H100-80G GPU.
 stage_args:
   - stage_id: 0
     runtime:
@@ -26,9 +28,6 @@ stage_args:
       seed: 42
       detokenize: True
       repetition_penalty: 1.1
-    # Distributed connector configuration (optional)
-    output_connectors:
-      to_stage_1: connector_thinker2talker # thinker -> talker: use local mooncake
   - stage_id: 1
     runtime:
       process: true
@@ -55,22 +54,17 @@ stage_args:
       detokenize: True
       repetition_penalty: 1.05
       stop_token_ids: [8294]
-    # Distributed connector configuration (optional)
-    input_connectors:
-      from_stage_0: connector_thinker2talker # talker <- thinker: use local mooncake
-    output_connectors:
-      to_stage_2: connector_talker2generator # talker -> code2wav: use remote mooncake (example)
   - stage_id: 2
     runtime:
       process: true
-      devices: "2"            # Example: use a different GPU than the previous stage; use "0" if single GPU
+      devices: "0"            # Example: use a different GPU than the previous stage; use "0" if single GPU
       max_batch_size: 1
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_cls: vllm_omni.worker.gpu_diffusion_worker.GPUDiffusionWorker
       scheduler_cls: vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler
-      gpu_memory_utilization: 0.3
+      gpu_memory_utilization: 0.15
       enforce_eager: true
       trust_remote_code: true
       enable_prefix_caching: false
@@ -86,9 +80,6 @@ stage_args:
       seed: 42
       detokenize: True
       repetition_penalty: 1.1
-    # Distributed connector configuration (optional)
-    input_connectors:
-      from_stage_1: connector_talker2generator # code2wav <- talker: use remote mooncake
 
 # Top-level runtime config (concise): default windows and stage edges
 runtime:
@@ -96,31 +87,6 @@ runtime:
   defaults:
     window_size: -1             # Simplified: trigger downstream only after full upstream completion
     max_inflight: 1             # Simplified: process serially within each stage
-
-  # Distributed connectors configuration (optional)
-  connectors:
-    # Local mooncake for same-node communication
-    connector_thinker2talker:
-      name: MooncakeConnector
-      extra:
-        host: "127.0.0.1"
-        metadata_server: "http://10.90.67.86:8080/metadata"
-        master: "10.90.67.86:50051"
-        segment: 512000000    # 512MB
-        localbuf: 64000000     # 64MB
-        proto: "tcp"
-
-    # Remote mooncake for cross-node communication (example)
-    connector_talker2generator:
-      name: MooncakeConnector
-      extra:
-        host: "127.0.0.1"
-        metadata_server: "http://10.90.67.86:8080/metadata"
-        master: "10.90.67.86:50051"
-        segment: 1073741824    # 1GB
-        localbuf: 134217728     # 128MB
-        proto: "tcp"
-
   edges:
     - from: 0                   # thinker → talker: trigger only after receiving full input (-1)
       to: 1

vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml

@@ -3,24 +3,25 @@
 # Stage 1: Talker (text embeddings → 8-layer RVQ codec codes)
 # Stage 2: Code2Wav (8-layer RVQ codes → audio waveform)
 
+# The following config has been verified on 2x H100-80G GPUs.
 stage_args:
   - stage_id: 0
     runtime:
-      devices: "0"
+      devices: "0,1"
       max_batch_size: 1
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
       scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
-      gpu_memory_utilization: 0.8
+      gpu_memory_utilization: 0.6
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: latent  # Output hidden states for talker
       distributed_executor_backend: "mp"
       enable_prefix_caching: false
       hf_config_name: thinker_config
-      # tensor_parallel_size: 2  # Commented out - using single GPU for thinker
+      tensor_parallel_size: 2
     final_output: true
     final_output_type: text
     is_comprehension: true
@@ -42,7 +43,7 @@ stage_args:
        model_arch: Qwen3OmniMoeForConditionalGeneration
        worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
        scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
-       gpu_memory_utilization: 0.5
+       gpu_memory_utilization: 0.3
        enforce_eager: true
        trust_remote_code: true
        engine_output_type: latent  # Output codec codes for code2wav
@@ -65,7 +66,7 @@ stage_args:
 
   - stage_id: 2
     runtime:
-      devices: "1"
+      devices: "0"
       max_batch_size: 1
     engine_args:
       model_stage: code2wav

vllm_omni/patch.py

@@ -15,6 +15,9 @@ from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding
 from vllm_omni.request import OmniRequest
 
 for module_name, module in sys.modules.items():
+    # only do patch on module of vllm, pass others
+    if "vllm" not in module_name:
+        continue
     if hasattr(module, "EngineCoreOutput") and module.EngineCoreOutput == _OriginalEngineCoreOutput:
         module.EngineCoreOutput = OmniEngineCoreOutput
     if hasattr(module, "EngineCoreOutputs") and module.EngineCoreOutputs == _OriginalEngineCoreOutputs: