p

Signed-off-by: Kevin H. Luu <khluu000@gmail.com>
p
2025-12-04 02:17:07 -08:00 · 2025-12-04 02:08:14 -08:00 · 2025-12-04 02:04:51 -08:00 · 2025-12-04 01:59:29 -08:00 · 2025-12-04 01:55:18 -08:00 · 2025-12-04 00:04:18 -08:00
106 changed files with 6833 additions and 1257 deletions
@@ -1,6 +1,23 @@
 steps:
+  - label: ":docker: Build image"
+    key: image-build
+    commands:
+      - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
+      - "docker build --file docker/Dockerfile.ci -t vllm-omni-ci ."
+      - "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:omni-latest"
+      - "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:omni-latest"
+    agents:
+      queue: "cpu_queue_premerge_us_east_1"
+
  - label: "Simple Unit Test"
+    depends_on: image-build
    commands:
      - ".buildkite/scripts/simple_test.sh"
    agents:
      queue: "cpu_queue_premerge"
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:omni-latest
+          always-pull: true
+          propagate-environment: true
+
@@ -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.
@@ -156,15 +156,9 @@ cython_debug/
 # cursor
 .cursor/

-# docker
-docker/
-
 # scripts
 /scripts/

-# tests
-tests/
-
 # PyCharm
 #  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
 #  be added to the global gitignore or merged into this project gitignore.  For a PyCharm
@@ -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/)
@@ -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://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>

@@ -57,8 +57,6 @@ vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, in

 Visit our [documentation](https://vllm-omni.readthedocs.io/en/latest/) to learn more.

-Visit our documentation to learn more.
-
 - [Installation](https://vllm-omni.readthedocs.io/en/latest/getting_started/installation/)
 - [Quickstart](https://vllm-omni.readthedocs.io/en/latest/getting_started/quickstart/)
 - [List of Supported Models](https://vllm-omni.readthedocs.io/en/latest/models/supported_models/)
@@ -68,10 +66,8 @@ Visit our documentation to learn more.
 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

@@ -0,0 +1,14 @@
+ARG VLLM_BASE_IMAGE=vllm/vllm-openai
+ARG VLLM_BASE_TAG=v0.11.0
+FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
+ARG APP_DIR=/workspace/vllm-omni
+WORKDIR ${APP_DIR}
+
+COPY . .
+
+# Install vllm-omni into the same uv-managed Python environment used by the base image.
+RUN uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]"
+
+RUN ln -sf /usr/bin/python3 /usr/bin/python
+
+ENTRYPOINT []
@@ -1,44 +1,45 @@
 nav:
-  - Home: README.md
-  - User Guide:
-    - Getting Started:
-      - getting_started/quickstart.md
-      - getting_started/installation
-    - Examples:
-      - examples/README.md
-      - Offline Inference:
-        - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
-      - Online Serving:
-        - Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
-        - Qwen2.5-Image: user_guide/examples/offline_inference/qwen_image.md
-    - General:
-      - usage/*
-    - Configuration:
-      - configuration/README.md
-      - configuration/*
-    - Models:
-      - models/supported_models.md
-  - Developer Guide:
-    - General:
-      - contributing/README.md
-      - glob: contributing/*
-        flatten_single_child_sections: true
-    - Model Implementation:
-      - contributing/model/README.md
-    - CI: contributing/ci
-    - Design Documents:
-      - design/index.md
-      - 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
-    - api/vllm_omni
-  - CLI Reference: cli
-  - Community:
-    - community/*
-    - Slack: https://slack.vllm.ai
-    - Blog: https://blog.vllm.ai
-    - Forum: https://discuss.vllm.ai
+- Home: README.md
+- User Guide:
+  - Getting Started:
+    - getting_started/quickstart.md
+    - getting_started/installation
+  - 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
+      - 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 Qwen3-omni: user_guide/examples/online_serving/qwen3_omni.md
+  - General:
+    - usage/*
+  - Configuration:
+    - configuration/README.md
+    - configuration/*
+  - Models:
+    - models/supported_models.md
+- Developer Guide:
+  - General:
+    - contributing/README.md
+    - glob: contributing/*
+      flatten_single_child_sections: true
+  - Model Implementation:
+    - contributing/model/README.md
+  - CI: contributing/ci
+  - Design Documents:
+    - design/index.md
+    - design/architecture_overview.md
+    - design/vllm_omni_design.md
+    - design/mrs_design.md
+  - Docs Guide: contributing/DOCS_GUIDE.md
+- API Reference:
+  - api/README.md
+  - api/vllm_omni
+- CLI Reference: cli
+- Community:
+  - community/*
+  - Slack: https://slack.vllm.ai
+  - Blog: https://blog.vllm.ai
+  - Forum: https://discuss.vllm.ai
@@ -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)
@@ -4,8 +4,8 @@

 Main entry points for vLLM-Omni inference and serving.

- [vllm_omni.entrypoints.async_omni_llm.AsyncOmniLLM][]
- [vllm_omni.entrypoints.async_omni_llm.AsyncOmniStageLLM][]
+- [vllm_omni.entrypoints.async_omni.AsyncOmni][]
+- [vllm_omni.entrypoints.async_omni.AsyncOmniStageLLM][]
 - [vllm_omni.entrypoints.chat_utils.OmniAsyncMultiModalContentParser][]
 - [vllm_omni.entrypoints.chat_utils.OmniAsyncMultiModalItemTracker][]
 - [vllm_omni.entrypoints.chat_utils.parse_chat_messages_futures][]
@@ -34,6 +34,8 @@ Engine classes for offline and online inference.
 - [vllm_omni.diffusion.diffusion_engine.DiffusionEngine][]
 - [vllm_omni.engine.AdditionalInformationEntry][]
 - [vllm_omni.engine.AdditionalInformationPayload][]
+- [vllm_omni.engine.OmniEngineCoreOutput][]
+- [vllm_omni.engine.OmniEngineCoreOutputs][]
 - [vllm_omni.engine.OmniEngineCoreRequest][]
 - [vllm_omni.engine.PromptEmbedsPayload][]
 - [vllm_omni.engine.arg_utils.AsyncOmniEngineArgs][]
@@ -48,6 +50,7 @@ Core scheduling and caching components.

 - [vllm_omni.core.dit_cache_manager.DiTCacheManager][]
 - [vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler][]
+- [vllm_omni.core.sched.generation_scheduler.GenerationScheduler][]
 - [vllm_omni.core.sched.output.OmniNewRequestData][]
 - [vllm_omni.core.sched.scheduler.OmniScheduler][]

@@ -55,7 +58,7 @@ Core scheduling and caching components.

 Model execution components.

- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni.OmniOutput][]
+- [vllm_omni.model_executor.models.output_templates.OmniOutput][]
 - [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni.Qwen2_5OmniForConditionalGeneration][]
 - [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_talker.Qwen2_5OmniTalkerForConditionalGeneration][]
 - [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker.Qwen2_5OmniAudioFeatureInputs][]
@@ -72,6 +75,22 @@ Model execution components.
 - [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_old.Qwen2EmbeddingModel][]
 - [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_old.Qwen2ForCausalLM][]
 - [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_old.Qwen2Model][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_moe.Qwen3MoeForCausalLM][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni.Qwen3OmniMoeForConditionalGeneration][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav.Qwen3OmniMoeCode2Wav][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mtp.Qwen3OmniCodePredictorBaseModel][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mtp.Qwen3OmniMoeTalkerCodePredictor][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_talker.Qwen3OmniMoeModel][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_talker.Qwen3OmniMoeTalkerForConditionalGeneration][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3MoeLLMForCausalLM][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3MoeLLMModel][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeConditionalGenerationMixin][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeThinkerForConditionalGeneration][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeThinkerMultiModalProcessor][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeThinkerProcessingInfo][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3Omni_VisionTransformer][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3_VisionPatchEmbed][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3_VisionPatchMerger][]

 ## Configuration

@@ -85,8 +104,17 @@ 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][]
 - [vllm_omni.worker.gpu_diffusion_worker.GPUDiffusionWorker][]
+- [vllm_omni.worker.gpu_generation_model_runner.GPUGenerationModelRunner][]
+- [vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker][]
 - [vllm_omni.worker.gpu_model_runner.OmniGPUModelRunner][]
+- [vllm_omni.worker.npu.npu_ar_model_runner.NPUARModelRunner][]
+- [vllm_omni.worker.npu.npu_ar_worker.NPUARWorker][]
+- [vllm_omni.worker.npu.npu_diffusion_model_runner.NPUDiffusionModelRunner][]
+- [vllm_omni.worker.npu.npu_diffusion_worker.NPUDiffusionWorker][]
+- [vllm_omni.worker.npu.npu_model_runner.OmniNPUModelRunner][]
@@ -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
@@ -0,0 +1,12 @@
+# Volunteers for Bugfix and CI
+
+We encourage you to check current docs and [issues](https://github.com/vllm-project/vllm-omni/issues) to find possible solutions for your questions. If non of these can solve it, please propose an issue to describe your questions about bug or CI problems for developing.
+
+If you have urgent need for locating and solving bugfix or CI problems, please find community volunteers below.
+
+| Dec 4-Dec 12 | Dec 15-Dec 19 | Dec 22-Dec 26 | Dec 29- Jan 2, 2026| Jan 5-Jan 9 | Jan 12-Jan 16 |
+|----------|----------|----------|----------|----------|----------|
+| <a href="https://github.com/congw729">Conw729</a> | <a href="https://github.com/yinpeiqi">yinpeiqi</a> | <a href="https://github.com/tzhouam">tzhouam</a> | <a href="https://github.com/SamitHuang">SamitHuang</a> | <a href="https://github.com/gcanlin">gcanlin</a> | <a href="https://github.com/natureofnature">natureofnature</a> |
+| <a href="https://github.com/david6666666">david6666666</a> | <a href="https://github.com/R2-Y">R2-Y</a> | <a href="https://github.com/hsliuustc0106">hsliuustc0106</a> | <a href="https://github.com/Gaohan123">Gaohan123</a> | <a href="https://github.com/ZJY0516">ZJY0516</a> | <a href="https://github.com/qibaoyuan">qibaoyuan</a> |
+
+We kindly welcome more contributors to fix bugs and contribute new features!
@@ -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`.
@@ -44,7 +47,7 @@ stage_args:
      trust_remote_code: true # Needed by huggingface config parsing
      engine_output_type: latent  # It claims that the stage will input latent hiddenstates besides token ids
      enable_prefix_caching: false # For request with hiddenstates output, the prefix caching is not supported now
-    is_comprehension: true # If the stage is a text or multimodal comprehension module. If it is, the AsyncOmniLLM will use its tokenizer as default
+    is_comprehension: true # If the stage is a text or multimodal comprehension module. If it is, the AsyncOmni will use its tokenizer as default
    final_output: true # If the stage has output as part of final outputs. If it is false, which means that the stage only works as a intermediate role.
    final_output_type: text # What is the final output type. It can be text and audio now.
    default_sampling_params: # sampling parameters for the stage. Their meaning aligns with vLLM.
@@ -206,7 +209,7 @@ Default: `false`

 ### `is_comprehension`

-Whether this stage is a text or multimodal comprehension module. When set to `true`, the stage acts as a comprehension module that processes input text or multimodal content. If this is the first comprehension stage, `AsyncOmniLLM` will use its tokenizer as the default tokenizer for the entire pipeline.
+Whether this stage is a text or multimodal comprehension module. When set to `true`, the stage acts as a comprehension module that processes input text or multimodal content. If this is the first comprehension stage, `AsyncOmni` will use its tokenizer as the default tokenizer for the entire pipeline.

 Default: `true`

@@ -9,7 +9,7 @@ This directory contains the source files for the vLLM-Omni documentation.
 Install documentation dependencies:

 ```bash
-pip install -e ".[docs]"
+uv pip install -e ".[docs]"
 ```

 ### Build and Serve Documentation
@@ -84,7 +84,7 @@ The documentation is automatically deployed to GitHub Pages using GitHub Actions
   ```

 3. **Documentation will be available at**:
-   - `https://vllm-project.github.io/vllm-omni/`
+   - `https://vllm-omni.readthedocs.io`

 The GitHub Actions workflow (`.github/workflows/docs.yml`) will automatically:
 - Build the documentation when you push to `main` branch
@@ -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]"
 ```

@@ -1,143 +0,0 @@
-# API Design Documentation
-
-This document contains comprehensive API design documentation for all core modules in vLLM-Omni. These templates provide a standardized structure for designing and implementing the core, engine, executor, and worker modules.
-
-## 📋 Module API
-
-### Core Module API
-**Core module** provides fundamental scheduling, caching, and resource management functionality.
-
-**Key Components:**
- Request scheduling and prioritization
- DiT cache management for diffusion models
- Resource allocation and coordination
- Inter-module communication
-
-### Engine Module API
-**Engine module** handles model loading, inference execution, and output processing.
-
-**Key Components:**
- Model loading and initialization
- Inference execution for AR and diffusion models
- Input preprocessing for various modalities
- Output postprocessing and formatting
-
-### Executor Module API
-**Executor module** coordinates and manages request execution across different engines and workers.
-
-**Key Components:**
- Request routing and load balancing
- Execution pipeline coordination
- Worker management and task distribution
- Error handling and recovery
-
-### Worker Module API
-**Worker module** provides the actual execution environment for model inference.
-
-**Key Components:**
- Model execution and inference
- GPU resource management
- Request batching and processing
- Performance optimization
-
-
-## 🔧 Implementation Guidelines
-
-### Code Standards
- Use type hints for all methods
- Include comprehensive docstrings
- Follow PEP 8 style guidelines
- Use async/await for I/O operations
-
-### Error Handling
- Define specific exception types
- Provide meaningful error messages
- Include error recovery strategies
- Log errors appropriately
-
-### Configuration
- Use dataclasses for configuration
- Provide sensible defaults
- Include validation methods
- Support environment variables
-
-### Testing
- Write unit tests for all public methods
- Include integration tests
- Test error conditions
- Validate configuration options
-
-## 🚀 Getting Started
-
-### For Developers
-
-1. **Choose a Module**
-   - Pick the module you want to work on
-   - Read the corresponding API documentation
-   - Understand the responsibilities and interfaces
-
-2. **Set Up Development Environment**
-   ```bash
-   # Clone the repository
-   git clone https://github.com/vllm-project/vllm-omni.git
-   cd vllm-omni
-
-   # Install dependencies
-   pip install -e ".[dev]"
-
-   # Set up pre-commit hooks
-   pre-commit install
-   ```
-
-3. **Start Implementation**
-   - Create the module directory
-   - Implement base classes first
-   - Add concrete implementations
-   - Write tests as you go
-
-### For Contributors
-
-1. **Review API Documentation**
-   - Read through all module APIs
-   - Understand the overall architecture
-   - Identify areas for improvement
-
-2. **Propose Changes**
-   - Create issues for API changes
-   - Discuss in pull requests
-   - Update documentation accordingly
-
-3. **Submit Contributions**
-   - Follow the coding standards
-   - Include tests for new features
-   - Update documentation
-   - Submit pull requests
-
-## 🤝 Contributing
-
-We welcome contributions to improve the API documentation and implementation. Please:
-
-1. **Follow the Template Structure**
-   - Use the [api template](https://github.com/vllm-project/vllm-omni/docs/contributing/design_documents/api_design_template.md) as a guide
-   - Maintain consistency across modules
-   - Include all required sections
-
-2. **Provide Working Examples**
-   - Test all code examples
-   - Include both basic and advanced usage
-   - Show integration patterns
-
-3. **Keep Documentation Updated**
-   - Update docs when APIs change
-   - Version control documentation changes
-   - Maintain backward compatibility notes
-
-## 📝 Notes
-
- All API documentation should be kept up-to-date with implementation
- Code examples should be tested and working
- Configuration options should be validated
- Error handling should be comprehensive
- Performance considerations should be documented
-
-For questions or suggestions about the API documentation, please open an issue or start a discussion in the repository.
@@ -1,336 +0,0 @@
-# API Design Template for vLLM-Omni Modules
-
-This template provides a standardized structure for designing APIs for core, engine, executor, and worker modules in vLLM-Omni. Use this template to ensure consistency and completeness across all modules.
-
-## Template Structure
-
-### 1. Module Overview
- **Purpose**: What this module does
- **Responsibilities**: Key responsibilities of the module
- **Dependencies**: What other modules this depends on
- **Integration Points**: How it integrates with other modules
-
-### 2. Core Classes/Interfaces
- **Base Classes**: Abstract base classes and interfaces
- **Implementation Classes**: Concrete implementations
- **Data Structures**: Key data structures and models
-
-### 3. Public API Methods
- **Initialization**: Constructor and setup methods
- **Core Operations**: Main functionality methods
- **Configuration**: Configuration and parameter methods
- **Lifecycle Management**: Start, stop, cleanup methods
- **Monitoring**: Status, metrics, and debugging methods
-
-### 4. Configuration
- **Configuration Classes**: Dataclasses or config objects
- **Required Parameters**: Must-have configuration
- **Optional Parameters**: Optional configuration with defaults
- **Validation**: Parameter validation rules
-
-### 5. Error Handling
- **Custom Exceptions**: Module-specific exceptions
- **Error Codes**: Standardized error codes
- **Recovery Strategies**: How to handle and recover from errors
-
-### 6. Examples
- **Basic Usage**: Simple usage examples
- **Advanced Usage**: Complex scenarios
- **Integration Examples**: How to use with other modules
-
---
-
-## Example: Core Module Template
-
-### 1. Module Overview
-
-**Purpose**: The core module provides fundamental scheduling and caching functionality for vLLM-Omni.
-
-**Responsibilities**:
- Request scheduling and prioritization
- DiT cache management
- Resource allocation and coordination
- Inter-module communication
-
-**Dependencies**:
- `vllm_omni.request` - Request handling
- `vllm_omni.config` - Configuration management
- `vllm_omni.utils` - Utility functions
-
-**Integration Points**:
- Receives requests from entrypoints
- Coordinates with engine modules
- Manages worker allocation
- Provides status to monitoring systems
-
-### 2. Core Classes/Interfaces
-
-```python
-from abc import ABC, abstractmethod
-from typing import List, Optional, Dict, Any
-from dataclasses import dataclass
-from enum import Enum
-
-class DiTCacheManager:
-    """Manages DiT cache for diffusion models."""
-
-    def __init__(self, config: DiTCacheConfig):
-        self.config = config
-        self.cache = {}
-        self.cache_stats = {}
-
-    async def get_cache(self, cache_key: str) -> Optional[Any]:
-        """Retrieve cached data."""
-        pass
-
-    async def set_cache(self, cache_key: str, data: Any) -> None:
-        """Store data in cache."""
-        pass
-
-    async def invalidate_cache(self, cache_key: str) -> None:
-        """Invalidate cached data."""
-        pass
-```
-
-### 3. Public API Methods
-
-#### Initialization
-```python
-class OmniScheduler():
-    def __init__(self, config: SchedulerConfig):
-        """Initialize the core scheduler."""
-        self.config = config
-        self.scheduler = self._create_scheduler()
-        self.cache_manager = DiTCacheManager(config.dit_cache_config)
-        self._running = False
-
-    def _create_scheduler(self) -> BaseScheduler:
-        """Factory method to create appropriate scheduler."""
-        if self.config.scheduler_type == SchedulerType.FIFO:
-            return FIFOScheduler(self.config)
-        elif self.config.scheduler_type == SchedulerType.PRIORITY:
-            return PriorityScheduler(self.config)
-        # ... other scheduler types
-```
-
-#### Core Operations
-```python
-    async def schedule(self, request: OmniRequest) -> bool:
-        """
-        Schedule a request for processing.
-
-        Args:
-            request: The request to schedule
-
-        Returns:
-            bool: True if successfully scheduled, False otherwise
-
-        Raises:
-            SchedulerError: If scheduling fails
-            QueueFullError: If queue is at capacity
-        """
-        pass
-```
-
-#### Configuration
-```python
-    def update_config(self, new_config: SchedulerConfig) -> None:
-        """Update scheduler configuration."""
-        pass
-
-    def get_config(self) -> SchedulerConfig:
-        """Get current configuration."""
-        pass
-```
-
-#### Lifecycle Management
-```python
-    async def start(self) -> None:
-        """Start the scheduler."""
-        self._running = True
-        # Start background tasks
-
-    async def stop(self) -> None:
-        """Stop the scheduler gracefully."""
-        self._running = False
-        # Cleanup resources
-
-    async def shutdown(self) -> None:
-        """Force shutdown the scheduler."""
-        # Immediate cleanup
-```
-
-#### Monitoring
-```python
-    def get_status(self) -> Dict[str, Any]:
-        """Get current scheduler status."""
-        return {
-            "running": self._running,
-            "queue_size": self.scheduler.queue_size(),
-            "processed_requests": self.scheduler.processed_count(),
-            "cache_hit_rate": self.cache_manager.get_hit_rate()
-        }
-
-    def get_metrics(self) -> Dict[str, float]:
-        """Get performance metrics."""
-        pass
-```
-
-### 4. Configuration
-
-```python
-@dataclass
-class OmniConfig:
-    """Configuration for the core module."""
-
-    # Scheduler configuration
-    scheduler: SchedulerConfig = field(default_factory=SchedulerConfig)
-
-    # Cache configuration
-    dit_cache: DiTCacheConfig = field(default_factory=DiTCacheConfig)
-
-    # Resource limits
-    max_memory_gb: float = 16.0
-    max_gpu_utilization: float = 0.8
-
-    # Timeouts
-    request_timeout: int = 300
-    worker_timeout: int = 60
-
-    def validate(self) -> None:
-        """Validate configuration parameters."""
-        if self.max_memory_gb <= 0:
-            raise ValueError("max_memory_gb must be positive")
-        if not 0 < self.max_gpu_utilization <= 1:
-            raise ValueError("max_gpu_utilization must be between 0 and 1")
-```
-
-### 5. Error Handling
-
-```python
-class CoreModuleError(Exception):
-    """Base exception for core module errors."""
-    pass
-
-class SchedulerError(CoreModuleError):
-    """Scheduler-related errors."""
-    pass
-
-class QueueFullError(SchedulerError):
-    """Queue is at capacity."""
-    pass
-
-class CacheError(CoreModuleError):
-    """Cache-related errors."""
-    pass
-
-class ResourceError(CoreModuleError):
-    """Resource allocation errors."""
-    pass
-```
-
-### 6. Examples
-
-#### Basic Usage
-```python
-from vllm_omni.core import CoreScheduler, CoreModuleConfig
-from vllm_omni.request import create_text_request
-
-# Create configuration
-config = CoreModuleConfig(
-    scheduler=SchedulerConfig(scheduler_type=SchedulerType.FIFO),
-    max_memory_gb=8.0
-)
-
-# Initialize scheduler
-scheduler = CoreScheduler(config)
-await scheduler.start()
-
-# Create and schedule a request
-request = create_text_request(
-    request_id="req_001",
-    prompt="Hello, world!",
-    sampling_params=sampling_params
-)
-
-success = await scheduler.schedule_request(request)
-if success:
-    result = await scheduler.process_request(request)
-    print(f"Result: {result}")
-```
-
-#### Advanced Usage
-```python
-# Custom scheduler with priority handling
-config = CoreModuleConfig(
-    scheduler=SchedulerConfig(
-        scheduler_type=SchedulerType.PRIORITY,
-        priority_weights={"high": 2.0, "normal": 1.0, "low": 0.5}
-    )
-)
-
-scheduler = CoreScheduler(config)
-await scheduler.start()
-
-# Monitor scheduler status
-status = scheduler.get_status()
-print(f"Queue size: {status['queue_size']}")
-print(f"Cache hit rate: {status['cache_hit_rate']:.2%}")
-```
-
---
-
-## Module-Specific Guidelines
-
-### Core Module
- Focus on scheduling, caching, and resource management
- Provide clear interfaces for other modules
- Handle concurrency and thread safety
- Implement comprehensive monitoring
-
-### Engine Module
- Handle model loading and inference
- Support both AR and diffusion models
- Provide unified interface for different model types
- Implement efficient memory management
-
-### Executor Module
- Coordinate between different engines
- Handle request routing and load balancing
- Manage execution pipelines
- Provide error recovery mechanisms
-
-### Worker Module
- Handle actual model execution
- Manage GPU resources
- Implement batching strategies
- Provide performance optimization
-
---
-
-## Checklist for API Design
-
- [ ] **Clear Purpose**: Module purpose is well-defined
- [ ] **Complete Interface**: All public methods are documented
- [ ] **Error Handling**: Comprehensive error handling strategy
- [ ] **Configuration**: Flexible configuration system
- [ ] **Examples**: Basic and advanced usage examples
- [ ] **Type Hints**: All methods have proper type hints
- [ ] **Documentation**: Comprehensive docstrings
- [ ] **Testing**: Unit tests for all public methods
- [ ] **Integration**: Clear integration points with other modules
- [ ] **Performance**: Performance considerations documented
-
---
-
-## Submission Guidelines
-
-1. Create a new file: `docs/api/[module_name]_api.md`
-2. Follow the template structure exactly
-3. Include all required sections
-4. Provide working code examples
-5. Ensure all methods have proper docstrings
-6. Include error handling strategies
-7. Submit for review before implementation
-
-This template ensures consistency and completeness across all vLLM-Omni modules.
@@ -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)
@@ -14,7 +14,7 @@ vLLM-Omni is a multi-modality extension for vLLM that supports non-autoregressiv

 ## Key Data Flow

-API Server --> OmniLLM/AsyncOmniLLM (New, including multi engines) --> LLMEngine/AsyncLLM --> Engine Core
+API Server --> OmniLLM/AsyncOmni (New, including multi engines) --> LLMEngine/AsyncLLM --> Engine Core
 --> Scheduler (New one for DiT) --> Executor (New one for diffusers) --> Worker (New one for DiT)
 --> ModelRunner (New one for AR hiddenstate, New one for DiT) --> RequestState --> OutputProcessoer (New one for final multimodal output)

@@ -40,7 +40,7 @@ graph TD
    B --> C{Detect --omni flag}
    C -->|Yes| D[Parse OmniConfig]
    C -->|No| E[Forward to vLLM CLI]
-    D --> F[Initialize AsyncOmniLLM]
+    D --> F[Initialize AsyncOmni]
    F --> G[Start omni Server]
    G --> H[Multi-stage Processing]
    E --> I[Standard vLLM Pipeline]
@@ -209,7 +209,7 @@ Similar to OmniLLM in offline inference, add some asynchronous processing, refer
 from vllm.v1.engine.async_llm import AsyncLLM


-class AsyncOmniLLM(AsyncLLM):
+class AsyncOmni(AsyncLLM):
    """Extended AsyncLLM supporting multiple engines and stage-based processing"""

    def __init__(self, stage_configs: List[StageConfig]):
@@ -449,15 +449,15 @@ AsyncLLM.add_request()
 ↓
 OutputProcessor.add_request()
 ↓
-RequestState.from_new_request() → 创建请求状态
+RequestState.from_new_request() → create request state

-AsyncLLM.__init__() / AsyncLLM.generate() / AsyncLLM.encode() →  创建一个Background loop 持续从EngineCore获取输出
+AsyncLLM.__init__() / AsyncLLM.generate() / AsyncLLM.encode() →  create a background loop that continuously pulls output from EngineCore
 ↓
-OutputProcessor.process_outputs() → 更新状态并处理输出
+OutputProcessor.process_outputs() → update state and process outputs
 ↓
-RequestState.make_request_output() → 转换为最终输出,格式为RequestOutput或者PoolingRequestOutput
+RequestState.make_request_output() → convert to the final output, formatted as RequestOutput or PoolingRequestOutput
 ↓
-RequestOutputCollector.put() → 推送到队列（AsyncLLM）
+RequestOutputCollector.put() → push to the queue (AsyncLLM)
 ```

 Need to add implementation for an existing method
@@ -4,4 +4,4 @@ vLLM supports the following hardware platforms:

 - [GPU](gpu.md)
    - [NVIDIA CUDA](gpu.md)
- [Ascend NPU](npu.md)
+- [NPU](npu.md)
@@ -37,4 +37,7 @@ 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!
+
+=== "NVIDIA CUDA"
+
+    --8<-- "docs/getting_started/installation/gpu/cuda.inc.md:pre-built-images"
@@ -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 .
 ```

@@ -73,3 +73,24 @@ uv pip install --no-build-isolation --editable .
 </details>

 # --8<-- [end:build-wheel-from-source]
+
+
+# --8<-- [start:pre-built-images]
+
+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.
+
+Here's an example deployment command that has been verified on 2 x H100's:
+```bash
+docker run --runtime nvidia --gpus 2 \
+    -v ~/.cache/huggingface:/root/.cache/huggingface \
+    --env "HF_TOKEN=$HF_TOKEN" \
+    -p 8091:8091 \
+    --ipc=host \
+    vllm/vllm-omni:v0.11.0rc1 \
+    --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
+```
+
+!!! tip
+    You can use this docker image to serve models the same way you would with in vLLM! To do so, make sure you overwrite the default entrypoint (`vllm serve --omni`) which works only for models supported in the vLLM-Omni project.
+
+# --8<-- [end:pre-built-images]
@@ -1,3 +1,23 @@
-# Ascend-NPU
+# NPU

-vLLM-Omni is a Python library that supports the following NPU variants. Select your NPU type to see vendor specific instructions:
+vLLM-Omni supports NPU through the vLLM Ascend Plugin (vllm-ascend). This is a community maintained hardware plugin for running vLLM on NPU.
+
+## Requirements
+
+- OS: Linux
+- Python: 3.12
+
+!!! note
+    vLLM-Omni is currently not natively supported on Windows.
+
+=== "NPU"
+
+    --8<-- "docs/getting_started/installation/npu/npu.inc.md:requirements"
+
+## Installation
+
+### Recommended
+
+=== "NPU"
+
+    --8<-- "docs/getting_started/installation/npu/npu.inc.md:installation"
@@ -1,5 +1,49 @@
+# --8<-- [start:requirements]
+
+For detailed hardware and software requirements, please refer to the [vllm-ascend installation documentation](https://docs.vllm.ai/projects/ascend/en/latest/installation.html).
+
+# --8<-- [end:requirements]
 # --8<-- [start:installation]

-vLLM-Omni mainly contains python implementations for framework and models.
+The recommended way to use vLLM-Omni on NPU is through the vllm-ascend pre-built Docker images:
+
+```bash
+# Update DEVICE according to your NPUs (/dev/davinci[0-7])
+export DEVICE0=/dev/davinci0
+export DEVICE1=/dev/davinci1
+# Update the vllm-ascend image
+# Atlas A2:
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.11.0rc2
+# Atlas A3:
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.11.0rc2-a3
+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 \
+    --device /dev/devmm_svm \
+    --device /dev/hisi_hdc \
+    -v /usr/local/dcmi:/usr/local/dcmi \
+    -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
+    -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
+    -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
+    -v /etc/ascend_install.info:/etc/ascend_install.info \
+    -v /root/.cache:/root/.cache \
+    -p 8000:8000 \
+    -it $IMAGE bash
+
+# Inside the container, install vLLM-Omni from source
+cd /vllm-workspace
+git clone https://github.com/vllm-project/vllm-omni.git
+cd vllm-omni
+pip install -v -e .
+export VLLM_WORKER_MULTIPROC_METHOD=spawn
+```
+
+The default workdir is `/workspace`, with vLLM, vLLM-Ascend and vLLM-Omni code placed in `/vllm-workspace` installed in development mode.
+
+For other installation methods (pip installation, building from source, custom Docker builds), please refer to the [vllm-ascend installation guide](https://docs.vllm.ai/projects/ascend/en/latest/installation.html).

 # --8<-- [end:installation]
@@ -0,0 +1,336 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import itertools
+import logging
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Literal
+
+import regex as re
+import yaml
+
+logger = logging.getLogger("mkdocs")
+
+ROOT_DIR = Path(__file__).parent.parent.parent.parent
+ROOT_DIR_RELATIVE = "../../../../.."
+EXAMPLE_DIR = ROOT_DIR / "examples"
+EXAMPLE_DOC_DIR = ROOT_DIR / "docs/user_guide/examples"
+NAV_FILE = ROOT_DIR / "docs/.nav.yml"
+
+
+def fix_case(text: str) -> str:
+    subs = {
+        "api": "API",
+        "cli": "CLI",
+        "cpu": "CPU",
+        "llm": "LLM",
+        "mae": "MAE",
+        "tpu": "TPU",
+        "gguf": "GGUF",
+        "lora": "LoRA",
+        "rlhf": "RLHF",
+        "vllm": "vLLM",
+        "openai": "OpenAI",
+        "lmcache": "LMCache",
+        "multilora": "MultiLoRA",
+        "mlpspeculator": "MLPSpeculator",
+        r"fp\d+": lambda x: x.group(0).upper(),  # e.g. fp16, fp32
+        r"int\d+": lambda x: x.group(0).upper(),  # e.g. int8, int16
+    }
+    for pattern, repl in subs.items():
+        text = re.sub(rf"\b{pattern}\b", repl, text, flags=re.IGNORECASE)
+    return text
+
+
+@dataclass
+class Example:
+    """
+    Example class for generating documentation content from a given path.
+
+    Attributes:
+        path (Path): The path to the main directory or file.
+        category (str): The category of the document.
+        main_file (Path): The main file in the directory.
+        other_files (list[Path]): list of other files in the directory.
+        title (str): The title of the document.
+
+    Methods:
+        __post_init__(): Initializes the main_file, other_files, and title attributes.
+        determine_main_file() -> Path: Determines the main file in the given path.
+        determine_other_files() -> list[Path]: Determines other files in the directory excluding the main file.
+        determine_title() -> str: Determines the title of the document.
+        generate() -> str: Generates the documentation content.
+    """  # noqa: E501
+
+    path: Path
+    category: str = None
+    main_file: Path = field(init=False)
+    other_files: list[Path] = field(init=False)
+    title: str = field(init=False)
+
+    def __post_init__(self):
+        self.main_file = self.determine_main_file()
+        self.other_files = self.determine_other_files()
+        self.title = self.determine_title()
+
+    @property
+    def is_code(self) -> bool:
+        return self.main_file.suffix != ".md"
+
+    def determine_main_file(self) -> Path:
+        """
+        Determines the main file in the given path.
+        If the path is a file, it returns the path itself. Otherwise, it searches
+        for Markdown files (*.md) in the directory and returns the first one found.
+        Returns:
+            Path: The main file path, either the original path if it's a file or the first
+            Markdown file found in the directory.
+        Raises:
+            IndexError: If no Markdown files are found in the directory.
+        """  # noqa: E501
+        return self.path if self.path.is_file() else list(self.path.glob("*.md")).pop()
+
+    def determine_other_files(self) -> list[Path]:
+        """
+        Determine other files in the directory excluding the main file.
+
+        This method checks if the given path is a file. If it is, it returns an empty list.
+        Otherwise, it recursively searches through the directory and returns a list of all
+        files that are not the main file.
+
+        Returns:
+            list[Path]: A list of Path objects representing the other files in the directory.
+        """  # noqa: E501
+        if self.path.is_file():
+            return []
+        # Binary file extensions to exclude
+        binary_extensions = {
+            ".wav",
+            ".mp3",
+            ".mp4",
+            ".avi",
+            ".mov",
+            ".mkv",  # Audio/Video
+            ".png",
+            ".jpg",
+            ".jpeg",
+            ".gif",
+            ".bmp",
+            ".ico",
+            ".svg",  # Images
+            ".pdf",
+            ".zip",
+            ".tar",
+            ".gz",
+            ".bz2",
+            ".xz",  # Archives/Documents
+            ".exe",
+            ".so",
+            ".dll",
+            ".dylib",  # Binaries
+            ".bin",
+            ".dat",
+            ".db",
+            ".sqlite",  # Data files
+        }
+
+        def is_other_file(file: Path) -> bool:
+            return file.is_file() and file != self.main_file and file.suffix.lower() not in binary_extensions
+
+        return [file for file in self.path.rglob("*") if is_other_file(file)]
+
+    def determine_title(self) -> str:
+        if not self.is_code:
+            # Specify encoding for building on Windows
+            with open(self.main_file, encoding="utf-8") as f:
+                first_line = f.readline().strip()
+            match = re.match(r"^#\s+(?P<title>.+)$", first_line)
+            if match:
+                return match.group("title")
+        return fix_case(self.path.stem.replace("_", " ").title())
+
+    def fix_relative_links(self, content: str) -> str:
+        """
+        Fix relative links in markdown content by converting them to gh-file
+        format.
+
+        Args:
+            content (str): The markdown content to process
+
+        Returns:
+            str: Content with relative links converted to gh-file format
+        """
+        # Regex to match markdown links [text](relative_path)
+        # This matches links that don't start with http, https, ftp, or #
+        link_pattern = r"\[([^\]]*)\]\((?!(?:https?|ftp)://|#)([^)]+)\)"
+
+        def replace_link(match):
+            link_text = match.group(1)
+            relative_path = match.group(2)
+
+            # Make relative to repo root
+            gh_file = (self.main_file.parent / relative_path).resolve()
+            gh_file = gh_file.relative_to(ROOT_DIR)
+
+            # Make GitHub URL
+            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}"
+
+            return f"[{link_text}]({gh_url})"
+
+        return re.sub(link_pattern, replace_link, content)
+
+    def generate(self) -> str:
+        content = f"# {self.title}\n\n"
+        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"
+
+        # Use long code fence to avoid issues with
+        # included files containing code fences too
+        code_fence = "``````"
+
+        if self.is_code:
+            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
+                main_content = f.readlines()[1:]
+            content += self.fix_relative_links("".join(main_content))
+        content += "\n"
+
+        if not self.other_files:
+            return content
+
+        content += "## Example materials\n\n"
+        for file in sorted(self.other_files):
+            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.relative_to(ROOT_DIR)}"\n'
+            if file.suffix != ".md":
+                content += f"    {code_fence}\n"
+
+        return content
+
+
+def update_nav_file(examples: list[Example]):
+    """
+    Update the .nav.yml file to include all generated examples.
+    This function completely regenerates the examples section based on the actual
+    folder structure, ensuring consistency between the examples folder and nav file.
+
+    Args:
+        examples: List of Example objects that have been generated
+    """
+    if not NAV_FILE.exists():
+        logger.warning("Navigation file not found: %s", NAV_FILE)
+        return
+
+    # Read the current nav file
+    with open(NAV_FILE, encoding="utf-8") as f:
+        nav_data = yaml.safe_load(f) or {}
+
+    nav_list = nav_data.get("nav", [])
+
+    # Find the "User Guide" section
+    user_guide_idx = None
+    examples_idx = None
+    for i, item in enumerate(nav_list):
+        if isinstance(item, dict) and "User Guide" in item:
+            user_guide_idx = i
+            user_guide_content = item["User Guide"]
+            # Find the "Examples" subsection
+            for j, subitem in enumerate(user_guide_content):
+                if isinstance(subitem, dict) and "Examples" in subitem:
+                    examples_idx = j
+                    break
+            break
+
+    if user_guide_idx is None or examples_idx is None:
+        logger.warning("Could not find 'User Guide' -> 'Examples' section in nav file")
+        return
+
+    # Get existing Examples section to preserve non-example items (like README.md)
+    existing_examples_content = nav_list[user_guide_idx]["User Guide"][examples_idx]["Examples"]
+
+    # Preserve string items (like "examples/README.md") that are not example categories
+    preserved_items = [
+        item
+        for item in existing_examples_content
+        if isinstance(item, str) and not item.startswith("user_guide/examples/")
+    ]
+
+    # Group examples by category
+    examples_by_category = {}
+    for example in examples:
+        category = example.category
+        if category not in examples_by_category:
+            examples_by_category[category] = []
+        examples_by_category[category].append(example)
+
+    # Build the new Examples section - start with preserved items
+    examples_section = preserved_items.copy()
+
+    # Add examples grouped by category, sorted by category name
+    for category in sorted(examples_by_category.keys()):
+        category_examples = sorted(examples_by_category[category], key=lambda e: e.path.stem)
+        category_items = []
+        for example in category_examples:
+            doc_path = EXAMPLE_DOC_DIR / example.category / f"{example.path.stem}.md"
+            rel_path = doc_path.relative_to(ROOT_DIR / "docs")
+            category_items.append({example.title: str(rel_path)})
+
+        if category_items:
+            # Format category name (e.g., "offline_inference" -> "Offline Inference")
+            category_title = fix_case(category.replace("_", " ").title())
+            examples_section.append({category_title: category_items})
+
+    # Update the nav structure
+    nav_list[user_guide_idx]["User Guide"][examples_idx]["Examples"] = examples_section
+
+    # Write back to file
+    nav_data["nav"] = nav_list
+    with open(NAV_FILE, "w", encoding="utf-8") as f:
+        yaml.dump(nav_data, f, default_flow_style=False, sort_keys=False, allow_unicode=True)
+    logger.info("Updated navigation file: %s", NAV_FILE.relative_to(ROOT_DIR))
+
+
+def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool):
+    logger.info("Generating example documentation")
+    logger.debug("Root directory: %s", ROOT_DIR.resolve())
+    logger.debug("Example directory: %s", EXAMPLE_DIR.resolve())
+    logger.debug("Example document directory: %s", EXAMPLE_DOC_DIR.resolve())
+
+    # Create the EXAMPLE_DOC_DIR if it doesn't exist
+    if not EXAMPLE_DOC_DIR.exists():
+        EXAMPLE_DOC_DIR.mkdir(parents=True)
+
+    categories = sorted(p for p in EXAMPLE_DIR.iterdir() if p.is_dir())
+
+    examples = []
+    glob_patterns = ["*.py", "*.md", "*.sh"]
+    # Find categorised examples
+    for category in categories:
+        globs = [category.glob(pattern) for pattern in glob_patterns]
+        for path in itertools.chain(*globs):
+            examples.append(Example(path, category.stem))
+        # Find examples in subdirectories
+        for path in category.glob("*/*.md"):
+            examples.append(Example(path.parent, category.stem))
+
+    # Generate the example documentation
+    for example in sorted(examples, key=lambda e: e.path.stem):
+        example_name = f"{example.path.stem}.md"
+        doc_path = EXAMPLE_DOC_DIR / example.category / example_name
+        if not doc_path.parent.exists():
+            doc_path.parent.mkdir(parents=True)
+        # Specify encoding for building on Windows
+        with open(doc_path, "w+", encoding="utf-8") as f:
+            f.write(example.generate())
+        logger.debug("Example generated: %s", doc_path.relative_to(ROOT_DIR))
+
+    # Update the navigation file
+    update_nav_file(examples)
@@ -25,8 +25,11 @@ a:not(:has(svg)):not(.md-icon):not(.autorefs-external) {

 a[href*="localhost"]::after,
 a[href*="127.0.0.1"]::after,
-a[href*="org.readthedocs.build"]::after,
-a[href*="docs.vllm.ai"]::after {
+
+/* Hide external link icons for all links */
+a[href^="//"]::after,
+a[href^="http://"]::after,
+a[href^="https://"]::after {
    display: none !important;
 }

@@ -17,11 +17,13 @@ th {

 | Architecture | Models | Example HF Models |
 |--------------|--------|-------------------|
+| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` |
 | `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` |
 | `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |
+|`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |


-## List of Supported Models for Ascend NPU
+## List of Supported Models for NPU

 <style>
 th {
@@ -33,3 +35,4 @@ th {
 | Architecture | Models | Example HF Models |
 |--------------|--------|-------------------|
 | `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B`|
+| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |
@@ -11,3 +11,15 @@ A: If you encounter error about backend of librosa, try to install ffmpeg with c
 sudo apt update
 sudo apt install ffmpeg
 ```
+
+> Q: I encounter some bugs or CI problems, which is urgent. How can I solve it?
+
+A: At first, you can check current [issues](https://github.com/vllm-project/vllm-omni/issues) to find possible solutions. If non of these satisfy your demand and it is urgent, please find these [volunteers](https://docs.vllm.ai/projects/vllm-omni/en/latest/community/volunteers/) for help.
+
+> Q: Does vLLM-Omni support AWQ or any other quantization?
+
+A: vLLM-Omni partitions model into several stages. For AR stages, it will reuse main logic of LLMEngine in vLLM. So current quantization supported in vLLM should be also supported in vLLM-Omni for them. But systematic verification is ongoing. For quantization for DiffusionEngine, we are working on it. Please stay tuned and welcome contribution!
+
+> Q: Does vLLM-Omni support multimodal streaming input and output?
+
+A: Not yet. We already put it on the [Roadmap](https://github.com/vllm-project/vllm-omni/issues/165). Please stay tuned!
@@ -1,55 +1,70 @@
-# Offline Inference Example of vLLM-Omni for Qwen2.5-Omni
+# Offline Example of vLLM-Omni for Qwen2.5-Omni

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


 ## 🛠️ Installation

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

-## Run Offline inference with Qwen2.5-Omni.
-First, navigate to the example folder
-```bash
-cd examples/offline_inference/qwen2_5_omni
-```
-Inside the directory, `end2end.py` is a comprehensive demo suite for initializing a model instance of `Qwen/Qwen2.5-Omni-7B` and use it for various offline inference tasks.
-??? abstract "end2end.py"
-    ``````py
-    --8<-- "examples/offline_inference/qwen2_5_omni/end2end.py"
-    ``````
-
-Below we also provide simple bash scripts to execute this file.
-### Single Prompt
-```bash
-bash run_single_prompt.sh
-```
-??? abstract "run_single_prompt.sh"
-    ``````sh
-    --8<-- "examples/offline_inference/qwen2_5_omni/run_single_prompt.sh"
-    ``````
+## Run examples (Qwen2.5-Omni)

 ### Multiple Prompts
-Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit) and extract the prompts with `extract_prompts.py`
-??? abstract "extract_prompts.py"
-    ``````py
-    --8<-- "examples/offline_inference/qwen2_5_omni/extract_prompts.py"
-    ``````
+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 10
+  --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
 ```
-??? abstract "run_multiple_prompts.sh"
-    ``````sh
-    --8<-- "examples/offline_inference/qwen2_5_omni/run_multiple_prompts.sh"
-    ``````
+
+### Single Prompt
+Get into the example folder
+```bash
+cd examples/offline_inference/qwen2_5_omni
+```
+Then run the command below.
+```bash
+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

@@ -58,3 +73,22 @@ If you encounter error about backend of librosa, try to install ffmpeg with comm
 sudo apt update
 sudo apt install ffmpeg
 ```
+
+## Example materials
+
+??? abstract "end2end.py"
+    ``````py
+    --8<-- "examples/offline_inference/qwen2_5_omni/end2end.py"
+    ``````
+??? abstract "extract_prompts.py"
+    ``````py
+    --8<-- "examples/offline_inference/qwen2_5_omni/extract_prompts.py"
+    ``````
+??? abstract "run_multiple_prompts.sh"
+    ``````sh
+    --8<-- "examples/offline_inference/qwen2_5_omni/run_multiple_prompts.sh"
+    ``````
+??? abstract "run_single_prompt.sh"
+    ``````sh
+    --8<-- "examples/offline_inference/qwen2_5_omni/run_single_prompt.sh"
+    ``````
@@ -0,0 +1,81 @@
+# Offline Example of vLLM-Omni for 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-omni/tree/main/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](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
+```
+Then run the command below.
+```bash
+bash run_multiple_prompts.sh
+```
+### Single Prompt
+Get into the example folder
+```bash
+cd examples/offline_inference/qwen3_omni
+```
+Then run the command below.
+```bash
+bash run_single_prompt.sh
+```
+If you have not enough memory, you can set thinker with tensor parallel. Just run the command below.
+```bash
+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.
+```
+sudo apt update
+sudo apt install ffmpeg
+```
+
+## Example materials
+
+??? abstract "end2end.py"
+    ``````py
+    --8<-- "examples/offline_inference/qwen3_omni/end2end.py"
+    ``````
+??? abstract "run_multiple_prompts.sh"
+    ``````sh
+    --8<-- "examples/offline_inference/qwen3_omni/run_multiple_prompts.sh"
+    ``````
+??? abstract "run_single_prompt.sh"
+    ``````sh
+    --8<-- "examples/offline_inference/qwen3_omni/run_single_prompt.sh"
+    ``````
+??? abstract "run_single_prompt_tp.sh"
+    ``````sh
+    --8<-- "examples/offline_inference/qwen3_omni/run_single_prompt_tp.sh"
+    ``````
@@ -0,0 +1,57 @@
+# Qwen-Image Offline Inference
+
+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:
+
+- `text_to_image.py`: command-line script for single image generation.
+- `web_demo.py`: lightweight Gradio UI for interactive prompt/seed/CFG exploration.
+
+
+## Local CLI Usage
+
+```bash
+python text_to_image.py \
+  --prompt "a cup of coffee on the table" \
+  --seed 42 \
+  --cfg_scale 4.0 \
+  --num_images_per_prompt 1 \
+  --num_inference_steps 50 \
+  --height 1024 \
+  --width 1024 \
+  --output outputs/coffee.png
+```
+
+Key arguments:
+
+- `--prompt`: text description (string).
+- `--seed`: integer seed for deterministic sampling.
+- `--cfg_scale`: true CFG scale (model-specific guidance strength).
+- `--num_images_per_prompt`: number of images to generate per prompt (saves as `output`, `output_1`, ...).
+- `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower).
+- `--height/--width`: output resolution (defaults 1024x1024).
+- `--output`: path to save the generated PNG.
+
+> ℹ️ Qwen-Image currently publishes best-effort presets at `1328x1328`, `1664x928`, `928x1664`, `1472x1140`, `1140x1472`, `1584x1056`, and `1056x1584`. Adjust `--height/--width` accordingly for the most reliable outcomes.
+
+## Web UI Demo
+
+Launch the gradio demo:
+
+```bash
+python gradio_demo.py --port 7862
+```
+
+Then open `http://localhost:7862/` on your local browser to interact with the web UI.
+
+## Example materials
+
+??? abstract "gradio_demo.py"
+    ``````py
+    --8<-- "examples/offline_inference/qwen_image/gradio_demo.py"
+    ``````
+??? abstract "text_to_image.py"
+    ``````py
+    --8<-- "examples/offline_inference/qwen_image/text_to_image.py"
+    ``````
@@ -1,43 +1,99 @@
-# 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-omni/blob/main/examples/online_serving/README.md>.
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/qwen2_5_omni>.


 ## 🛠️ Installation

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

-## Deploy Qwen/Qwen2.5-Omni-7B
+## Run examples (Qwen2.5-Omni)
+
+### Launch the Server

-First launch the OpenAI-compatible inference server
 ```bash
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
 ```
+
 If you have custom stage configs file, launch the server with command below
 ```bash
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
 ```
-## Query the model
-Navigate to the example folder
+
+### Send Multi-modal Request
+
+Get into the example folder
 ```bash
-cd examples/online_serving
+cd examples/online_serving/qwen2_5_omni
 ```
-Query the model server with OpenAI Python API client:
+
+####  Send request via python
+
 ```bash
 python openai_chat_completion_client_for_multimodal_generation.py --query-type mixed_modalities
 ```
-??? abstract "openai_chat_completion_client_for_multimodal_generation.py"
-    ``````py
-    --8<-- "examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py"
-    ``````
-You can also query the model with `curl` command:
+
+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
 ```
-??? abstract "run_curl_multimodal_generation.sh"
-    ``````py
-    --8<-- "examples/online_serving/run_curl_multimodal_generation.sh"
-    ``````
+
+## Run Local Web UI Demo
+
+This Web UI demo allows users to interact with the model through a web browser.
+
+### Running Gradio Demo
+
+Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmni by
+
+```bash
+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:
+
+- `--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)

 ### FAQ

@@ -47,34 +103,17 @@ sudo apt update
 sudo apt install ffmpeg
 ```

-## Run Local Web UI Demo
+## Example materials

-You can also deploy a Gradio Web UI that allows users to interact with the model through a web browser. Below is an example on how to do so with `Qwen/Qwen2.5-Omni-7B`.
-
-### Running Gradio Demo
-
-Install gradio with `uv pip install "gradio>=5.49.1,<6.0.0"`, then you can launch the web service built on AsyncOmni by
-
-```bash
-python gradio_demo.py  --model Qwen/Qwen2.5-Omni-7B --port 7861
-```
-Now you can interact with model via the web UI at `http://localhost:7861/` on your local browser.
-
-
-### Options
-
-You can customize its basic launch parameters:
-
-```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
-```
-
- `--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.
+??? abstract "gradio_demo.py"
+    ``````py
+    --8<-- "examples/online_serving/qwen2_5_omni/gradio_demo.py"
+    ``````
+??? abstract "openai_chat_completion_client_for_multimodal_generation.py"
+    ``````py
+    --8<-- "examples/online_serving/qwen2_5_omni/openai_chat_completion_client_for_multimodal_generation.py"
+    ``````
+??? abstract "run_curl_multimodal_generation.sh"
+    ``````sh
+    --8<-- "examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh"
+    ``````
@@ -0,0 +1,120 @@
+# Online serving Example of vLLM-Omni for 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-omni/tree/main/README.md)
+
+## Run examples (Qwen3-Omni)
+
+### Launch the Server
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
+```
+
+If you have custom stage configs file, launch the server with command below
+```bash
+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
+
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py --query-type use_image
+```
+
+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.
+```
+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.
+
+### Running Gradio Demo
+
+Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmni by
+
+```bash
+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.
+
+
+The gradio script supports the following arguments:
+
+- `--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)
+
+## Example materials
+
+??? abstract "gradio_demo.py"
+    ``````py
+    --8<-- "examples/online_serving/qwen3_omni/gradio_demo.py"
+    ``````
+??? abstract "openai_chat_completion_client_for_multimodal_generation.py"
+    ``````py
+    --8<-- "examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py"
+    ``````
+??? abstract "run_curl_multimodal_generation.sh"
+    ``````sh
+    --8<-- "examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh"
+    ``````
@@ -1,10 +1,11 @@
-# Offline Example of vLLM-Omni for Qwen2.5-omni
+# Offline Example of vLLM-Omni for Qwen2.5-Omni

 ## 🛠️ Installation

 Please refer to [README.md](../../../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
@@ -23,7 +24,8 @@ Then run the command below.
 ```bash
 bash run_multiple_prompts.sh
 ```
-### Single Prompts
+
+### Single Prompt
 Get into the example folder
 ```bash
 cd examples/offline_inference/qwen2_5_omni
@@ -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.
@@ -6,17 +6,20 @@ 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

-from vllm_omni.entrypoints.omni_llm import OmniLLM
+from vllm_omni.entrypoints.omni import Omni

 SEED = 42

@@ -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,19 +190,137 @@ 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]()

-    omni_llm = OmniLLM(
+    # 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,
        log_stats=args.enable_stats,
        log_file=("omni_llm_pipeline.log" if args.enable_stats else None),
@@ -283,6 +463,39 @@ def parse_args():
        default=None,
        help="Path to a .txt file with one prompt per line (preferred).",
    )
+    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()

@@ -4,7 +4,7 @@ from typing import Optional


 def extract_prompt(line: str) -> Optional[str]:
-    # 提取第一个 '|' 与第二个 '|' 之间的内容
+    # Extract the content between the first '|' and the second '|'
    i = line.find("|")
    if i == -1:
        return None
@@ -1,2 +1,2 @@
 python end2end.py --output-wav output_audio \
-                  --query-type use_audio_in_video
+                  --query-type mixed_modalities
@@ -1,12 +1,13 @@
-# Offline Example of vLLM-Omni for Qwen2.5-omni
+# Offline Example of vLLM-Omni for Qwen3-Omni

 ## 🛠️ Installation

 Please refer to [README.md](../../../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](../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](../qwen2_5_omni/README.md)
 Get into the example folder
 ```bash
 cd examples/offline_inference/qwen3_omni
@@ -15,7 +16,7 @@ Then run the command below.
 ```bash
 bash run_multiple_prompts.sh
 ```
-### Single Prompts
+### Single Prompt
 Get into the example folder
 ```bash
 cd examples/offline_inference/qwen3_omni
@@ -24,6 +25,30 @@ Then run the command below.
 ```bash
 bash run_single_prompt.sh
 ```
+If you have not enough memory, you can set thinker with tensor parallel. Just run the command below.
+```bash
+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

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

-from vllm_omni.entrypoints.omni_llm import OmniLLM
+from vllm_omni.entrypoints.omni import Omni

 SEED = 42

@@ -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,10 +153,26 @@ query_map = {

 def main(args):
    model_name = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
-    query_result = query_map[args.query_type]()

-    omni_llm = OmniLLM(
+    # 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,
+        stage_configs_path=args.stage_configs_path,
    )

    thinker_sampling_params = SamplingParams(
@@ -276,6 +320,45 @@ def parse_args():
        default=None,
        help="Path to a .txt file with one prompt per line (preferred).",
    )
+    parser.add_argument(
+        "--stage-configs-path",
+        type=str,
+        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()

@@ -283,7 +366,3 @@ def parse_args():
 if __name__ == "__main__":
    args = parse_args()
    main(args)
-
-    # use examples:
-    # python end2end.py --model Qwen/Qwen3-Omni-30B-A3B-Instruct --query-type text
-    # python end2end.py --model /custom_path/Qwen3-Omni-30B-A3B-Instruct --query-type use_video_only
@@ -0,0 +1,5 @@
+python end2end.py --output-wav output_audio \
+                  --query-type use_audio \
+                  --init-sleep-seconds 90
+
+# init-sleep-seconds works to avoid two vLLM stages initialized at the same time within a card.
@@ -5,6 +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, is_npu

 ASPECT_RATIOS: dict[str, tuple[int, int]] = {
    "1:1": (1328, 1328),
@@ -58,11 +59,18 @@ 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:
-    device = "cuda"
+    device = detect_device_type()
    omni = get_omni(args.model)

    def run_inference(
@@ -7,6 +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, is_npu


 def parse_args() -> argparse.Namespace:
@@ -45,10 +46,18 @@ def parse_args() -> argparse.Namespace:

 def main():
    args = parse_args()
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+    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,
@@ -1,12 +1,13 @@
-# Online serving Example of vLLM-Omni for Qwen2.5-omni
+# Online serving Example of vLLM-Omni for Qwen2.5-Omni

 ## 🛠️ Installation

-Please refer to [README.md](../../README.md)
+Please refer to [README.md](../../../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
 ```
@@ -16,36 +17,66 @@ 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.

 ### Running Gradio Demo

-Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmniLLM by
+Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmni by

 ```bash
 python gradio_demo.py  --model Qwen/Qwen2.5-Omni-7B --port 7861
@@ -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.
@@ -15,7 +15,7 @@ from PIL import Image
 from vllm.assets.video import video_get_metadata, video_to_ndarrays
 from vllm.sampling_params import SamplingParams

-from vllm_omni.entrypoints.async_omni_llm import AsyncOmniLLM
+from vllm_omni.entrypoints.async_omni import AsyncOmni

 # Import utils from offline inference example
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../offline_inference/qwen2_5_omni"))
@@ -236,8 +236,8 @@ def process_video_file(
    return frames, metadata, audio_tuple


-async def run_inference_async_omni_llm(
-    omni_llm: AsyncOmniLLM,
+async def run_inference_async_omni(
+    omni: AsyncOmni,
    sampling_params: list[SamplingParams],
    prompt_args_template: SimpleNamespace,
    user_prompt: str,
@@ -246,7 +246,7 @@ async def run_inference_async_omni_llm(
    video_file: Optional[str] = None,
    use_audio_in_video: bool = False,
 ):
-    """Run inference using AsyncOmniLLM directly with multimodal support."""
+    """Run inference using AsyncOmni directly with multimodal support."""
    if not user_prompt.strip() and not audio_file and not image_file and not video_file:
        return "Please provide at least a text prompt or multimodal input.", None

@@ -328,7 +328,7 @@ async def run_inference_async_omni_llm(
        text_outputs: list[str] = []
        audio_output = None

-        async for stage_outputs in omni_llm.generate(
+        async for stage_outputs in omni.generate(
            prompt=omni_prompt,
            request_id=request_id,
            sampling_params_list=sampling_params,
@@ -362,12 +362,12 @@ async def run_inference_async_omni_llm(


 def build_interface(
-    omni_llm: AsyncOmniLLM,
+    omni: AsyncOmni,
    sampling_params: list[SamplingParams],
    prompt_args_template: SimpleNamespace,
    model: str,
 ):
-    """Build Gradio interface for AsyncOmniLLM mode."""
+    """Build Gradio interface for AsyncOmni mode."""

    async def run_inference(
        user_prompt: str,
@@ -376,8 +376,8 @@ def build_interface(
        video_file: Optional[str],
        use_audio_in_video: bool,
    ):
-        return await run_inference_async_omni_llm(
-            omni_llm,
+        return await run_inference_async_omni(
+            omni,
            sampling_params,
            prompt_args_template,
            user_prompt,
@@ -475,7 +475,7 @@ def build_interface(

 def main():
    args = parse_args()
-    omni_llm = None
+    omni = None

    model_name = "/".join(args.model.split("/")[-2:])
    assert model_name in SUPPORTED_MODELS, (
@@ -485,9 +485,9 @@ def main():
    # Register signal handlers for graceful shutdown
    def signal_handler(sig, frame):
        print("\nReceived interrupt signal, shutting down...")
-        if omni_llm is not None:
+        if omni is not None:
            try:
-                omni_llm.shutdown()
+                omni.shutdown()
            except Exception as e:
                print(f"Error during shutdown: {e}")
        sys.exit(0)
@@ -495,21 +495,21 @@ def main():
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

-    print(f"Initializing AsyncOmniLLM with model: {args.model}")
+    print(f"Initializing AsyncOmni with model: {args.model}")
    if args.stage_configs_path:
        print(f"Using custom stage configs: {args.stage_configs_path}")

    sampling_params = build_sampling_params(SEED, model_name)
-    omni_llm = AsyncOmniLLM(
+    omni = AsyncOmni(
        model=args.model,
        stage_configs_path=args.stage_configs_path,
        init_timeout=ASYNC_INIT_TIMEOUT,
    )
-    print("✓ AsyncOmniLLM initialized successfully")
+    print("✓ AsyncOmni initialized successfully")
    prompt_args_template = create_prompt_args(args)

    demo = build_interface(
-        omni_llm,
+        omni,
        sampling_params,
        prompt_args_template,
        args.model,
@@ -524,9 +524,9 @@ def main():
        print("\nShutting down...")
    finally:
        # Cleanup
-        if omni_llm is not None:
+        if omni is not None:
            try:
-                omni_llm.shutdown()
+                omni.shutdown()
            except Exception as e:
                print(f"Error during cleanup: {e}")

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

@@ -1,12 +1,13 @@
-# Online serving Example of vLLM-omni for Qwen2.5-omni
+# Online serving Example of vLLM-Omni for Qwen3-omni

 ## 🛠️ Installation

-Please refer to [README.md](../../README.md)
+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
+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 mixed_modalities
+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
-bash run_curl_multimodal_generation.sh mixed_modalities
+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.
@@ -45,7 +84,7 @@ This Web UI demo allows users to interact with the model through a web browser.

 ### Running Gradio Demo

-Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmniLLM by
+Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmni by

 ```bash
 python gradio_demo.py  --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 7861
@@ -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 AsyncOmniLLM 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)
@@ -15,7 +15,7 @@ from PIL import Image
 from vllm.assets.video import video_get_metadata, video_to_ndarrays
 from vllm.sampling_params import SamplingParams

-from vllm_omni.entrypoints.async_omni_llm import AsyncOmniLLM
+from vllm_omni.entrypoints.async_omni import AsyncOmni

 # Import utils from offline inference example
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../offline_inference/qwen3_omni"))
@@ -42,19 +42,20 @@ SUPPORTED_MODELS: dict[str, dict[str, Any]] = {
                "seed": SEED,
            },
            "talker": {
-                "temperature": 0.0,
-                "top_p": 1.0,
-                "top_k": -1,
+                "temperature": 0.9,
+                "top_k": 50,
                "max_tokens": 4096,
-                "detokenize": True,
-                "repetition_penalty": 1.1,
+                "seed": SEED,
+                "detokenize": False,
+                "repetition_penalty": 1.05,
                "stop_token_ids": [2150],
            },
            "code2wav": {
                "temperature": 0.0,
                "top_p": 1.0,
                "top_k": -1,
-                "max_tokens": 2048,
+                "max_tokens": 4096 * 16,
+                "seed": SEED,
                "detokenize": True,
                "repetition_penalty": 1.1,
            },
@@ -238,8 +239,8 @@ def process_video_file(
    return frames, metadata, audio_tuple


-async def run_inference_async_omni_llm(
-    omni_llm: AsyncOmniLLM,
+async def run_inference_async_omni(
+    omni: AsyncOmni,
    sampling_params: list[SamplingParams],
    prompt_args_template: SimpleNamespace,
    user_prompt: str,
@@ -248,7 +249,7 @@ async def run_inference_async_omni_llm(
    video_file: Optional[str] = None,
    use_audio_in_video: bool = False,
 ):
-    """Run inference using AsyncOmniLLM directly with multimodal support."""
+    """Run inference using AsyncOmni directly with multimodal support."""
    if not user_prompt.strip() and not audio_file and not image_file and not video_file:
        return "Please provide at least a text prompt or multimodal input.", None

@@ -263,7 +264,7 @@ async def run_inference_async_omni_llm(
        # 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)
@@ -281,7 +282,9 @@ async def run_inference_async_omni_llm(
                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
@@ -310,11 +313,11 @@ async def run_inference_async_omni_llm(
            )
            prompt = f"<|im_start|>system\n{default_system}<|im_end|>\n<|im_start|>user\n"
            if audio_data:
-                prompt += "<|audio_bos|><|AUDIO|><|audio_eos|>"
+                prompt += "<|audio_start|><|audio_pad|><|audio_end|>"
            if image_data:
-                prompt += "<|vision_bos|><|IMAGE|><|vision_eos|>"
+                prompt += "<|vision_start|><|image_pad|><|vision_end|>"
            if video_payload is not None:
-                prompt += "<|vision_bos|><|VIDEO|><|vision_eos|>"
+                prompt += "<|vision_start|><|video_pad|><|vision_end|>"
            if user_prompt.strip():
                prompt += f"{user_prompt}"
            prompt += "<|im_end|>\n<|im_start|>assistant\n"
@@ -330,7 +333,7 @@ async def run_inference_async_omni_llm(
        text_outputs: list[str] = []
        audio_output = None

-        async for stage_outputs in omni_llm.generate(
+        async for stage_outputs in omni.generate(
            prompt=omni_prompt,
            request_id=request_id,
            sampling_params_list=sampling_params,
@@ -344,14 +347,15 @@ async def run_inference_async_omni_llm(
                            text_outputs.append(output.text)
            elif stage_outputs.final_output_type == "audio":
                # multimodal_output is on the RequestOutput object
-                # See vllm_omni/entrypoints/openai/serving_chat.py:680 for reference
+                # See vllm_omni/entrypoints/openai/serving_chat.py:733 for reference
                if hasattr(request_output, "multimodal_output") and request_output.multimodal_output:
                    if "audio" in request_output.multimodal_output:
                        audio_tensor = request_output.multimodal_output["audio"]
                        # Ensure audio is 1D (flatten if needed)
                        if hasattr(audio_tensor, "ndim") and audio_tensor.ndim > 1:
                            audio_tensor = audio_tensor.flatten()
-                        audio_np = audio_tensor.detach().cpu().numpy()
+                        # Convert BFloat16 to Float32 before numpy conversion (numpy doesn't support BFloat16)
+                        audio_np = audio_tensor.float().detach().cpu().numpy()
                        audio_output = (
                            24000,  # sampling rate in Hz
                            audio_np,
@@ -364,12 +368,12 @@ async def run_inference_async_omni_llm(


 def build_interface(
-    omni_llm: AsyncOmniLLM,
+    omni: AsyncOmni,
    sampling_params: list[SamplingParams],
    prompt_args_template: SimpleNamespace,
    model: str,
 ):
-    """Build Gradio interface for AsyncOmniLLM mode."""
+    """Build Gradio interface for AsyncOmni mode."""

    async def run_inference(
        user_prompt: str,
@@ -378,8 +382,8 @@ def build_interface(
        video_file: Optional[str],
        use_audio_in_video: bool,
    ):
-        return await run_inference_async_omni_llm(
-            omni_llm,
+        return await run_inference_async_omni(
+            omni,
            sampling_params,
            prompt_args_template,
            user_prompt,
@@ -477,7 +481,7 @@ def build_interface(

 def main():
    args = parse_args()
-    omni_llm = None
+    omni = None

    model_name = "/".join(args.model.split("/")[-2:])
    assert model_name in SUPPORTED_MODELS, (
@@ -487,9 +491,9 @@ def main():
    # Register signal handlers for graceful shutdown
    def signal_handler(sig, frame):
        print("\nReceived interrupt signal, shutting down...")
-        if omni_llm is not None:
+        if omni is not None:
            try:
-                omni_llm.shutdown()
+                omni.shutdown()
            except Exception as e:
                print(f"Error during shutdown: {e}")
        sys.exit(0)
@@ -497,21 +501,21 @@ def main():
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

-    print(f"Initializing AsyncOmniLLM with model: {args.model}")
+    print(f"Initializing AsyncOmni with model: {args.model}")
    if args.stage_configs_path:
        print(f"Using custom stage configs: {args.stage_configs_path}")

    sampling_params = build_sampling_params(SEED, model_name)
-    omni_llm = AsyncOmniLLM(
+    omni = AsyncOmni(
        model=args.model,
        stage_configs_path=args.stage_configs_path,
        init_timeout=ASYNC_INIT_TIMEOUT,
    )
-    print("✓ AsyncOmniLLM initialized successfully")
+    print("✓ AsyncOmni initialized successfully")
    prompt_args_template = create_prompt_args(args)

    demo = build_interface(
-        omni_llm,
+        omni,
        sampling_params,
        prompt_args_template,
        args.model,
@@ -526,9 +530,9 @@ def main():
        print("\nShutting down...")
    finally:
        # Cleanup
-        if omni_llm is not None:
+        if omni is not None:
            try:
-                omni_llm.shutdown()
+                omni.shutdown()
            except Exception as e:
                print(f"Error during cleanup: {e}")

@@ -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.
    }
@@ -197,7 +340,6 @@ def run_multimodal_generation(args) -> None:

    count = 0
    for choice in chat_completion.choices:
-        print(f"choice: {choice}")
        if choice.message.audio:
            audio_data = base64.b64decode(choice.message.audio.data)
            audio_file_path = f"audio_{count}.wav"
@@ -226,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()

@@ -2,16 +2,16 @@
 set -euo pipefail

 # Default query type
-QUERY_TYPE="${1:-mixed_modalities}"
+QUERY_TYPE="${1:-use_video}"

 # Validate query type
-if [[ ! "$QUERY_TYPE" =~ ^(mixed_modalities|use_audio_in_video|multi_audios|text)$ ]]; then
+if [[ ! "$QUERY_TYPE" =~ ^(text|use_audio|use_image|use_video)$ ]]; then
    echo "Error: Invalid query type '$QUERY_TYPE'"
-    echo "Usage: $0 [mixed_modalities|use_audio_in_video|multi_audios|text]"
-    echo "  mixed_modalities: Audio + Image + Video + Text query"
-    echo "  use_audio_in_video: Video + Text query (with audio extraction from video)"
-    echo "  multi_audios: Two audio clips + Text query"
+    echo "Usage: $0 [text|use_audio|use_image|use_video]"
    echo "  text: Text query"
+    echo "  use_audio: Audio + Text query"
+    echo "  use_image: Image + Text query"
+    echo "  use_video: Video + Text query"
    exit 1
 fi

@@ -23,7 +23,6 @@ thinker_sampling_params='{
  "top_k": 1,
  "max_tokens": 16384,
  "seed": 42,
-  "detokenize": true,
  "repetition_penalty": 1.05,
  "stop_token_ids": [151645]
 }'
@@ -33,7 +32,7 @@ talker_sampling_params='{
  "top_k": 50,
  "max_tokens": 4096,
  "seed": 42,
-  "detokenize": true,
+  "detokenize": false,
  "repetition_penalty": 1.05,
  "stop_token_ids": [2150]
 }'
@@ -42,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
@@ -51,7 +50,6 @@ code2wav_sampling_params='{

 # Define URLs for assets
 MARY_HAD_LAMB_AUDIO_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/mary_had_lamb.ogg"
-WINNING_CALL_AUDIO_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/winning_call.ogg"
 CHERRY_BLOSSOM_IMAGE_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
 SAMPLE_VIDEO_URL="https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"

@@ -71,7 +69,7 @@ case "$QUERY_TYPE" in
    ]'
    mm_processor_kwargs="{}"
    ;;
-  mixed_modalities)
+  use_audio)
    user_content='[
        {
          "type": "audio_url",
@@ -79,21 +77,9 @@ case "$QUERY_TYPE" in
            "url": "'"$MARY_HAD_LAMB_AUDIO_URL"'"
          }
        },
-        {
-          "type": "image_url",
-          "image_url": {
-            "url": "'"$CHERRY_BLOSSOM_IMAGE_URL"'"
-          }
-        },
-        {
-          "type": "video_url",
-          "video_url": {
-            "url": "'"$SAMPLE_VIDEO_URL"'"
-          }
-        },
        {
          "type": "text",
-          "text": "What is recited in the audio? What is the content of this image? Why is this video funny?"
+          "text": "What is the content of this audio?"
        }
      ]'
    sampling_params_list='[
@@ -103,7 +89,27 @@ case "$QUERY_TYPE" in
    ]'
    mm_processor_kwargs="{}"
    ;;
-  use_audio_in_video)
+  use_image)
+    user_content='[
+        {
+          "type": "image_url",
+          "image_url": {
+            "url": "'"$CHERRY_BLOSSOM_IMAGE_URL"'"
+          }
+        },
+        {
+          "type": "text",
+          "text": "What is the content of this image?"
+        }
+      ]'
+    sampling_params_list='[
+      '"$thinker_sampling_params"',
+      '"$talker_sampling_params"',
+      '"$code2wav_sampling_params"'
+    ]'
+    mm_processor_kwargs="{}"
+    ;;
+  use_video)
    user_content='[
        {
          "type": "video_url",
@@ -113,35 +119,7 @@ case "$QUERY_TYPE" in
        },
        {
          "type": "text",
-          "text": "Describe the content of the video, then convert what the baby say into text."
-        }
-      ]'
-    sampling_params_list='[
-      '"$thinker_sampling_params"',
-      '"$talker_sampling_params"',
-      '"$code2wav_sampling_params"'
-    ]'
-    mm_processor_kwargs='{
-      "use_audio_in_video": true
-    }'
-    ;;
-  multi_audios)
-    user_content='[
-        {
-          "type": "audio_url",
-          "audio_url": {
-            "url": "'"$MARY_HAD_LAMB_AUDIO_URL"'"
-          }
-        },
-        {
-          "type": "audio_url",
-          "audio_url": {
-            "url": "'"$WINNING_CALL_AUDIO_URL"'"
-          }
-        },
-        {
-          "type": "text",
-          "text": "Are these two audio clips the same?"
+          "text": "Why is this video funny?"
        }
      ]'
    sampling_params_list='[
@@ -161,7 +139,7 @@ output=$(curl -sS -X POST http://localhost:8091/v1/chat/completions \
    -H "Content-Type: application/json" \
    -d @- <<EOF
 {
-  "model": "Qwen/Qwen2.5-Omni-7B",
+  "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
  "sampling_params_list": $sampling_params_list,
  "mm_processor_kwargs": $mm_processor_kwargs,
  "messages": [
@@ -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/

@@ -58,6 +58,7 @@ theme:
 hooks:
  - docs/mkdocs/hooks/generate_api_readme.py
  - docs/mkdocs/hooks/url_schemes.py
+  - docs/mkdocs/hooks/generate_examples.py

 # Plugins
 plugins:
@@ -87,6 +88,7 @@ plugins:
      show_full_namespace: false  # Show only module name, not full path
      exclude:
        - "re:vllm_omni\\._.*"  # Internal modules
+        - "vllm_omni.diffusion.models.qwen_image"  # avoid importing vllm in mkdocs building
  - mkdocstrings:
      handlers:
        python:
@@ -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",
@@ -1,6 +1,6 @@
 """
-Test suite for vLLM-omni.
+Test suite for vLLM-Omni.

 This package contains unit tests, integration tests, and benchmarks
-for vLLM-omni.
+for vLLM-Omni.
 """
@@ -342,14 +342,14 @@ def _setup_log_mocks(monkeypatch):
 def mock_get_config(monkeypatch):
    """Auto-mock get_config and related model loading functions to avoid model path validation."""
    # CRITICAL: Mock tokenizer-related imports FIRST, before any module imports
-    # This prevents ImportError when async_omni_llm is imported (which happens via omni_stage)
+    # This prevents ImportError when async_omni is imported (which happens via omni_stage)
    import sys

    fake_tokenizer = MagicMock()
    fake_tokenizer.encode = MagicMock(return_value=[1, 2, 3])
    fake_tokenizer.decode = MagicMock(return_value="test")

-    # Mock init_tokenizer_from_configs (used in async_omni_llm)
+    # Mock init_tokenizer_from_configs (used in async_omni)
    def _mock_init_tokenizer_from_configs(model_config=None, **kwargs):
        return fake_tokenizer

@@ -367,8 +367,9 @@ def mock_get_config(monkeypatch):
        tokenizer_module = sys.modules[tokenizer_module_path]
        setattr(tokenizer_module, "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs)

-    # CRITICAL: Mock length_from_prompt_token_ids_or_embeds BEFORE trying to mock async_omni_llm
-    # This is because async_omni_llm imports processor.py, which imports this function at module level
+    # CRITICAL: Mock length_from_prompt_token_ids_or_embeds BEFORE trying to mock async_omni
+
+    # This is because async_omni imports processor.py, which imports this function at module level
    # Mock length_from_prompt_token_ids_or_embeds (used in processor.py)
    def _mock_length_from_prompt_token_ids_or_embeds(prompt_token_ids=None, prompt_embeds=None):
        # Return a reasonable default length
@@ -402,19 +403,19 @@ def mock_get_config(monkeypatch):
            processor_module, "length_from_prompt_token_ids_or_embeds", _mock_length_from_prompt_token_ids_or_embeds
        )

-    # Strategy 3: Now mock async_omni_llm AFTER length_from_prompt_token_ids_or_embeds is mocked
-    # This prevents ImportError when async_omni_llm imports processor.py
+    # Strategy 3: Now mock async_omni AFTER length_from_prompt_token_ids_or_embeds is mocked
+    # This prevents ImportError when async_omni imports processor.py
    monkeypatch.setattr(
-        "vllm_omni.entrypoints.async_omni_llm.init_tokenizer_from_configs",
+        "vllm_omni.entrypoints.async_omni.init_tokenizer_from_configs",
        _mock_init_tokenizer_from_configs,
        raising=False,
    )

-    # Strategy 4: If async_omni_llm is already imported, patch it directly
-    async_omni_llm_path = "vllm_omni.entrypoints.async_omni_llm"
-    if async_omni_llm_path in sys.modules:
-        async_omni_llm_module = sys.modules[async_omni_llm_path]
-        setattr(async_omni_llm_module, "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs)
+    # Strategy 4: If async_omni is already imported, patch it directly
+    async_omni_path = "vllm_omni.entrypoints.async_omni"
+    if async_omni_path in sys.modules:
+        async_omni_module = sys.modules[async_omni_path]
+        setattr(async_omni_module, "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs)

    # Now mock get_config and other functions
    fake_hf_config = MagicMock()
@@ -21,10 +21,10 @@ except ModuleNotFoundError as exc:  # pragma: no cover - optional dependency
    patch = None  # type: ignore

 from .config import OmniModelConfig
-from .entrypoints.async_omni_llm import AsyncOmniLLM
+from .entrypoints.async_omni import AsyncOmni

 # Main entry points
-# from .entrypoints.omni import Omni
+from .entrypoints.omni import Omni

 from .version import __version__, __version_tuple__  # isort:skip

@@ -33,8 +33,8 @@ __all__ = [
    "__version__",
    "__version_tuple__",
    # Main components
-    # "Omni",
-    "AsyncOmniLLM",
+    "Omni",
+    "AsyncOmni",
    # Configuration
    "OmniModelConfig",
    # All other components are available through their respective modules
@@ -4,9 +4,9 @@ from typing import Optional

 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
 from vllm.v1.core.sched.request_queue import create_request_queue
-from vllm.v1.core.sched.scheduler import EngineCoreOutputs, Request, RequestStatus, SchedulerOutput, SpecDecodingStats
+from vllm.v1.core.sched.scheduler import Request, RequestStatus, SchedulerOutput, SpecDecodingStats
 from vllm.v1.core.sched.utils import remove_all
-from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput
+from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs

 from vllm_omni.core.sched.output import OmniNewRequestData
 from vllm_omni.core.sched.scheduler import OmniScheduler
@@ -5,14 +5,13 @@ from typing import Optional
 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
 from vllm.v1.core.sched.request_queue import create_request_queue
 from vllm.v1.core.sched.scheduler import (
-    EngineCoreOutputs,
    Request,
    RequestStatus,
    SchedulerOutput,
    SpecDecodingStats,
 )
 from vllm.v1.core.sched.utils import remove_all
-from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput
+from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs

 from vllm_omni.core.sched.output import OmniNewRequestData
 from vllm_omni.core.sched.scheduler import OmniScheduler
@@ -1,8 +1,15 @@
 from __future__ import annotations

+from collections import defaultdict
+
 from vllm.logger import init_logger
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
+from vllm.v1.core.sched.utils import check_stop, remove_all
+from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs
+from vllm.v1.outputs import ModelRunnerOutput
+from vllm.v1.request import Request, RequestStatus
+from vllm.v1.spec_decode.metrics import SpecDecodingStats


 class OmniScheduler(VLLMScheduler):
@@ -50,3 +57,177 @@ class OmniScheduler(VLLMScheduler):
            init_logger(__name__).exception("Failed to wrap scheduled_new_reqs with OmniNewRequestData")

        return scheduler_output
+
+    def update_from_output(
+        self,
+        scheduler_output: SchedulerOutput,
+        model_runner_output: ModelRunnerOutput,
+    ) -> dict[int, EngineCoreOutputs]:
+        """Override update_from_output to fix check_stop bug when output_token_ids is empty.
+
+        The original vLLM implementation calls check_stop when pooler_outputs exist,
+        but doesn't check if output_token_ids is empty. This causes IndexError during
+        prefill phase when pooler_outputs exist but output_token_ids is still empty.
+
+        This override adds a safety check before calling check_stop.
+        """
+        sampled_token_ids = model_runner_output.sampled_token_ids
+        logprobs = model_runner_output.logprobs
+        prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict
+        num_scheduled_tokens = scheduler_output.num_scheduled_tokens
+        pooler_outputs = model_runner_output.pooler_output
+        num_nans_in_logits = model_runner_output.num_nans_in_logits
+        kv_connector_output = model_runner_output.kv_connector_output
+
+        outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
+        spec_decoding_stats: SpecDecodingStats | None = None
+        kv_connector_stats = kv_connector_output.kv_connector_stats if kv_connector_output else None
+
+        # NOTE(woosuk): As len(num_scheduled_tokens) can be up to 1K or more,
+        # the below loop can be a performance bottleneck. We should do our best
+        # to avoid expensive operations inside the loop.
+        stopped_running_reqs: set[Request] = set()
+        stopped_preempted_reqs: set[Request] = set()
+        for req_id, num_tokens_scheduled in num_scheduled_tokens.items():
+            assert num_tokens_scheduled > 0
+            request = self.requests.get(req_id)
+            if request is None:
+                # The request is already finished. This can happen if the
+                # request is aborted while the model is executing it (e.g.,
+                # in pipeline parallelism).
+                continue
+
+            req_index = model_runner_output.req_id_to_index[req_id]
+            generated_token_ids = sampled_token_ids[req_index] if sampled_token_ids else []
+
+            scheduled_spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(req_id)
+            if scheduled_spec_token_ids:
+                num_draft_tokens = len(scheduled_spec_token_ids)
+                num_accepted = len(generated_token_ids) - 1
+                num_rejected = num_draft_tokens - num_accepted
+                # num_computed_tokens represents the number of tokens
+                # processed in the current step, considering scheduled
+                # tokens and rejections. If some tokens are rejected,
+                # num_computed_tokens is decreased by the number of rejected
+                # tokens.
+                request.num_computed_tokens -= num_rejected
+                spec_decoding_stats = self.make_spec_decoding_stats(
+                    spec_decoding_stats, num_draft_tokens=num_draft_tokens, num_accepted_tokens=num_accepted
+                )
+
+            stopped = False
+            new_logprobs = None
+            new_token_ids = generated_token_ids
+            kv_transfer_params = None
+            status_before_stop = request.status
+
+            # Check for stop and update request status.
+            if new_token_ids:
+                new_token_ids, stopped = self._update_request_with_output(request, new_token_ids)
+
+            """
+            # Stop checking for pooler models.
+            pooler_output = None
+            if pooler_outputs:
+                pooler_output = pooler_outputs[req_index]
+                stopped = check_stop(request, self.max_model_len,
+                                     pooler_output)
+            """
+
+            # Stop checking for pooler models.
+            # FIX: Check if output_token_ids is empty before calling check_stop
+            # to prevent IndexError during prefill phase when pooler_outputs exist
+            # but output_token_ids is still empty.
+            pooler_output = None
+            if pooler_outputs:
+                pooler_output = pooler_outputs[req_index]
+                # Only call check_stop if output_token_ids is not empty
+                # This prevents IndexError when accessing output_token_ids[-1]
+                if request.output_token_ids:
+                    stopped = check_stop(request, self.max_model_len, pooler_output)
+                # If output_token_ids is empty, we don't call check_stop
+                # and keep stopped=False (don't stop during prefill phase)
+
+            if stopped:
+                kv_transfer_params = self._free_request(request)
+                if status_before_stop == RequestStatus.RUNNING:
+                    stopped_running_reqs.add(request)
+                else:
+                    stopped_preempted_reqs.add(request)
+
+            # Extract sample logprobs if needed.
+            if request.sampling_params is not None and request.sampling_params.logprobs is not None and logprobs:
+                # NOTE: once we support N tokens per step (spec decode),
+                # the outer lists can be of length > 1.
+                new_logprobs = logprobs.slice(req_index, req_index + 1)
+
+            if new_token_ids and self.structured_output_manager.should_advance(request):
+                # NOTE: structured_output_request
+                # should not be None if use_structured_output, we have
+                # checked above, so safe to ignore type warning
+                request.structured_output_request.grammar.accept_tokens(  # type: ignore[union-attr]
+                    req_id, new_token_ids
+                )
+
+            if num_nans_in_logits is not None and req_id in num_nans_in_logits:
+                request.num_nans_in_logits = num_nans_in_logits[req_id]
+
+            # Get prompt logprobs for this request.
+            prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
+            if new_token_ids or pooler_output is not None or kv_transfer_params:
+                # Add EngineCoreOutput for this Request.
+                outputs[request.client_index].append(
+                    EngineCoreOutput(
+                        request_id=req_id,
+                        new_token_ids=new_token_ids,
+                        finish_reason=request.get_finished_reason(),
+                        new_logprobs=new_logprobs,
+                        new_prompt_logprobs_tensors=prompt_logprobs_tensors,
+                        pooling_output=pooler_output,
+                        stop_reason=request.stop_reason,
+                        events=request.take_events(),
+                        kv_transfer_params=kv_transfer_params,
+                        trace_headers=request.trace_headers,
+                        num_cached_tokens=request.num_cached_tokens,
+                    )
+                )
+            else:
+                # Invariant: EngineCore returns no partial prefill outputs.
+                assert not prompt_logprobs_tensors
+
+        # Remove the stopped requests from the running and waiting queues.
+        if stopped_running_reqs:
+            self.running = remove_all(self.running, stopped_running_reqs)
+        if stopped_preempted_reqs:
+            # This is a rare case and unlikely to impact performance.
+            self.waiting.remove_requests(stopped_preempted_reqs)
+
+        # KV Connector: update state for finished KV Transfers.
+        if model_runner_output.kv_connector_output:
+            self._update_from_kv_xfer_finished(model_runner_output.kv_connector_output)
+
+        # Create EngineCoreOutputs for all clients that have requests with
+        # outputs in this step.
+        engine_core_outputs = {client_index: EngineCoreOutputs(outputs=outs) for client_index, outs in outputs.items()}
+
+        finished_req_ids = self.finished_req_ids_dict
+        if finished_req_ids:
+            # Include ids of requests that finished since last outputs
+            # were sent.
+            for client_index, finished_set in finished_req_ids.items():
+                # Set finished request set in EngineCoreOutputs for this client.
+                if (eco := engine_core_outputs.get(client_index)) is not None:
+                    eco.finished_requests = finished_set
+                else:
+                    engine_core_outputs[client_index] = EngineCoreOutputs(finished_requests=finished_set)
+            finished_req_ids.clear()
+
+        if (stats := self.make_stats(spec_decoding_stats, kv_connector_stats)) is not None:
+            # Return stats to only one of the front-ends.
+            if (eco := next(iter(engine_core_outputs.values()), None)) is None:
+                # We must return the stats even if there are no request
+                # outputs this step.
+                engine_core_outputs[0] = eco = EngineCoreOutputs()
+            eco.scheduler_stats = stats
+
+        return engine_core_outputs
@@ -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
@@ -9,7 +9,7 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.registry import get_diffusion_post_process_func
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.scheduler import scheduler
-from vllm_omni.diffusion.worker.gpu_worker import WorkerProc
+from vllm_omni.utils.platform_utils import get_diffusion_worker_class

 logger = init_logger(__name__)

@@ -76,6 +76,9 @@ class DiffusionEngine:
        mp.set_start_method("spawn", force=True)
        processes = []

+        # Get the appropriate worker class for current device
+        worker_proc = get_diffusion_worker_class()
+
        # Launch all worker processes
        scheduler_pipe_readers = []
        scheduler_pipe_writers = []
@@ -84,7 +87,7 @@ class DiffusionEngine:
            reader, writer = mp.Pipe(duplex=False)
            scheduler_pipe_writers.append(writer)
            process = mp.Process(
-                target=WorkerProc.worker_main,
+                target=worker_proc.worker_main,
                args=(
                    i,  # rank
                    od_config,
@@ -5,7 +5,11 @@ import os

 import torch

+from vllm_omni.utils.platform_utils import detect_device_type
+

 def get_local_device() -> torch.device:
-    """Return the torch device for the current rank."""
-    return torch.device(f"cuda:{os.environ.get('LOCAL_RANK', 0)}")
+    """Return the torch device for the current rank based on detected device type."""
+    device_type = detect_device_type()
+    local_rank = os.environ.get("LOCAL_RANK", 0)
+    return torch.device(f"{device_type}:{local_rank}")
@@ -0,0 +1,3 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Diffusion model implementations."""
@@ -0,0 +1,17 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Qwen Image diffusion model components."""
+
+from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import (
+    QwenImagePipeline,
+    get_qwen_image_post_process_func,
+)
+from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
+    QwenImageTransformer2DModel,
+)
+
+__all__ = [
+    "QwenImagePipeline",
+    "QwenImageTransformer2DModel",
+    "get_qwen_image_post_process_func",
+]
@@ -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,24 +235,28 @@ class QwenImagePipeline(
    def __init__(
        self,
        *,
-        od_config: OmniDiffusionConfig = None,
+        od_config: OmniDiffusionConfig,
        prefix: str = "",
    ):
        super().__init__()
+        self.od_config = od_config
        self.device = get_local_device()
        model = od_config.model
-        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(model, subfolder="scheduler")
-        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")
-            logger.info("Loaded Qwen-Image text encoder successfully")
-            self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae").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")
-            logger.info("Loaded Qwen-Image tokenizer successfully.")
+        # Check if model is a local path
+        local_files_only = os.path.exists(model)
+
+        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+            model, subfolder="scheduler", local_files_only=local_files_only
+        )
+        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

@@ -392,8 +395,6 @@ class QwenImagePipeline(
        Args:
            prompt (`str` or `list[str]`, *optional*):
                prompt to be encoded
-            device: (`torch.device`):
-                torch device
            num_images_per_prompt (`int`):
                number of images that should be generated per prompt
            prompt_embeds (`torch.Tensor`, *optional*):
@@ -765,10 +766,14 @@ class QwenImagePipeline(
                    for name, tensor in state_dict.items():
                        yield name, tensor

-        model_name = "Qwen/Qwen-Image"
        try:
-            # Use the arguments as verified by the user
-            model_path = download_weights_from_hf_specific(model_name, None, ["*"])
+            # Get model path from config or download from HF
+            model_name = self.od_config.model if hasattr(self, "od_config") else "Qwen/Qwen-Image"
+            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))

@@ -113,10 +113,12 @@ class QwenEmbedRope(nn.Module):
        # DO NOT USING REGISTER BUFFER HERE, IT WILL CAUSE COMPLEX NUMBERS LOSE ITS IMAGINARY PART
        self.scale_rope = scale_rope

-    def rope_params(self, index, dim, theta=10000):
+    def rope_params(self, index: torch.Tensor, dim: int, theta: int = 10000):
        """
        Args:
-            index: [0, 1, 2, 3] 1D Tensor representing the position index of the token
+            index (`torch.Tensor`): [0, 1, 2, 3] 1D Tensor representing the position index of the token
+            dim (`int`): Dimension for the rope parameters
+            theta (`int`): Theta parameter for rope
        """
        assert dim % 2 == 0
        freqs = torch.outer(
@@ -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
@@ -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.")
@@ -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",
 }


@@ -1,12 +1,13 @@
 from functools import lru_cache

-from diffusers.pipelines.pipeline_utils import DiffusionPipeline
 from vllm.logger import init_logger

 logger = init_logger(__name__)


 def load_diffusers_config(model_name) -> dict:
+    from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+
    config = DiffusionPipeline.load_config(model_name)
    return config

@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Worker classes for diffusion models."""
+
+from vllm_omni.diffusion.worker.gpu_worker import GPUWorker, WorkerProc
+
+__all__ = ["GPUWorker", "WorkerProc"]
@@ -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:
@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""NPU worker classes for diffusion models."""
+
+from vllm_omni.diffusion.worker.npu.npu_worker import NPUWorker, NPUWorkerProc
+
+__all__ = ["NPUWorker", "NPUWorkerProc"]
@@ -0,0 +1,204 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import multiprocessing as mp
+import os
+
+import torch
+import zmq
+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
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+
+logger = init_logger(__name__)
+
+
+class NPUWorker:
+    """
+    A worker that executes the model on a single NPU.
+    """
+
+    def __init__(
+        self,
+        local_rank: int,
+        rank: int,
+        od_config: OmniDiffusionConfig,
+    ):
+        self.local_rank = local_rank
+        self.rank = rank
+        self.od_config = od_config
+        self.pipeline = None
+
+        self.init_device_and_model()
+
+    def init_device_and_model(self) -> None:
+        """Initialize the device and load the model."""
+        world_size = self.od_config.num_gpus
+        rank = self.rank
+        # Set environment variables for distributed initialization
+        os.environ["MASTER_ADDR"] = "localhost"
+        os.environ["MASTER_PORT"] = str(self.od_config.master_port)
+        os.environ["LOCAL_RANK"] = str(self.local_rank)
+        os.environ["RANK"] = str(rank)
+        os.environ["WORLD_SIZE"] = str(world_size)
+
+        device = torch.device(f"npu:{rank}")
+        torch.npu.set_device(device)
+
+        # hack
+        # 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)
+
+        init_distributed_environment(world_size=world_size, rank=rank)
+        initialize_model_parallel(tensor_model_parallel_size=world_size)
+
+        with device:
+            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.")
+
+    @torch.inference_mode()
+    def execute_model(self, reqs: list[OmniDiffusionRequest], od_config: OmniDiffusionConfig) -> DiffusionOutput:
+        """
+        Execute a forward pass.
+        """
+        assert self.pipeline is not None
+        # TODO: dealing with first req for now
+        req = reqs[0]
+        output = self.pipeline.forward(req)
+        return output
+
+
+class NPUWorkerProc:
+    """Wrapper that runs one Worker in a separate process."""
+
+    def __init__(
+        self,
+        od_config: OmniDiffusionConfig,
+        gpu_id: int,
+        broadcast_handle,
+    ):
+        self.od_config = od_config
+
+        # Inter-process Communication
+        self.context = zmq.Context(io_threads=2)
+
+        # Initialize MessageQueue reader from handle
+        self.mq = MessageQueue.create_from_handle(broadcast_handle, gpu_id)
+
+        self.result_mq = None
+        self.result_mq_handle = None
+
+        # Setup result sender (only for rank 0 for now, or whoever needs to reply)
+        # Assuming only rank 0 replies to scheduler as per original logic
+        if gpu_id == 0:
+            # Create MessageQueue for results (1 writer -> 1 reader)
+            # We assume the reader (SyncScheduler) will act as rank 0
+            self.result_mq = MessageQueue(n_reader=1, n_local_reader=1, local_reader_ranks=[0])
+            self.result_mq_handle = self.result_mq.export_handle()
+            logger.info(f"Worker {gpu_id} created result MessageQueue")
+
+        assert od_config.master_port is not None
+        worker = NPUWorker(
+            local_rank=gpu_id,
+            rank=gpu_id,
+            od_config=od_config,
+        )
+        self.worker = worker
+        self.gpu_id = gpu_id
+        self._running = True
+
+    def return_result(self, output: DiffusionOutput):
+        """
+        replies to client, only on rank 0
+        """
+        if self.result_mq is not None:
+            self.result_mq.enqueue(output)
+
+    def recv_reqs(self):
+        """
+        Receive requests from broadcast queue
+        """
+        return self.mq.dequeue(indefinite=True)
+
+    # TODO: queueing, cancellation
+    def worker_busy_loop(self) -> None:
+        """Main busy loop for Multiprocessing Workers"""
+
+        logger.info(f"Worker {self.gpu_id} ready to receive requests via shared memory")
+
+        while self._running:
+            reqs = None
+            # 1: receive requests
+            try:
+                reqs = self.recv_reqs()
+            except Exception as e:
+                logger.error(
+                    f"Error receiving requests in scheduler event loop: {e}",
+                    exc_info=True,
+                )
+                continue
+
+            # 2: execute, make sure a reply is always sent
+            try:
+                output = self.worker.execute_model(reqs, self.od_config)
+            except Exception as e:
+                logger.error(
+                    f"Error executing forward in event loop: {e}",
+                    exc_info=True,
+                )
+                output = DiffusionOutput(error=str(e))
+
+            try:
+                self.return_result(output)
+            except zmq.ZMQError as e:
+                # Reply failed; log and keep loop alive to accept future requests
+                logger.error(f"ZMQ error sending reply: {e}")
+                continue
+
+        logger.info("event loop terminated.")
+        # if self.result_sender is not None:
+        #     self.result_sender.close()
+        self.context.term()
+
+    @staticmethod
+    def worker_main(
+        rank: int,
+        od_config: OmniDiffusionConfig,
+        pipe_writer: mp.connection.Connection,
+        broadcast_handle,
+    ) -> None:
+        """Worker initialization and execution loops."""
+
+        worker_proc = NPUWorkerProc(
+            od_config,
+            gpu_id=rank,
+            broadcast_handle=broadcast_handle,
+        )
+        logger.info(f"Worker {rank}: Scheduler loop started.")
+        pipe_writer.send(
+            {
+                "status": "ready",
+                "result_handle": worker_proc.result_mq_handle if rank == 0 else None,
+            }
+        )
+        worker_proc.worker_busy_loop()
+        logger.info(f"Worker {rank}: Shutdown complete.")
@@ -51,7 +51,7 @@ from vllm_omni.outputs import OmniRequestOutput
 logger = init_logger(__name__)


-class AsyncOmniLLM(EngineClient):
+class AsyncOmni(EngineClient):
    """Async entry point for vLLM-Omni inference.

    This class provides an asynchronous interface for running multi-modal
@@ -75,7 +75,7 @@ class AsyncOmniLLM(EngineClient):
        **kwargs: Additional keyword arguments passed to stage engines

    Example:
-        >>> async_llm = AsyncOmniLLM(model="Qwen/Qwen2.5-Omni-7B")
+        >>> async_llm = AsyncOmni(model="Qwen/Qwen2.5-Omni-7B")
        >>> async for output in async_llm.generate(
        ...     prompt="Hello",
        ...     request_id="req-1",
@@ -90,11 +90,11 @@ class AsyncOmniLLM(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,
-        **kwargs,
+        **kwargs: Any,
    ):
        self.batch_timeout = batch_timeout
        self._enable_stats: bool = bool(log_stats)
@@ -175,7 +175,7 @@ class AsyncOmniLLM(EngineClient):

        Sends shutdown signals to all stage input queues and stops
        all stage worker processes. This method should be called
-        when done using the AsyncOmniLLM instance.
+        when done using the AsyncOmni instance.
        """
        for q in self._stage_in_queues:
            try:
@@ -193,7 +193,7 @@ class AsyncOmniLLM(EngineClient):
                logger.warning("[Orchestrator] Failed to stop stage worker: %s", e)

    def __del__(self) -> None:  # best-effort
-        print("[AsyncOmniLLM] __del__ close()", flush=True)
+        print("[AsyncOmni] __del__ close()", flush=True)
        raise Exception("test")
        try:
            self.close()
@@ -611,13 +611,13 @@ class AsyncOmniLLM(EngineClient):
        **kwargs,
    ):
        """Generate outputs for a request from a pooling model."""
-        raise NotImplementedError("encode() is not implemented for AsyncOmniLLM")
+        raise NotImplementedError("encode() is not implemented for AsyncOmni")

    async def start_profile(self) -> None:
-        raise NotImplementedError("start_profile() is not implemented for AsyncOmniLLM")
+        raise NotImplementedError("start_profile() is not implemented for AsyncOmni")

    async def stop_profile(self) -> None:
-        raise NotImplementedError("stop_profile() is not implemented for AsyncOmniLLM")
+        raise NotImplementedError("stop_profile() is not implemented for AsyncOmni")


 class AsyncOmniStageLLM(AsyncLLM):
@@ -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:
@@ -534,8 +534,9 @@ class OmniStageLLM(LLM):
        self,
        model: str,
        compilation_config: Optional[Union[int, dict[str, Any], CompilationConfig]] = None,
+        hf_overrides: Optional[dict[str, Any]] = None,
        structured_outputs_config: Optional[Union[dict[str, Any], StructuredOutputsConfig]] = None,
-        **kwargs,
+        **kwargs: Any,
    ):
        """LLM constructor."""
        if "disable_log_stats" not in kwargs:
@@ -32,7 +32,7 @@ from vllm_omni.entrypoints.stage_utils import (
    _to_dict,
    maybe_dump_to_shm,
    maybe_load_from_ipc_with_metrics,
-    set_stage_gpu_devices,
+    set_stage_devices,
 )
 from vllm_omni.inputs.data import OmniTokensPrompt

@@ -295,6 +295,8 @@ class OmniStage:
                raise ValueError("engine_input_source is empty")
            source_stage_id = self.engine_input_source[0]
            source_outputs = stage_list[source_stage_id].engine_outputs
+            if not isinstance(prompt, list):
+                prompt = [prompt]
            multi_modal_data = {
                source_output.request_id: p.get("multi_modal_data", None)
                for source_output, p in zip(source_outputs, prompt)
@@ -425,7 +427,10 @@ def _stage_worker(

    # Device mapping
    try:
-        set_stage_gpu_devices(stage_id, runtime_cfg.get("devices"))
+        from vllm_omni.utils import detect_device_type
+
+        device_type = detect_device_type()
+        set_stage_devices(stage_id, runtime_cfg.get("devices"), device_type=device_type)
    except Exception as e:
        _logging.getLogger(__name__).warning("[Stage-%s] Device setup failed: %s", stage_id, e)

@@ -452,9 +457,9 @@ def _stage_worker(
        max_batch_size = int(runtime_cfg.get("max_batch_size", 1) or 1)
        print(f"[Stage-{stage_id}] Max batch size: {max_batch_size}")
        batch_tasks: list[dict[str, Any]] = [task]
+        start_time = _time.time()
        if max_batch_size > 1:
            while len(batch_tasks) < max_batch_size:
-                start_time = _time.time()
                if not in_q.empty():
                    extra = in_q.get_nowait()
                    if extra is None:
@@ -666,7 +671,7 @@ async def _stage_worker_async(
    import logging as _logging
    import time as _time

-    from vllm_omni.entrypoints.async_omni_llm import AsyncOmniStageLLM
+    from vllm_omni.entrypoints.async_omni import AsyncOmniStageLLM
    from vllm_omni.entrypoints.log_utils import (
        compute_and_log_stage_request_stats,
        count_tokens_from_outputs,
@@ -717,7 +722,10 @@ async def _stage_worker_async(

    # Device mapping
    try:
-        set_stage_gpu_devices(stage_id, runtime_cfg.get("devices"))
+        from vllm_omni.utils import detect_device_type
+
+        device_type = detect_device_type()
+        set_stage_devices(stage_id, runtime_cfg.get("devices"), device_type=device_type)
    except Exception as e:
        _logging.getLogger(__name__).warning("[Stage-%s] Device setup failed: %s", stage_id, e)

@@ -38,7 +38,7 @@ from vllm.logger import init_logger
 from vllm.transformers_utils.tokenizer import MistralTokenizer
 from vllm.utils import decorate_logs

-from vllm_omni.entrypoints.async_omni_llm import AsyncOmniLLM
+from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat

 logger = init_logger(__name__)
@@ -65,7 +65,7 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
    if log_config is not None:
        uvicorn_kwargs["log_config"] = log_config

-    async with build_async_omni_llm(
+    async with build_async_omni(
        args,
        client_config=client_config,
    ) as engine_client:
@@ -108,15 +108,15 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,


@asynccontextmanager
-async def build_async_omni_llm(
+async def build_async_omni(
    args: Namespace,
    *,
    disable_frontend_multiprocessing: Optional[bool] = None,
    client_config: Optional[dict[str, Any]] = None,
 ) -> AsyncIterator[EngineClient]:
-    """Build an AsyncOmniLLM instance from command-line arguments.
+    """Build an AsyncOmni instance from command-line arguments.

-    Creates an async context manager that yields an AsyncOmniLLM instance
+    Creates an async context manager that yields an AsyncOmni instance
    configured from the provided arguments. Handles forkserver setup if
    needed and ensures proper cleanup on exit.

@@ -127,7 +127,7 @@ async def build_async_omni_llm(
        client_config: Optional client configuration dictionary

    Yields:
-        EngineClient instance (AsyncOmniLLM) ready for use
+        EngineClient instance (AsyncOmni) ready for use
    """
    if os.getenv("VLLM_WORKER_MULTIPROC_METHOD") == "forkserver":
        # The executor is expected to be mp.
@@ -138,25 +138,25 @@ async def build_async_omni_llm(
        forkserver.ensure_running()
        logger.debug("Forkserver setup complete!")

-    # Context manager to handle async_omni_llm lifecycle
+    # Context manager to handle async_omni lifecycle
    # Ensures everything is shutdown and cleaned up on error/exit
-    async with build_async_omni_llm_from_stage_config(
+    async with build_async_omni_from_stage_config(
        args,
        disable_frontend_multiprocessing=disable_frontend_multiprocessing,
-    ) as async_omni_llm:
-        yield async_omni_llm
+    ) as async_omni:
+        yield async_omni


@asynccontextmanager
-async def build_async_omni_llm_from_stage_config(
+async def build_async_omni_from_stage_config(
    args: Namespace,
    *,
    disable_frontend_multiprocessing: bool = False,
    client_config: Optional[dict[str, Any]] = None,
 ) -> AsyncIterator[EngineClient]:
-    """Create AsyncOmniLLM from stage configuration.
+    """Create AsyncOmni from stage configuration.

-    Creates an AsyncOmniLLM instance either in-process or using multiprocess
+    Creates an AsyncOmni instance either in-process or using multiprocess
    RPC. Loads stage configurations from the model or from a specified path.

    Args:
@@ -166,7 +166,7 @@ async def build_async_omni_llm_from_stage_config(
        client_config: Optional client configuration dictionary

    Yields:
-        EngineClient instance (AsyncOmniLLM) ready for use
+        EngineClient instance (AsyncOmni) ready for use

    Note:
        Stage configurations are loaded from args.stage_configs_path if provided,
@@ -182,21 +182,21 @@ async def build_async_omni_llm_from_stage_config(
            "To disable frontend multiprocessing, set VLLM_USE_V1=0."
        )

-    async_omni_llm: Optional[EngineClient] = None
+    async_omni: Optional[EngineClient] = None

    try:
        if getattr(args, "stage_configs_path", None):
-            async_omni_llm = AsyncOmniLLM(model=args.model, stage_configs_path=args.stage_configs_path)
+            async_omni = AsyncOmni(model=args.model, stage_configs_path=args.stage_configs_path)
        else:
-            async_omni_llm = AsyncOmniLLM(model=args.model)
+            async_omni = AsyncOmni(model=args.model)

        # # Don't keep the dummy data in memory
        # await async_llm.reset_mm_cache()

-        yield async_omni_llm
+        yield async_omni
    finally:
-        if async_omni_llm:
-            async_omni_llm.shutdown()
+        if async_omni:
+            async_omni.shutdown()


 async def omni_init_app_state(
@@ -211,7 +211,7 @@ async def omni_init_app_state(
    and other server configuration needed for handling API requests.

    Args:
-        engine_client: Engine client instance (AsyncOmniLLM)
+        engine_client: Engine client instance (AsyncOmni)
        vllm_config: vLLM configuration object
        state: FastAPI application state object to initialize
        args: Parsed command-line arguments
@@ -13,100 +13,143 @@ from omegaconf import OmegaConf
 logger = logging.getLogger(__name__)


-def set_stage_gpu_devices(stage_id: int, devices: str | int | None) -> None:
-    """Configure per-stage CUDA visibility and current device.
+def set_stage_devices(
+    stage_id: int,
+    devices: str | int | None,
+    device_type: str | None = None,
+) -> None:
+    """Configure per-stage device visibility and current device (CUDA or NPU).

-    Behavior
-    - Comma-separated string (e.g. "2,5,7"): set CUDA_VISIBLE_DEVICES exactly
-      to this list; logical index 0 is used as current device.
-    - Integer or digit-string: treat as logical index (0-based) into the current
-      CUDA_VISIBLE_DEVICES mapping; map to the physical device, and then set
-      CUDA_VISIBLE_DEVICES to this single device.
-    - None/"cpu": keep default visibility.
-    - Otherwise: set CUDA_VISIBLE_DEVICES to the provided single device string.
+    This function sets environment variables that control which devices are visible
+    to the process, and sets the current device. It must be called BEFORE worker
+    initialization so that workers see the correct devices.
+
+    Args:
+        stage_id: Stage identifier for logging
+        devices: Device specification:
+            - Comma-separated string (e.g. "2,5,7"): set device visibility env var
+              exactly to this list; logical index 0 is used as current device.
+            - Integer or digit-string: treat as logical index (0-based) into the
+              current device visibility mapping; map to physical device, then set
+              env var to this single device.
+            - None/"cpu": keep default visibility.
+            - Otherwise: set env var to the provided single device string.
+        device_type: Device type ("cuda" or "npu"). If None, auto-detects.
+
+    Behavior:
+        - CUDA: Sets CUDA_VISIBLE_DEVICES and calls torch.cuda.set_device()
+        - NPU: Sets ASCEND_RT_VISIBLE_DEVICES and calls torch.npu.set_device()
    """
+    from vllm_omni.utils import detect_device_type, get_device_control_env_var
+
+    if device_type is None:
+        device_type = detect_device_type()
+
+    env_var = get_device_control_env_var()
+
+    # Select device-specific torch functions
+    if device_type == "npu":
+        try:
+            import torch.npu  # type: ignore[import-untyped]
+        except ImportError:
+            logger.debug("[Stage-%s] torch.npu not available, skipping NPU device setup", stage_id)
+            return
+
+        is_available_fn = torch.npu.is_available
+        set_device_fn = torch.npu.set_device
+        device_count_fn = torch.npu.device_count
+        get_device_properties_fn = torch.npu.get_device_properties
+        mem_get_info_fn = torch.npu.mem_get_info
+        get_device_name_fn = torch.npu.get_device_name
+        device_type_label = "NPU"
+    elif device_type == "cuda":
+        import torch  # noqa: WPS433
+
+        is_available_fn = torch.cuda.is_available
+        set_device_fn = torch.cuda.set_device
+        device_count_fn = torch.cuda.device_count
+        get_device_properties_fn = torch.cuda.get_device_properties
+        mem_get_info_fn = torch.cuda.mem_get_info
+        get_device_name_fn = torch.cuda.get_device_name
+        device_type_label = "CUDA"
+    else:
+        logger.debug("[Stage-%s] Unsupported device type: %s", stage_id, device_type)
+        return
+
    try:
        selected_physical: int | None = None
        logical_idx: int | None = None

        if isinstance(devices, str) and "," in devices:
-            os.environ["CUDA_VISIBLE_DEVICES"] = devices
+            os.environ[env_var] = devices
            toks = [t.strip() for t in devices.split(",") if t.strip() != ""]
            if toks:
                try:
                    selected_physical = int(toks[0])
                    logger.debug(
-                        "[Stage-%s] Set CUDA_VISIBLE_DEVICES to %s; logical 0 -> physical %s",
+                        "[Stage-%s] Set %s to %s; logical 0 -> physical %s",
                        stage_id,
+                        env_var,
                        devices,
                        selected_physical,
                    )
                except Exception as e:
-                    logger.debug("[Stage-%s] Failed to parse first CUDA device: %s", stage_id, e)
+                    logger.debug("[Stage-%s] Failed to parse first %s device: %s", stage_id, device_type_label, e)
                    selected_physical = None
        elif isinstance(devices, (int, str)) and (isinstance(devices, int) or str(devices).isdigit()):
            logical_idx = max(0, int(devices))
-            vis = os.environ.get("CUDA_VISIBLE_DEVICES")
+            vis = os.environ.get(env_var)
            if vis:
                try:
                    mapping = [int(x) for x in vis.split(",") if x.strip() != ""]
                    if 0 <= logical_idx < len(mapping):
                        selected_physical = mapping[logical_idx]
                except Exception as e:
-                    logger.debug("[Stage-%s] Failed to map logical index via CUDA_VISIBLE_DEVICES: %s", stage_id, e)
+                    logger.debug("[Stage-%s] Failed to map logical index via %s: %s", stage_id, env_var, e)
                    selected_physical = None
            if selected_physical is None:
                selected_physical = int(logical_idx)
-            os.environ["CUDA_VISIBLE_DEVICES"] = str(selected_physical)
+            os.environ[env_var] = str(selected_physical)
            logger.debug(
-                "[Stage-%s] Logical index %d -> physical %s; set CUDA_VISIBLE_DEVICES to single device",
+                "[Stage-%s] Logical index %d -> physical %s; set %s to single device",
                stage_id,
                logical_idx + 1,
                selected_physical,
+                env_var,
            )
        elif devices in (None, "cpu"):
            logger.debug("[Stage-%s] Using default device visibility (devices=%s)", stage_id, devices)
        else:
            selected_physical = int(str(devices))
-            os.environ["CUDA_VISIBLE_DEVICES"] = str(selected_physical)
-            logger.debug(
-                "[Stage-%s] Set CUDA_VISIBLE_DEVICES to single device %s (fallback)", stage_id, selected_physical
-            )
+            os.environ[env_var] = str(selected_physical)
+            logger.debug("[Stage-%s] Set %s to single device %s (fallback)", stage_id, env_var, selected_physical)

        try:
-            import torch
+            import torch  # noqa: WPS433

-            if torch.cuda.is_available():
+            if is_available_fn():
                try:
-                    torch.cuda.set_device(0)
+                    set_device_fn(0)
                except Exception as e:
                    logger.debug(
-                        "[Stage-%s] torch.cuda.set_device(0) failed: %s",
-                        stage_id,
-                        e,
-                        exc_info=True,
+                        "[Stage-%s] %s set_device(0) failed: %s", stage_id, device_type_label, e, exc_info=True
                    )
-                num = torch.cuda.device_count()
+                num = device_count_fn()
                info = []
                for i in range(num):
-                    total = torch.cuda.get_device_properties(i).total_memory
-                    free, _ = torch.cuda.mem_get_info(i)
+                    total = get_device_properties_fn(i).total_memory
+                    free, _ = mem_get_info_fn(i)
                    info.append(
                        {
                            "idx": i,
-                            "name": torch.cuda.get_device_name(i),
+                            "name": get_device_name_fn(i),
                            "total": int(total),
                            "free": int(free),
                        }
                    )
-                logger.debug("[Stage-%s] CUDA devices visible=%s info=%s", stage_id, num, info)
+                logger.debug("[Stage-%s] %s devices visible=%s info=%s", stage_id, device_type_label, num, info)
        except Exception as e:
-            logger.debug(
-                "[Stage-%s] Failed to query CUDA devices: %s",
-                stage_id,
-                e,
-                exc_info=True,
-            )
+            logger.debug("[Stage-%s] Failed to query %s devices: %s", stage_id, device_type_label, e, exc_info=True)
    except Exception as e:
        logger.warning("Failed to interpret devices for stage %s: %s", stage_id, e)

@@ -1,13 +1,10 @@
-from __future__ import annotations
-
-import logging
 import os
 from pathlib import Path

 from omegaconf import OmegaConf
 from vllm.transformers_utils.config import get_config

-logger = logging.getLogger(__name__)
+from vllm_omni.utils import detect_device_type

 # Get the project root directory (2 levels up from this file)
 PROJECT_ROOT = Path(__file__).parent.parent.parent
@@ -16,9 +13,9 @@ PROJECT_ROOT = Path(__file__).parent.parent.parent
 def load_stage_configs_from_model(model: str) -> list:
    """Load stage configurations from model's default config file.

-    Loads stage configurations based on the model type. Looks for a
-    YAML configuration file in the stage_configs directory matching
-    the model's model_type.
+    Loads stage configurations based on the model type and device type.
+    First tries to load 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)
@@ -31,6 +28,17 @@ def load_stage_configs_from_model(model: str) -> list:
    """
    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):
@@ -35,7 +35,7 @@ from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
 )
 from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, split_list_into_ranges
 from vllm_omni.model_executor.models.vision import get_llm_pos_ids_for_vision
-from vllm_omni.model_executor.preprocess_mixin import PreprocessMixin
+from vllm_omni.utils.platform_utils import is_npu

 TALKER_CODEC_EOS_TOKEN_ID = 8294
 TALKER_CODEC_BOS_TOKEN_ID = 8293
@@ -50,12 +50,10 @@ logger = init_logger(__name__)
    dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder,
 )
 class Qwen2_5OmniForConditionalGeneration(
-    nn.Module, SupportsMultiModal, SupportsPP, SupportsMRoPE,
-    Qwen2_5OmniConditionalGenerationMixin, PreprocessMixin
+    nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin, SupportsMRoPE
 ):
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
-        self.has_preprocess = False
        self.have_multimodal_outputs = True
        config: Qwen2_5OmniConfig = vllm_config.model_config.hf_config
        multimodal_config = vllm_config.model_config.multimodal_config
@@ -86,9 +84,6 @@ class Qwen2_5OmniForConditionalGeneration(
            self.token2wav = None

        elif self.model_stage == "talker":
-            # register the process function for the talker stage
-            self.has_preprocess = True
-            self.set_preprocess(self.process_fn)
            self.thinker = None
            # Initialize talker model wrapper (handles projection + LM)
            self.talker = init_vllm_registered_model(
@@ -207,6 +202,7 @@ class Qwen2_5OmniForConditionalGeneration(
        """
        if self.model_stage == "thinker":
            # Normalize to batched inputs if caller provides 1D/2D unbatched tensors
+            # TODO: Remove this hack when NPU supports batched inputs properly
            added_batch_dim = False
            if input_ids is not None and input_ids.ndim == 1:
                input_ids = input_ids.unsqueeze(0)
@@ -234,15 +230,27 @@ class Qwen2_5OmniForConditionalGeneration(
                positions = positions.to(thinker_dev)
            if inputs_embeds is not None and inputs_embeds.device != thinker_dev:
                inputs_embeds = inputs_embeds.to(thinker_dev)
-            # Run thinker
-            with torch.inference_mode():
-                thinker_output = self.thinker(
-                    input_ids=input_ids,
-                    positions=positions[0],
-                    intermediate_tensors=intermediate_tensors,
-                    inputs_embeds=inputs_embeds,
-                    **kwargs,
+
+            if is_npu():
+                # TODO: remove this hack when NPU supports batched inputs properly
+                thinker_input_ids = input_ids[0] if input_ids is not None and added_batch_dim else input_ids
+                thinker_positions = positions[0] if positions.ndim > 1 else positions
+                thinker_inputs_embeds = (
+                    inputs_embeds[0] if inputs_embeds is not None and added_batch_dim else inputs_embeds
                )
+            else:
+                thinker_input_ids = input_ids
+                thinker_positions = positions[0]
+                thinker_inputs_embeds = inputs_embeds
+
+            # Run thinker
+            thinker_output = self.thinker(
+                input_ids=thinker_input_ids,
+                positions=thinker_positions,
+                intermediate_tensors=intermediate_tensors,
+                inputs_embeds=thinker_inputs_embeds,
+                **kwargs,
+            )

            if isinstance(thinker_output, tuple):
                embeds, text_hidden_states = thinker_output
@@ -257,13 +265,151 @@ class Qwen2_5OmniForConditionalGeneration(

        # 2) Talker (if codec not provided)
        if self.model_stage == "talker":
-            # mock data for profile
+            # Mixed-mode support: In a single step, both Prefill*n and Decode*n are supported.
+            # Rules:
+            # - Prefill segments are wrapped with special tokens: [BOS][PAD...][EOS]
+            # - Decode segments consist of a single non-special token.
+            # - If additional_information is provided
+            # (can be a list split by request or a concatenated tensor plus a list of shapes),
+            #   then for each request, reconstruct the thinker→talker input embeddings for the Prefill segments;
+            # - For Decode segments, if per-request auxiliary decode embeddings are provided (optional), add them;
+            # otherwise, keep the original embedding.
+
            if input_ids is None and additional_information is None:
                input_ids = torch.zeros(
-                    inputs_embeds.shape[0], dtype=torch.long, device=inputs_embeds.device
+                    inputs_embeds.shape[0],
+                    dtype=torch.long,
+                    device=inputs_embeds.device,
                )
                additional_information = {}
                self.thinker_reply_part = torch.zeros_like(inputs_embeds)
+                is_profile = True
+            else:
+                is_profile = False
+
+            # Ensure we have base embeddings when only ids are provided
+            if inputs_embeds is None and input_ids is not None:
+                inputs_embeds = self.talker.get_input_embeddings(input_ids)
+
+            # ------- Request-scoped additional information (no cross-request concat) -------
+            request_ids: Optional[list[str]] = kwargs.get("request_ids")  # ordered
+            request_token_spans: Optional[list[tuple[int, int]]] = kwargs.get("request_token_spans")
+            addi_by_req: Optional[dict] = kwargs.get("additional_information_by_req_id")
+            runtime_addi = kwargs.get("runtime_additional_information")
+
+            # Normalize runtime_addi into a mapping by request_id for convenience
+            runtime_addi_by_req: dict[str, dict] = {}
+            if (
+                isinstance(request_ids, list)
+                and isinstance(runtime_addi, list)
+                and len(runtime_addi) == len(request_ids)
+            ):
+                for i, rid in enumerate(request_ids):
+                    if isinstance(rid, str) and isinstance(runtime_addi[i], dict):
+                        runtime_addi_by_req[rid] = runtime_addi[i]
+            elif isinstance(request_ids, list) and isinstance(runtime_addi, dict):
+                for rid in request_ids:
+                    if isinstance(rid, str) and isinstance(runtime_addi.get(rid), dict):
+                        runtime_addi_by_req[rid] = runtime_addi[rid]
+
+            # Containers to return per-request updates (e.g., thinker_reply_part_per_request)
+            update_by_req_id: dict[str, dict] = {}
+
+            # ------- Prefill: span_len > 1 -------
+            if (
+                not is_profile
+                and isinstance(request_ids, list)
+                and isinstance(request_token_spans, list)
+                and isinstance(addi_by_req, dict)
+            ):
+                for idx_req, rid in enumerate(request_ids):
+                    s, e = request_token_spans[idx_req]
+                    span_len = int(e) - int(s)
+                    if span_len <= 1:
+                        continue
+                    info = addi_by_req.get(rid, {}) if isinstance(rid, str) else {}
+                    if not isinstance(info, dict):
+                        info = {}
+                    pe = info.get("prompt_embeds")  # Tensor [P,H]
+                    tr = info.get("thinker_result")  # Tensor [K,H]
+                    ptoks = info.get("prompt_token_ids")  # list[int]
+                    otoks = info.get("thinker_output_token_ids")  # list[int]
+
+                    if not isinstance(pe, torch.Tensor):
+                        pe = torch.zeros(
+                            0,
+                            self.talker.config.hidden_size,
+                            dtype=inputs_embeds.dtype,
+                            device=self._module_device(self.model),
+                        )
+                    if not isinstance(tr, torch.Tensor):
+                        tr = torch.zeros(
+                            0,
+                            self.talker.config.hidden_size,
+                            dtype=inputs_embeds.dtype,
+                            device=self._module_device(self.model),
+                        )
+                    if not isinstance(ptoks, (list, torch.Tensor)):
+                        ptoks = []
+                    if not isinstance(otoks, (list, torch.Tensor)):
+                        otoks = []
+
+                    req_input_ids, req_embeds = self._thinker_to_talker_prefill(
+                        voice_type=voice_type,
+                        output_prompt_embeds=tr.to(inputs_embeds.dtype).to(self._module_device(self.model)),
+                        output_token_ids=otoks,
+                        thinker_prompt_embeds=pe.to(inputs_embeds.dtype).to(self._module_device(self.model)),
+                        prompt_token_ids=ptoks,
+                    )
+                    seg_len = min(span_len, req_embeds.shape[0])
+                    inputs_embeds[s : s + seg_len] = req_embeds[:seg_len]
+                    if isinstance(req_input_ids, torch.Tensor) and req_input_ids.numel() == seg_len:
+                        input_ids[s : s + seg_len] = req_input_ids
+
+                    # Prepare per-request reply queue for subsequent decode: drop first row
+                    if tr.ndim == 2 and tr.shape[0] > 0:
+                        update_by_req_id.setdefault(rid, {})["thinker_reply_part_per_request"] = (
+                            tr[1:].detach().to("cpu").contiguous()
+                        )
+
+            # ------- Decode: span_len == 1 -------
+            if not is_profile and isinstance(request_ids, list) and isinstance(request_token_spans, list):
+                for idx_req, rid in enumerate(request_ids):
+                    s, e = request_token_spans[idx_req]
+                    if (int(e) - int(s)) != 1:
+                        continue
+                    # choose step vector in priority order
+                    step_vec = None
+                    # A) runtime queue
+                    q = None
+                    if isinstance(rid, str):
+                        q = runtime_addi_by_req.get(rid, {}).get("thinker_reply_part_per_request")
+                    if isinstance(q, torch.Tensor) and q.numel() > 0:
+                        step_vec = q[0:1]
+                        new_q = q[1:].detach().to("cpu").contiguous()
+                        update_by_req_id.setdefault(rid, {})["thinker_reply_part_per_request"] = new_q
+                    else:
+                        # B) per-request provided decode vector (optional)
+                        info = addi_by_req.get(rid, {}) if isinstance(addi_by_req, dict) else {}
+                        dv = info.get("decode_output_prompt_embeds") if isinstance(info, dict) else None
+                        if isinstance(dv, torch.Tensor) and dv.numel() > 0:
+                            step_vec = dv[0:1] if dv.ndim == 2 else dv.view(1, -1)
+                        elif (
+                            hasattr(self, "thinker_reply_part")
+                            and isinstance(self.thinker_reply_part, torch.Tensor)
+                            and self.thinker_reply_part.numel() > 0
+                        ):
+                            # C) fallback shared pool
+                            step_vec = self.thinker_reply_part[0:1]
+                            self.thinker_reply_part = self.thinker_reply_part[1:]
+
+                    if isinstance(step_vec, torch.Tensor) and step_vec.numel() > 0:
+                        one_id = input_ids[s : s + 1]
+                        _, one_embed = self._thinker_to_talker_decode_one_step(
+                            output_prompt_embeds=step_vec.to(inputs_embeds.dtype).to(self._module_device(self.model)),
+                            output_token_ids=one_id,
+                        )
+                        inputs_embeds[s] = one_embed[0]

            with torch.inference_mode():
                talker_hidden = self.talker(
@@ -271,6 +417,10 @@ class Qwen2_5OmniForConditionalGeneration(
                    positions=positions[0],
                    inputs_embeds=inputs_embeds,
                )
+            multimodal_outputs: dict = None
+            # Return updates if any
+            if update_by_req_id:
+                multimodal_outputs = {"additional_information_update_by_req_id": update_by_req_id}

            if sampling_metadata is not None:
                # the padding token id is set to text model's pad token id,
@@ -279,7 +429,7 @@ class Qwen2_5OmniForConditionalGeneration(

            return OmniOutput(
                text_hidden_states=talker_hidden,
-                multimodal_outputs=None,
+                multimodal_outputs=multimodal_outputs,
            )

        if self.model_stage == "code2wav":
@@ -591,68 +741,6 @@ class Qwen2_5OmniForConditionalGeneration(
            return talker_hf_config.tts_text_start_token_id
        return self.tts_text_spk_token_ids[voice_type]

-    def process_fn(self,
-        input_ids: torch.Tensor,
-        inputs_embeds: torch.Tensor,
-        **info_dict: object,
-    ):
-        # Mixed-mode support: In a single step, both Prefill*n and Decode*n are supported.
-        # Rules:
-        # - Prefill segments are wrapped with special tokens: [BOS][PAD...][EOS]
-        # - Decode segments consist of a single non-special token.
-        # - If additional_information is provided (can be a list split by request or a 
-        #   concatenated tensor plus a list of shapes), then for each request, reconstruct 
-        #   the thinker→talker input embeddings for the Prefill segments;
-        # - For Decode segments, if per-request auxiliary decode embeddings are provided (optional),
-        #   add them; otherwise, keep the original embedding.
-
-        # Ensure we have base embeddings when only ids are provided
-        if inputs_embeds is None and input_ids is not None:
-            inputs_embeds = self.talker.get_input_embeddings(input_ids)
-
-        span_len = input_ids.shape[0]
-        if span_len > 1:
-            # prefill
-            return self.thinker_to_talker_process(input_ids, inputs_embeds, **info_dict)
-        else:
-            # decode
-            return self.thinker_to_talker_decode_one_step(input_ids, inputs_embeds, **info_dict)
-
-
-    def thinker_to_talker_process(self, 
-        input_ids: torch.Tensor,
-        inputs_embeds: torch.Tensor,
-        **info_dict: object,
-    ):
-        update_dict = {}
-
-        prompt_embeds = info_dict.get("prompt_embeds")  # Tensor [P,H]
-        thinker_result = info_dict.get("thinker_result")  # Tensor [K,H]
-        prompt_token_ids = info_dict.get("prompt_token_ids")  # list[int]
-        thinker_output_token_ids = info_dict.get("thinker_output_token_ids")  # list[int]
-
-        if not isinstance(prompt_embeds, torch.Tensor):
-            prompt_embeds = torch.zeros(0, self.talker.config.hidden_size, dtype=inputs_embeds.dtype, device=self._module_device(self.model))
-        if not isinstance(thinker_result, torch.Tensor):
-            thinker_result = torch.zeros(0, self.talker.config.hidden_size, dtype=inputs_embeds.dtype, device=self._module_device(self.model))
-        if not isinstance(prompt_token_ids, (list, torch.Tensor)):
-            prompt_token_ids = []
-        if not isinstance(thinker_output_token_ids, (list, torch.Tensor)):
-            thinker_output_token_ids = []
-        
-        req_input_ids, req_embeds = self._thinker_to_talker_prefill(
-            voice_type=info_dict.get("voice_type"),
-            output_prompt_embeds=thinker_result.to(inputs_embeds.dtype).to(self._module_device(self.model)),
-            output_token_ids=thinker_output_token_ids,
-            thinker_prompt_embeds=prompt_embeds.to(inputs_embeds.dtype).to(self._module_device(self.model)),
-            prompt_token_ids=prompt_token_ids,
-        )
-        
-        if thinker_result.ndim == 2 and thinker_result.shape[0] > 0:
-            update_dict["thinker_reply_part"] = thinker_result[1:].detach().to("cpu").contiguous()
-
-        return req_input_ids, req_embeds, update_dict
-
    def _thinker_to_talker_prefill(
        self,
        voice_type: str,
@@ -698,38 +786,6 @@ class Qwen2_5OmniForConditionalGeneration(
            )
        return prompt_token_ids_processed, prompt_embeds

-    def thinker_to_talker_decode_one_step(self,
-        input_ids,
-        inputs_embeds,
-        **info_dict
-    ):
-        update_dict = {}
-        # choose step vector in priority order
-        step_vec = None
-        q = info_dict.get("thinker_reply_part", None)
-        if isinstance(q, torch.Tensor) and q.numel() > 0:
-            step_vec = q[0:1]
-            new_q = q[1:].detach().to("cpu").contiguous()
-            update_dict["thinker_reply_part"] = new_q
-        else:
-            # B) per-request provided decode vector (optional)
-            dv = info_dict.get("decode_output_prompt_embeds") if isinstance(info_dict, dict) else None
-            if isinstance(dv, torch.Tensor) and dv.numel() > 0:
-                step_vec = dv[0:1] if dv.ndim == 2 else dv.view(1, -1)
-            elif hasattr(self, "thinker_reply_part") and isinstance(self.thinker_reply_part, torch.Tensor) and self.thinker_reply_part.numel() > 0:
-                # C) fallback shared pool
-                step_vec = self.thinker_reply_part[0:1]
-                self.thinker_reply_part = self.thinker_reply_part[1:]
-
-        if isinstance(step_vec, torch.Tensor) and step_vec.numel() > 0:
-            one_id = input_ids[0:1]
-            _, one_embed = self._thinker_to_talker_decode_one_step(
-                output_prompt_embeds=step_vec.to(inputs_embeds.dtype).to(self._module_device(self.model)),
-                output_token_ids=one_id,
-            )
-            inputs_embeds[0] = one_embed[0]
-        return input_ids[0:1], inputs_embeds[0:1], update_dict
-
    def _thinker_to_talker_decode_one_step(
        self,
        output_prompt_embeds,
@@ -810,7 +866,7 @@ class Qwen2_5OmniForConditionalGeneration(
        __init__."""
        if self.token2wav is None or self.token2wav_config is None:
            return
-        device = "cuda" if torch.cuda.is_available() else "cpu"
+        device = self._module_device(self.token2wav)
        # optional speaker resources
        conds = getattr(self.token2wav_config, "conds", None)
        ref_mels = getattr(self.token2wav_config, "ref_mels", None)
@@ -33,6 +33,8 @@ from vllm.v1.outputs import SamplerOutput
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.sampler import Sampler

+from vllm_omni.utils.platform_utils import is_npu
+

 # Provide a no-op auto_docstring decorator to satisfy annotations if missing
 def auto_docstring(func=None, **_kwargs):
@@ -724,7 +726,13 @@ def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) ->
    else:
        beta = 0.0

-    kaiser_window = torch.kaiser_window(kernel_size, beta=beta, periodic=False, dtype=torch.float32)
+    # TODO: When torch.kaiser_window supports NPU, remove the device="cpu" argument
+    if is_npu():
+        kaiser_window = torch.kaiser_window(
+            kernel_size, beta=beta, periodic=False, dtype=torch.float32, device="cpu"
+        ).to("npu")
+    else:
+        kaiser_window = torch.kaiser_window(kernel_size, beta=beta, periodic=False, dtype=torch.float32)

    # Compute time indices
    if is_even:
@@ -734,7 +742,7 @@ def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) ->

    # Compute sinc filter
    if cutoff == 0:
-        return torch.zeros((1, 1, kernel_size), dtype=torch.float32)  # Ensures correct shape
+        return torch.zeros((1, 1, kernel_size), dtype=torch.float32)

    sinc_filter = torch.sinc(2 * cutoff * time_indices)
    normalized_filter = 2 * cutoff * kaiser_window * sinc_filter
@@ -745,6 +753,29 @@ def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) ->
    return normalized_filter.view(1, 1, kernel_size)


+def replication_pad_1d(hidden_states: torch.Tensor, pad_left: int, pad_right: int) -> torch.Tensor:
+    """
+    Manual replicate padding to avoid replication_pad1d kernel limits on NPU.
+    TODO: remove when F.pad supports replicate mode on NPU.
+    """
+    # NOTE: a immature implementation for running in NPU. Need to discuss.
+    if pad_left == 0 and pad_right == 0:
+        return hidden_states
+
+    segments = []
+    if pad_left > 0:
+        left = hidden_states[..., :1].expand(*hidden_states.shape[:-1], pad_left)
+        segments.append(left)
+
+    segments.append(hidden_states)
+
+    if pad_right > 0:
+        right = hidden_states[..., -1:].expand(*hidden_states.shape[:-1], pad_right)
+        segments.append(right)
+
+    return torch.cat(segments, dim=-1)
+
+
 class UpSample1d(nn.Module):
    def __init__(self, ratio=2, kernel_size=None):
        super().__init__()
@@ -760,14 +791,28 @@ class UpSample1d(nn.Module):

    def forward(self, hidden_states):
        channels = hidden_states.shape[1]
-        hidden_states_dtype = hidden_states.dtype
-        hidden_states = F.pad(hidden_states, (self.pad, self.pad), mode="replicate").to(self.filter.dtype)
-        hidden_states = self.ratio * F.conv_transpose1d(
-            hidden_states,
-            self.filter.expand(channels, -1, -1),
-            stride=self.stride,
-            groups=channels,
-        ).to(hidden_states_dtype)
+        if is_npu():
+            # TODO: When F.pad supports replicate mode on NPU, remove this branch
+            input_dtype = hidden_states.dtype
+            # F.pad in NPU doesn't support BF16 when mode is replicate.
+            # To ensure the accuracy, manually pad the input tensor.
+            hidden_states = replication_pad_1d(hidden_states.to(self.filter.dtype), self.pad, self.pad)
+            filter_convert_dtype = self.filter.to(hidden_states.dtype)
+            hidden_states = self.ratio * F.conv_transpose1d(
+                hidden_states,
+                filter_convert_dtype.expand(channels, -1, -1),
+                stride=self.stride,
+                groups=channels,
+            ).to(input_dtype)
+        else:
+            hidden_states_dtype = hidden_states.dtype
+            hidden_states = F.pad(hidden_states, (self.pad, self.pad), mode="replicate").to(self.filter.dtype)
+            hidden_states = self.ratio * F.conv_transpose1d(
+                hidden_states,
+                self.filter.expand(channels, -1, -1),
+                stride=self.stride,
+                groups=channels,
+            ).to(hidden_states_dtype)
        hidden_states = hidden_states[..., self.pad_left : -self.pad_right]

        return hidden_states
@@ -793,14 +838,29 @@ class DownSample1d(nn.Module):

    def forward(self, hidden_states):
        channels = hidden_states.shape[1]
-        hidden_states_dtype = hidden_states.dtype
-        hidden_states = F.pad(hidden_states, (self.pad_left, self.pad_right), mode="replicate").to(self.filter.dtype)
-        out = F.conv1d(
-            hidden_states,
-            self.filter.expand(channels, -1, -1),
-            stride=self.stride,
-            groups=channels,
-        ).to(hidden_states_dtype)
+        if is_npu():
+            input_dtype = hidden_states.dtype
+            # F.pad in NPU doesn't support BF16 when mode is replicate.
+            # To ensure the accuracy, manually pad the input tensor.
+            hidden_states = replication_pad_1d(hidden_states.to(self.filter.dtype), self.pad_left, self.pad_right)
+            filter_on_device = self.filter.to(device=hidden_states.device, dtype=hidden_states.dtype)
+            out = F.conv1d(
+                hidden_states,
+                filter_on_device.expand(channels, -1, -1),
+                stride=self.stride,
+                groups=channels,
+            ).to(input_dtype)
+        else:
+            hidden_states_dtype = hidden_states.dtype
+            hidden_states = F.pad(hidden_states, (self.pad_left, self.pad_right), mode="replicate").to(
+                self.filter.dtype
+            )
+            out = F.conv1d(
+                hidden_states,
+                self.filter.expand(channels, -1, -1),
+                stride=self.stride,
+                groups=channels,
+            ).to(hidden_states_dtype)
        return out


@@ -519,9 +519,9 @@ class Qwen3OmniMoeForConditionalGeneration(
                        tts_eos_thinker=tts_eos_thinker,
                        tts_pad_thinker=tts_pad_thinker,
                    )
-                    # req_embeds: [S,D], req_input_ids: [S]  →  map到当前 span
+                    # req_embeds: [S,D], req_input_ids: [S]  →  map onto the current span
                    seg_len = min(span_len, req_embeds.shape[0])
-                    # 强制对齐 dtype/device/内存布局，避免精度或布局差异导致比较不相等
+                    # Force-align dtype/device/memory layout to avoid mismatches from precision or layout differences
                    inputs_embeds[s : s + seg_len] = (
                        req_embeds[:seg_len].to(device=inputs_embeds.device, dtype=inputs_embeds.dtype).contiguous()
                    )
@@ -6,6 +6,7 @@ The code predictor generates residual RVQ (Residual Vector Quantization) codes
 autoregressively, predicting layers 1 to N based on layer-0 codes from the talker.
 """

+from collections import namedtuple
 from typing import Any, Optional

 import torch
@@ -164,21 +165,48 @@ class Qwen3OmniCodePredictorAttention(nn.Module):

        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS

-        attention_interface = ALL_ATTENTION_FUNCTIONS["flash_attention_2"]
-        attn_output, _ = attention_interface(
-            self,
-            q_heads,
-            k_heads,
-            v_heads,
-            None,
-            dropout=0.0 if not self.training else self.attention_dropout,
-            scaling=self.head_dim**-0.5,
-            sliding_window=None,
-            use_cache=use_cache,
-            position_ids=position_ids,
-            output_hidden_states=True,
-            output_attentions=False,
-        )
+        # Try attention backends in order of preference, with runtime error handling
+        # This handles cases where the backend is registered but not actually available
+        attention_backends = ["flash_attention_2", "xformers", "eager", "sdpa"]
+        attn_output = None
+        last_error = None
+
+        for backend_name in attention_backends:
+            if backend_name not in ALL_ATTENTION_FUNCTIONS:
+                continue
+
+            try:
+                attention_interface = ALL_ATTENTION_FUNCTIONS[backend_name]
+                attn_output, _ = attention_interface(
+                    self,
+                    q_heads,
+                    k_heads,
+                    v_heads,
+                    None,
+                    dropout=0.0 if not self.training else getattr(self, "attention_dropout", 0.0),
+                    scaling=self.head_dim**-0.5,
+                    sliding_window=None,
+                    use_cache=use_cache,
+                    position_ids=position_ids,
+                    output_hidden_states=True,
+                    output_attentions=False,
+                )
+                # Success - log fallback if not using flash_attention_2
+                if backend_name != "flash_attention_2":
+                    logger.warning_once(
+                        f"Using {backend_name} attention backend (flash_attention_2 not available or failed)"
+                    )
+                break
+            except (ValueError, ImportError, RuntimeError, AttributeError) as e:
+                # Store error and try next backend
+                last_error = e
+                continue
+
+        if attn_output is None:
+            raise RuntimeError(
+                f"All attention backends failed. Last error: {last_error}. "
+                "Please install flash-attn, or ensure PyTorch's scaled_dot_product_attention is available."
+            )
        attn_output = attn_output.reshape(*(hidden_states.shape[:-1]), -1).contiguous()

        attn_output = self.o_proj(attn_output)
@@ -373,16 +401,22 @@ class Qwen3OmniCodePredictorBaseModel(nn.Module):
        past_key_values: Optional[Any] = None,
        use_cache: Optional[bool] = False,
        cache_position: Optional[torch.LongTensor] = None,
-        **kwargs,
-    ):
+        **kwargs: Any,
+    ) -> Any:
        """
        Forward pass matching HF structure.

        Args:
            inputs_embeds: [batch, seq_len, hidden_size]
+            attention_mask: Optional attention mask tensor
+            position_ids: Optional position IDs tensor
+            past_key_values: Optional cached key-value pairs
+            use_cache: Whether to use cache
+            cache_position: Optional cache position tensor
+            **kwargs: Additional keyword arguments

        Returns:
-            Object with .last_hidden_state attribute
+            Named tuple with .last_hidden_state and .past_key_values attributes
        """
        batch_size, seq_len, _ = inputs_embeds.shape

@@ -421,8 +455,6 @@ class Qwen3OmniCodePredictorBaseModel(nn.Module):
        hidden_states = self.norm(hidden_states)

        # Return in HF-compatible format
-        from collections import namedtuple
-
        Output = namedtuple("Output", ["last_hidden_state", "past_key_values"])
        return Output(last_hidden_state=hidden_states, past_key_values=None)  # [batch, num_code_groups-1, hidden_size]

@@ -1,4 +1,5 @@
 from collections.abc import Iterable
+from typing import Any

 import torch
 import torch.nn as nn
@@ -284,7 +285,7 @@ class Qwen3OmniMoeTalkerForConditionalGeneration(

        return result_codes, summed_embeddings

-    def init_multi_modal(self, thinker_config):
+    def init_multi_modal(self, thinker_config: Any) -> None:
        """
        Initialize multimodal components from the thinker.

@@ -1,27 +0,0 @@
-from collections.abc import Callable
-import torch
-
-class PreprocessMixin:
-    """
-    Mixin class for all stages in the Omni model.
-    """
-    def set_preprocess(self, preprocess_fn: Callable) -> None:
-        """
-        Set a preprocess function for the stage.
-        Args:
-            preprocess_fn: The preprocess function to register.
-        """
-        self.preprocess = preprocess_fn
-
-    def preprocess(self,
-                input_ids: torch.Tensor,
-                inputs_embeds: torch.Tensor,
-                **input_dict: object
-    ) -> tuple[torch.Tensor, torch.Tensor, dict]:
-        """
-        Process the input_ids and inputs_embeds for the given input_dict.
-        Returns the processed input_ids, inputs_embeds, and the input_dict.
-
-        If the stage don't applicable, return the original input_ids, inputs_embeds, and an empty dict.
-        """
-        return input_ids, inputs_embeds, {}
@@ -0,0 +1,94 @@
+# stage config for running qwen2.5-omni with architecture of OmniLLM.
+stage_args:
+  - stage_id: 0
+    runtime:
+      process: true            # Run this stage in a separate process
+      devices: "0"            # Visible devices for this stage
+      max_batch_size: 1
+    engine_args:
+      model_stage: thinker
+      model_arch: Qwen2_5OmniForConditionalGeneration
+      worker_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker
+      scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
+      gpu_memory_utilization: 0.8
+      enforce_eager: true  # Now we only support eager mode
+      trust_remote_code: true
+      engine_output_type: latent
+      enable_prefix_caching: false
+    is_comprehension: true
+    final_output: true
+    final_output_type: text
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.1
+  - stage_id: 1
+    runtime:
+      process: true
+      devices: "1"
+      max_batch_size: 1
+    engine_args:
+      model_stage: talker
+      model_arch: Qwen2_5OmniForConditionalGeneration
+      worker_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker
+      scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
+      gpu_memory_utilization: 0.8
+      enforce_eager: true
+      trust_remote_code: true
+      enable_prefix_caching: false
+      engine_output_type: latent
+    engine_input_source: [0]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker
+    default_sampling_params:
+      temperature: 0.9
+      top_p: 0.8
+      top_k: 40
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
+      stop_token_ids: [8294]
+  - stage_id: 2
+    runtime:
+      process: true
+      devices: "0"            # Example: use a different NPU than the previous stage; use "0" if single NPU
+      max_batch_size: 1
+    engine_args:
+      model_stage: code2wav
+      model_arch: Qwen2_5OmniForConditionalGeneration
+      worker_cls: vllm_omni.worker.npu.npu_diffusion_worker.NPUDiffusionWorker
+      scheduler_cls: vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler
+      gpu_memory_utilization: 0.15
+      enforce_eager: true
+      trust_remote_code: true
+      enable_prefix_caching: false
+      engine_output_type: audio
+    engine_input_source: [1]
+    final_output: true
+    final_output_type: audio
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.1
+
+# Top-level runtime config (concise): default windows and stage edges
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1             # Simplified: trigger downstream only after full upstream completion
+    max_inflight: 1             # Simplified: process serially within each stage
+  edges:
+    - from: 0                   # thinker → talker: trigger only after receiving full input (-1)
+      to: 1
+      window_size: -1
+    - from: 1                   # talker → code2wav: trigger only after receiving full input (-1)
+      to: 2
+      window_size: -1
@@ -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:
@@ -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
@@ -3,7 +3,7 @@
 # Copyright 2025 The Qwen team.
 """Stage input processor for Qwen3 Omni MoE: Thinker → Talker transition."""

-from typing import Union
+from typing import Any, Union

 import torch
 from vllm.inputs import TextPrompt
@@ -47,11 +47,11 @@ def _compute_talker_prompt_ids_length(info):


 def thinker2talker(
-    stage_list,
-    engine_input_source,
-    prompt: Union[OmniTokensPrompt, TextPrompt] = None,
+    stage_list: list[Any],
+    engine_input_source: list[int],
+    prompt: Union[OmniTokensPrompt, TextPrompt, None] = None,
    requires_multimodal_data: bool = False,
-):
+) -> list[OmniTokensPrompt]:
    """
    Process thinker outputs to create talker inputs.

@@ -64,6 +64,7 @@ def thinker2talker(
        stage_list: List of stage objects
        engine_input_source: Source stage IDs (typically [0] for thinker)
        prompt: Original prompt data
+        requires_multimodal_data: Whether multimodal data is required

    Returns:
        List of OmniTokensPrompt for talker stage
@@ -111,11 +112,11 @@ def thinker2talker(


 def talker2code2wav(
-    stage_list,
-    engine_input_source,
-    prompt: Union[OmniTokensPrompt, TextPrompt] = None,
+    stage_list: list[Any],
+    engine_input_source: list[int],
+    prompt: Union[OmniTokensPrompt, TextPrompt, None] = None,
    requires_multimodal_data: bool = False,
-):
+) -> list[OmniTokensPrompt]:
    """
    Process talker outputs to create code2wav inputs.

@@ -128,6 +129,7 @@ def talker2code2wav(
        stage_list: List of stage objects
        engine_input_source: Source stage IDs (typically [1] for talker)
        prompt: Original prompt data
+        requires_multimodal_data: Whether multimodal data is required

    Returns:
        List of OmniTokensPrompt for code2wav stage
@@ -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:
@@ -0,0 +1,11 @@
+from vllm_omni.utils.platform_utils import (
+    detect_device_type,
+    get_device_control_env_var,
+    is_npu,
+)
+
+__all__ = [
+    "detect_device_type",
+    "get_device_control_env_var",
+    "is_npu",
+]
@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+import torch
+from vllm.platforms import current_platform
+
+
+def detect_device_type() -> str:
+    device_type = getattr(current_platform, "device_type", None)
+    if isinstance(device_type, str) and device_type:
+        return device_type.lower()
+    if torch.cuda.is_available():
+        return "cuda"
+    if hasattr(torch, "npu") and torch.npu.is_available():  # type: ignore[attr-defined]
+        return "npu"
+    return "cpu"
+
+
+def is_npu() -> bool:
+    return detect_device_type() == "npu"
+
+
+def get_device_control_env_var() -> str:
+    """Return the environment variable name for device visibility control."""
+    if hasattr(current_platform, "device_control_env_var"):
+        env_var = getattr(current_platform, "device_control_env_var", None)
+        if isinstance(env_var, str) and env_var:
+            return env_var
+
+    device_type = detect_device_type()
+    if device_type == "npu":
+        return "ASCEND_RT_VISIBLE_DEVICES"
+    return "CUDA_VISIBLE_DEVICES"  # fallback
+
+
+def get_diffusion_worker_class() -> type:
+    """Get the appropriate diffusion WorkerProc class based on current device type.
+
+    Returns:
+        The WorkerProc class for the detected device type (either NPUWorkerProc or WorkerProc).
+
+    Raises:
+        ImportError: If the worker module for the detected device is not available.
+    """
+    device_type = detect_device_type()
+
+    if device_type == "npu":
+        from vllm_omni.diffusion.worker.npu.npu_worker import NPUWorkerProc
+
+        return NPUWorkerProc
+    else:
+        # Default to GPU worker for cuda and other devices
+        from vllm_omni.diffusion.worker.gpu_worker import WorkerProc
+
+        return WorkerProc
@@ -356,4 +356,4 @@ class GPUARModelRunner(OmniGPUModelRunner):
                merged[k] = v
        setattr(req_state, "additional_information_cpu", merged)

-    # ===== Helper functions extracted for clarity and reuse =====
+    # ===== Helper functions extracted for clarity and reuse =====
@@ -1,4 +1,4 @@
-"""Code2Wav GPU Model Runner for vLLM-omni.
+"""Code2Wav GPU Model Runner for vLLM-Omni.

 Handles direct conversion from codec codes to audio waveforms for Qwen3 Omni MoE Code2Wav.
 This is a non-autoregressive model that doesn't require sampling or logits computation.
@@ -37,7 +37,7 @@ logger = logging.getLogger(__name__)


 class GPUGenerationModelRunner(OmniGPUModelRunner):
-    """Generation model runner for vLLM-omni (non-autoregressive).
+    """Generation model runner for vLLM-Omni (non-autoregressive).

    - Reuses GPUModelRunner preparation, multimodal handling, and TP/PP/DP glue.
    - Does not compute logits or perform token sampling.
@@ -6,6 +6,8 @@ from vllm.config import CUDAGraphMode
 from vllm.distributed.parallel_state import get_pp_group
 from vllm.forward_context import BatchDescriptor, set_forward_context
 from vllm.logger import init_logger
+from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
+from vllm.model_executor.models.interfaces import supports_mrope
 from vllm.model_executor.models.interfaces_base import VllmModelForPooling
 from vllm.sampling_params import SamplingType
 from vllm.utils import LazyLoader, cdiv
@@ -33,6 +35,63 @@ logger = init_logger(__name__)


 class OmniGPUModelRunner(GPUModelRunner):
+    def _init_mrope_positions(self, req_state: CachedRequestState):
+        """Initialize M-RoPE positions for multimodal inputs.
+
+        Extracts multimodal feature metadata (image grids, video grids,
+        audio features) and computes M-RoPE positions for proper positional
+        encoding of multimodal tokens.
+
+        Args:
+            req_state: Cached request state containing multimodal features
+
+        Raises:
+            AssertionError: If the model does not support M-RoPE
+        """
+        image_grid_thw = []
+        video_grid_thw = []
+        second_per_grid_ts = []
+        audio_feature_lengths = []
+        use_audio_in_video = False
+        for mm_feature in req_state.mm_features:
+            mm_item = mm_feature.data
+            if mm_item is None:
+                continue
+            mm_input = mm_item.get_data()
+            if (t := mm_input.get("image_grid_thw")) is not None:
+                image_grid_thw.append(t.tolist())
+            if (t := mm_input.get("video_grid_thw")) is not None:
+                video_grid_thw.append(t.tolist())
+            if (t := mm_input.get("second_per_grid_ts")) is not None:
+                second_per_grid_ts.append(t)
+            if (t := mm_input.get("audio_feature_lengths")) is not None:
+                audio_feature_lengths.append(t)
+            # Check for use_audio_in_video
+            use_audio_in_video_value = mm_input.get("use_audio_in_video")
+            if use_audio_in_video_value is not None:
+                use_audio_in_video = bool(use_audio_in_video_value.item())
+
+        if supports_mrope(self.model):
+            req_state.mrope_positions, req_state.mrope_position_delta = self.model.get_mrope_input_positions(
+                req_state.prompt_token_ids,
+                hf_config=self.model_config.hf_config,
+                image_grid_thw=image_grid_thw,
+                video_grid_thw=video_grid_thw,
+                second_per_grid_ts=second_per_grid_ts,
+                audio_feature_lengths=audio_feature_lengths,
+                use_audio_in_video=use_audio_in_video,
+            )
+        else:
+            req_state.mrope_positions, req_state.mrope_position_delta = MRotaryEmbedding.get_input_positions_tensor(
+                req_state.prompt_token_ids,
+                hf_config=self.model_config.hf_config,
+                image_grid_thw=image_grid_thw,
+                video_grid_thw=video_grid_thw,
+                second_per_grid_ts=second_per_grid_ts,
+                audio_feature_lengths=audio_feature_lengths,
+                use_audio_in_video=use_audio_in_video,
+            )
+
    def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
        """Update the cached states and the persistent batch with the scheduler
        output.
@@ -639,34 +698,6 @@ class OmniGPUModelRunner(GPUModelRunner):
            encoder_inputs = self._extract_encoder_inputs(scheduler_output)
            model_kwargs.update(encoder_inputs)

-        if self.model.has_preprocess:
-            # Overlay custom prompt_embeds per request for the prompt portion;
-            # collect additional_information (tensor/list) for prefill portion only
-            for req_index, req_id in enumerate(self.input_batch.req_ids):
-                req_state = self.requests.get(req_id)
-                req_infos = (getattr(req_state, "additional_information_cpu", None)
-                             if req_state is not None else None)
-
-                start_offset = int(self.query_start_loc.cpu[req_index])
-                sched_tokens = int(num_scheduled_tokens_np[req_index])
-                s, e = start_offset, start_offset + sched_tokens
-                span_len = int(e) - int(s)
-
-                # call the custom process function
-                req_input_ids, req_embeds, update_dict = self.model.preprocess(
-                    input_ids=input_ids[s:e],
-                    inputs_embeds=inputs_embeds[s:e],
-                    **req_infos
-                )
-                # TODO(Peiqi): the merge stage could move out from the critical path
-                self._merge_additional_information_update(req_id, update_dict)
-
-                # update the inputs_embeds and input_ids
-                seg_len = min(span_len, req_embeds.shape[0])
-                inputs_embeds[s : s + seg_len] = req_embeds[:seg_len]
-                if isinstance(req_input_ids, torch.Tensor) and req_input_ids.numel() == seg_len:
-                    input_ids[s : s + seg_len] = req_input_ids
-
        return (
            num_scheduled_tokens,
            num_input_tokens,
Author	SHA1	Message	Date
Kevin H. Luu	edda01380d	p Signed-off-by: Kevin H. Luu <khluu000@gmail.com>	2025-12-04 02:17:07 -08:00
Kevin H. Luu	205238f7e3	p Signed-off-by: Kevin H. Luu <khluu000@gmail.com>	2025-12-04 02:08:14 -08:00
Kevin H. Luu	9dc7a70ca6	p Signed-off-by: Kevin H. Luu <khluu000@gmail.com>	2025-12-04 02:04:51 -08:00
Kevin H. Luu	ebd17001c4	p Signed-off-by: Kevin H. Luu <khluu000@gmail.com>	2025-12-04 01:59:29 -08:00
Kevin H. Luu	e63c151f86	p Signed-off-by: Kevin H. Luu <khluu000@gmail.com>	2025-12-04 01:55:18 -08:00
AliciaandRoger Wang	c0b88a87ba	[Doc] Add installation part for pre built docker. (#141 ) Signed-off-by: WANG Cong <115451386+congw729@users.noreply.github.com> Signed-off-by: Alicia Wang <115451386+congw729@users.noreply.github.com> Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com> Signed-off-by: Roger Wang <hey@rogerw.io> Co-authored-by: Roger Wang <hey@rogerw.io>	2025-12-04 00:04:18 -08:00
Canlin Guo	863cdb82b6	[Docs] remove Ascend word to make docs general (#190 ) Signed-off-by: gcanlin <canlinguosdu@gmail.com>	2025-12-03 23:21:11 -08:00
knlnguyen1802andn00909098	f73791d246	[Bugfix] Init early toch cuda (#185 ) Signed-off-by: n00909098 <nguyen.kha.long@huawei.com> Co-authored-by: n00909098 <nguyen.kha.long@huawei.com>	2025-12-04 11:38:52 +08:00
Gao Han	0e07a0ae9e	[Docs] Supplement volunteers and faq docs (#182 ) Signed-off-by: Gaohan123 <hgaoaf@connect.ust.hk>	2025-12-04 10:37:00 +08:00
汪志鹏andZhipeng Wang	598d26580e	[Doc] change `vllm_omni` to `vllm-omni` (#177 ) Signed-off-by: Zhipeng Wang <zhipeng.wang@rakuten.com> Co-authored-by: Zhipeng Wang <zhipeng.wang@rakuten.com>	2025-12-03 17:43:59 +08:00
Cyrus Leung	574e1fb68d	[Doc] Update Security Advisory link (#173 ) Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>	2025-12-03 12:39:08 +08:00
wjcwjc77 wangjiangcheng Gao Han	8ed9716d5b	[Doc] Fix wrong examples URLs (#166 ) Signed-off-by: wangjiangcheng <wangjiangcheng@kuaishou.com> Co-authored-by: wangjiangcheng <wangjiangcheng@kuaishou.com> Co-authored-by: Gao Han <gaohan19@huawei.com>	2025-12-03 03:25:58 +00:00
Jiangyun ZhuandSamit	7a6d2fe3f6	[diffusion] z-image support (#149 ) Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: Samit <285365963@qq.com>	2025-12-03 02:16:36 +00:00
Samit	5a0d7882c5	[Misc] Qwen-Omni support offline inference with local files (#167 ) Signed-off-by: samithuang <285365963@qq.com>	2025-12-02 19:05:45 +08:00
Roger Wang	55ea6ef7e9	[Misc] Update link in issue template (#155 ) Signed-off-by: Roger Wang <hey@rogerw.io>	2025-12-02 08:20:37 +08:00
Roger Wang	9fe730a44b	[Misc] update image reference for PyPI (#150 ) Signed-off-by: Roger Wang <hey@rogerw.io>	2025-12-01 09:37:29 -08:00
Mengqing Cao	00b1a9a1a0	[Misc] Temporarily support downloading models from ModelScope by snapshot download (#132 ) Signed-off-by: MengqingCao <cmq0113@163.com>	2025-12-01 17:13:23 +00:00
Canlin Guo	7e60c49ce4	[Bugfix] Add Fake VllmConfig on NPU and add slicing/tiling args in Qwen-Image (#145 ) Signed-off-by: gcanlin <canlinguosdu@gmail.com>	2025-12-01 16:40:01 +00:00
Samit	98d94e0e3d	[Bugfix] Fix Qwen3-Omni gradio audio input bug (#147 ) Signed-off-by: samithuang <285365963@qq.com>	2025-12-01 13:52:31 +00:00
Roger Wang	510aeb7da4	[Misc] Misc description updates (#146 ) Signed-off-by: Roger Wang <hey@rogerw.io>	2025-12-01 12:09:24 +00:00
Roger Wang	5939089c98	[Doc] Logo (#143 ) Signed-off-by: Roger Wang <hey@rogerw.io>	2025-12-01 02:26:39 -08:00
Samit	1a23d70a97	[Misc][WIP] Support qwen-omni online inference with local video/audio/image path (#131 ) Signed-off-by: samithuang <285365963@qq.com>	2025-12-01 00:30:22 -08:00
Roger Wang	a6c5a6aca4	[Doc] Update developer & user channel (#138 ) Signed-off-by: Roger Wang <hey@rogerw.io>	2025-12-01 00:25:44 -08:00
Gao Han	d53155f45c	[Bugfix] Fix the curl bug of qwen3-omni and doc errors (#135 ) Signed-off-by: Gaohan123 <gaohan19@huawei.com>	2025-12-01 08:01:54 +00:00
youkaichao	e8fb85d38c	Update README with vllm-omni blogpost link (#137 ) Signed-off-by: youkaichao <youkaichao@gmail.com>	2025-11-30 23:31:32 -08:00
Samit	f378f80430	[Bugfix] Fix redundant shm broadcast warnings in diffusion workers (#133 ) Signed-off-by: samithuang <285365963@qq.com>	2025-11-30 23:17:56 -08:00
Roger Wang	a1914e1176	[Doc] Fix arch pic reference (#136 ) Signed-off-by: Roger Wang <hey@rogerw.io>	2025-11-30 23:17:21 -08:00
Roger Wang	3b12ed9cf4	[Doc] Cleanup reference to deleted files (#134 ) Signed-off-by: Roger Wang <hey@rogerw.io>	2025-11-30 22:50:21 -08:00
Roger Wang	c852267a1a	[Misc] Update default stage config for qwen3-omni (#124 ) Signed-off-by: Roger Wang <hey@rogerw.io>	2025-11-30 20:01:20 -08:00
Cyrus Leung	2c5d0481b3	[Doc] Fix outdated CONTRIBUTING link (#127 ) Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>	2025-12-01 11:11:33 +08:00
Roger Wang	7391db62d1	[Doc] Cleanup non-english comments (#125 ) Signed-off-by: Roger Wang <hey@rogerw.io>	2025-11-30 16:45:30 -08:00
Yikun Jiang	90ba189232	Refresh supported models and address nits in doc (#119 ) Signed-off-by: Yikun Jiang <yikunkero@gmail.com>	2025-11-30 16:36:04 -08:00
Alicia	a33bbcedea	[Fix] Resolve comments & update vLLM-Omni name usages. (#122 ) Signed-off-by: WANG Cong <115451386+congw729@users.noreply.github.com>	2025-12-01 08:05:24 +08:00
Samit	04987ffa7a	[Bugfix] Fix Qwen-omni Online Inference Bug caused by check_stop and long sequence (#112 ) Signed-off-by: samithuang <285365963@qq.com>	2025-11-30 23:23:45 +08:00
Alicia	3427fb4774	[Feat] Rename AsyncOmniLLM -> AsyncOmni (#103 ) Signed-off-by: WANG Cong <115451386+congw729@users.noreply.github.com>	2025-11-30 22:26:33 +08:00
Hongsheng LiuandRoger Wang	0c245dd1f9	[Doc] Fix readme typos (#108 ) Signed-off-by: hsliu <liuhongsheng4@huawei.com> Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com> Co-authored-by: Roger Wang <hey@rogerw.io>	2025-11-30 21:32:04 +08:00
Gao Han	bfd3cad5ee	[Bugfix] Fix bug of online serving for qwen2.5-omni (#118 ) Signed-off-by: Gaohan123 <gaohan19@huawei.com>	2025-11-30 21:30:17 +08:00
Canlin Guo	0d110312df	[Feat] Enable loading local Qwen-Image model (#117 ) Signed-off-by: gcanlin <canlinguosdu@gmail.com>	2025-11-30 21:07:53 +08:00
Samit	3ea7fd280b	[Feature] Support Gradio Demo for Qwen3-Omni (#116 ) Signed-off-by: samithuang <285365963@qq.com>	2025-11-30 20:44:04 +08:00
Canlin Guo	ec9e04027d	[Feat] Add NPU Backend support for vLLM-Omni (#89 ) Signed-off-by: gcanlin <canlinguosdu@gmail.com>	2025-11-30 20:39:06 +08:00
Gao Han	88c00fb271	[Bugfix] Remove mandatory flash-attn dependency and optimzie docs (#113 ) Signed-off-by: Gaohan123 <gaohan19@huawei.com>	2025-11-30 18:20:44 +08:00
Zhou Taichang	228f4e5070	[Bug-fix] Fix Bugs in Qwen3/Qwen2.5 Omni Rebased Support (#114 ) Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2025-11-30 18:15:10 +08:00
Samit	f77e44d5bf	[Docs] Add qwen image missing doc in user guide (#111 ) Signed-off-by: samithuang <285365963@qq.com>	2025-11-30 17:46:39 +08:00