push

Signed-off-by: Kevin H. Luu <khluu000@gmail.com>
push
2025-12-08 16:44:19 -08:00 · 2025-12-08 16:41:04 -08:00 · 2025-12-08 16:39:58 -08:00 · 2025-12-08 16:35:25 -08:00 · 2025-12-08 16:04:10 -08:00 · 2025-12-08 15:34:04 -08:00
141 changed files with 5802 additions and 4279 deletions
@@ -0,0 +1,10 @@
+name: vllm_omni_ci
+github_repo_name: vllm-project/vllm-omni
+job_dirs:
+  - ".buildkite/jobs"
+run_all_patterns: []
+run_all_exclude_patterns: []
+registries: public.ecr.aws/q9t5s3a7
+repositories:
+  main: "vllm-ci-postmerge-repo"
+  premerge: "vllm-ci-test-repo"
@@ -0,0 +1,10 @@
+group: Build
+steps:
+  - label: ":docker: build vllm-omni-ci image"
+    key: image-build
+    depends_on: []
+    commands:
+      - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
+      - "docker build --file docker/Dockerfile.ci -t vllm-omni-ci ."
+      - "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
+      - "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
@@ -0,0 +1,30 @@
+group: Tests
+depends_on:
+  - image-build
+steps:
+- label: "Simple Unit Test"
+  commands:
+    - ".buildkite/scripts/simple_test.sh"
+  no_gpu: true
+  no_plugin: true
+
+- label: "Diffusion Model Test"
+  timeout_in_minutes: 15
+  commands:
+    - pytest -s -v tests/single_stage/test_diffusion_model.py
+
+- label: "Omni Model Test"
+  timeout_in_minutes: 15
+  num_gpus: 4
+  commands:
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - pytest -s -v tests/multi_stages/
+
+- label: "Omni Model Test with H100"
+  timeout_in_minutes: 20
+  gpu: h100
+  num_gpus: 2
+  commands:
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - pytest -s -v tests/multi_stages_h100/
@@ -1,6 +0,0 @@
-steps:
-  - label: "Simple Unit Test"
-    commands:
-      - ".buildkite/scripts/simple_test.sh"
-    agents:
-      queue: "cpu_queue_premerge"
@@ -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,9 +156,6 @@ cython_debug/
 # cursor
 .cursor/

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

@@ -232,4 +229,3 @@ configs/development.yaml
 # Docker
 .dockerignore
 Dockerfile.dev
-discussion
@@ -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://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> |
+| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> |
 </p>

 ---
@@ -30,7 +30,7 @@ Easy, fast, and cheap omni-modality model serving for everyone

 <p align="center">
  <picture>
-    <img alt="vllm-omni" src="./docs/source/architecture/omni-modality-model-architecture.png" width=55%>
+    <img alt="vllm-omni" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/omni-modality-model-architecture.png" width=55%>
  </picture>
 </p>

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

-## Weekly Meeting
-
- vLLM Omni Weekly Meeting: https://tinyurl.com/vllm-omni-meeting
- Wednesday, 11:30 - 12:30 (UTC+8, [Convert to your timezone](https://dateful.com/convert/gmt8?t=15))
+## Join the Community
+Feel free to ask questions, provide feedbacks and discuss with fellow users of vLLM-Omni in `#sig-omni` slack channel at [slack.vllm.ai](https://slack.vllm.ai) or vLLM user forum at [discuss.vllm.ai](https://discuss.vllm.ai).

 ## License

@@ -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 []
@@ -7,11 +7,11 @@ nav:
  - Examples:
    - examples/README.md
    - Offline Inference:
-      - Offline Example of vLLM-Omni for Qwen2.5-omni: user_guide/examples/offline_inference/qwen2_5_omni.md
-      - Offline Example of vLLM-Omni for Qwen3-omni: user_guide/examples/offline_inference/qwen3_omni.md
+      - Offline Example of vLLM-Omni for Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
+      - Offline Example of vLLM-Omni for Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
      - Qwen-Image Offline Inference: user_guide/examples/offline_inference/qwen_image.md
    - Online Serving:
-      - Online serving Example of vLLM-Omni for Qwen2.5-omni: user_guide/examples/online_serving/qwen2_5_omni.md
+      - Online serving Example of vLLM-Omni for Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
      - Online serving Example of vLLM-Omni for Qwen3-omni: user_guide/examples/online_serving/qwen3_omni.md
  - General:
    - usage/*
@@ -32,8 +32,13 @@ nav:
    - design/index.md
    - design/architecture_overview.md
    - design/vllm_omni_design.md
-    - design/mrs_design.md
-    - design/api_design_doc.md
+    - Feature Design:
+      - design/feature/multi_request_streaming_design.md
+      - design/feature/omni_connector_design.md
+    - Module Design:
+      - design/module/ar_module_design.md
+      - design/module/dit_module_design.md
+      - design/module/entrypoint_module_design.md
  - Docs Guide: contributing/DOCS_GUIDE.md
 - API Reference:
  - api/README.md
@@ -8,8 +8,8 @@ hide:

 <p align="center">
  <picture>
-    <source media="(prefers-color-scheme: dark)" src="./source/logos/vllm-omni-logo-text-dark.png">
-    <img alt="vllm-omni" src="./source/logos/vllm-omni-logo-text-dark.png" width=55%>
+    <source media="(prefers-color-scheme: dark)" src="./source/logos/vllm-omni-logo.png">
+    <img alt="vllm-omni" src="./source/logos/vllm-omni-logo.png" width=55%>
  </picture>
 </p>
 <h3 align="center">
@@ -34,7 +34,8 @@ Easy, fast, and cheap omni-modality model serving for everyone

 <p align="center">
  <picture>
-    <img alt="vllm-omni" src="./source/architecture/omni-modality model architecture.png" width=55%>
+    <source media="(prefers-color-scheme: dark)" src="./source/architecture/omni-modality-model-architecture.png">
+    <img alt="vllm-omni-arch" src="./source/architecture/omni-modality-model-architecture.png" width=55%>
  </picture>
 </p>

@@ -58,4 +59,6 @@ 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 architecture design and recent roadmaps](https://docs.google.com/presentation/d/1Y7t2Zm3BIISPN-_X_sDpxZMNCWir2rAx/edit?usp=drive_link&ouid=110473603432222024453&rtpof=true&sd=true)
+- [vllm-omni announcement blogpost](https://blog.vllm.ai/2025/11/30/vllm-omni.html)
@@ -49,10 +49,9 @@ Engine classes for offline and online inference.
 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.omni_ar_scheduler.OmniARScheduler][]
+- [vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler][]
 - [vllm_omni.core.sched.output.OmniNewRequestData][]
- [vllm_omni.core.sched.scheduler.OmniScheduler][]

 ## Model Executor

@@ -102,15 +101,17 @@ Configuration classes.

 Worker classes and model runners for distributed inference.

+- [vllm_omni.diffusion.worker.gpu_worker.GPUWorker][]
+- [vllm_omni.diffusion.worker.gpu_worker.WorkerProc][]
+- [vllm_omni.diffusion.worker.npu.npu_worker.NPUWorker][]
+- [vllm_omni.diffusion.worker.npu.npu_worker.NPUWorkerProc][]
 - [vllm_omni.worker.gpu_ar_model_runner.GPUARModelRunner][]
 - [vllm_omni.worker.gpu_ar_worker.GPUARWorker][]
- [vllm_omni.worker.gpu_diffusion_model_runner.GPUDiffusionModelRunner][]
- [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_generation_model_runner.NPUGenerationModelRunner][]
+- [vllm_omni.worker.npu.npu_generation_worker.NPUGenerationWorker][]
 - [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`.
@@ -9,7 +9,7 @@ stage_args:
      model_stage: thinker
      model_arch: Qwen2_5OmniForConditionalGeneration
      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
-      scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
      gpu_memory_utilization: 0.8
      enforce_eager: true  # Now we only support eager mode
      trust_remote_code: true
@@ -35,7 +35,7 @@ stage_args:
      model_stage: talker
      model_arch: Qwen2_5OmniForConditionalGeneration
      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
-      scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
      gpu_memory_utilization: 0.8
      enforce_eager: true
      trust_remote_code: true
@@ -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]"
 ```

@@ -0,0 +1 @@
+## OmniConnector for fully disaggregated serving (update soon)
@@ -6,8 +6,14 @@ 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
+## Feature Design Documents

- [vLLM-Omni API Design Documentation](api_design_doc.md)
+- [Multi-Request Streaming](feature/multi_request_streaming_design.md)
+- [Omni Connector Design](feature/omni_connector_design.md)
+
+## Module Design Documents
+
+- [AR Module Design](module/ar_module_design.md)
+- [DIT Module Design](module/dit_module_design.md)
+- [Entrypoint Module Design](module/entrypoint_module_design.md)
@@ -0,0 +1 @@
+Architecture design of the AR module (update soon)
@@ -0,0 +1 @@
+Architecture design of the DiT module (update soon)
@@ -0,0 +1 @@
+Architecture design of the entrypoint (update soon)
@@ -16,7 +16,7 @@ vLLM-Omni is a multi-modality extension for vLLM that supports non-autoregressiv

 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)
+ --> ModelRunner (New one for AR hiddenstate, New one for DiT) --> RequestState --> OutputProcessor (New one for final multimodal output)

 ## Core Components (Ordered by Data Flow)

@@ -137,7 +137,7 @@ class OmniLLM(LLM):
            prompt_str, engine_request, tokenization_kwargs = self._process_stage_inputs(stage_config, **stage_args)

            # Add inputs to Engine
-            stage_engine.add_request(requesy_id, prompt_str, tokenization_kwargs)
+            stage_engine.add_request(request_id, prompt_str, tokenization_kwargs)
            # Run Engine
            stage_output = stage_engine.step()

@@ -157,7 +157,7 @@ class OmniLLM(LLM):
            return self._process_ar_inputs(**stage_args)
        elif stage_config.engine_type == "DiT":
            return self._process_dit_inputs(**stage_args)
-        else
+        else:
            raise NotImplementedError

    def _process_dit_inputs(**stage_args)-> Any:
@@ -1,6 +1,6 @@
 # Installation

-vLLM supports the following hardware platforms:
+vLLM-Omni supports the following hardware platforms:

 - [GPU](gpu.md)
    - [NVIDIA CUDA](gpu.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]
@@ -19,6 +19,7 @@ export DEVICE1=/dev/davinci1
 export IMAGE=quay.io/ascend/vllm-ascend:v0.11.0rc2
 docker run --rm \
    --name vllm-omni-npu \
+    --shm-size=1g \
    --device $DEVICE0 \
    --device $DEVICE1 \
    --device /dev/davinci_manager \
@@ -16,7 +16,19 @@ Please refer to [installation](installation/README.md)

 ## Offline Inference

-Please refer to [offline inference](../user_guide/examples/offline_inference/qwen2_5_omni.md)
+Text-to-image generation quickstart with vLLM-Omni:
+
+```python
+from vllm_omni.entrypoints.omni import Omni
+
+if __name__ == "__main__":
+    omni = Omni(model="Tongyi-MAI/Z-Image-Turbo")
+    prompt = "a cup of coffee on the table"
+    images = omni.generate(prompt)
+    images[0].save("coffee.png")
+```
+
+For more usages, please refer to [offline inference](../user_guide/examples/offline_inference/qwen2_5_omni.md)

 ## Online Serving with OpenAI-Completions API

@@ -173,7 +173,7 @@ class Example:
            gh_file = gh_file.relative_to(ROOT_DIR)

            # Make GitHub URL
-            url = "https://github.com/vllm-project/vllm/"
+            url = "https://github.com/vllm-project/vllm-omni/"
            url += "tree/main" if self.path.is_dir() else "blob/main"
            gh_url = f"{url}/{gh_file}"

@@ -183,7 +183,7 @@ class Example:

    def generate(self) -> str:
        content = f"# {self.title}\n\n"
-        url = "https://github.com/vllm-project/vllm/"
+        url = "https://github.com/vllm-project/vllm-omni/"
        url += "tree/main" if self.path.is_dir() else "blob/main"
        content += f"Source <{url}/{self.path.relative_to(ROOT_DIR)}>.\n\n"

@@ -192,7 +192,8 @@ class Example:
        code_fence = "``````"

        if self.is_code:
-            content += f'{code_fence}{self.main_file.suffix[1:]}\n--8<-- "{self.main_file}"\n{code_fence}\n'
+            main_file_rel = self.main_file.relative_to(ROOT_DIR)
+            content += f'{code_fence}{self.main_file.suffix[1:]}\n--8<-- "{main_file_rel}"\n{code_fence}\n'
        else:
            with open(self.main_file) as f:
                # Skip the title from md snippets as it's been included above
@@ -208,7 +209,7 @@ class Example:
            content += f'??? abstract "{file.relative_to(self.path)}"\n'
            if file.suffix != ".md":
                content += f"    {code_fence}{file.suffix[1:]}\n"
-            content += f'    --8<-- "{file}"\n'
+            content += f'    --8<-- "{file.relative_to(ROOT_DIR)}"\n'
            if file.suffix != ".md":
                content += f"    {code_fence}\n"

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

@@ -4,7 +4,7 @@ vLLM-Omni supports unified multimodal comprehension and generation models across

 ## Model Implementation

-If vLLM-Omni natively supports a model, its implementation can be found in <gh-file:vllm-omni/model_executor/models> and <gh-file:vllm_omni/diffusion/models>.
+If vLLM-Omni natively supports a model, its implementation can be found in <gh-file:vllm_omni/model_executor/models> and <gh-file:vllm_omni/diffusion/models>.

 ## List of Supported Models for Nvidia GPU

@@ -20,9 +20,11 @@ th {
 | `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` |
 | `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` |
 | `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |
+| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` |
+|`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |


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

 <style>
 th {
@@ -2,7 +2,7 @@

 > Q: How many chips do I need to infer a model in vLLM-Omni?

-A: Now, we support natively disaggregated deployment for different model stages within a model. There is a restriction that one chip can only has one AutoRegressive model stage. This is because the unified KV cache management of vLLM. Stages of other types can coexist with in a chip. The restriction will be resolved in later version.
+A: Now, we support natively disaggregated deployment for different model stages within a model. There is a restriction that one chip can only have one AutoRegressive model stage. This is because the unified KV cache management of vLLM. Stages of other types can coexist within a chip. The restriction will be resolved in later version.

 > Q: When trying to run examples, I encounter error about backend of librosa or soundfile. How to solve it?

@@ -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 none 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,13 +1,14 @@
-# Offline Example of vLLM-Omni for Qwen2.5-omni
+# Offline Example of vLLM-Omni for Qwen2.5-Omni

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


 ## 🛠️ Installation

-Please refer to [README.md](https://github.com/vllm-project/vllm/tree/main/README.md)
+Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
+
+## Run examples (Qwen2.5-Omni)

-## Run examples (Qwen2.5-omni)
 ### Multiple Prompts
 Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). To get the prompt, you can:
 ```bash
@@ -26,6 +27,7 @@ Then run the command below.
 ```bash
 bash run_multiple_prompts.sh
 ```
+
 ### Single Prompt
 Get into the example folder
 ```bash
@@ -36,6 +38,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.
@@ -48,17 +78,17 @@ sudo apt install ffmpeg

 ??? abstract "end2end.py"
    ``````py
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen2_5_omni/end2end.py"
+    --8<-- "examples/offline_inference/qwen2_5_omni/end2end.py"
    ``````
 ??? abstract "extract_prompts.py"
    ``````py
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen2_5_omni/extract_prompts.py"
+    --8<-- "examples/offline_inference/qwen2_5_omni/extract_prompts.py"
    ``````
 ??? abstract "run_multiple_prompts.sh"
    ``````sh
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen2_5_omni/run_multiple_prompts.sh"
+    --8<-- "examples/offline_inference/qwen2_5_omni/run_multiple_prompts.sh"
    ``````
 ??? abstract "run_single_prompt.sh"
    ``````sh
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen2_5_omni/run_single_prompt.sh"
+    --8<-- "examples/offline_inference/qwen2_5_omni/run_single_prompt.sh"
    ``````
@@ -1,15 +1,16 @@
-# Offline Example of vLLM-Omni for Qwen3-omni
+# Offline Example of vLLM-Omni for Qwen3-Omni

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


 ## 🛠️ Installation

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

-## Run examples (Qwen3-omni)
 ### Multiple Prompts
-Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). For processing dataset please refer to [Qwen2.5-omni README.md](https://github.com/vllm-project/vllm/tree/main/examples/offline_inference/qwen2_5_omni/README.md)
+Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). For processing dataset please refer to [Qwen2.5-Omni README.md](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen2_5_omni/README.md)
 Get into the example folder
 ```bash
 cd examples/offline_inference/qwen3_omni
@@ -32,6 +33,26 @@ If you have not enough memory, you can set thinker with tensor parallel. Just ru
 bash run_single_prompt_tp.sh
 ```

+#### Using Local Media Files
+The `end2end.py` script supports local media files (audio, video, image) via command-line arguments:
+
+```bash
+# Use local video file
+python end2end.py --query-type use_video --video-path /path/to/video.mp4
+
+# Use local image file
+python end2end.py --query-type use_image --image-path /path/to/image.jpg
+
+# Use local audio file
+python end2end.py --query-type use_audio --audio-path /path/to/audio.wav
+```
+
+If media file paths are not provided, the script will use default assets. Supported query types:
+- `use_video`: Video input
+- `use_image`: Image input
+- `use_audio`: Audio input
+- `text`: Text-only query
+
 ### FAQ

 If you encounter error about backend of librosa, try to install ffmpeg with command below.
@@ -44,21 +65,17 @@ sudo apt install ffmpeg

 ??? abstract "end2end.py"
    ``````py
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen3_omni/end2end.py"
-    ``````
-??? abstract "qwen3_omni_moe_tp.yaml"
-    ``````yaml
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen3_omni/qwen3_omni_moe_tp.yaml"
+    --8<-- "examples/offline_inference/qwen3_omni/end2end.py"
    ``````
 ??? abstract "run_multiple_prompts.sh"
    ``````sh
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen3_omni/run_multiple_prompts.sh"
+    --8<-- "examples/offline_inference/qwen3_omni/run_multiple_prompts.sh"
    ``````
 ??? abstract "run_single_prompt.sh"
    ``````sh
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen3_omni/run_single_prompt.sh"
+    --8<-- "examples/offline_inference/qwen3_omni/run_single_prompt.sh"
    ``````
 ??? abstract "run_single_prompt_tp.sh"
    ``````sh
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen3_omni/run_single_prompt_tp.sh"
+    --8<-- "examples/offline_inference/qwen3_omni/run_single_prompt_tp.sh"
    ``````
@@ -1,6 +1,6 @@
 # Qwen-Image Offline Inference

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


 This folder provides two simple entrypoints for experimenting with `Qwen/Qwen-Image` using vLLM-Omni:
@@ -49,9 +49,9 @@ Then open `http://localhost:7862/` on your local browser to interact with the we

 ??? abstract "gradio_demo.py"
    ``````py
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen_image/gradio_demo.py"
+    --8<-- "examples/offline_inference/qwen_image/gradio_demo.py"
    ``````
 ??? abstract "text_to_image.py"
    ``````py
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen_image/text_to_image.py"
+    --8<-- "examples/offline_inference/qwen_image/text_to_image.py"
    ``````
@@ -1,15 +1,16 @@
-# Online serving Example of vLLM-Omni for Qwen2.5-omni
+# Online serving Example of vLLM-Omni for Qwen2.5-Omni

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


 ## 🛠️ Installation

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

-## Run examples (Qwen2.5-omni)
+## Run examples (Qwen2.5-Omni)
+
+### Launch the Server

-Launch the server
 ```bash
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
 ```
@@ -19,29 +20,59 @@ If you have custom stage configs file, launch the server with command below
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
 ```

+### Send Multi-modal Request
+
 Get into the example folder
 ```bash
-cd examples/online_serving
+cd examples/online_serving/qwen2_5_omni
 ```

-Send request via python
+####  Send request via python
+
 ```bash
 python openai_chat_completion_client_for_multimodal_generation.py --query-type mixed_modalities
 ```

-Send request via curl
+The Python client supports the following command-line arguments:
+
+- `--query-type` (or `-q`): Query type (default: `mixed_modalities`)
+  - Options: `mixed_modalities`, `use_audio_in_video`, `multi_audios`, `text`
+- `--video-path` (or `-v`): Path to local video file or URL
+  - If not provided and query-type uses video, uses default video URL
+  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
+  - Example: `--video-path /path/to/video.mp4` or `--video-path https://example.com/video.mp4`
+- `--image-path` (or `-i`): Path to local image file or URL
+  - If not provided and query-type uses image, uses default image URL
+  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
+  - Supports common image formats: JPEG, PNG, GIF, WebP
+  - Example: `--image-path /path/to/image.jpg` or `--image-path https://example.com/image.png`
+- `--audio-path` (or `-a`): Path to local audio file or URL
+  - If not provided and query-type uses audio, uses default audio URL
+  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
+  - Supports common audio formats: MP3, WAV, OGG, FLAC, M4A
+  - Example: `--audio-path /path/to/audio.wav` or `--audio-path https://example.com/audio.mp3`
+- `--prompt` (or `-p`): Custom text prompt/question
+  - If not provided, uses default prompt for the selected query type
+  - Example: `--prompt "What are the main activities shown in this video?"`
+
+
+For example, to use mixed modalities with all local files:
+
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py \
+    --query-type mixed_modalities \
+    --video-path /path/to/your/video.mp4 \
+    --image-path /path/to/your/image.jpg \
+    --audio-path /path/to/your/audio.wav \
+    --prompt "Analyze all the media content and provide a comprehensive summary."
+```
+
+####  Send request via curl
+
 ```bash
 bash run_curl_multimodal_generation.sh mixed_modalities
 ```

-### FAQ
-
-If you encounter error about backend of librosa, try to install ffmpeg with command below.
-```
-sudo apt update
-sudo apt install ffmpeg
-```
-
 ## Run Local Web UI Demo

 This Web UI demo allows users to interact with the model through a web browser.
@@ -56,36 +87,33 @@ python gradio_demo.py  --model Qwen/Qwen2.5-Omni-7B --port 7861

 Then open `http://localhost:7861/` on your local browser to interact with the web UI.

+The gradio script supports the following arguments:

-### Options
+- `--model`: Model name
+- `--ip`: Host/IP for Gradio server (default: 127.0.0.1)
+- `--port`: Port for Gradio server (default: 7861)
+- `--stage-configs-path`: Path to custom stage configs YAML file (optional)
+- `--share`: Share the Gradio demo publicly (creates a public link)

-You can customize its basic launch parameters:
+### FAQ

-```bash
-python gradio_demo.py \
-    --model Qwen/Qwen2.5-Omni-7B \
-    --ip 127.0.0.1 \
-    --port 7861 \
-    --stage-configs-path /path/to/stage_configs.yaml
+If you encounter error about backend of librosa, try to install ffmpeg with command below.
+```
+sudo apt update
+sudo apt install ffmpeg
 ```
-
- `--model`: Local model checkpoint to load (default `Qwen/Qwen2.5-Omni-7B`).
- `--ip`: Host/IP for the Gradio server (default `127.0.0.1`).
- `--port`: Port for the Gradio server (default `7861`).
- `--stage-configs-path`: Optional path to custom stage configs YAML.
- `--share`: Set to expose a temporary public link via Gradio.

 ## Example materials

 ??? abstract "gradio_demo.py"
    ``````py
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/online_serving/qwen2_5_omni/gradio_demo.py"
+    --8<-- "examples/online_serving/qwen2_5_omni/gradio_demo.py"
    ``````
 ??? abstract "openai_chat_completion_client_for_multimodal_generation.py"
    ``````py
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/online_serving/qwen2_5_omni/openai_chat_completion_client_for_multimodal_generation.py"
+    --8<-- "examples/online_serving/qwen2_5_omni/openai_chat_completion_client_for_multimodal_generation.py"
    ``````
 ??? abstract "run_curl_multimodal_generation.sh"
    ``````sh
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh"
+    --8<-- "examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh"
    ``````
@@ -1,15 +1,16 @@
 # Online serving Example of vLLM-Omni for Qwen3-omni

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


 ## 🛠️ Installation

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

 ## Run examples (Qwen3-Omni)

-Launch the server
+### Launch the Server
+
 ```bash
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
@@ -19,21 +20,59 @@ If you have custom stage configs file, launch the server with command below
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
 ```

+### Send Multi-modal Request
+
 Get into the example folder
 ```bash
-cd examples/online_serving
+cd examples/online_serving/qwen3_omni
 ```

-Send request via python
+####  Send request via python
+
 ```bash
 python openai_chat_completion_client_for_multimodal_generation.py --query-type use_image
 ```

-Send request via curl
+The Python client supports the following command-line arguments:
+
+- `--query-type` (or `-q`): Query type (default: `use_video`)
+  - Options: `text`, `use_audio`, `use_image`, `use_video`
+- `--model` (or `-m`): Model name/path (default: `Qwen/Qwen3-Omni-30B-A3B-Instruct`)
+- `--video-path` (or `-v`): Path to local video file or URL
+  - If not provided and query-type is `use_video`, uses default video URL
+  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
+  - Example: `--video-path /path/to/video.mp4` or `--video-path https://example.com/video.mp4`
+- `--image-path` (or `-i`): Path to local image file or URL
+  - If not provided and query-type is `use_image`, uses default image URL
+  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
+  - Supports common image formats: JPEG, PNG, GIF, WebP
+  - Example: `--image-path /path/to/image.jpg` or `--image-path https://example.com/image.png`
+- `--audio-path` (or `-a`): Path to local audio file or URL
+  - If not provided and query-type is `use_audio`, uses default audio URL
+  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
+  - Supports common audio formats: MP3, WAV, OGG, FLAC, M4A
+  - Example: `--audio-path /path/to/audio.wav` or `--audio-path https://example.com/audio.mp3`
+- `--prompt` (or `-p`): Custom text prompt/question
+  - If not provided, uses default prompt for the selected query type
+  - Example: `--prompt "What are the main activities shown in this video?"`
+
+
+For example, to use a local video file with custom prompt:
+
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py \
+    --query-type use_video \
+    --video-path /path/to/your/video.mp4 \
+    --prompt "What are the main activities shown in this video?"
+```
+
+####  Send request via curl
+
 ```bash
 bash run_curl_multimodal_generation.sh use_image
 ```

+
 ### FAQ

 If you encounter error about backend of librosa, try to install ffmpeg with command below.
@@ -57,23 +96,9 @@ python gradio_demo.py  --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 7861
 Then open `http://localhost:7861/` on your local browser to interact with the web UI.


-### Options
-
-The gradio demo also supports running with an existing API server and can be customized with the following arguments.
-
-
-```bash
-python gradio_demo.py \
-    --model Qwen/Qwen3-Omni-30B-A3B-Instruct \
-    --use-api-server \
-    --api-base http://localhost:8091/v1 \
-    --ip 127.0.0.1 \
-    --port 7861
-```
+The gradio script supports the following arguments:

 - `--model`: Model name
- `--use-api-server`: If set, connect to an existing vLLM HTTP API server instead of running AsyncOmni locally.
- `--api-base`: Base URL for vllm serve (only used when `use-api-server` is set, default: http://localhost:8091/v1)
 - `--ip`: Host/IP for Gradio server (default: 127.0.0.1)
 - `--port`: Port for Gradio server (default: 7861)
 - `--stage-configs-path`: Path to custom stage configs YAML file (optional)
@@ -83,13 +108,17 @@ python gradio_demo.py \

 ??? abstract "gradio_demo.py"
    ``````py
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/online_serving/qwen3_omni/gradio_demo.py"
+    --8<-- "examples/online_serving/qwen3_omni/gradio_demo.py"
    ``````
 ??? abstract "openai_chat_completion_client_for_multimodal_generation.py"
    ``````py
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py"
+    --8<-- "examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py"
+    ``````
+??? abstract "qwen3_omni_moe_thinking.yaml"
+    ``````yaml
+    --8<-- "examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml"
    ``````
 ??? abstract "run_curl_multimodal_generation.sh"
    ``````sh
-    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh"
+    --8<-- "examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh"
    ``````
@@ -5,6 +5,7 @@
 Please refer to [README.md](../../../README.md)

 ## Run examples (Qwen2.5-Omni)
+
 ### Multiple Prompts
 Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). To get the prompt, you can:
 ```bash
@@ -23,6 +24,7 @@ Then run the command below.
 ```bash
 bash run_multiple_prompts.sh
 ```
+
 ### Single Prompt
 Get into the example folder
 ```bash
@@ -33,6 +35,34 @@ Then run the command below.
 bash run_single_prompt.sh
 ```

+#### Using Local Media Files
+The `end2end.py` script supports local media files (audio, video, image) via CLI arguments:
+
+```bash
+# Use single local media files
+python end2end.py --query-type use_image --image-path /path/to/image.jpg
+python end2end.py --query-type use_video --video-path /path/to/video.mp4
+python end2end.py --query-type use_audio --audio-path /path/to/audio.wav
+
+# Combine multiple local media files
+python end2end.py --query-type mixed_modalities \
+    --video-path /path/to/video.mp4 \
+    --image-path /path/to/image.jpg \
+    --audio-path /path/to/audio.wav
+
+# Use audio from video file
+python end2end.py --query-type use_audio_in_video --video-path /path/to/video.mp4
+
+```
+
+If media file paths are not provided, the script will use default assets. Supported query types:
+- `use_image`: Image input only
+- `use_video`: Video input only
+- `use_audio`: Audio input only
+- `mixed_modalities`: Audio + image + video
+- `use_audio_in_video`: Extract audio from video
+- `text`: Text-only query
+
 ### FAQ

 If you encounter error about backend of librosa, try to install ffmpeg with command below.
@@ -6,12 +6,15 @@ with the correct prompt format on Qwen2.5-Omni
 """

 import os
-from typing import NamedTuple
+from typing import NamedTuple, Optional

+import librosa
+import numpy as np
 import soundfile as sf
+from PIL import Image
 from vllm.assets.audio import AudioAsset
 from vllm.assets.image import ImageAsset
-from vllm.assets.video import VideoAsset
+from vllm.assets.video import VideoAsset, video_to_ndarrays
 from vllm.multimodal.image import convert_image_mode
 from vllm.sampling_params import SamplingParams
 from vllm.utils import FlexibleArgumentParser
@@ -54,7 +57,13 @@ def get_text_query(question: str = None) -> QueryResult:
    )


-def get_mixed_modalities_query() -> QueryResult:
+def get_mixed_modalities_query(
+    video_path: Optional[str] = None,
+    image_path: Optional[str] = None,
+    audio_path: Optional[str] = None,
+    num_frames: int = 16,
+    sampling_rate: int = 16000,
+) -> QueryResult:
    question = "What is recited in the audio? What is the content of this image? Why is this video funny?"
    prompt = (
        f"<|im_start|>system\n{default_system}<|im_end|>\n"
@@ -64,20 +73,49 @@ def get_mixed_modalities_query() -> QueryResult:
        f"{question}<|im_end|>\n"
        f"<|im_start|>assistant\n"
    )
+
+    # Load video
+    if video_path:
+        if not os.path.exists(video_path):
+            raise FileNotFoundError(f"Video file not found: {video_path}")
+        video_frames = video_to_ndarrays(video_path, num_frames=num_frames)
+    else:
+        video_frames = VideoAsset(name="baby_reading", num_frames=num_frames).np_ndarrays
+
+    # Load image
+    if image_path:
+        if not os.path.exists(image_path):
+            raise FileNotFoundError(f"Image file not found: {image_path}")
+        pil_image = Image.open(image_path)
+        image_data = convert_image_mode(pil_image, "RGB")
+    else:
+        image_data = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB")
+
+    # Load audio
+    if audio_path:
+        if not os.path.exists(audio_path):
+            raise FileNotFoundError(f"Audio file not found: {audio_path}")
+        audio_signal, sr = librosa.load(audio_path, sr=sampling_rate)
+        audio_data = (audio_signal.astype(np.float32), sr)
+    else:
+        audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
+
    return QueryResult(
        inputs={
            "prompt": prompt,
            "multi_modal_data": {
-                "audio": AudioAsset("mary_had_lamb").audio_and_sample_rate,
-                "image": convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB"),
-                "video": VideoAsset(name="baby_reading", num_frames=16).np_ndarrays,
+                "audio": audio_data,
+                "image": image_data,
+                "video": video_frames,
            },
        },
        limit_mm_per_prompt={"audio": 1, "image": 1, "video": 1},
    )


-def get_use_audio_in_video_query() -> QueryResult:
+def get_use_audio_in_video_query(
+    video_path: Optional[str] = None, num_frames: int = 16, sampling_rate: int = 16000
+) -> QueryResult:
    question = "Describe the content of the video, then convert what the baby say into text."
    prompt = (
        f"<|im_start|>system\n{default_system}<|im_end|>\n"
@@ -85,14 +123,24 @@ def get_use_audio_in_video_query() -> QueryResult:
        f"{question}<|im_end|>\n"
        f"<|im_start|>assistant\n"
    )
-    asset = VideoAsset(name="baby_reading", num_frames=16)
-    audio = asset.get_audio(sampling_rate=16000)
+
+    if video_path:
+        if not os.path.exists(video_path):
+            raise FileNotFoundError(f"Video file not found: {video_path}")
+        video_frames = video_to_ndarrays(video_path, num_frames=num_frames)
+        # Extract audio from video file
+        audio_signal, sr = librosa.load(video_path, sr=sampling_rate)
+        audio = (audio_signal.astype(np.float32), sr)
+    else:
+        asset = VideoAsset(name="baby_reading", num_frames=num_frames)
+        video_frames = asset.np_ndarrays
+        audio = asset.get_audio(sampling_rate=sampling_rate)

    return QueryResult(
        inputs={
            "prompt": prompt,
            "multi_modal_data": {
-                "video": asset.np_ndarrays,
+                "video": video_frames,
                "audio": audio,
            },
            "mm_processor_kwargs": {
@@ -103,7 +151,7 @@ def get_use_audio_in_video_query() -> QueryResult:
    )


-def get_multi_audios_query() -> QueryResult:
+def get_multi_audios_query(audio_path: Optional[str] = None, sampling_rate: int = 16000) -> QueryResult:
    question = "Are these two audio clips the same?"
    prompt = (
        f"<|im_start|>system\n{default_system}<|im_end|>\n"
@@ -112,14 +160,28 @@ def get_multi_audios_query() -> QueryResult:
        f"{question}<|im_end|>\n"
        f"<|im_start|>assistant\n"
    )
+
+    if audio_path:
+        if not os.path.exists(audio_path):
+            raise FileNotFoundError(f"Audio file not found: {audio_path}")
+        audio_signal, sr = librosa.load(audio_path, sr=sampling_rate)
+        audio_data = (audio_signal.astype(np.float32), sr)
+        # Use the provided audio as the first audio, default as second
+        audio_list = [
+            audio_data,
+            AudioAsset("mary_had_lamb").audio_and_sample_rate,
+        ]
+    else:
+        audio_list = [
+            AudioAsset("winning_call").audio_and_sample_rate,
+            AudioAsset("mary_had_lamb").audio_and_sample_rate,
+        ]
+
    return QueryResult(
        inputs={
            "prompt": prompt,
            "multi_modal_data": {
-                "audio": [
-                    AudioAsset("winning_call").audio_and_sample_rate,
-                    AudioAsset("mary_had_lamb").audio_and_sample_rate,
-                ],
+                "audio": audio_list,
            },
        },
        limit_mm_per_prompt={
@@ -128,17 +190,135 @@ def get_multi_audios_query() -> QueryResult:
    )


+def get_image_query(question: str = None, image_path: Optional[str] = None) -> QueryResult:
+    if question is None:
+        question = "What is the content of this image?"
+    prompt = (
+        f"<|im_start|>system\n{default_system}<|im_end|>\n"
+        "<|im_start|>user\n<|vision_bos|><|IMAGE|><|vision_eos|>"
+        f"{question}<|im_end|>\n"
+        f"<|im_start|>assistant\n"
+    )
+
+    if image_path:
+        if not os.path.exists(image_path):
+            raise FileNotFoundError(f"Image file not found: {image_path}")
+        pil_image = Image.open(image_path)
+        image_data = convert_image_mode(pil_image, "RGB")
+    else:
+        image_data = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB")
+
+    return QueryResult(
+        inputs={
+            "prompt": prompt,
+            "multi_modal_data": {
+                "image": image_data,
+            },
+        },
+        limit_mm_per_prompt={"image": 1},
+    )
+
+
+def get_video_query(question: str = None, video_path: Optional[str] = None, num_frames: int = 16) -> QueryResult:
+    if question is None:
+        question = "Why is this video funny?"
+    prompt = (
+        f"<|im_start|>system\n{default_system}<|im_end|>\n"
+        "<|im_start|>user\n<|vision_bos|><|VIDEO|><|vision_eos|>"
+        f"{question}<|im_end|>\n"
+        f"<|im_start|>assistant\n"
+    )
+
+    if video_path:
+        if not os.path.exists(video_path):
+            raise FileNotFoundError(f"Video file not found: {video_path}")
+        video_frames = video_to_ndarrays(video_path, num_frames=num_frames)
+    else:
+        video_frames = VideoAsset(name="baby_reading", num_frames=num_frames).np_ndarrays
+
+    return QueryResult(
+        inputs={
+            "prompt": prompt,
+            "multi_modal_data": {
+                "video": video_frames,
+            },
+        },
+        limit_mm_per_prompt={"video": 1},
+    )
+
+
+def get_audio_query(question: str = None, audio_path: Optional[str] = None, sampling_rate: int = 16000) -> QueryResult:
+    if question is None:
+        question = "What is the content of this audio?"
+    prompt = (
+        f"<|im_start|>system\n{default_system}<|im_end|>\n"
+        "<|im_start|>user\n<|audio_bos|><|AUDIO|><|audio_eos|>"
+        f"{question}<|im_end|>\n"
+        f"<|im_start|>assistant\n"
+    )
+
+    if audio_path:
+        if not os.path.exists(audio_path):
+            raise FileNotFoundError(f"Audio file not found: {audio_path}")
+        audio_signal, sr = librosa.load(audio_path, sr=sampling_rate)
+        audio_data = (audio_signal.astype(np.float32), sr)
+    else:
+        audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
+
+    return QueryResult(
+        inputs={
+            "prompt": prompt,
+            "multi_modal_data": {
+                "audio": audio_data,
+            },
+        },
+        limit_mm_per_prompt={"audio": 1},
+    )
+
+
 query_map = {
    "mixed_modalities": get_mixed_modalities_query,
    "use_audio_in_video": get_use_audio_in_video_query,
    "multi_audios": get_multi_audios_query,
+    "use_image": get_image_query,
+    "use_video": get_video_query,
+    "use_audio": get_audio_query,
    "text": get_text_query,
 }


 def main(args):
    model_name = "Qwen/Qwen2.5-Omni-7B"
-    query_result = query_map[args.query_type]()
+
+    # Get paths from args
+    video_path = getattr(args, "video_path", None)
+    image_path = getattr(args, "image_path", None)
+    audio_path = getattr(args, "audio_path", None)
+    num_frames = getattr(args, "num_frames", 16)
+    sampling_rate = getattr(args, "sampling_rate", 16000)
+
+    # Get the query function and call it with appropriate parameters
+    query_func = query_map[args.query_type]
+    if args.query_type == "mixed_modalities":
+        query_result = query_func(
+            video_path=video_path,
+            image_path=image_path,
+            audio_path=audio_path,
+            num_frames=num_frames,
+            sampling_rate=sampling_rate,
+        )
+    elif args.query_type == "use_audio_in_video":
+        query_result = query_func(video_path=video_path, num_frames=num_frames, sampling_rate=sampling_rate)
+    elif args.query_type == "multi_audios":
+        query_result = query_func(audio_path=audio_path, sampling_rate=sampling_rate)
+    elif args.query_type == "use_image":
+        query_result = query_func(image_path=image_path)
+    elif args.query_type == "use_video":
+        query_result = query_func(video_path=video_path, num_frames=num_frames)
+    elif args.query_type == "use_audio":
+        query_result = query_func(audio_path=audio_path, sampling_rate=sampling_rate)
+    else:
+        query_result = query_func()

    omni_llm = Omni(
        model=model_name,
@@ -284,18 +464,38 @@ def parse_args():
        help="Path to a .txt file with one prompt per line (preferred).",
    )
    parser.add_argument(
-        "--worker-backend",
-        type=str,
-        default="process",
-        choices=["process","ray"],
-        help="backend"
-    )
-    parser.add_argument(
-        "--ray-address",
+        "--video-path",
+        "-v",
        type=str,
        default=None,
-        help="Path to a .txt file with one prompt per line (preferred).",
-
+        help="Path to local video file. If not provided, uses default video asset.",
+    )
+    parser.add_argument(
+        "--image-path",
+        "-i",
+        type=str,
+        default=None,
+        help="Path to local image file. If not provided, uses default image asset.",
+    )
+    parser.add_argument(
+        "--audio-path",
+        "-a",
+        type=str,
+        default=None,
+        help="Path to local audio file. If not provided, uses default audio asset.",
+    )
+    parser.add_argument(
+        "--num-frames",
+        type=int,
+        default=16,
+        help="Number of frames to extract from video (default: 16).",
+    )
+    parser.add_argument(
+        "--sampling-rate",
+        type=int,
+        default=16000,
+        help="Sampling rate for audio loading (default: 16000).",
+    )

    return parser.parse_args()

@@ -5,6 +5,7 @@
 Please refer to [README.md](../../../README.md)

 ## Run examples (Qwen3-Omni)
+
 ### Multiple Prompts
 Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). For processing dataset please refer to [Qwen2.5-Omni README.md](../qwen2_5_omni/README.md)
 Get into the example folder
@@ -29,6 +30,26 @@ If you have not enough memory, you can set thinker with tensor parallel. Just ru
 bash run_single_prompt_tp.sh
 ```

+#### Using Local Media Files
+The `end2end.py` script supports local media files (audio, video, image) via command-line arguments:
+
+```bash
+# Use local video file
+python end2end.py --query-type use_video --video-path /path/to/video.mp4
+
+# Use local image file
+python end2end.py --query-type use_image --image-path /path/to/image.jpg
+
+# Use local audio file
+python end2end.py --query-type use_audio --audio-path /path/to/audio.wav
+```
+
+If media file paths are not provided, the script will use default assets. Supported query types:
+- `use_video`: Video input
+- `use_image`: Image input
+- `use_audio`: Audio input
+- `text`: Text-only query
+
 ### FAQ

 If you encounter error about backend of librosa, try to install ffmpeg with command below.
@@ -6,13 +6,16 @@ with the correct prompt format on Qwen3-Omni (thinker only).
 """

 import os
-from typing import NamedTuple
+from typing import NamedTuple, Optional

+import librosa
+import numpy as np
 import soundfile as sf
+from PIL import Image
 from vllm import SamplingParams
 from vllm.assets.audio import AudioAsset
 from vllm.assets.image import ImageAsset
-from vllm.assets.video import VideoAsset
+from vllm.assets.video import VideoAsset, video_to_ndarrays
 from vllm.multimodal.image import convert_image_mode
 from vllm.utils import FlexibleArgumentParser

@@ -54,7 +57,7 @@ def get_text_query(question: str = None) -> QueryResult:
    )


-def get_video_query(question: str = None) -> QueryResult:
+def get_video_query(question: str = None, video_path: Optional[str] = None, num_frames: int = 16) -> QueryResult:
    if question is None:
        question = "Why is this video funny?"
    prompt = (
@@ -64,18 +67,25 @@ def get_video_query(question: str = None) -> QueryResult:
        f"<|im_start|>assistant\n"
    )

+    if video_path:
+        if not os.path.exists(video_path):
+            raise FileNotFoundError(f"Video file not found: {video_path}")
+        video_frames = video_to_ndarrays(video_path, num_frames=num_frames)
+    else:
+        video_frames = VideoAsset(name="baby_reading", num_frames=num_frames).np_ndarrays
+
    return QueryResult(
        inputs={
            "prompt": prompt,
            "multi_modal_data": {
-                "video": VideoAsset(name="baby_reading", num_frames=16).np_ndarrays,
+                "video": video_frames,
            },
        },
        limit_mm_per_prompt={"video": 1},
    )


-def get_image_query(question: str = None) -> QueryResult:
+def get_image_query(question: str = None, image_path: Optional[str] = None) -> QueryResult:
    if question is None:
        question = "What is the content of this image?"
    prompt = (
@@ -84,18 +94,27 @@ def get_image_query(question: str = None) -> QueryResult:
        f"{question}<|im_end|>\n"
        f"<|im_start|>assistant\n"
    )
+
+    if image_path:
+        if not os.path.exists(image_path):
+            raise FileNotFoundError(f"Image file not found: {image_path}")
+        pil_image = Image.open(image_path)
+        image_data = convert_image_mode(pil_image, "RGB")
+    else:
+        image_data = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB")
+
    return QueryResult(
        inputs={
            "prompt": prompt,
            "multi_modal_data": {
-                "image": convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB"),
+                "image": image_data,
            },
        },
        limit_mm_per_prompt={"image": 1},
    )


-def get_audio_query(question: str = None) -> QueryResult:
+def get_audio_query(question: str = None, audio_path: Optional[str] = None, sampling_rate: int = 16000) -> QueryResult:
    if question is None:
        question = "What is the content of this audio?"
    prompt = (
@@ -104,11 +123,20 @@ def get_audio_query(question: str = None) -> QueryResult:
        f"{question}<|im_end|>\n"
        f"<|im_start|>assistant\n"
    )
+
+    if audio_path:
+        if not os.path.exists(audio_path):
+            raise FileNotFoundError(f"Audio file not found: {audio_path}")
+        audio_signal, sr = librosa.load(audio_path, sr=sampling_rate)
+        audio_data = (audio_signal.astype(np.float32), sr)
+    else:
+        audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
+
    return QueryResult(
        inputs={
            "prompt": prompt,
            "multi_modal_data": {
-                "audio": AudioAsset("mary_had_lamb").audio_and_sample_rate,
+                "audio": audio_data,
            },
        },
        limit_mm_per_prompt={"audio": 1},
@@ -125,7 +153,22 @@ query_map = {

 def main(args):
    model_name = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
-    query_result = query_map[args.query_type]()
+
+    # Get paths from args
+    video_path = getattr(args, "video_path", None)
+    image_path = getattr(args, "image_path", None)
+    audio_path = getattr(args, "audio_path", None)
+
+    # Get the query function and call it with appropriate parameters
+    query_func = query_map[args.query_type]
+    if args.query_type == "use_video":
+        query_result = query_func(video_path=video_path, num_frames=getattr(args, "num_frames", 16))
+    elif args.query_type == "use_image":
+        query_result = query_func(image_path=image_path)
+    elif args.query_type == "use_audio":
+        query_result = query_func(audio_path=audio_path, sampling_rate=getattr(args, "sampling_rate", 16000))
+    else:
+        query_result = query_func()

    omni_llm = Omni(
        model=model_name,
@@ -283,6 +326,39 @@ def parse_args():
        default=None,
        help="Path to a stage configs file.",
    )
+    parser.add_argument(
+        "--video-path",
+        "-v",
+        type=str,
+        default=None,
+        help="Path to local video file. If not provided, uses default video asset.",
+    )
+    parser.add_argument(
+        "--image-path",
+        "-i",
+        type=str,
+        default=None,
+        help="Path to local image file. If not provided, uses default image asset.",
+    )
+    parser.add_argument(
+        "--audio-path",
+        "-a",
+        type=str,
+        default=None,
+        help="Path to local audio file. If not provided, uses default audio asset.",
+    )
+    parser.add_argument(
+        "--num-frames",
+        type=int,
+        default=16,
+        help="Number of frames to extract from video (default: 16).",
+    )
+    parser.add_argument(
+        "--sampling-rate",
+        type=int,
+        default=16000,
+        help="Sampling rate for audio loading (default: 16000).",
+    )

    return parser.parse_args()

@@ -1,6 +1,5 @@
 python end2end.py --output-wav output_audio \
                  --query-type use_audio \
-                  --init-sleep-seconds 90 \
-                  --stage-configs-path qwen3_omni_moe_tp.yaml
+                  --init-sleep-seconds 90

 # init-sleep-seconds works to avoid two vLLM stages initialized at the same time within a card.
@@ -1,10 +1,21 @@
 # Qwen-Image Offline Inference

-This folder provides two simple entrypoints for experimenting with `Qwen/Qwen-Image` using vLLM-Omni:
+This folder provides several entrypoints for experimenting with `Qwen/Qwen-Image` using vLLM-Omni:

- `text_to_image.py`: command-line script for single image generation.
+- `text_to_image.py`: command-line script for single image generation with advanced options.
 - `web_demo.py`: lightweight Gradio UI for interactive prompt/seed/CFG exploration.

+## Basic Usage
+
+```python
+from vllm_omni.entrypoints.omni import Omni
+
+if __name__ == "__main__":
+    omni = Omni(model="Qwen/Qwen-Image")
+    prompt = "a cup of coffee on the table"
+    images = omni.generate(prompt)
+    images[0].save("coffee.png")
+```

 ## Local CLI Usage

@@ -5,7 +5,7 @@ import gradio as gr
 import torch

 from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.utils.platform_utils import detect_device_type
+from vllm_omni.utils.platform_utils import detect_device_type, is_npu

 ASPECT_RATIOS: dict[str, tuple[int, int]] = {
    "1:1": (1328, 1328),
@@ -59,7 +59,14 @@ def parse_args() -> argparse.Namespace:

@lru_cache(maxsize=1)
 def get_omni(model_name: str) -> Omni:
-    return Omni(model=model_name)
+    # Enable VAE memory optimizations on NPU
+    vae_use_slicing = is_npu()
+    vae_use_tiling = is_npu()
+    return Omni(
+        model=model_name,
+        vae_use_slicing=vae_use_slicing,
+        vae_use_tiling=vae_use_tiling,
+    )


 def build_demo(args: argparse.Namespace) -> gr.Blocks:
@@ -98,7 +105,6 @@ def build_demo(args: argparse.Namespace) -> gr.Blocks:
            generator=generator,
            true_cfg_scale=float(cfg_scale_value),
            num_inference_steps=num_steps,
-            num_images_per_prompt=num_images,
            num_outputs_per_prompt=num_images,
        )
        return [img for img in images[:num_images]]
@@ -0,0 +1,143 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Example script for image editing with Qwen-Image-Edit.
+
+Usage:
+    python image_edit.py \
+        --image input.png \
+        --prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'" \
+        --output output_image_edit.png \
+        --num_inference_steps 50 \
+        --cfg_scale 4.0
+
+For more options, run:
+    python image_edit.py --help
+"""
+
+import argparse
+import os
+from pathlib import Path
+
+import torch
+from PIL import Image
+
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.utils.platform_utils import detect_device_type, is_npu
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Edit an image with Qwen-Image-Edit.")
+    parser.add_argument(
+        "--model",
+        default="Qwen/Qwen-Image-Edit",
+        help="Diffusion model name or local path.",
+    )
+    parser.add_argument(
+        "--image",
+        type=str,
+        required=True,
+        help="Path to input image file (PNG, JPG, etc.).",
+    )
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        required=True,
+        help="Text prompt describing the edit to make to the image.",
+    )
+    parser.add_argument(
+        "--negative_prompt",
+        type=str,
+        default=" ",
+        required=False,
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=0,
+        help="Random seed for deterministic results.",
+    )
+    parser.add_argument(
+        "--cfg_scale",
+        type=float,
+        default=4.0,
+        help="True classifier-free guidance scale specific to Qwen-Image-Edit.",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="output_image_edit.png",
+        help="Path to save the edited image (PNG).",
+    )
+    parser.add_argument(
+        "--num_outputs_per_prompt",
+        type=int,
+        default=1,
+        help="Number of images to generate for the given prompt.",
+    )
+    parser.add_argument(
+        "--num_inference_steps",
+        type=int,
+        default=50,
+        help="Number of denoising steps for the diffusion sampler.",
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    # Validate input image exists
+    if not os.path.exists(args.image):
+        raise FileNotFoundError(f"Input image not found: {args.image}")
+
+    # Load input image
+    input_image = Image.open(args.image).convert("RGB")
+    print(f"Loaded input image from {args.image} (size: {input_image.size})")
+
+    device = detect_device_type()
+    generator = torch.Generator(device=device).manual_seed(args.seed)
+
+    # Enable VAE memory optimizations on NPU
+    vae_use_slicing = is_npu()
+    vae_use_tiling = is_npu()
+
+    # Initialize Omni with QwenImageEditPipeline
+    omni = Omni(
+        model=args.model,
+        model_class_name="QwenImageEditPipeline",
+        vae_use_slicing=vae_use_slicing,
+        vae_use_tiling=vae_use_tiling,
+    )
+    print("Pipeline loaded")
+
+    # Generate edited image
+    images = omni.generate(
+        prompt=args.prompt,
+        pil_image=input_image,
+        negative_prompt=args.negative_prompt,
+        generator=generator,
+        true_cfg_scale=args.cfg_scale,
+        num_inference_steps=args.num_inference_steps,
+        num_outputs_per_prompt=args.num_outputs_per_prompt,
+    )
+
+    # Save output image(s)
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    suffix = output_path.suffix or ".png"
+    stem = output_path.stem or "output_image_edit"
+
+    if args.num_outputs_per_prompt <= 1:
+        images[0].save(output_path)
+        print(f"Saved edited image to {os.path.abspath(output_path)}")
+    else:
+        for idx, img in enumerate(images):
+            save_path = output_path.parent / f"{stem}_{idx}{suffix}"
+            img.save(save_path)
+            print(f"Saved edited image to {os.path.abspath(save_path)}")
+
+
+if __name__ == "__main__":
+    main()
@@ -7,7 +7,7 @@ from pathlib import Path
 import torch

 from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.utils.platform_utils import detect_device_type
+from vllm_omni.utils.platform_utils import detect_device_type, is_npu


 def parse_args() -> argparse.Namespace:
@@ -49,7 +49,15 @@ def main():
    device = detect_device_type()
    generator = torch.Generator(device=device).manual_seed(args.seed)

-    omni = Omni(model=args.model)
+    # Enable VAE memory optimizations on NPU
+    vae_use_slicing = is_npu()
+    vae_use_tiling = is_npu()
+
+    omni = Omni(
+        model=args.model,
+        vae_use_slicing=vae_use_slicing,
+        vae_use_tiling=vae_use_tiling,
+    )
    images = omni.generate(
        args.prompt,
        height=args.height,
@@ -57,7 +65,6 @@ def main():
        generator=generator,
        true_cfg_scale=args.cfg_scale,
        num_inference_steps=args.num_inference_steps,
-        num_images_per_prompt=args.num_images_per_prompt,
        num_outputs_per_prompt=args.num_images_per_prompt,
    )

@@ -1,34 +0,0 @@
-# Offline Example of vLLM-omni for Qwen2.5-omni
-
-## 🛠️ Installation
-
-Please refer to [README.md](../../../README.md)
-
-## Run examples (Qwen2.5-omni)
-### Multiple Prompts
-Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). To get the prompt, you can:
-```bash
-tar -xf <Your Download Path>/seedtts_testset.tar
-cp seedtts_testset/en/meta.lst examples/offline_inference/qwen2_5_omni/meta.lst
-python3 examples/offline_inference/qwen2_5_omni/extract_prompts.py \
-  --input examples/offline_inference/qwen2_5_omni/meta.lst \
-  --output examples/offline_inference/qwen2_5_omni/top100.txt \
-  --topk 100
-```
-Get into the example folder
-```bash
-cd examples/offline_inference/qwen2_5_omni
-```
-Then run the command below.
-```bash
-bash run_multiple_prompts.sh
-```
-### Single Prompts
-Get into the example folder
-```bash
-cd examples/offline_inference/qwen2_5_omni
-```
-Then run the command below.
-```bash
-bash run_single_prompt.sh
-```
@@ -1,264 +0,0 @@
-import argparse
-import os
-import os as _os_env_toggle
-import random
-
-import numpy as np
-import soundfile as sf
-import torch
-from utils import make_omni_prompt
-from vllm.sampling_params import SamplingParams
-
-from vllm_omni.entrypoints.omni_llm import OmniLLM
-
-_os_env_toggle.environ["VLLM_USE_V1"] = "1"
-
-SEED = 42
-# Set all random seeds
-random.seed(SEED)
-np.random.seed(SEED)
-torch.manual_seed(SEED)
-torch.cuda.manual_seed(SEED)
-torch.cuda.manual_seed_all(SEED)
-
-# Make PyTorch deterministic
-torch.backends.cudnn.deterministic = True
-torch.backends.cudnn.benchmark = False
-
-# Set environment variables for deterministic behavior
-os.environ["PYTHONHASHSEED"] = str(SEED)
-os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
-
-
-def parse_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--model",
-        required=True,
-        help="Path to merged model directory (will be created if downloading).",
-    )
-    parser.add_argument("--thinker-model", type=str, default=None)
-    parser.add_argument("--talker-model", type=str, default=None)
-    parser.add_argument("--code2wav-model", type=str, default=None)
-    parser.add_argument(
-        "--hf-hub-id",
-        default="Qwen/Qwen2.5-Omni-7B",
-        help="Hugging Face repo id to download if needed.",
-    )
-    parser.add_argument("--hf-revision", default=None, help="Optional HF revision (branch/tag/commit).")
-    parser.add_argument("--prompts", nargs="+", default=None, help="Input text prompts.")
-    parser.add_argument("--voice-type", default="default", help="Voice type, e.g., m02, f030, default.")
-    parser.add_argument(
-        "--code2wav-dir",
-        default=None,
-        help="Path to code2wav folder (contains spk_dict.pt).",
-    )
-    parser.add_argument("--dit-ckpt", default=None, help="Path to DiT checkpoint file (e.g., dit.pt).")
-    parser.add_argument("--bigvgan-ckpt", default=None, help="Path to BigVGAN checkpoint file.")
-    parser.add_argument("--dtype", default="bfloat16", choices=["float16", "bfloat16", "float32"])
-    parser.add_argument("--max-model-len", type=int, default=32768)
-    parser.add_argument(
-        "--init-sleep-seconds",
-        type=int,
-        default=20,
-        help="Sleep seconds after starting each stage process to allow initialization (default: 20)",
-    )
-
-    parser.add_argument("--thinker-only", action="store_true")
-    parser.add_argument("--text-only", action="store_true")
-    parser.add_argument("--do-wave", action="store_true")
-    parser.add_argument(
-        "--prompt_type",
-        choices=[
-            "text",
-            "audio",
-            "audio-long",
-            "audio-long-chunks",
-            "audio-long-expand-chunks",
-            "image",
-            "video",
-            "video-frames",
-            "audio-in-video",
-            "audio-in-video-v2",
-            "audio-multi-round",
-            "badcase-vl",
-            "badcase-text",
-            "badcase-image-early-stop",
-            "badcase-two-audios",
-            "badcase-two-videos",
-            "badcase-multi-round",
-            "badcase-voice-type",
-            "badcase-voice-type-v2",
-            "badcase-audio-tower-1",
-            "badcase-audio-only",
-        ],
-        default="text",
-    )
-    parser.add_argument("--use-torchvision", action="store_true")
-    parser.add_argument("--tokenize", action="store_true")
-    parser.add_argument(
-        "--output-wav",
-        default="output.wav",
-        help="[Deprecated] Output wav directory (use --output-dir).",
-    )
-    parser.add_argument(
-        "--output-dir",
-        default="outputs",
-        help="Output directory to save text and wav files together.",
-    )
-    parser.add_argument(
-        "--thinker-hidden-states-dir",
-        default="thinker_hidden_states",
-        help="Path to thinker hidden states directory.",
-    )
-    parser.add_argument(
-        "--batch-timeout",
-        type=int,
-        default=5,
-        help="Timeout for batching in seconds (default: 5)",
-    )
-    parser.add_argument(
-        "--init-timeout",
-        type=int,
-        default=300,
-        help="Timeout for initializing stages in seconds (default: 300)",
-    )
-    parser.add_argument(
-        "--shm-threshold-bytes",
-        type=int,
-        default=65536,
-        help="Threshold for using shared memory in bytes (default: 65536)",
-    )
-    parser.add_argument(
-        "--enable-stats",
-        action="store_true",
-        default=False,
-        help="Enable writing detailed statistics (default: disabled)",
-    )
-    parser.add_argument(
-        "--txt-prompts",
-        type=str,
-        default=None,
-        help="Path to a .txt file with one prompt per line (preferred).",
-    )
-    parser.add_argument(
-        "--worker-backend",
-        type=str,
-        default="process",
-        choices=["process","ray"],
-        help="backend"
-    )
-    parser.add_argument(
-        "--ray-address",
-        type=str,
-        default=None,
-        help="Path to a .txt file with one prompt per line (preferred)."
-        )
-
-    args = parser.parse_args()
-    return args
-
-
-def main():
-    args = parse_args()
-    model_name = args.model
-    try:
-        # Preferred: load from txt file (one prompt per line)
-        if getattr(args, "txt_prompts", None) and args.prompt_type == "text":
-            with open(args.txt_prompts, encoding="utf-8") as f:
-                lines = [ln.strip() for ln in f.readlines()]
-            args.prompts = [ln for ln in lines if ln != ""]
-            print(f"[Info] Loaded {len(args.prompts)} prompts from {args.txt_prompts}")
-    except Exception as e:
-        print(f"[Error] Failed to load prompts: {e}")
-        raise
-
-    if args.prompts is None:
-        raise ValueError("No prompts provided. Use --prompts ... or --txt-prompts <file.txt> (with --prompt_type text)")
-    omni_llm = OmniLLM(
-        model=model_name,
-        log_stats=args.enable_stats,
-        log_file=("omni_llm_pipeline.log" if args.enable_stats else None),
-        init_sleep_seconds=args.init_sleep_seconds,
-        batch_timeout=args.batch_timeout,
-        init_timeout=args.init_timeout,
-        shm_threshold_bytes=args.shm_threshold_bytes,
-        worker_backend=args.worker_backend,
-        ray_address=args.ray_address
-    )
-    thinker_sampling_params = SamplingParams(
-        temperature=0.0,  # Deterministic - no randomness
-        top_p=1.0,  # Disable nucleus sampling
-        top_k=-1,  # Disable top-k sampling
-        max_tokens=2048,
-        seed=SEED,  # Fixed seed for sampling
-        detokenize=True,
-        repetition_penalty=1.1,
-    )
-    talker_sampling_params = SamplingParams(
-        temperature=0.9,
-        top_p=0.8,
-        top_k=40,
-        max_tokens=2048,
-        seed=SEED,  # Fixed seed for sampling
-        detokenize=True,
-        repetition_penalty=1.05,
-        stop_token_ids=[8294],
-    )
-    code2wav_sampling_params = SamplingParams(
-        temperature=0.0,  # Deterministic - no randomness
-        top_p=1.0,  # Disable nucleus sampling
-        top_k=-1,  # Disable top-k sampling
-        max_tokens=2048,
-        seed=SEED,  # Fixed seed for sampling
-        detokenize=True,
-        repetition_penalty=1.1,
-    )
-
-    sampling_params_list = [
-        thinker_sampling_params,
-        talker_sampling_params,
-        code2wav_sampling_params,
-    ]
-    import time
-    for i in range(1):
-        t1 = time.time()
-        prompt = [make_omni_prompt(args, prompt) for prompt in args.prompts]
-        print(f"prompt:{prompt}")
-        omni_outputs = omni_llm.generate(prompt, sampling_params_list)
-        t2 = time.time()
-        print(f"==========> time:{t2-t1}")
-
-    # Determine output directory: prefer --output-dir; fallback to --output-wav
-    output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
-    os.makedirs(output_dir, exist_ok=True)
-    for stage_outputs in omni_outputs:
-        if stage_outputs.final_output_type == "text":
-            for output in stage_outputs.request_output:
-                request_id = int(output.request_id)
-                text_output = output.outputs[0].text
-                # Save aligned text file per request
-                prompt_text = args.prompts[request_id]
-                out_txt = os.path.join(output_dir, f"{request_id:05d}.txt")
-                lines = []
-                lines.append("Prompt:\n")
-                lines.append(str(prompt_text) + "\n")
-                lines.append("vllm_text_output:\n")
-                lines.append(str(text_output).strip() + "\n")
-                try:
-                    with open(out_txt, "w", encoding="utf-8") as f:
-                        f.writelines(lines)
-                except Exception as e:
-                    print(f"[Warn] Failed writing text file {out_txt}: {e}")
-                print(f"Request ID: {request_id}, Text saved to {out_txt}")
-        elif stage_outputs.final_output_type == "audio":
-            for output in stage_outputs.request_output:
-                request_id = int(output.request_id)
-                audio_tensor = output.multimodal_output["audio"]
-                output_wav = os.path.join(output_dir, f"output_{output.request_id}.wav")
-                sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000)
-                print(f"Request ID: {request_id}, Saved audio to {output_wav}")
-
-
-if __name__ == "__main__":
-    main()
@@ -1,45 +0,0 @@
-#!/usr/bin/env python3
-import argparse
-from typing import Optional
-
-
-def extract_prompt(line: str) -> Optional[str]:
-    # 提取第一个 '|' 与第二个 '|' 之间的内容
-    i = line.find("|")
-    if i == -1:
-        return None
-    j = line.find("|", i + 1)
-    if j == -1:
-        return None
-    return line[i + 1 : j].strip()
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--input", "-i", required=True, help="Input .lst file path")
-    parser.add_argument("--output", "-o", required=True, help="Output file path")
-    parser.add_argument(
-        "--topk",
-        "-k",
-        type=int,
-        default=100,
-        help="Extract the top K prompts (default: 100)",
-    )
-    args = parser.parse_args()
-
-    prompts = []
-    with open(args.input, encoding="utf-8", errors="ignore") as f:
-        for line in f:
-            if len(prompts) >= args.topk:
-                break
-            p = extract_prompt(line.rstrip("\n"))
-            if p:
-                prompts.append(p)
-
-    with open(args.output, "w", encoding="utf-8") as f:
-        for p in prompts:
-            f.write(p + "\n")
-
-
-if __name__ == "__main__":
-    main()
@@ -1,367 +0,0 @@
-from __future__ import annotations
-
-import base64
-import logging
-import math
-import os
-import time
-import warnings
-from functools import lru_cache
-from io import BytesIO
-
-import requests
-import torch
-import torchvision
-from packaging import version
-from PIL import Image
-from torchvision import io, transforms
-from torchvision.transforms import InterpolationMode
-
-logger = logging.getLogger(__name__)
-
-IMAGE_FACTOR = 28
-MIN_PIXELS = 4 * 28 * 28
-MAX_PIXELS = 16384 * 28 * 28
-MAX_RATIO = 200
-
-VIDEO_MIN_PIXELS = 128 * 28 * 28
-VIDEO_MAX_PIXELS = 768 * 28 * 28
-VIDEO_TOTAL_PIXELS = 24576 * 28 * 28
-FRAME_FACTOR = 2
-FPS = 2.0
-FPS_MIN_FRAMES = 4
-FPS_MAX_FRAMES = 768
-
-temporal_patch_size = 2
-spatial_patch_size = 14
-spatial_merge_size = 2
-
-
-def round_by_factor(number: int, factor: int) -> int:
-    """Returns the closest integer to 'number' that is divisible by 'factor'."""
-    return round(number / factor) * factor
-
-
-def ceil_by_factor(number: int, factor: int) -> int:
-    """Returns the smallest integer greater than or equal to 'number' that is
-    divisible by 'factor'."""
-    return math.ceil(number / factor) * factor
-
-
-def floor_by_factor(number: int, factor: int) -> int:
-    """Returns the largest integer less than or equal to 'number' that is
-    divisible by 'factor'."""
-    return math.floor(number / factor) * factor
-
-
-def smart_resize(
-    height: int,
-    width: int,
-    factor: int = IMAGE_FACTOR,
-    min_pixels: int = MIN_PIXELS,
-    max_pixels: int = MAX_PIXELS,
-) -> tuple[int, int]:
-    """
-    Rescales the image so that the following conditions are met:
-
-    1. Both dimensions (height and width) are divisible by 'factor'.
-
-    2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
-
-    3. The aspect ratio of the image is maintained as closely as possible.
-    """
-    if max(height, width) / min(height, width) > MAX_RATIO:
-        raise ValueError(
-            f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}"
-        )
-    h_bar = max(factor, round_by_factor(height, factor))
-    w_bar = max(factor, round_by_factor(width, factor))
-    if h_bar * w_bar > max_pixels:
-        beta = math.sqrt((height * width) / max_pixels)
-        h_bar = floor_by_factor(height / beta, factor)
-        w_bar = floor_by_factor(width / beta, factor)
-    elif h_bar * w_bar < min_pixels:
-        beta = math.sqrt(min_pixels / (height * width))
-        h_bar = ceil_by_factor(height * beta, factor)
-        w_bar = ceil_by_factor(width * beta, factor)
-    return h_bar, w_bar
-
-
-def fetch_image(ele: dict[str, str | Image.Image], size_factor: int = IMAGE_FACTOR) -> Image.Image:
-    if "image" in ele:
-        image = ele["image"]
-    else:
-        image = ele["image_url"]
-    image_obj = None
-    if isinstance(image, Image.Image):
-        image_obj = image
-    elif image.startswith("http://") or image.startswith("https://"):
-        image_obj = Image.open(requests.get(image, stream=True).raw)
-    elif image.startswith("file://"):
-        image_obj = Image.open(image[7:])
-    elif image.startswith("data:image"):
-        if "base64," in image:
-            _, base64_data = image.split("base64,", 1)
-            data = base64.b64decode(base64_data)
-            image_obj = Image.open(BytesIO(data))
-    else:
-        image_obj = Image.open(image)
-    if image_obj is None:
-        raise ValueError(f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}")
-    image = image_obj.convert("RGB")
-    # resize
-    if "resized_height" in ele and "resized_width" in ele:
-        resized_height, resized_width = smart_resize(
-            ele["resized_height"],
-            ele["resized_width"],
-            factor=size_factor,
-        )
-    else:
-        width, height = image.size
-        min_pixels = ele.get("min_pixels", MIN_PIXELS)
-        max_pixels = ele.get("max_pixels", MAX_PIXELS)
-        resized_height, resized_width = smart_resize(
-            height,
-            width,
-            factor=size_factor,
-            min_pixels=min_pixels,
-            max_pixels=max_pixels,
-        )
-    image = image.resize((resized_width, resized_height))
-
-    return image
-
-
-def smart_nframes(
-    ele: dict,
-    total_frames: int,
-    video_fps: int | float,
-) -> int:
-    """calculate the number of frames for video used for model inputs.
-
-    Args:
-        ele (dict): a dict contains the configuration of video.
-            support either `fps` or `nframes`:
-                - nframes: the number of frames to extract for model inputs.
-                - fps: the fps to extract frames for model inputs.
-                    - min_frames: the minimum number of frames of the video,
-                      only used when fps is provided.
-                    - max_frames: the maximum number of frames of the video,
-                      only used when fps is provided.
-        total_frames (int): the original total number of frames of the video.
-        video_fps (int | float): the original fps of the video.
-
-    Raises:
-        ValueError: nframes should in interval [FRAME_FACTOR, total_frames].
-
-    Returns:
-        int: the number of frames for video used for model inputs.
-    """
-    assert not ("fps" in ele and "nframes" in ele), "Only accept either `fps` or `nframes`"
-    if "nframes" in ele:
-        nframes = round_by_factor(ele["nframes"], FRAME_FACTOR)
-    else:
-        fps = ele.get("fps", FPS)
-        min_frames = ceil_by_factor(ele.get("min_frames", FPS_MIN_FRAMES), FRAME_FACTOR)
-        max_frames = floor_by_factor(
-            ele.get("max_frames", min(FPS_MAX_FRAMES, total_frames)),
-            FRAME_FACTOR,
-        )
-        nframes = total_frames / video_fps * fps
-        nframes = min(max(nframes, min_frames), max_frames)
-        nframes = round_by_factor(nframes, FRAME_FACTOR)
-    if not (FRAME_FACTOR <= nframes and nframes <= total_frames):
-        raise ValueError(f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}.")
-    return nframes
-
-
-def _read_video_torchvision(
-    ele: dict,
-) -> torch.Tensor:
-    """read video using torchvision.io.read_video
-
-    Args:
-        ele (dict): a dict contains the configuration of video.
-        support keys:
-            - video: the path of video. support "file://", "http://",
-              "https://" and local path.
-            - video_start: the start time of video.
-            - video_end: the end time of video.
-    Returns:
-        torch.Tensor: the video tensor with shape (T, C, H, W).
-    """
-    video_path = ele["video"]
-    if version.parse(torchvision.__version__) < version.parse("0.19.0"):
-        if "http://" in video_path or "https://" in video_path:
-            warnings.warn("torchvision < 0.19.0 does not support http/https video path, please upgrade to 0.19.0.")
-        if "file://" in video_path:
-            video_path = video_path[7:]
-    st = time.time()
-    video, audio, info = io.read_video(
-        video_path,
-        start_pts=ele.get("video_start", 0.0),
-        end_pts=ele.get("video_end", None),
-        pts_unit="sec",
-        output_format="TCHW",
-    )
-    total_frames, video_fps = video.size(0), info["video_fps"]
-    total_duration = round(total_frames / video_fps, 3)
-    logger.info(
-        f"torchvision:  {video_path=}, {total_frames=}, {video_fps=}, duration={total_duration}s, time={time.time() - st:.3f}s"
-    )
-    nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
-    idx = torch.linspace(0, total_frames - 1, nframes).round().long()
-    video = video[idx]
-    return video, total_duration, nframes
-
-
-def is_decord_available() -> bool:
-    import importlib.util
-
-    return importlib.util.find_spec("decord") is not None
-
-
-def _read_video_decord(
-    ele: dict,
-) -> torch.Tensor:
-    """read video using decord.VideoReader
-
-    Args:
-        ele (dict): a dict contains the configuration of video.
-        support keys:
-            - video: the path of video. support "file://", "http://",
-              "https://" and local path.
-            - video_start: the start time of video.
-            - video_end: the end time of video.
-    Returns:
-        torch.Tensor: the video tensor with shape (T, C, H, W).
-    """
-    import decord
-
-    video_path = ele["video"]
-    st = time.time()
-    vr = decord.VideoReader(video_path)
-    # TODO: support start_pts and end_pts
-    if "video_start" in ele or "video_end" in ele:
-        raise NotImplementedError("not support start_pts and end_pts in decord for now.")
-    total_frames, video_fps = len(vr), vr.get_avg_fps()
-    total_duration = round(total_frames / video_fps, 3)
-    logger.info(f"decord:  {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s")
-    nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
-    idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist()
-    video = vr.get_batch(idx).asnumpy()
-    video = torch.tensor(video).permute(0, 3, 1, 2)  # Convert to TCHW format
-    return video, total_duration, nframes
-
-
-VIDEO_READER_BACKENDS = {
-    "decord": _read_video_decord,
-    "torchvision": _read_video_torchvision,
-}
-
-FORCE_QWENVL_VIDEO_READER = os.getenv("FORCE_QWENVL_VIDEO_READER", None)
-
-
-@lru_cache(maxsize=1)
-def get_video_reader_backend() -> str:
-    if FORCE_QWENVL_VIDEO_READER is not None:
-        video_reader_backend = FORCE_QWENVL_VIDEO_READER
-    elif is_decord_available():
-        video_reader_backend = "decord"
-    else:
-        video_reader_backend = "torchvision"
-    # print(f"qwen-vl-utils using {video_reader_backend} to read video.",
-    #       file=sys.stderr)
-    return video_reader_backend
-
-
-def fetch_video(ele: dict, image_factor: int = IMAGE_FACTOR) -> torch.Tensor | list[Image.Image]:
-    if isinstance(ele["video"], str):
-        video_reader_backend = get_video_reader_backend()
-        video, total_dur, nframes = VIDEO_READER_BACKENDS[video_reader_backend](ele)
-        frame_timestamps = total_dur * torch.arange(1, nframes + 1) / nframes
-        grid_timestamps = frame_timestamps[::FRAME_FACTOR]
-        second_per_grid = grid_timestamps[1] - grid_timestamps[0]
-        nframes, _, height, width = video.shape
-        min_pixels = ele.get("min_pixels", VIDEO_MIN_PIXELS)
-        total_pixels = ele.get("total_pixels", VIDEO_TOTAL_PIXELS)
-        max_pixels = max(
-            min(VIDEO_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR),
-            int(min_pixels * 1.05),
-        )
-        max_pixels = ele.get("max_pixels", max_pixels)
-        # min_pixels = (factor ** 2) * 52
-        # max_pixels = (factor ** 2) * min(768, (16384 / nframes * temporal_patch_size))
-        if "resized_height" in ele and "resized_width" in ele:
-            resized_height, resized_width = smart_resize(
-                ele["resized_height"],
-                ele["resized_width"],
-                factor=image_factor,
-            )
-        else:
-            resized_height, resized_width = smart_resize(
-                height,
-                width,
-                factor=image_factor,
-                min_pixels=min_pixels,
-                max_pixels=max_pixels,
-            )
-        video = transforms.functional.resize(
-            video,
-            [resized_height, resized_width],
-            interpolation=InterpolationMode.BICUBIC,
-            antialias=True,
-        ).float()
-        return video, total_dur, nframes, second_per_grid
-    else:
-        assert isinstance(ele["video"], (list, tuple))
-        process_info = ele.copy()
-        process_info.pop("type", None)
-        process_info.pop("video", None)
-        images = [
-            fetch_image({"image": video_element, **process_info}, size_factor=image_factor)
-            for video_element in ele["video"]
-        ]
-        nframes = ceil_by_factor(len(images), FRAME_FACTOR)
-        if len(images) < nframes:
-            images.extend([images[-1]] * (nframes - len(images)))
-        return images, None, None, None
-
-
-def extract_vision_info(conversations: list[dict] | list[list[dict]]) -> list[dict]:
-    vision_infos = []
-    if isinstance(conversations[0], dict):
-        conversations = [conversations]
-    for conversation in conversations:
-        for message in conversation:
-            if isinstance(message["content"], list):
-                for ele in message["content"]:
-                    if (
-                        "image" in ele
-                        or "image_url" in ele
-                        or "video" in ele
-                        or ele["type"] in ("image", "image_url", "video")
-                    ):
-                        vision_infos.append(ele)
-    return vision_infos
-
-
-def process_vision_info(
-    conversations: list[dict] | list[list[dict]],
-) -> tuple[list[Image.Image] | None, list[torch.Tensor | list[Image.Image]] | None]:
-    vision_infos = extract_vision_info(conversations)
-    # Read images or videos
-    image_inputs = []
-    video_inputs = []
-    for vision_info in vision_infos:
-        if "image" in vision_info or "image_url" in vision_info:
-            image_inputs.append(fetch_image(vision_info))
-        elif "video" in vision_info:
-            video_inputs.append(fetch_video(vision_info))
-        else:
-            raise ValueError("image, image_url or video should in content.")
-    if len(image_inputs) == 0:
-        image_inputs = None
-    if len(video_inputs) == 0:
-        video_inputs = None
-    return image_inputs, video_inputs
@@ -1 +0,0 @@
-ray start --head --port=6399 --num-gpus=1
@@ -1 +0,0 @@
-ray start --address=10.90.67.86:6399
@@ -1,8 +0,0 @@
-python end2end.py --model Qwen/Qwen2.5-Omni-7B \
-                                 --voice-type "m02" \
-                                 --dit-ckpt none \
-                                 --bigvgan-ckpt none \
-                                 --output-wav output_audio \
-                                 --prompt_type text \
-                                 --init-sleep-seconds 0 \
-                                 --txt-prompts top100.txt
@@ -1,11 +0,0 @@
-export PYTHONPATH=/workspace/omni/vllm-omni/:$PYTHONPATH
-CUDA_VISIBLE_DEVICES=4 python end2end.py --model /workspace/Qwen2.5-Omni-7B/ \
-                                 --voice-type "m02" \
-                                 --dit-ckpt none \
-                                 --bigvgan-ckpt none \
-                                 --output-wav output_audio \
-                                 --prompt_type text \
-                                 --init-sleep-seconds 0 \
-				 --worker-backend ray \
-				 --ray-address auto \
-                                 --prompts "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
@@ -1,312 +0,0 @@
-import tempfile
-from typing import Optional, Union
-from urllib.request import urlopen
-
-import librosa
-import requests
-import resampy
-import soundfile as sf
-import torch
-import torchvision.io
-from processing_omni import fetch_image, fetch_video
-from transformers import AutoConfig, AutoProcessor
-from vllm.inputs import TextPrompt
-
-from vllm_omni.inputs.data import OmniTokensPrompt
-
-# Simple caches to avoid repeated heavy HF loads per prompt
-_PROCESSOR_CACHE: dict[str, "AutoProcessor"] = {}
-_CONFIG_CACHE: dict[str, "AutoConfig"] = {}
-
-
-def get_system_prompt():
-    return {
-        "role": "system",
-        "content": [
-            {
-                "type": "text",
-                "text": (
-                    "You are Qwen, a virtual human developed by the Qwen Team, "
-                    "Alibaba Group, capable of perceiving auditory and visual inputs, "
-                    "as well as generating text and speech."
-                ),
-            }
-        ],
-    }
-
-
-def resample_wav_to_16khz(input_filepath):
-    data, original_sample_rate = sf.read(input_filepath)
-    # Only use the first channel
-    if len(data.shape) > 1:
-        data = data[:, 0]
-    # resample to 16kHz
-    data_resampled = resampy.resample(data, sr_orig=original_sample_rate, sr_new=16000)
-    return data_resampled
-
-
-def fetch_and_read_video(args, video_url: str, fps=2):
-    def read_video_with_torchvision(video_file_name: str):
-        video, audio, info = torchvision.io.read_video(
-            video_file_name,
-            start_pts=0.0,
-            end_pts=None,
-            pts_unit="sec",
-            output_format="TCHW",
-        )
-
-        total_frames, video_fps = video.size(0), info["video_fps"]
-        total_duration = round(total_frames / video_fps, 3)
-        nframes = int(total_frames / video_fps * fps)
-
-        frame_timestamps = total_duration * torch.arange(1, nframes + 1) / nframes
-        grid_timestamps = frame_timestamps[::2]
-        second_per_grid = grid_timestamps[1] - grid_timestamps[0]
-
-        idx = torch.linspace(0, video.size(0) - 1, nframes).round().long()
-        video = video[idx]
-
-        if args.legacy_omni_video:
-            return [video, total_duration, nframes, second_per_grid.item()]
-        else:
-            return video
-
-    def read_video_with_transformers(video_file_name: Union[str, list[str]]):
-        video, total_duration, nframes, second_per_grid = fetch_video({"video": video_file_name})
-        if total_duration is None and nframes is None:
-            nframes = len(video)
-            total_duration = 0.5 * nframes
-            second_per_grid = 1.0
-        if args.legacy_omni_video:
-            return [video, total_duration, nframes, second_per_grid]
-        else:
-            return video
-
-    def read_video(video_file_name: str):
-        if args.use_torchvision:
-            return read_video_with_torchvision(video_file_name)
-        else:
-            return read_video_with_transformers(video_file_name)
-
-    if isinstance(video_url, str) and video_url.startswith("http"):
-        with tempfile.NamedTemporaryFile(delete=True) as temp_video_file:
-            resp = requests.get(video_url)
-            assert resp.status_code == requests.codes.ok, (
-                f"Failed to fetch video from {video_url}, status_code:{resp.status_code}, resp:{resp}"
-            )
-
-            temp_video_file.write(urlopen(video_url).read())
-            temp_video_file_path = temp_video_file.name
-            video_file_name = temp_video_file_path
-            return read_video(video_file_name)
-    else:
-        video_file_name = video_url
-        return read_video(video_file_name)
-
-
-def make_inputs_qwen2_omni(
-    args,
-    messages: list[dict[str, Union[str, list[dict[str, str]]]]],
-    use_audio_in_video: Optional[bool] = False,
-    tokenize: bool = False,
-) -> Union[OmniTokensPrompt, TextPrompt]:
-    from transformers import AutoConfig, AutoProcessor
-
-    # Cached processor/config to prevent per-prompt reloading and repeated warnings
-    if args.model not in _PROCESSOR_CACHE:
-        _PROCESSOR_CACHE[args.model] = AutoProcessor.from_pretrained(args.model)
-    processor = _PROCESSOR_CACHE[args.model]
-
-    config = _CONFIG_CACHE.get(args.model)
-    if config is None:
-        try:
-            config = AutoConfig.from_pretrained(args.model)
-        except Exception:
-            config = None
-        _CONFIG_CACHE[args.model] = config  # cache even if None to avoid retry storms
-
-    # Decide legacy flag only once based on config (default True if unknown)
-    if getattr(args, "legacy_omni_video", None) is None:
-        if config is not None and hasattr(config, "architectures"):
-            args.legacy_omni_video = "Qwen2_5OmniModel" not in config.architectures
-        else:
-            args.legacy_omni_video = True
-
-    audios, images, videos = [], [], []
-    for message in messages:
-        if not isinstance(message["content"], list):
-            message["content"] = [
-                {
-                    "type": "text",
-                    "text": message["content"],
-                }
-            ]
-        index, num_contents = 0, len(message["content"])
-        while index < num_contents:
-            ele = message["content"][index]
-            if "type" not in ele:
-                if "text" in ele:
-                    ele["type"] = "text"
-                elif "audio" in ele:
-                    ele["type"] = "audio"
-                elif "audio_url" in ele:
-                    ele["type"] = "audio_url"
-                elif "image" in ele:
-                    ele["type"] = "image"
-                elif "image_url" in ele:
-                    ele["type"] = "image_url"
-                elif "video" in ele:
-                    ele["type"] = "video"
-                elif "video_url" in ele:
-                    ele["type"] = "video_url"
-                else:
-                    raise ValueError(f"Unknown ele: {ele}")
-
-            if ele["type"] == "audio" or ele["type"] == "audio_url":
-                if "audio_url" in ele:
-                    audio_key = "audio_url"
-                    with tempfile.NamedTemporaryFile(delete=True) as temp_audio_file:
-                        temp_audio_file.write(urlopen(ele[audio_key]).read())
-                        temp_audio_file_path = temp_audio_file.name
-                        audios.append(resample_wav_to_16khz(temp_audio_file_path))
-                        ele["audio"] = temp_audio_file_path
-                elif "audio" in ele:
-                    audio_key = "audio"
-                    audios.append(resample_wav_to_16khz(ele[audio_key]))
-                else:
-                    raise ValueError(f"Unknown ele {ele}")
-            elif use_audio_in_video and (ele["type"] == "video" or ele["type"] == "video_url"):
-                # use video as audio as well
-                if "video_url" in ele:
-                    audio_key = "video_url"
-                    with tempfile.NamedTemporaryFile(delete=True) as temp_video_file:
-                        temp_video_file.write(urlopen(ele[audio_key]).read())
-                        temp_video_file_path = temp_video_file.name
-                        ele[audio_key] = temp_video_file_path
-                        audios.append(librosa.load(temp_video_file_path, sr=16000)[0])
-                        videos.append(fetch_and_read_video(args, temp_video_file_path))
-                        ele["video"] = temp_video_file_path
-                elif "video" in ele:
-                    audio_key = "video"
-                    audios.append(librosa.load(ele[audio_key], sr=16000)[0])
-                    videos.append(fetch_and_read_video(args, audio_key))
-                else:
-                    raise ValueError(f"Unknown ele {ele}")
-                # insert a audio after the video
-                message["content"].insert(
-                    index + 1,
-                    {
-                        "type": "audio",
-                        "audio": ele[audio_key],
-                    },
-                )
-                # no need to load the added audio again
-                index += 1
-            elif ele["type"] == "video" or ele["type"] == "video_url":
-                if "video_url" in ele:
-                    video_key = "video_url"
-                    with tempfile.NamedTemporaryFile(delete=True) as temp_video_file:
-                        temp_video_file.write(urlopen(ele["video_url"]).read())
-                        temp_video_file_path = temp_video_file.name
-                        videos.append(fetch_and_read_video(args, temp_video_file))
-                        ele["video"] = temp_video_file_path
-                else:
-                    video_key = "video"
-                    videos.append(fetch_and_read_video(args, ele[video_key]))
-            elif ele["type"] == "image" or ele["type"] == "image_url":
-                images.append(fetch_image(ele))
-
-            # move to the next content
-            index += 1
-
-    prompt = processor.apply_chat_template(
-        messages,
-        tokenize=tokenize,
-        add_generation_prompt=True,
-        add_vision_id=True,
-    )
-
-    audios = audios if len(audios) > 0 else None
-    images = images if len(images) > 0 else None
-    videos = videos if len(videos) > 0 else None
-
-    multi_modal_data = {}
-    if audios:
-        multi_modal_data["audio"] = audios
-    if images:
-        multi_modal_data["image"] = images
-    if videos:
-        multi_modal_data["video"] = videos
-
-    if isinstance(prompt, list) and isinstance(prompt[0], (list, str)):
-        prompt = prompt[0]
-
-    if tokenize:
-        return OmniTokensPrompt(
-            prompt_token_ids=prompt,
-            multi_modal_data=multi_modal_data,
-        )
-    else:
-        return TextPrompt(
-            prompt=prompt,
-            multi_modal_data=multi_modal_data,
-        )
-
-
-def make_text_prompt(args, prompt):
-    messages = [
-        get_system_prompt(),
-        {
-            "role": "user",
-            "content": [
-                {"type": "text", "text": prompt},
-            ],
-        },
-    ]
-
-    prompt = make_inputs_qwen2_omni(args, messages, tokenize=args.tokenize)
-    return prompt
-
-
-def make_audio_in_video_v2_prompt(args):
-    messages = [
-        {
-            "role": "system",
-            "content": [
-                {
-                    "type": "text",
-                    "text": (
-                        "You are Qwen, a virtual human developed by the Qwen Team, "
-                        "Alibaba Group, capable of perceiving auditory and visual "
-                        "inputs, as well as generating text and speech."
-                    ),
-                }
-            ],
-        },
-        {
-            "role": "user",
-            "content": [
-                {
-                    "type": "video_url",
-                    "video_url": ("https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-Omni/draw_small.mp4"),
-                },
-            ],
-        },
-    ]
-    prompt = make_inputs_qwen2_omni(
-        args,
-        messages,
-        use_audio_in_video=True,
-        tokenize=args.tokenize,
-    )
-    return prompt
-
-
-def make_omni_prompt(args, prompt=None) -> Union[OmniTokensPrompt, list[OmniTokensPrompt]]:
-    if args.prompt_type == "text":
-        prompt = make_text_prompt(args, prompt)
-    elif args.prompt_type == "audio-in-video-v2":
-        prompt = make_audio_in_video_v2_prompt(args)
-    else:
-        raise ValueError(f"Unsupported prompt type: {args.prompt_type}")
-    return prompt
@@ -6,7 +6,8 @@ Please refer to [README.md](../../../README.md)

 ## Run examples (Qwen2.5-Omni)

-Launch the server
+### Launch the Server
+
 ```bash
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
 ```
@@ -16,29 +17,59 @@ If you have custom stage configs file, launch the server with command below
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
 ```

+### Send Multi-modal Request
+
 Get into the example folder
 ```bash
 cd examples/online_serving/qwen2_5_omni
 ```

-Send request via python
+####  Send request via python
+
 ```bash
 python openai_chat_completion_client_for_multimodal_generation.py --query-type mixed_modalities
 ```

-Send request via curl
+The Python client supports the following command-line arguments:
+
+- `--query-type` (or `-q`): Query type (default: `mixed_modalities`)
+  - Options: `mixed_modalities`, `use_audio_in_video`, `multi_audios`, `text`
+- `--video-path` (or `-v`): Path to local video file or URL
+  - If not provided and query-type uses video, uses default video URL
+  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
+  - Example: `--video-path /path/to/video.mp4` or `--video-path https://example.com/video.mp4`
+- `--image-path` (or `-i`): Path to local image file or URL
+  - If not provided and query-type uses image, uses default image URL
+  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
+  - Supports common image formats: JPEG, PNG, GIF, WebP
+  - Example: `--image-path /path/to/image.jpg` or `--image-path https://example.com/image.png`
+- `--audio-path` (or `-a`): Path to local audio file or URL
+  - If not provided and query-type uses audio, uses default audio URL
+  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
+  - Supports common audio formats: MP3, WAV, OGG, FLAC, M4A
+  - Example: `--audio-path /path/to/audio.wav` or `--audio-path https://example.com/audio.mp3`
+- `--prompt` (or `-p`): Custom text prompt/question
+  - If not provided, uses default prompt for the selected query type
+  - Example: `--prompt "What are the main activities shown in this video?"`
+
+
+For example, to use mixed modalities with all local files:
+
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py \
+    --query-type mixed_modalities \
+    --video-path /path/to/your/video.mp4 \
+    --image-path /path/to/your/image.jpg \
+    --audio-path /path/to/your/audio.wav \
+    --prompt "Analyze all the media content and provide a comprehensive summary."
+```
+
+####  Send request via curl
+
 ```bash
 bash run_curl_multimodal_generation.sh mixed_modalities
 ```

-### FAQ
-
-If you encounter error about backend of librosa, try to install ffmpeg with command below.
-```
-sudo apt update
-sudo apt install ffmpeg
-```
-
 ## Run Local Web UI Demo

 This Web UI demo allows users to interact with the model through a web browser.
@@ -53,21 +84,18 @@ python gradio_demo.py  --model Qwen/Qwen2.5-Omni-7B --port 7861

 Then open `http://localhost:7861/` on your local browser to interact with the web UI.

+The gradio script supports the following arguments:

-### Options
+- `--model`: Model name
+- `--ip`: Host/IP for Gradio server (default: 127.0.0.1)
+- `--port`: Port for Gradio server (default: 7861)
+- `--stage-configs-path`: Path to custom stage configs YAML file (optional)
+- `--share`: Share the Gradio demo publicly (creates a public link)

-You can customize its basic launch parameters:
+### FAQ

-```bash
-python gradio_demo.py \
-    --model Qwen/Qwen2.5-Omni-7B \
-    --ip 127.0.0.1 \
-    --port 7861 \
-    --stage-configs-path /path/to/stage_configs.yaml
+If you encounter error about backend of librosa, try to install ffmpeg with command below.
+```
+sudo apt update
+sudo apt install ffmpeg
 ```
-
- `--model`: Local model checkpoint to load (default `Qwen/Qwen2.5-Omni-7B`).
- `--ip`: Host/IP for the Gradio server (default `127.0.0.1`).
- `--port`: Port for the Gradio server (default `7861`).
- `--stage-configs-path`: Optional path to custom stage configs YAML.
- `--share`: Set to expose a temporary public link via Gradio.
@@ -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()

@@ -6,7 +6,8 @@ Please refer to [README.md](../../../README.md)

 ## Run examples (Qwen3-Omni)

-Launch the server
+### Launch the Server
+
 ```bash
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
@@ -16,21 +17,59 @@ If you have custom stage configs file, launch the server with command below
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
 ```

+### Send Multi-modal Request
+
 Get into the example folder
 ```bash
 cd examples/online_serving/qwen3_omni
 ```

-Send request via python
+####  Send request via python
+
 ```bash
 python openai_chat_completion_client_for_multimodal_generation.py --query-type use_image
 ```

-Send request via curl
+The Python client supports the following command-line arguments:
+
+- `--query-type` (or `-q`): Query type (default: `use_video`)
+  - Options: `text`, `use_audio`, `use_image`, `use_video`
+- `--model` (or `-m`): Model name/path (default: `Qwen/Qwen3-Omni-30B-A3B-Instruct`)
+- `--video-path` (or `-v`): Path to local video file or URL
+  - If not provided and query-type is `use_video`, uses default video URL
+  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
+  - Example: `--video-path /path/to/video.mp4` or `--video-path https://example.com/video.mp4`
+- `--image-path` (or `-i`): Path to local image file or URL
+  - If not provided and query-type is `use_image`, uses default image URL
+  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
+  - Supports common image formats: JPEG, PNG, GIF, WebP
+  - Example: `--image-path /path/to/image.jpg` or `--image-path https://example.com/image.png`
+- `--audio-path` (or `-a`): Path to local audio file or URL
+  - If not provided and query-type is `use_audio`, uses default audio URL
+  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
+  - Supports common audio formats: MP3, WAV, OGG, FLAC, M4A
+  - Example: `--audio-path /path/to/audio.wav` or `--audio-path https://example.com/audio.mp3`
+- `--prompt` (or `-p`): Custom text prompt/question
+  - If not provided, uses default prompt for the selected query type
+  - Example: `--prompt "What are the main activities shown in this video?"`
+
+
+For example, to use a local video file with custom prompt:
+
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py \
+    --query-type use_video \
+    --video-path /path/to/your/video.mp4 \
+    --prompt "What are the main activities shown in this video?"
+```
+
+####  Send request via curl
+
 ```bash
 bash run_curl_multimodal_generation.sh use_image
 ```

+
 ### FAQ

 If you encounter error about backend of librosa, try to install ffmpeg with command below.
@@ -54,23 +93,9 @@ python gradio_demo.py  --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 7861
 Then open `http://localhost:7861/` on your local browser to interact with the web UI.


-### Options
-
-The gradio demo also supports running with an existing API server and can be customized with the following arguments.
-
-
-```bash
-python gradio_demo.py \
-    --model Qwen/Qwen3-Omni-30B-A3B-Instruct \
-    --use-api-server \
-    --api-base http://localhost:8091/v1 \
-    --ip 127.0.0.1 \
-    --port 7861
-```
+The gradio script supports the following arguments:

 - `--model`: Model name
- `--use-api-server`: If set, connect to an existing vLLM HTTP API server instead of running AsyncOmni locally.
- `--api-base`: Base URL for vllm serve (only used when `use-api-server` is set, default: http://localhost:8091/v1)
 - `--ip`: Host/IP for Gradio server (default: 127.0.0.1)
 - `--port`: Port for Gradio server (default: 7861)
 - `--stage-configs-path`: Path to custom stage configs YAML file (optional)
@@ -264,7 +264,7 @@ async def run_inference_async_omni(
        # Process audio
        audio_data = process_audio_file(audio_file)
        if audio_data is not None:
-            multi_modal_data["audio"] = audio_data
+            multi_modal_data["audio"] = [audio_data] if not isinstance(audio_data, list) else audio_data

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

        # Build the prompt input
@@ -1,5 +1,6 @@
 import base64
-from typing import NamedTuple
+import os
+from typing import NamedTuple, Optional

 import requests
 from openai import OpenAI
@@ -34,6 +35,129 @@ def encode_base64_content_from_url(content_url: str) -> str:
    return result


+def encode_base64_content_from_file(file_path: str) -> str:
+    """Encode a local file to base64 format."""
+    with open(file_path, "rb") as f:
+        content = f.read()
+        result = base64.b64encode(content).decode("utf-8")
+    return result
+
+
+def get_video_url_from_path(video_path: Optional[str]) -> str:
+    """Convert a video path (local file or URL) to a video URL format for the API.
+
+    If video_path is None or empty, returns the default URL.
+    If video_path is a local file path, encodes it to base64 data URL.
+    If video_path is a URL, returns it as-is.
+    """
+    if not video_path:
+        # Default video URL
+        return "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
+
+    # Check if it's a URL (starts with http:// or https://)
+    if video_path.startswith(("http://", "https://")):
+        return video_path
+
+    # Otherwise, treat it as a local file path
+    if not os.path.exists(video_path):
+        raise FileNotFoundError(f"Video file not found: {video_path}")
+
+    # Detect video MIME type from file extension
+    video_path_lower = video_path.lower()
+    if video_path_lower.endswith(".mp4"):
+        mime_type = "video/mp4"
+    elif video_path_lower.endswith(".webm"):
+        mime_type = "video/webm"
+    elif video_path_lower.endswith(".mov"):
+        mime_type = "video/quicktime"
+    elif video_path_lower.endswith(".avi"):
+        mime_type = "video/x-msvideo"
+    elif video_path_lower.endswith(".mkv"):
+        mime_type = "video/x-matroska"
+    else:
+        # Default to mp4 if extension is unknown
+        mime_type = "video/mp4"
+
+    video_base64 = encode_base64_content_from_file(video_path)
+    return f"data:{mime_type};base64,{video_base64}"
+
+
+def get_image_url_from_path(image_path: Optional[str]) -> str:
+    """Convert an image path (local file or URL) to an image URL format for the API.
+
+    If image_path is None or empty, returns the default URL.
+    If image_path is a local file path, encodes it to base64 data URL.
+    If image_path is a URL, returns it as-is.
+    """
+    if not image_path:
+        # Default image URL
+        return "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
+
+    # Check if it's a URL (starts with http:// or https://)
+    if image_path.startswith(("http://", "https://")):
+        return image_path
+
+    # Otherwise, treat it as a local file path
+    if not os.path.exists(image_path):
+        raise FileNotFoundError(f"Image file not found: {image_path}")
+
+    # Detect image MIME type from file extension
+    image_path_lower = image_path.lower()
+    if image_path_lower.endswith((".jpg", ".jpeg")):
+        mime_type = "image/jpeg"
+    elif image_path_lower.endswith(".png"):
+        mime_type = "image/png"
+    elif image_path_lower.endswith(".gif"):
+        mime_type = "image/gif"
+    elif image_path_lower.endswith(".webp"):
+        mime_type = "image/webp"
+    else:
+        # Default to jpeg if extension is unknown
+        mime_type = "image/jpeg"
+
+    image_base64 = encode_base64_content_from_file(image_path)
+    return f"data:{mime_type};base64,{image_base64}"
+
+
+def get_audio_url_from_path(audio_path: Optional[str]) -> str:
+    """Convert an audio path (local file or URL) to an audio URL format for the API.
+
+    If audio_path is None or empty, returns the default URL.
+    If audio_path is a local file path, encodes it to base64 data URL.
+    If audio_path is a URL, returns it as-is.
+    """
+    if not audio_path:
+        # Default audio URL
+        return AudioAsset("mary_had_lamb").url
+
+    # Check if it's a URL (starts with http:// or https://)
+    if audio_path.startswith(("http://", "https://")):
+        return audio_path
+
+    # Otherwise, treat it as a local file path
+    if not os.path.exists(audio_path):
+        raise FileNotFoundError(f"Audio file not found: {audio_path}")
+
+    # Detect audio MIME type from file extension
+    audio_path_lower = audio_path.lower()
+    if audio_path_lower.endswith((".mp3", ".mpeg")):
+        mime_type = "audio/mpeg"
+    elif audio_path_lower.endswith(".wav"):
+        mime_type = "audio/wav"
+    elif audio_path_lower.endswith(".ogg"):
+        mime_type = "audio/ogg"
+    elif audio_path_lower.endswith(".flac"):
+        mime_type = "audio/flac"
+    elif audio_path_lower.endswith(".m4a"):
+        mime_type = "audio/mp4"
+    else:
+        # Default to wav if extension is unknown
+        mime_type = "audio/wav"
+
+    audio_base64 = encode_base64_content_from_file(audio_path)
+    return f"data:{mime_type};base64,{audio_base64}"
+
+
 def get_system_prompt():
    return {
        "role": "system",
@@ -50,8 +174,10 @@ def get_system_prompt():
    }


-def get_text_query():
-    question = "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
+def get_text_query(custom_prompt: Optional[str] = None):
+    question = (
+        custom_prompt or "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
+    )
    prompt = {
        "role": "user",
        "content": [
@@ -71,16 +197,15 @@ default_system = (
 )


-def get_video_query() -> QueryResult:
-    question = "Why is this video funny?"
+def get_video_query(video_path: Optional[str] = None, custom_prompt: Optional[str] = None):
+    question = custom_prompt or "Why is this video funny?"
+    video_url = get_video_url_from_path(video_path)
    prompt = {
        "role": "user",
        "content": [
            {
                "type": "video_url",
-                "video_url": {
-                    "url": "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
-                },
+                "video_url": {"url": video_url},
            },
            {
                "type": "text",
@@ -91,16 +216,15 @@ def get_video_query() -> QueryResult:
    return prompt


-def get_image_query() -> QueryResult:
-    question = "What is the content of this image?"
+def get_image_query(image_path: Optional[str] = None, custom_prompt: Optional[str] = None):
+    question = custom_prompt or "What is the content of this image?"
+    image_url = get_image_url_from_path(image_path)
    prompt = {
        "role": "user",
        "content": [
            {
                "type": "image_url",
-                "image_url": {
-                    "url": "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
-                },
+                "image_url": {"url": image_url},
            },
            {
                "type": "text",
@@ -111,14 +235,15 @@ def get_image_query() -> QueryResult:
    return prompt


-def get_audio_query() -> QueryResult:
-    question = "What is the content of this audio?"
+def get_audio_query(audio_path: Optional[str] = None, custom_prompt: Optional[str] = None):
+    question = custom_prompt or "What is the content of this audio?"
+    audio_url = get_audio_url_from_path(audio_path)
    prompt = {
        "role": "user",
        "content": [
            {
                "type": "audio_url",
-                "audio_url": {"url": AudioAsset("mary_had_lamb").url},
+                "audio_url": {"url": audio_url},
            },
            {
                "type": "text",
@@ -178,7 +303,25 @@ def run_multimodal_generation(args) -> None:
        code2wav_sampling_params,
    ]

-    prompt = query_map[args.query_type]()
+    # Get paths and custom prompt from args
+    video_path = getattr(args, "video_path", None)
+    image_path = getattr(args, "image_path", None)
+    audio_path = getattr(args, "audio_path", None)
+    custom_prompt = getattr(args, "prompt", None)
+
+    # Get the query function and call it with appropriate parameters
+    query_func = query_map[args.query_type]
+    if args.query_type == "use_video":
+        prompt = query_func(video_path=video_path, custom_prompt=custom_prompt)
+    elif args.query_type == "use_image":
+        prompt = query_func(image_path=image_path, custom_prompt=custom_prompt)
+    elif args.query_type == "use_audio":
+        prompt = query_func(audio_path=audio_path, custom_prompt=custom_prompt)
+    elif args.query_type == "text":
+        prompt = query_func(custom_prompt=custom_prompt)
+    else:
+        prompt = query_func()
+
    extra_body = {
        "sampling_params_list": sampling_params_list  # Optional, it has a default setting in stage_configs of the corresponding model.
    }
@@ -225,6 +368,34 @@ def parse_args():
        default="Qwen/Qwen3-Omni-30B-A3B-Instruct",
        help="Model Name / Path",
    )
+    parser.add_argument(
+        "--video-path",
+        "-v",
+        type=str,
+        default=None,
+        help="Path to local video file or URL. If not provided and query-type is 'use_video', uses default video URL.",
+    )
+    parser.add_argument(
+        "--image-path",
+        "-i",
+        type=str,
+        default=None,
+        help="Path to local image file or URL. If not provided and query-type is 'use_image', uses default image URL.",
+    )
+    parser.add_argument(
+        "--audio-path",
+        "-a",
+        type=str,
+        default=None,
+        help="Path to local audio file or URL. If not provided and query-type is 'use_audio', uses default audio URL.",
+    )
+    parser.add_argument(
+        "--prompt",
+        "-p",
+        type=str,
+        default=None,
+        help="Custom text prompt/question to use instead of the default prompt for the selected query type.",
+    )

    return parser.parse_args()

@@ -0,0 +1,36 @@
+# Stage config for running Qwen3-Omni-MoE-Thinking (text-only output)
+# This config is for models like Qwen3-Omni-30B-A3B-Thinking that only have the
+# thinker component and do not support audio output.
+#
+# Single stage: Thinker (multimodal understanding + text generation)
+
+# The following config has been verified on 2x H100-80G GPUs.
+stage_args:
+  - stage_id: 0
+    runtime:
+      devices: "0,1"
+      max_batch_size: 1
+    engine_args:
+      model_stage: thinker
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.9
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: text
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      hf_config_name: thinker_config
+      tensor_parallel_size: 2
+    final_output: true
+    final_output_type: text
+    is_comprehension: true
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
@@ -41,7 +41,7 @@ code2wav_sampling_params='{
  "temperature": 0.0,
  "top_p": 1.0,
  "top_k": -1,
-  "max_tokens": 4096 * 16,
+  "max_tokens": 65536,
  "seed": 42,
  "detokenize": true,
  "repetition_penalty": 1.1
@@ -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/

@@ -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",
@@ -155,6 +156,7 @@ markers = [
    "integration: Integration tests",
    "benchmark: Benchmark tests",
    "slow: Slow tests",
+    "core_model: enable this model test in each PR instead of only nightly",
 ]

 [tool.typos.default]
@@ -0,0 +1,36 @@
+import os
+
+import pytest
+import torch
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@pytest.fixture(autouse=True)
+def clean_gpu_memory_between_tests():
+    if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1":
+        yield
+        return
+
+    # Wait for GPU memory to be cleared before starting the test
+    import gc
+
+    from tests.utils import wait_for_gpu_memory_to_clear
+
+    num_gpus = torch.cuda.device_count()
+    if num_gpus > 0:
+        try:
+            wait_for_gpu_memory_to_clear(
+                devices=list(range(num_gpus)),
+                threshold_ratio=0.1,
+            )
+        except ValueError as e:
+            logger.info("Failed to clean GPU memory: %s", e)
+
+    yield
+
+    # Clean up GPU memory after the test
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        gc.collect()
@@ -0,0 +1,341 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Pytest configuration and fixtures for vllm-omni tests.
+"""
+
+from typing import Any
+
+import pytest
+from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.entrypoints.omni import Omni
+
+PromptAudioInput = list[tuple[Any, int]] | tuple[Any, int] | None
+PromptImageInput = list[Any] | Any | None
+PromptVideoInput = list[Any] | Any | None
+
+
+class OmniRunner:
+    """
+    Test runner for Omni models.
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        seed: int = 42,
+        init_sleep_seconds: int = 20,
+        batch_timeout: int = 10,
+        init_timeout: int = 300,
+        shm_threshold_bytes: int = 65536,
+        log_stats: bool = False,
+        stage_configs_path: str | None = None,
+        **kwargs,
+    ) -> None:
+        """
+        Initialize an OmniRunner for testing.
+
+        Args:
+            model_name: The model name or path
+            seed: Random seed for reproducibility
+            init_sleep_seconds: Sleep time after starting each stage
+            batch_timeout: Timeout for batching in seconds
+            init_timeout: Timeout for initializing stages in seconds
+            shm_threshold_bytes: Threshold for using shared memory
+            log_stats: Enable detailed statistics logging
+            stage_configs_path: Optional path to YAML stage config file
+            **kwargs: Additional arguments passed to Omni
+        """
+        self.model_name = model_name
+        self.seed = seed
+
+        self.omni = Omni(
+            model=model_name,
+            log_stats=log_stats,
+            init_sleep_seconds=init_sleep_seconds,
+            batch_timeout=batch_timeout,
+            init_timeout=init_timeout,
+            shm_threshold_bytes=shm_threshold_bytes,
+            stage_configs_path=stage_configs_path,
+            **kwargs,
+        )
+
+    def get_default_sampling_params_list(self) -> list[SamplingParams]:
+        """
+        Get a list of default sampling parameters for all stages.
+
+        Returns:
+            List of SamplingParams with default decoding for each stage
+        """
+        return [st.default_sampling_params for st in self.omni.instance.stage_list]
+
+    def get_omni_inputs(
+        self,
+        prompts: list[str] | str,
+        system_prompt: str | None = None,
+        audios: PromptAudioInput = None,
+        images: PromptImageInput = None,
+        videos: PromptVideoInput = None,
+        mm_processor_kwargs: dict[str, Any] | None = None,
+    ) -> list[dict[str, Any]]:
+        """
+        Construct Omni input format from prompts and multimodal data.
+
+        Args:
+            prompts: Text prompt(s) - either a single string or list of strings
+            system_prompt: Optional system prompt (defaults to Qwen system prompt)
+            audios: Audio input(s) - tuple of (audio_array, sample_rate) or list of tuples
+            images: Image input(s) - PIL Image or list of PIL Images
+            videos: Video input(s) - numpy array or list of numpy arrays
+            mm_processor_kwargs: Optional processor kwargs (e.g., use_audio_in_video)
+
+        Returns:
+            List of prompt dictionaries suitable for Omni.generate()
+        """
+        if system_prompt is None:
+            system_prompt = (
+                "You are Qwen, a virtual human developed by the Qwen Team, Alibaba "
+                "Group, capable of perceiving auditory and visual inputs, as well as "
+                "generating text and speech."
+            )
+
+        video_padding_token = "<|VIDEO|>"
+        image_padding_token = "<|IMAGE|>"
+        audio_padding_token = "<|AUDIO|>"
+
+        if self.model_name == "Qwen/Qwen3-Omni-30B-A3B-Instruct":
+            video_padding_token = "<|video_pad|>"
+            image_padding_token = "<|image_pad|>"
+            audio_padding_token = "<|audio_pad|>"
+
+        if isinstance(prompts, str):
+            prompts = [prompts]
+
+        def _normalize_mm_input(mm_input, num_prompts):
+            if mm_input is None:
+                return [None] * num_prompts
+            if isinstance(mm_input, list):
+                if len(mm_input) != num_prompts:
+                    raise ValueError(
+                        f"Multimodal input list length ({len(mm_input)}) must match prompts length ({num_prompts})"
+                    )
+                return mm_input
+            return [mm_input] * num_prompts
+
+        num_prompts = len(prompts)
+        audios_list = _normalize_mm_input(audios, num_prompts)
+        images_list = _normalize_mm_input(images, num_prompts)
+        videos_list = _normalize_mm_input(videos, num_prompts)
+
+        omni_inputs = []
+        for i, prompt_text in enumerate(prompts):
+            user_content = ""
+            multi_modal_data = {}
+
+            audio = audios_list[i]
+            if audio is not None:
+                if isinstance(audio, list):
+                    for _ in audio:
+                        user_content += f"<|audio_bos|>{audio_padding_token}<|audio_eos|>"
+                    multi_modal_data["audio"] = audio
+                else:
+                    user_content += f"<|audio_bos|>{audio_padding_token}<|audio_eos|>"
+                    multi_modal_data["audio"] = audio
+
+            image = images_list[i]
+            if image is not None:
+                if isinstance(image, list):
+                    for _ in image:
+                        user_content += f"<|vision_bos|>{image_padding_token}<|vision_eos|>"
+                    multi_modal_data["image"] = image
+                else:
+                    user_content += f"<|vision_bos|>{image_padding_token}<|vision_eos|>"
+                    multi_modal_data["image"] = image
+
+            video = videos_list[i]
+            if video is not None:
+                if isinstance(video, list):
+                    for _ in video:
+                        user_content += f"<|vision_bos|>{video_padding_token}<|vision_eos|>"
+                    multi_modal_data["video"] = video
+                else:
+                    user_content += f"<|vision_bos|>{video_padding_token}<|vision_eos|>"
+                    multi_modal_data["video"] = video
+
+            user_content += prompt_text
+
+            full_prompt = (
+                f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
+                f"<|im_start|>user\n{user_content}<|im_end|>\n"
+                f"<|im_start|>assistant\n"
+            )
+
+            input_dict: dict[str, Any] = {"prompt": full_prompt}
+            if multi_modal_data:
+                input_dict["multi_modal_data"] = multi_modal_data
+            if mm_processor_kwargs:
+                input_dict["mm_processor_kwargs"] = mm_processor_kwargs
+
+            omni_inputs.append(input_dict)
+
+        return omni_inputs
+
+    def generate(
+        self,
+        prompts: list[dict[str, Any]],
+        sampling_params_list: list[SamplingParams] | None = None,
+    ) -> list[Any]:
+        """
+        Generate outputs for the given prompts.
+
+        Args:
+            prompts: List of prompt dictionaries with 'prompt' and optionally
+                    'multi_modal_data' keys
+            sampling_params_list: List of sampling parameters for each stage.
+                                 If None, uses default parameters.
+
+        Returns:
+            List of OmniRequestOutput objects from stages with final_output=True
+        """
+        if sampling_params_list is None:
+            sampling_params_list = self.get_default_sampling_params_list()
+
+        return self.omni.generate(prompts, sampling_params_list)
+
+    def generate_multimodal(
+        self,
+        prompts: list[str] | str,
+        sampling_params_list: list[SamplingParams] | None = None,
+        system_prompt: str | None = None,
+        audios: PromptAudioInput = None,
+        images: PromptImageInput = None,
+        videos: PromptVideoInput = None,
+        mm_processor_kwargs: dict[str, Any] | None = None,
+    ) -> list[Any]:
+        """
+        Convenience method to generate with multimodal inputs.
+
+        Args:
+            prompts: Text prompt(s)
+            sampling_params_list: List of sampling parameters for each stage
+            system_prompt: Optional system prompt
+            audios: Audio input(s)
+            images: Image input(s)
+            videos: Video input(s)
+            mm_processor_kwargs: Optional processor kwargs
+
+        Returns:
+            List of OmniRequestOutput objects from stages with final_output=True
+        """
+        omni_inputs = self.get_omni_inputs(
+            prompts=prompts,
+            system_prompt=system_prompt,
+            audios=audios,
+            images=images,
+            videos=videos,
+            mm_processor_kwargs=mm_processor_kwargs,
+        )
+        return self.generate(omni_inputs, sampling_params_list)
+
+    def generate_audio(
+        self,
+        prompts: list[str] | str,
+        sampling_params_list: list[SamplingParams] | None = None,
+        system_prompt: str | None = None,
+        audios: PromptAudioInput = None,
+        mm_processor_kwargs: dict[str, Any] | None = None,
+    ) -> list[Any]:
+        """
+        Convenience method to generate with multimodal inputs.
+        Args:
+            prompts: Text prompt(s)
+            sampling_params_list: List of sampling parameters for each stage
+            system_prompt: Optional system prompt
+            audios: Audio input(s)
+            mm_processor_kwargs: Optional processor kwargs
+        Returns:
+            List of OmniRequestOutput objects from stages with final_output=True
+        """
+        omni_inputs = self.get_omni_inputs(
+            prompts=prompts,
+            system_prompt=system_prompt,
+            audios=audios,
+            mm_processor_kwargs=mm_processor_kwargs,
+        )
+        return self.generate(omni_inputs, sampling_params_list)
+
+    def generate_video(
+        self,
+        prompts: list[str] | str,
+        sampling_params_list: list[SamplingParams] | None = None,
+        system_prompt: str | None = None,
+        videos: PromptVideoInput = None,
+        mm_processor_kwargs: dict[str, Any] | None = None,
+    ) -> list[Any]:
+        """
+        Convenience method to generate with multimodal inputs.
+        Args:
+            prompts: Text prompt(s)
+            sampling_params_list: List of sampling parameters for each stage
+            system_prompt: Optional system prompt
+            videos: Video input(s)
+            mm_processor_kwargs: Optional processor kwargs
+        Returns:
+            List of OmniRequestOutput objects from stages with final_output=True
+        """
+        omni_inputs = self.get_omni_inputs(
+            prompts=prompts,
+            system_prompt=system_prompt,
+            videos=videos,
+            mm_processor_kwargs=mm_processor_kwargs,
+        )
+        return self.generate(omni_inputs, sampling_params_list)
+
+    def generate_image(
+        self,
+        prompts: list[str] | str,
+        sampling_params_list: list[SamplingParams] | None = None,
+        system_prompt: str | None = None,
+        images: PromptImageInput = None,
+        mm_processor_kwargs: dict[str, Any] | None = None,
+    ) -> list[Any]:
+        """
+        Convenience method to generate with multimodal inputs.
+        Args:
+            prompts: Text prompt(s)
+            sampling_params_list: List of sampling parameters for each stage
+            system_prompt: Optional system prompt
+            images: Image input(s)
+            mm_processor_kwargs: Optional processor kwargs
+        Returns:
+            List of OmniRequestOutput objects from stages with final_output=True
+        """
+        omni_inputs = self.get_omni_inputs(
+            prompts=prompts,
+            system_prompt=system_prompt,
+            images=images,
+            mm_processor_kwargs=mm_processor_kwargs,
+        )
+        return self.generate(omni_inputs, sampling_params_list)
+
+    def __enter__(self):
+        """Context manager entry."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit - cleanup resources."""
+        self.close()
+        del self.omni
+        cleanup_dist_env_and_memory()
+
+    def close(self):
+        """Close and cleanup the Omni instance."""
+        if hasattr(self.omni.instance, "close"):
+            self.omni.instance.close()
+
+
+@pytest.fixture(scope="session")
+def omni_runner():
+    return OmniRunner
@@ -0,0 +1,105 @@
+# stage config for running qwen2.5-omni with architecture of OmniLLM.
+
+# The following config has been verified on 1x 24GB GPU (L4/RTX3090).
+# This config is optimized for CI e2e tests.
+stage_args:
+  - stage_id: 0
+    runtime:
+      process: true            # Run this stage in a separate process
+      devices: "0"            # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
+      max_batch_size: 1
+    engine_args:
+      model_stage: thinker
+      model_arch: Qwen2_5OmniForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      max_model_len: 896
+      max_num_batched_tokens: 896
+      max_num_seqs: 1
+      gpu_memory_utilization: 0.8
+      skip_mm_profiling: true
+      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: 128
+      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.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      max_model_len: 896
+      max_num_batched_tokens: 896
+      max_num_seqs: 1
+      gpu_memory_utilization: 0.8
+      skip_mm_profiling: true
+      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: 128
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
+      stop_token_ids: [8294]
+  - stage_id: 2
+    runtime:
+      process: true
+      devices: "0"            # Example: use a different GPU than the previous stage; use "0" if single GPU
+      max_batch_size: 1
+    engine_args:
+      model_stage: code2wav
+      model_arch: Qwen2_5OmniForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      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: 128
+      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
@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E tests for Qwen2.5-Omni model with mixed modality inputs and audio output.
+"""
+
+from pathlib import Path
+
+import pytest
+from vllm.assets.audio import AudioAsset
+from vllm.assets.image import ImageAsset
+from vllm.assets.video import VideoAsset
+from vllm.multimodal.image import convert_image_mode
+
+from .conftest import OmniRunner
+from .utils import create_new_process_for_each_test
+
+models = ["Qwen/Qwen2.5-Omni-3B"]
+
+# CI stage config optimized for 24GB GPU (L4/RTX3090)
+CI_STAGE_CONFIG_PATH = str(Path(__file__).parent / "stage_configs" / "qwen2_5_omni_ci.yaml")
+
+
+@pytest.mark.core_model
+@pytest.mark.parametrize("model", models)
+@create_new_process_for_each_test()
+def test_mixed_modalities_to_audio(omni_runner: type[OmniRunner], model: str) -> None:
+    """Test processing audio, image, and video together, generating audio output."""
+    with omni_runner(model, seed=42, stage_configs_path=CI_STAGE_CONFIG_PATH) as runner:
+        # Prepare multimodal inputs
+        question = "What is recited in the audio? What is in this image? Describe the video briefly."
+        audio = AudioAsset("mary_had_lamb").audio_and_sample_rate
+        audio = (audio[0][: 16000 * 5], audio[1])  # Trim to first 5 seconds
+        image = convert_image_mode(ImageAsset("cherry_blossom").pil_image.resize((128, 128)), "RGB")
+        video = VideoAsset(name="baby_reading", num_frames=4).np_ndarrays
+
+        outputs = runner.generate_multimodal(
+            prompts=question,
+            audios=audio,
+            images=image,
+            videos=video,
+        )
+
+        # Verify we got outputs from multiple stages
+        assert len(outputs) > 0
+
+        # Find and verify text output (thinker stage)
+        text_output = None
+        for stage_output in outputs:
+            if stage_output.final_output_type == "text":
+                text_output = stage_output
+                break
+
+        assert text_output is not None
+        assert len(text_output.request_output) > 0
+        text_content = text_output.request_output[0].outputs[0].text
+        assert text_content is not None
+        assert len(text_content.strip()) > 0
+
+        # Find and verify audio output (code2wav stage)
+        audio_output = None
+        for stage_output in outputs:
+            if stage_output.final_output_type == "audio":
+                audio_output = stage_output
+                break
+
+        assert audio_output is not None
+        assert len(audio_output.request_output) > 0
+
+        # Verify audio tensor exists and has content
+        audio_tensor = audio_output.request_output[0].multimodal_output["audio"]
+        assert audio_tensor is not None
+        assert audio_tensor.numel() > 0
@@ -0,0 +1,206 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import contextlib
+import functools
+import os
+import signal
+import subprocess
+import sys
+import tempfile
+from collections.abc import Callable
+from contextlib import ExitStack, suppress
+from pathlib import Path
+from typing import Any, Literal
+
+import cloudpickle
+from typing_extensions import ParamSpec
+from vllm.platforms import current_platform
+
+VLLM_PATH = Path(__file__).parent.parent.parent
+"""Path to root of the vLLM repository."""
+
+
+_P = ParamSpec("_P")
+
+
+def fork_new_process_for_each_test(func: Callable[_P, None]) -> Callable[_P, None]:
+    """Decorator to fork a new process for each test function.
+    See https://github.com/vllm-project/vllm/issues/7053 for more details.
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
+        # Make the process the leader of its own process group
+        # to avoid sending SIGTERM to the parent process
+        os.setpgrp()
+        from _pytest.outcomes import Skipped
+
+        # Create a unique temporary file to store exception info from child
+        # process. Use test function name and process ID to avoid collisions.
+        with (
+            tempfile.NamedTemporaryFile(
+                delete=False,
+                mode="w+b",
+                prefix=f"vllm_test_{func.__name__}_{os.getpid()}_",
+                suffix=".exc",
+            ) as exc_file,
+            ExitStack() as delete_after,
+        ):
+            exc_file_path = exc_file.name
+            delete_after.callback(os.remove, exc_file_path)
+
+            pid = os.fork()
+            print(f"Fork a new process to run a test {pid}")
+            if pid == 0:
+                # Parent process responsible for deleting, don't delete
+                # in child.
+                delete_after.pop_all()
+                try:
+                    func(*args, **kwargs)
+                except Skipped as e:
+                    # convert Skipped to exit code 0
+                    print(str(e))
+                    os._exit(0)
+                except Exception as e:
+                    import traceback
+
+                    tb_string = traceback.format_exc()
+
+                    # Try to serialize the exception object first
+                    exc_to_serialize: dict[str, Any]
+                    try:
+                        # First, try to pickle the actual exception with
+                        # its traceback.
+                        exc_to_serialize = {"pickled_exception": e}
+                        # Test if it can be pickled
+                        cloudpickle.dumps(exc_to_serialize)
+                    except (Exception, KeyboardInterrupt):
+                        # Fall back to string-based approach.
+                        exc_to_serialize = {
+                            "exception_type": type(e).__name__,
+                            "exception_msg": str(e),
+                            "traceback": tb_string,
+                        }
+                    try:
+                        with open(exc_file_path, "wb") as f:
+                            cloudpickle.dump(exc_to_serialize, f)
+                    except Exception:
+                        # Fallback: just print the traceback.
+                        print(tb_string)
+                    os._exit(1)
+                else:
+                    os._exit(0)
+            else:
+                pgid = os.getpgid(pid)
+                _pid, _exitcode = os.waitpid(pid, 0)
+                # ignore SIGTERM signal itself
+                old_signal_handler = signal.signal(signal.SIGTERM, signal.SIG_IGN)
+                # kill all child processes
+                os.killpg(pgid, signal.SIGTERM)
+                # restore the signal handler
+                signal.signal(signal.SIGTERM, old_signal_handler)
+                if _exitcode != 0:
+                    # Try to read the exception from the child process
+                    exc_info = {}
+                    if os.path.exists(exc_file_path):
+                        with (
+                            contextlib.suppress(Exception),
+                            open(exc_file_path, "rb") as f,
+                        ):
+                            exc_info = cloudpickle.load(f)
+
+                    if (original_exception := exc_info.get("pickled_exception")) is not None:
+                        # Re-raise the actual exception object if it was
+                        # successfully pickled.
+                        assert isinstance(original_exception, Exception)
+                        raise original_exception
+
+                    if (original_tb := exc_info.get("traceback")) is not None:
+                        # Use string-based traceback for fallback case
+                        raise AssertionError(
+                            f"Test {func.__name__} failed when called with"
+                            f" args {args} and kwargs {kwargs}"
+                            f" (exit code: {_exitcode}):\n{original_tb}"
+                        ) from None
+
+                    # Fallback to the original generic error
+                    raise AssertionError(
+                        f"function {func.__name__} failed when called with"
+                        f" args {args} and kwargs {kwargs}"
+                        f" (exit code: {_exitcode})"
+                    ) from None
+
+    return wrapper
+
+
+def spawn_new_process_for_each_test(f: Callable[_P, None]) -> Callable[_P, None]:
+    """Decorator to spawn a new process for each test function."""
+
+    @functools.wraps(f)
+    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
+        # Check if we're already in a subprocess
+        if os.environ.get("RUNNING_IN_SUBPROCESS") == "1":
+            # If we are, just run the function directly
+            return f(*args, **kwargs)
+
+        import torch.multiprocessing as mp
+
+        with suppress(RuntimeError):
+            mp.set_start_method("spawn")
+
+        # Get the module
+        module_name = f.__module__
+
+        # Create a process with environment variable set
+        env = os.environ.copy()
+        env["RUNNING_IN_SUBPROCESS"] = "1"
+
+        with tempfile.TemporaryDirectory() as tempdir:
+            output_filepath = os.path.join(tempdir, "new_process.tmp")
+
+            # `cloudpickle` allows pickling complex functions directly
+            input_bytes = cloudpickle.dumps((f, output_filepath))
+
+            repo_root = str(VLLM_PATH.resolve())
+
+            env = dict(env or os.environ)
+            env["PYTHONPATH"] = repo_root + os.pathsep + env.get("PYTHONPATH", "")
+
+            cmd = [sys.executable, "-m", f"{module_name}"]
+
+            returned = subprocess.run(cmd, input=input_bytes, capture_output=True, env=env)
+
+            # check if the subprocess is successful
+            try:
+                returned.check_returncode()
+            except Exception as e:
+                # wrap raised exception to provide more information
+                raise RuntimeError(f"Error raised in subprocess:\n{returned.stderr.decode()}") from e
+
+    return wrapper
+
+
+def create_new_process_for_each_test(
+    method: Literal["spawn", "fork"] | None = None,
+) -> Callable[[Callable[_P, None]], Callable[_P, None]]:
+    """Creates a decorator that runs each test function in a new process.
+
+    Args:
+        method: The process creation method. Can be either "spawn" or "fork".
+               If not specified, it defaults to "spawn" on ROCm and XPU
+               platforms and "fork" otherwise.
+
+    Returns:
+        A decorator to run test functions in separate processes.
+    """
+    if method is None:
+        use_spawn = current_platform.is_rocm() or current_platform.is_xpu()
+        method = "spawn" if use_spawn else "fork"
+
+    assert method in ["spawn", "fork"], "Method must be either 'spawn' or 'fork'"
+
+    if method == "fork":
+        return fork_new_process_for_each_test
+
+    return spawn_new_process_for_each_test
@@ -0,0 +1,14 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Pytest configuration for multi_stages_h100 tests.
+
+This conftest imports fixtures from the parent multi_stages conftest.
+"""
+
+# Import the omni_runner fixture from the parent directory's conftest
+# This makes the fixture available to tests in this directory
+from ..multi_stages.conftest import omni_runner  # noqa: F401
+
+# Explicitly expose the fixture for pytest discovery
+__all__ = ["omni_runner"]
@@ -3,6 +3,7 @@
 # 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:
@@ -12,8 +13,8 @@ stage_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
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.6
      enforce_eager: true
      trust_remote_code: true
      engine_output_type: latent  # Output hidden states for talker
@@ -21,6 +22,7 @@ stage_args:
      enable_prefix_caching: false
      hf_config_name: thinker_config
      tensor_parallel_size: 2
+      load_format: dummy
    final_output: true
    final_output_type: text
    is_comprehension: true
@@ -28,7 +30,7 @@ stage_args:
      temperature: 0.4
      top_p: 0.9
      top_k: 1
-      max_tokens: 2048
+      max_tokens: 100
      seed: 42
      detokenize: True
      repetition_penalty: 1.05
@@ -41,7 +43,7 @@ stage_args:
       model_stage: talker
       model_arch: Qwen3OmniMoeForConditionalGeneration
       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
-       scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
+       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.5
       enforce_eager: true
       trust_remote_code: true
@@ -50,6 +52,7 @@ stage_args:
       enable_prefix_caching: false
       distributed_executor_backend: "mp"
       hf_config_name: talker_config
+       load_format: dummy
    engine_input_source: [0]
    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
    # final_output: true
@@ -57,7 +60,7 @@ stage_args:
    default_sampling_params:
      temperature: 0.9
      top_k: 50
-      max_tokens: 4096
+      max_tokens: 100
      seed: 42
      detokenize: False
      repetition_penalty: 1.05
@@ -65,13 +68,13 @@ stage_args:

  - stage_id: 2
    runtime:
-      devices: "1"
+      devices: "0"
      max_batch_size: 1
    engine_args:
      model_stage: code2wav
      model_arch: Qwen3OmniMoeForConditionalGeneration
      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
-      scheduler_cls: vllm_omni.core.sched.generation_scheduler.GenerationScheduler
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
      enforce_eager: true
      trust_remote_code: true
      enable_prefix_caching: false
@@ -80,6 +83,7 @@ stage_args:
      distributed_executor_backend: "mp"
      max_num_batched_tokens: 1000000
      hf_config_name: thinker_config
+      load_format: dummy
    engine_input_source: [1]
    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
    final_output: true
@@ -88,7 +92,7 @@ stage_args:
      temperature: 0.0
      top_p: 1.0
      top_k: -1
-      max_tokens: 65536
+      max_tokens: 200
      seed: 42
      detokenize: True
      repetition_penalty: 1.1
@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E tests for Qwen3-Omni model with audio/video/image input and audio output.
+"""
+
+import os
+from pathlib import Path
+
+import pytest
+from vllm.assets.video import VideoAsset
+
+from ..multi_stages.conftest import OmniRunner
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
+
+# CI stage config for 2xH100-80G GPUs
+CI_STAGE_CONFIG_PATH = str(Path(__file__).parent / "stage_configs" / "qwen3_omni_ci.yaml")
+
+
+@pytest.mark.parametrize("model", models)
+def test_video_to_audio(omni_runner: type[OmniRunner], model: str) -> None:
+    """Test processing video, generating audio output."""
+    with omni_runner(model, seed=42, stage_configs_path=CI_STAGE_CONFIG_PATH) as runner:
+        # Prepare inputs
+        question = "Describe the video briefly."
+        video = VideoAsset(name="baby_reading", num_frames=4).np_ndarrays
+
+        outputs = runner.generate_multimodal(
+            prompts=question,
+            videos=video,
+        )
+
+        # Verify we got outputs from multiple stages
+        assert len(outputs) > 0
+
+        # Find and verify text output (thinker stage)
+        text_output = None
+        for stage_output in outputs:
+            if stage_output.final_output_type == "text":
+                text_output = stage_output
+                break
+
+        assert text_output is not None
+        assert len(text_output.request_output) > 0
+        text_content = text_output.request_output[0].outputs[0].text
+        assert text_content is not None
+        assert len(text_content.strip()) > 0
+
+        # Find and verify audio output (code2wav stage)
+        audio_output = None
+        for stage_output in outputs:
+            if stage_output.final_output_type == "audio":
+                audio_output = stage_output
+                break
+
+        assert audio_output is not None
+        assert len(audio_output.request_output) > 0
+
+        # Verify audio tensor exists and has content
+        audio_tensor = audio_output.request_output[0].multimodal_output["audio"]
+        assert audio_tensor is not None
+        assert audio_tensor.numel() > 0
@@ -0,0 +1,39 @@
+import os
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+
+# ruff: noqa: E402
+REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from vllm_omni import Omni
+
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+
+models = ["Tongyi-MAI/Z-Image-Turbo", "riverclouds/qwen_image_random"]
+
+
+@pytest.mark.parametrize("model_name", models)
+def test_diffusion_model(model_name: str):
+    m = Omni(model=model_name)
+    # high resolution may cause OOM on L4
+    height = 256
+    width = 256
+    images = m.generate(
+        "a photo of a cat sitting on a laptop keyboard",
+        height=height,
+        width=width,
+        num_inference_steps=2,
+        guidance_scale=0.0,
+        generator=torch.Generator("cuda").manual_seed(42),
+        num_outputs_per_prompt=2,
+    )
+    assert len(images) == 2
+    # check image size
+    assert images[0].width == width
+    assert images[0].height == height
+    images[0].save("image_output.png")
@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+import time
+from contextlib import contextmanager
+
+from vllm.platforms import current_platform
+
+if current_platform.is_rocm():
+    from amdsmi import (
+        amdsmi_get_gpu_vram_usage,
+        amdsmi_get_processor_handles,
+        amdsmi_init,
+        amdsmi_shut_down,
+    )
+
+    @contextmanager
+    def _nvml():
+        try:
+            amdsmi_init()
+            yield
+        finally:
+            amdsmi_shut_down()
+elif current_platform.is_cuda():
+    from vllm.third_party.pynvml import (
+        nvmlDeviceGetHandleByIndex,
+        nvmlDeviceGetMemoryInfo,
+        nvmlInit,
+        nvmlShutdown,
+    )
+
+    @contextmanager
+    def _nvml():
+        try:
+            nvmlInit()
+            yield
+        finally:
+            nvmlShutdown()
+else:
+
+    @contextmanager
+    def _nvml():
+        yield
+
+
+def get_physical_device_indices(devices):
+    visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
+    if visible_devices is None:
+        return devices
+
+    visible_indices = [int(x) for x in visible_devices.split(",")]
+    index_mapping = {i: physical for i, physical in enumerate(visible_indices)}
+    return [index_mapping[i] for i in devices if i in index_mapping]
+
+
+@_nvml()
+def wait_for_gpu_memory_to_clear(
+    *,
+    devices: list[int],
+    threshold_bytes: int | None = None,
+    threshold_ratio: float | None = None,
+    timeout_s: float = 120,
+) -> None:
+    assert threshold_bytes is not None or threshold_ratio is not None
+    # Use nvml instead of pytorch to reduce measurement error from torch cuda
+    # context.
+    devices = get_physical_device_indices(devices)
+    start_time = time.time()
+    while True:
+        output: dict[int, str] = {}
+        output_raw: dict[int, tuple[float, float]] = {}
+        for device in devices:
+            if current_platform.is_rocm():
+                dev_handle = amdsmi_get_processor_handles()[device]
+                mem_info = amdsmi_get_gpu_vram_usage(dev_handle)
+                gb_used = mem_info["vram_used"] / 2**10
+                gb_total = mem_info["vram_total"] / 2**10
+            else:
+                dev_handle = nvmlDeviceGetHandleByIndex(device)
+                mem_info = nvmlDeviceGetMemoryInfo(dev_handle)
+                gb_used = mem_info.used / 2**30
+                gb_total = mem_info.total / 2**30
+            output_raw[device] = (gb_used, gb_total)
+            output[device] = f"{gb_used:.02f}/{gb_total:.02f}"
+
+        print("gpu memory used/total (GiB): ", end="")
+        for k, v in output.items():
+            print(f"{k}={v}; ", end="")
+        print("")
+
+        if threshold_bytes is not None:
+            is_free = lambda used, total: used <= threshold_bytes / 2**30  # noqa E731
+            threshold = f"{threshold_bytes / 2**30} GiB"
+        else:
+            is_free = lambda used, total: used / total <= threshold_ratio  # noqa E731
+            threshold = f"{threshold_ratio:.2f}"
+
+        dur_s = time.time() - start_time
+        if all(is_free(used, total) for used, total in output_raw.values()):
+            print(f"Done waiting for free GPU memory on devices {devices=} ({threshold=}) {dur_s=:.02f}")
+            break
+
+        if dur_s >= timeout_s:
+            raise ValueError(f"Memory of devices {devices=} not free after {dur_s=:.02f} ({threshold=})")
+
+        time.sleep(5)
@@ -2,12 +2,12 @@
 Scheduling components for vLLM-Omni.
 """

-from .diffusion_scheduler import DiffusionScheduler
+from .omni_ar_scheduler import OmniARScheduler
+from .omni_generation_scheduler import OmniGenerationScheduler
 from .output import OmniNewRequestData
-from .scheduler import OmniScheduler

 __all__ = [
-    "OmniScheduler",
-    "DiffusionScheduler",
+    "OmniARScheduler",
+    "OmniGenerationScheduler",
    "OmniNewRequestData",
 ]
@@ -1,315 +0,0 @@
-import time
-from collections import defaultdict
-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 Request, RequestStatus, SchedulerOutput, SpecDecodingStats
-from vllm.v1.core.sched.utils import remove_all
-from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs
-
-from vllm_omni.core.sched.output import OmniNewRequestData
-from vllm_omni.core.sched.scheduler import OmniScheduler
-from vllm_omni.outputs import OmniModelRunnerOutput
-
-
-class DiffusionScheduler(OmniScheduler):
-    """Scheduler for diffusion-based non-autoregressive stages.
-
-    Implements a fast-path scheduling strategy for diffusion models where
-    all input tokens are fed at once. Falls back to standard vLLM scheduling
-    if the token budget cannot accommodate all tokens in a single iteration.
-    """
-
-    def schedule(self) -> SchedulerOutput:
-        """Schedule requests for diffusion model execution.
-
-        Uses a fast-path approach that feeds all input tokens at once
-        (or 1 placeholder token if input is empty). Falls back to
-        standard scheduling if token budget is insufficient.
-
-        Returns:
-            SchedulerOutput containing scheduled requests and token allocations
-        """
-
-        # Select requests with zero prompt and using pooling
-        # (diffusion results returned via pooler_output)
-        token_budget = self.max_num_scheduled_tokens
-        capacity = self.max_num_running_reqs - len(self.running)
-        scheduled_timestamp = time.monotonic()
-
-        scheduled_new_reqs: list[Request] = []
-        scheduled_resumed_reqs: list[Request] = []
-        scheduled_running_reqs: list[Request] = []
-
-        req_to_new_blocks: dict[str, KVCacheBlocks] = {}
-        num_scheduled_tokens: dict[str, int] = {}
-        scheduled_spec_decode_tokens: dict[str, list[int]] = {}
-        scheduled_encoder_inputs: dict[str, list[int]] = {}
-
-        # Temporary queue: preserve waiting order, do not disturb non-diffusion requests
-        skipped_waiting_requests = create_request_queue(self.policy)
-
-        # Fast path selection and scheduling (treat all as diffusion requests,
-        # independent of pooling_params)
-        while self.waiting and token_budget > 0 and capacity > 0:
-            request = self.waiting.peek_request()
-            # Uniformly treat as diffusion. A feature flag can be added later
-            # via config or request tag.
-
-            # Allocate all input tokens for the request in one shot
-            # (allocate 1 placeholder if zero)
-            required_tokens = max(getattr(request, "num_prompt_tokens", 0), 1)
-            if required_tokens > token_budget:
-                # Insufficient budget to process all inputs at once;
-                # stop fast path attempt
-                break
-            num_new_tokens = required_tokens
-            new_blocks = self.kv_cache_manager.allocate_slots(
-                request,
-                num_new_tokens,
-                num_lookahead_tokens=self.num_lookahead_tokens,
-            )
-            if new_blocks is None:
-                # Allocation failed (e.g., VRAM pressure); stop fast path and
-                # fall back to default scheduling
-                # Put the current request back to the head of the waiting queue
-                # Note: the original queue order is preserved
-                break
-
-            # Officially schedule this request
-            request = self.waiting.pop_request()
-            self.running.append(request)
-            request.status = RequestStatus.RUNNING
-            if self.log_stats:
-                request.record_event(EngineCoreEventType.SCHEDULED, scheduled_timestamp)
-
-            req_to_new_blocks[request.request_id] = new_blocks
-            num_scheduled_tokens[request.request_id] = num_new_tokens
-            token_budget -= num_new_tokens
-            capacity -= 1
-            scheduled_new_reqs.append(request)
-
-        # Return skipped waiting requests
-        if skipped_waiting_requests:
-            self.waiting.prepend_requests(skipped_waiting_requests)
-
-        # If fast path scheduled none, fall back to the original scheduling
-        if not num_scheduled_tokens:
-            return super().schedule()
-
-        # Compute common prefix blocks (aligned with v1)
-        num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups)
-        if self.running:
-            any_request = self.running[0]
-            num_common_prefix_blocks = self.kv_cache_manager.get_num_common_prefix_blocks(
-                any_request, len(self.running)
-            )
-
-        # Assemble SchedulerOutput
-        new_reqs_data = [
-            OmniNewRequestData.from_request(req, req_to_new_blocks[req.request_id].get_block_ids())
-            for req in scheduled_new_reqs
-        ]
-        cached_reqs_data = self._make_cached_request_data(
-            scheduled_running_reqs,
-            scheduled_resumed_reqs,
-            num_scheduled_tokens,
-            scheduled_spec_decode_tokens,
-            req_to_new_blocks,
-        )
-        scheduled_requests = scheduled_new_reqs + scheduled_running_reqs + scheduled_resumed_reqs
-        structured_output_request_ids, grammar_bitmask = self.get_grammar_bitmask(
-            scheduled_requests, scheduled_spec_decode_tokens
-        )
-
-        total_num_scheduled_tokens = sum(num_scheduled_tokens.values())
-        scheduler_output = SchedulerOutput(
-            scheduled_new_reqs=new_reqs_data,
-            scheduled_cached_reqs=cached_reqs_data,
-            num_scheduled_tokens=num_scheduled_tokens,
-            total_num_scheduled_tokens=total_num_scheduled_tokens,
-            scheduled_spec_decode_tokens=scheduled_spec_decode_tokens,
-            scheduled_encoder_inputs=scheduled_encoder_inputs,
-            num_common_prefix_blocks=num_common_prefix_blocks,
-            finished_req_ids=self.finished_req_ids,
-            free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(),
-            structured_output_request_ids=structured_output_request_ids,
-            grammar_bitmask=grammar_bitmask,
-        )
-
-        # KVTransfer: package metadata
-        if self.connector is not None:
-            meta = self.connector.build_connector_meta(scheduler_output)
-            scheduler_output.kv_connector_metadata = meta
-
-        # Update internal state (advance num_computed_tokens, free encoder inputs,
-        # etc.)
-        self._update_after_schedule(scheduler_output)
-        return scheduler_output
-
-    """
-    Scheduler for the diffusion model.
-    This scheduler is modified to stop the request immediately for the diffusion model.
-    This is because the diffusion model can generate the final image/audio in one step.
-    Note: This is just a minimal modification to the original scheduler,
-    and there should be some further efforts to optimize the scheduler.
-    The original scheduler is still used for the AR model.
-    """
-
-    def update_from_output(
-        self,
-        scheduler_output: SchedulerOutput,
-        model_runner_output: OmniModelRunnerOutput,
-    ) -> dict[int, EngineCoreOutputs]:
-        """Update the scheduler state based on the model runner output.
-
-        This method is modified to stop the request immediately for the diffusion model.
-        """
-        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: Optional[SpecDecodingStats] = 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,
-                )
-
-            new_logprobs = None
-            new_token_ids = generated_token_ids
-            kv_transfer_params = None
-            status_before_stop = request.status
-            pooler_output = None
-            if pooler_outputs:
-                pooler_output = pooler_outputs[req_index]
-
-            # Diffusion request: completes in one step; mark finished and free resources
-            request.status = RequestStatus.FINISHED_STOPPED
-            # Optional: set a stop_reason for front-end clarity
-            # (does not affect protocol)
-            request.stop_reason = request.stop_reason or "diffusion_done"
-            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 check above, so safe to ignore
-                # type warning
-                request.structured_output_request.grammar.accept_tokens(req_id, new_token_ids)  # type: ignore[union-attr]  # noqa: E501
-
-            # spec_token_ids comes from the model runner output
-            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)
-            # Convert pooler_output tensor to dict format expected by OmniEngineCoreOutput
-            pooling_output_dict = None
-            if pooler_output is not None:
-                # Wrap tensor in dict to match OmniEngineCoreOutput.pooling_output type
-                # which expects Optional[dict[str, torch.Tensor]]
-                pooling_output_dict = {"model_outputs": pooler_output}
-            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=pooling_output_dict,
-                        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:
-            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
@@ -12,7 +12,7 @@ from vllm.v1.request import Request, RequestStatus
 from vllm.v1.spec_decode.metrics import SpecDecodingStats


-class OmniScheduler(VLLMScheduler):
+class OmniARScheduler(VLLMScheduler):
    """
    OmniScheduler: Scheduler for vLLM-Omni multimodal processing.

@@ -10,15 +10,15 @@ from vllm.v1.core.sched.scheduler import (
    SchedulerOutput,
    SpecDecodingStats,
 )
+from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 from vllm.v1.core.sched.utils import remove_all
 from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs

 from vllm_omni.core.sched.output import OmniNewRequestData
-from vllm_omni.core.sched.scheduler import OmniScheduler
 from vllm_omni.outputs import OmniModelRunnerOutput


-class GenerationScheduler(OmniScheduler):
+class OmniGenerationScheduler(VLLMScheduler):
    def schedule(self) -> SchedulerOutput:
        """Diffusion fast path:
        - Feed all input tokens of the request at once
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Generic, TypeVar
+
+import torch
+
+
+class AttentionBackend(ABC):
+    """Abstract class for diffusion attention backends."""
+
+    accept_output_buffer: bool = False
+
+    @staticmethod
+    @abstractmethod
+    def get_name() -> str:
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def get_impl_cls() -> type["AttentionImpl"]:
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def get_builder_cls():  # -> Type["AttentionMetadataBuilder"]:
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def get_supported_head_sizes() -> list[int]:
+        """Get the list of supported head sizes for this backend."""
+        raise NotImplementedError
+
+    @classmethod
+    def supports_head_size(cls, head_size: int) -> bool:
+        supported_head_sizes = cls.get_supported_head_sizes()
+        return (not supported_head_sizes) or head_size in supported_head_sizes
+
+
+@dataclass
+class AttentionMetadata:
+    attn_mask: torch.Tensor | None = None
+
+
+T = TypeVar("T", bound=AttentionMetadata)
+
+
+class AttentionImpl(ABC, Generic[T]):
+    @abstractmethod
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: T,
+    ) -> torch.Tensor:
+        raise NotImplementedError
@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+
+logger = init_logger(__name__)
+
+try:
+    # only tested with flash_attn v3
+    # from flash_attn_interface import flash_attn_func as flash_attn_3_func  # not available in flash-attn 2.8.1
+    from flash_attn import flash_attn_func  # can be FA2 or FA3
+except ImportError:
+    logger.warning(
+        "FlashAttentionBackend is not available. You may install flash-attn "
+        "by running `uv pip install flash-attn==2.8.1 --no-build-isolation`"
+        " or install pre-built flash-attn from https://github.com/Dao-AILab/flash-attention/releases"
+    )
+    flash_attn_3_func = None
+
+
+class FlashAttentionBackend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 96, 128, 192, 256]
+
+    @staticmethod
+    def get_name() -> str:
+        return "FLASH_ATTN"
+
+    @staticmethod
+    def get_impl_cls() -> type["FlashAttentionImpl"]:
+        return FlashAttentionImpl
+
+
+class FlashAttentionImpl(AttentionImpl):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.num_heads = num_heads
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata = None,
+    ) -> torch.Tensor:
+        # TODO: flash_attn_func does not support attn_mask.
+        out: torch.Tensor = flash_attn_func(
+            query,
+            key,
+            value,
+            causal=self.causal,
+            softmax_scale=self.softmax_scale,
+        )
+        return out
@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+
+logger = init_logger(__name__)
+
+
+class SDPABackend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [x for x in range(1024)]  # todo
+
+    @staticmethod
+    def get_name() -> str:
+        return "SDPA"
+
+    @staticmethod
+    def get_impl_cls() -> type["SDPAImpl"]:
+        return SDPAImpl
+
+
+class SDPAImpl(AttentionImpl):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata = None,
+    ) -> torch.Tensor:
+        query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value))
+        attention_mask = attn_metadata.attn_mask if attn_metadata else None
+
+        output = torch.nn.functional.scaled_dot_product_attention(
+            query,
+            key,
+            value,
+            attn_mask=attention_mask,
+            dropout_p=0.0,
+            is_causal=self.causal,
+            scale=self.softmax_scale,
+        )
+        out = output.permute(0, 2, 1, 3)
+        return out
@@ -0,0 +1,43 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+import torch.nn as nn
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionMetadata,
+)
+from vllm_omni.diffusion.attention.selector import get_attn_backend
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.attn_backend = get_attn_backend(-1)
+        self.attn_impl_cls = self.attn_backend.get_impl_cls()
+        self.attention = self.attn_impl_cls(
+            num_heads=num_heads,
+            head_size=head_size,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            num_kv_heads=num_kv_heads,
+        )
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata = None,
+    ) -> torch.Tensor:
+        # shape: (batch_size, seq_len, num_heads, head_size)
+        attn_output = self.attention.forward(query, key, value, attn_metadata)
+        return attn_output
@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+from functools import cache
+
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionBackend,
+)
+from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionBackend
+from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
+
+logger = init_logger(__name__)
+
+# environment variable value -> backend class
+SUPPORTED_BACKENDS = {
+    "FLASH_ATTN": FlashAttentionBackend,
+    "TORCH_SDPA": SDPABackend,
+}
+
+
+@cache
+def get_attn_backend(head_size: int) -> type[AttentionBackend]:
+    """
+    Get attention backend for diffusion models.
+
+    The backend is selected based on the following priority:
+    1. DIFFUSION_ATTENTION_BACKEND environment variable (if set, e.g. export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN)
+    2. Default backend (SDPA)
+
+    Args:
+        head_size: Head size (currently not used for selection, but kept for API compatibility)
+
+    Returns:
+        The selected attention backend class
+    """
+    # Check environment variable
+    backend_name: str | None = os.environ.get("DIFFUSION_ATTENTION_BACKEND")
+
+    if backend_name is not None:
+        backend_name_upper = backend_name.upper()
+        if backend_name_upper not in SUPPORTED_BACKENDS:
+            valid_backends = list(SUPPORTED_BACKENDS.keys())
+            raise ValueError(
+                f"Invalid attention backend for diffusion: '{backend_name}'. Valid backends are: {valid_backends}"
+            )
+        logger.info(f"Using attention backend '{backend_name_upper}' for diffusion")
+        return SUPPORTED_BACKENDS[backend_name_upper]
+
+    # Default to SDPA
+    return SDPABackend
@@ -14,6 +14,32 @@ from vllm_omni.diffusion.utils.network_utils import is_port_available
 logger = init_logger(__name__)


+@dataclass
+class TransformerConfig:
+    """Container for raw transformer configuration dictionaries."""
+
+    params: dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "TransformerConfig":
+        if not isinstance(data, dict):
+            raise TypeError(f"Expected transformer config dict, got {type(data)!r}")
+        return cls(params=dict(data))
+
+    def to_dict(self) -> dict[str, Any]:
+        return dict(self.params)
+
+    def get(self, key: str, default: Any | None = None) -> Any:
+        return self.params.get(key, default)
+
+    def __getattr__(self, item: str) -> Any:
+        params = object.__getattribute__(self, "params")
+        try:
+            return params[item]
+        except KeyError as exc:
+            raise AttributeError(item) from exc
+
+
@dataclass
 class OmniDiffusionConfig:
    # Model and path configuration (for convenience)
@@ -21,6 +47,10 @@ class OmniDiffusionConfig:

    model_class_name: str | None = None

+    dtype: torch.dtype = torch.bfloat16
+
+    tf_model_config: TransformerConfig = field(default_factory=TransformerConfig)
+
    # Attention
    # attention_backend: str = None

@@ -80,6 +110,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
@@ -208,3 +242,7 @@ class AttentionBackendEnum(enum.Enum):

    def __str__(self):
        return self.name.lower()
+
+
+# Special message broadcast via scheduler queues to signal worker shutdown.
+SHUTDOWN_MESSAGE = {"type": "shutdown"}
@@ -2,11 +2,12 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

 import multiprocessing as mp
+import time

 from vllm.logger import init_logger

-from vllm_omni.diffusion.data import OmniDiffusionConfig
-from vllm_omni.diffusion.registry import get_diffusion_post_process_func
+from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE, OmniDiffusionConfig
+from vllm_omni.diffusion.registry import get_diffusion_post_process_func, get_diffusion_pre_process_func
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.scheduler import scheduler
 from vllm_omni.utils.platform_utils import get_diffusion_worker_class
@@ -26,16 +27,32 @@ class DiffusionEngine:
        self.od_config = od_config

        self.post_process_func = get_diffusion_post_process_func(od_config)
+        self.pre_process_func = get_diffusion_pre_process_func(od_config)

+        self._processes: list[mp.Process] = []
+        self._closed = False
        self._make_client()

    def step(self, requests: list[OmniDiffusionRequest]):
        try:
+            # Apply pre-processing if available
+            if self.pre_process_func is not None:
+                preprocess_start_time = time.time()
+                requests = self.pre_process_func(requests)
+                preprocess_time = time.time() - preprocess_start_time
+                logger.info(f"Pre-processing completed in {preprocess_time:.4f} seconds")
+
            output = self.add_req_and_wait_for_response(requests)
            if output.error:
                raise Exception(f"{output.error}")
            logger.info("Generation completed successfully.")
-            return self.post_process_func(output.output)
+
+            postprocess_start_time = time.time()
+            result = self.post_process_func(output.output)
+            postprocess_time = time.time() - postprocess_start_time
+            logger.info(f"Post-processing completed in {postprocess_time:.4f} seconds")
+
+            return result
        except Exception as e:
            logger.error(f"Generation failed: {e}")
            return None
@@ -68,6 +85,8 @@ class DiffusionEngine:
        else:
            logger.error("Failed to get result queue handle from workers")

+        self._processes = processes
+
    def _launch_workers(self, broadcast_handle):
        od_config = self.od_config
        logger.info("Starting server...")
@@ -131,3 +150,32 @@ class DiffusionEngine:

    def add_req_and_wait_for_response(self, requests: list[OmniDiffusionRequest]):
        return scheduler.add_req(requests)
+
+    def close(self, *, timeout_s: float = 30.0) -> None:
+        if self._closed:
+            return
+        self._closed = True
+
+        # Send shutdown signal to worker processes via broadcast queue
+        try:
+            if getattr(scheduler, "mq", None) is not None:
+                for _ in range(self.od_config.num_gpus or 1):
+                    scheduler.mq.enqueue(SHUTDOWN_MESSAGE)
+        except Exception as exc:  # pragma: no cover - best effort cleanup
+            logger.warning("Failed to send shutdown signal: %s", exc)
+
+        # Join all worker processes, terminate if they refuse to exit
+        for proc in self._processes:
+            if not proc.is_alive():
+                continue
+            proc.join(timeout_s)
+            if proc.is_alive():
+                logger.warning("Terminating diffusion worker %s after timeout", proc.name)
+                proc.terminate()
+                proc.join(timeout_s)
+
+        scheduler.close()
+        self._processes = []
+
+    def __del__(self):  # pragma: no cover - best effort cleanup
+        self.close()
@@ -0,0 +1,228 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import dataclasses
+import glob
+import os
+import time
+from collections.abc import Generator, Iterable
+from pathlib import Path
+from typing import cast
+
+import torch
+from torch import nn
+from vllm.config import ModelConfig
+from vllm.config.load import LoadConfig
+from vllm.logger import init_logger
+from vllm.model_executor.model_loader.utils import set_default_torch_dtype
+from vllm.model_executor.model_loader.weight_utils import (
+    download_safetensors_index_file_from_hf,
+    download_weights_from_hf,
+    filter_duplicate_safetensors_files,
+    filter_files_not_needed_for_inference,
+    maybe_download_from_modelscope,
+    safetensors_weights_iterator,
+)
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.registry import initialize_model
+
+logger = init_logger(__name__)
+
+
+MODEL_INDEX = "model_index.json"
+DIFFUSION_MODEL_WEIGHTS_INDEX = "diffusion_pytorch_model.safetensors.index.json"
+
+
+class DiffusersPipelineLoader:
+    """Model loader that can load diffusers pipeline components from disk."""
+
+    # default number of thread when enable multithread weight loading
+    DEFAULT_NUM_THREADS = 8
+
+    @dataclasses.dataclass
+    class ComponentSource:
+        """A source for weights."""
+
+        model_or_path: str
+        """The model ID or path."""
+
+        subfolder: str | None
+        """The subfolder inside the model repo."""
+
+        revision: str | None
+        """The optional model revision."""
+
+        prefix: str = ""
+        """A prefix to prepend to all weights."""
+
+        fall_back_to_pt: bool = True
+        """Whether .pt weights can be used."""
+
+        allow_patterns_overrides: list[str] | None = None
+        """If defined, weights will load exclusively using these patterns."""
+
+    counter_before_loading_weights: float = 0.0
+    counter_after_loading_weights: float = 0.0
+
+    def __init__(self, load_config: LoadConfig):
+        self.load_config = load_config
+
+        # TODO(Isotr0py): Enable multithreaded weight loading
+        # extra_config = load_config.model_loader_extra_config
+        # allowed_keys = {"enable_multithread_load", "num_threads"}
+        # unexpected_keys = set(extra_config.keys()) - allowed_keys
+
+        # if unexpected_keys:
+        #     raise ValueError(
+        #         f"Unexpected extra config keys for load format {load_config.load_format}: {unexpected_keys}"
+        #     )
+
+    def _prepare_weights(
+        self,
+        model_name_or_path: Path,
+        subfolder: str | None,
+        revision: str | None,
+        fall_back_to_pt: bool,
+        allow_patterns_overrides: list[str] | None,
+    ) -> tuple[str, list[str], bool]:
+        """Prepare weights for the model.
+
+        If the model is not local, it will be downloaded."""
+        model_name_or_path = maybe_download_from_modelscope(model_name_or_path, revision) or model_name_or_path
+
+        is_local = os.path.isdir(model_name_or_path)
+        load_format = self.load_config.load_format
+        use_safetensors = False
+        index_file = DIFFUSION_MODEL_WEIGHTS_INDEX
+
+        # only hf is supported currently
+        if load_format == "auto":
+            load_format = "hf"
+
+        # Some quantized models use .pt files for storing the weights.
+        if load_format == "hf":
+            allow_patterns = ["*.safetensors", "*.bin"]
+        else:
+            raise ValueError(f"Unknown load_format: {load_format}")
+
+        if fall_back_to_pt:
+            allow_patterns += ["*.pt"]
+
+        if allow_patterns_overrides is not None:
+            allow_patterns = allow_patterns_overrides
+
+        if subfolder is not None:
+            allow_patterns = [f"{subfolder}/{pattern}" for pattern in allow_patterns]
+
+        if not is_local:
+            hf_folder = download_weights_from_hf(
+                model_name_or_path,
+                self.load_config.download_dir,
+                allow_patterns,
+                revision,
+                ignore_patterns=self.load_config.ignore_patterns,
+            )
+        else:
+            hf_folder = model_name_or_path
+
+        hf_weights_files: list[str] = []
+        for pattern in allow_patterns:
+            hf_weights_files += glob.glob(os.path.join(hf_folder, pattern))
+            if len(hf_weights_files) > 0:
+                if pattern == "*.safetensors":
+                    use_safetensors = True
+                break
+
+        if use_safetensors:
+            # For models like Mistral-7B-Instruct-v0.3
+            # there are both sharded safetensors files and a consolidated
+            # safetensors file. Using both breaks.
+            # Here, we download the `model.safetensors.index.json` and filter
+            # any files not found in the index.
+            if not is_local:
+                download_safetensors_index_file_from_hf(
+                    model_name_or_path,
+                    index_file,
+                    self.load_config.download_dir,
+                    revision,
+                )
+            hf_weights_files = filter_duplicate_safetensors_files(hf_weights_files, hf_folder, index_file)
+        else:
+            hf_weights_files = filter_files_not_needed_for_inference(hf_weights_files)
+
+        if len(hf_weights_files) == 0:
+            raise RuntimeError(f"Cannot find any model weights with `{model_name_or_path}`")
+
+        return hf_folder, hf_weights_files, use_safetensors
+
+    def _get_weights_iterator(self, source: "ComponentSource") -> Generator[tuple[str, torch.Tensor], None, None]:
+        """Get an iterator for the model weights based on the load format."""
+        hf_folder, hf_weights_files, use_safetensors = self._prepare_weights(
+            source.model_or_path,
+            source.subfolder,
+            source.revision,
+            source.fall_back_to_pt,
+            source.allow_patterns_overrides,
+        )
+        weights_iterator = safetensors_weights_iterator(
+            hf_weights_files,
+            self.load_config.use_tqdm_on_load,
+            self.load_config.safetensors_load_strategy,
+        )
+
+        if self.counter_before_loading_weights == 0.0:
+            self.counter_before_loading_weights = time.perf_counter()
+        # Apply the prefix.
+        return ((source.prefix + name, tensor) for (name, tensor) in weights_iterator)
+
+    def get_all_weights(
+        self,
+        model: nn.Module,
+    ) -> Generator[tuple[str, torch.Tensor], None, None]:
+        sources = cast(
+            Iterable[DiffusersPipelineLoader.ComponentSource],
+            getattr(model, "weights_sources", ()),
+        )
+        for source in sources:
+            yield from self._get_weights_iterator(source)
+
+    def download_model(self, model_config: ModelConfig) -> None:
+        self._prepare_weights(
+            model_config.model,
+            model_config.revision,
+            fall_back_to_pt=True,
+            allow_patterns_overrides=None,
+        )
+
+    def load_model(self, od_config: OmniDiffusionConfig, load_device: str) -> nn.Module:
+        """Load a model with the given configurations."""
+        target_device = torch.device(load_device)
+        with set_default_torch_dtype(od_config.dtype):
+            with target_device:
+                model = initialize_model(od_config)
+
+            logger.debug("Loading weights on %s ...", load_device)
+            # Quantization does not happen in `load_weights` but after it
+            self.load_weights(model)
+        return model.eval()
+
+    def load_weights(self, model: nn.Module) -> None:
+        weights_to_load = {name for name, _ in model.named_parameters()}
+        loaded_weights = model.load_weights(self.get_all_weights(model))
+
+        self.counter_after_loading_weights = time.perf_counter()
+        logger.info_once(
+            "Loading weights took %.2f seconds",
+            self.counter_after_loading_weights - self.counter_before_loading_weights,
+        )
+        # TODO(Isotr0py): Enable weights loading check after decoupling
+        # all components' weights loading (AutoModel.from_pretrained etc).
+        # We only enable strict check for non-quantized models
+        # that have loaded weights tracking currently.
+        if loaded_weights is not None:
+            _ = weights_to_load - loaded_weights
+        #     if weights_not_loaded:
+        #         raise ValueError(
+        #             "Following weights were not initialized from "
+        #             f"checkpoint: {weights_not_loaded}"
+        #         )
@@ -6,6 +6,7 @@ import json
 import logging
 import math
 import os
+from collections.abc import Iterable
 from typing import Any, Optional, Union

 import numpy as np
@@ -20,10 +21,11 @@ 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.model_executor.models.utils import AutoWeightsLoader

 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
    QwenImageTransformer2DModel,
 )
@@ -236,37 +238,37 @@ class QwenImagePipeline(
    def __init__(
        self,
        *,
-        od_config: OmniDiffusionConfig = None,
+        od_config: OmniDiffusionConfig,
        prefix: str = "",
    ):
        super().__init__()
        self.od_config = od_config
+        self.weights_sources = [
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=od_config.model,
+                subfolder="transformer",
+                revision=None,
+                prefix="transformer.",
+                fall_back_to_pt=True,
+            )
+        ]
+
        self.device = get_local_device()
        model = od_config.model
-
        # 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
        )
-        logger.info("Loaded Qwen-Image scheduler successfully")
-
-        with set_default_torch_dtype(torch.bfloat16):
-            self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-                model, subfolder="text_encoder", local_files_only=local_files_only
-            )
-            logger.info("Loaded Qwen-Image text encoder successfully")
-            self.vae = AutoencoderKLQwenImage.from_pretrained(
-                model, subfolder="vae", local_files_only=local_files_only
-            ).to(self.device)
-            logger.info("Loaded Qwen-Image VAE successfully")
-            self.transformer = QwenImageTransformer2DModel()
-            logger.info("Initialized Qwen-Image transformer successfully.")
-            self.tokenizer = Qwen2Tokenizer.from_pretrained(
-                model, subfolder="tokenizer", local_files_only=local_files_only
-            )
-            logger.info("Loaded Qwen-Image tokenizer successfully.")
+        self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+            model, subfolder="text_encoder", local_files_only=local_files_only
+        )
+        self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
+            self.device
+        )
+        self.transformer = QwenImageTransformer2DModel(od_config=od_config)
+        self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)

        self.stage = None

@@ -741,52 +743,6 @@ class QwenImagePipeline(

        return DiffusionOutput(output=image)

-    def load_weights(self):
-        self.load_transformer()
-
-    # TODO: find a better way to load
-    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 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))
-
-        except Exception as e:
-            print(f"An error occurred: {e}")
-            raise e
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights)
@@ -0,0 +1,829 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import inspect
+import json
+import logging
+import math
+import os
+from collections.abc import Iterable
+from typing import Any, Optional, Union
+
+import numpy as np
+import PIL.Image
+import torch
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.models.autoencoders.autoencoder_kl_qwenimage import (
+    AutoencoderKLQwenImage,
+)
+from diffusers.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler,
+)
+from diffusers.utils.torch_utils import randn_tensor
+from torch import nn
+from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer, Qwen2VLProcessor
+from vllm.model_executor.models.utils import AutoWeightsLoader
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import calculate_shift
+from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
+    QwenImageTransformer2DModel,
+)
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.model_executor.model_loader.weight_utils import (
+    download_weights_from_hf_specific,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def get_qwen_image_edit_pre_process_func(
+    od_config: OmniDiffusionConfig,
+):
+    """Pre-processing function for QwenImageEditPipeline."""
+    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["temporal_downsample"]) if "temporal_downsample" in vae_config else 8
+
+    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2)
+    latent_channels = vae_config.get("z_dim", 16)
+
+    def pre_process_func(
+        requests: list[OmniDiffusionRequest],
+    ):
+        """Pre-process requests for QwenImageEditPipeline."""
+        for req in requests:
+            image = req.pil_image
+
+            image_size = image[0].size if isinstance(image, list) else image.size
+            calculated_width, calculated_height = calculate_dimensions(1024 * 1024, image_size[0] / image_size[1])
+            height = req.height or calculated_height
+            width = req.width or calculated_width
+
+            # Ensure dimensions are multiples of vae_scale_factor * 2
+            multiple_of = vae_scale_factor * 2
+            height = height // multiple_of * multiple_of
+            width = width // multiple_of * multiple_of
+
+            # Store calculated dimensions in request
+            req.calculated_height = calculated_height
+            req.calculated_width = calculated_width
+            req.height = height
+            req.width = width
+
+            # Preprocess image
+            if image is not None and not (
+                isinstance(image, torch.Tensor) and len(image.shape) > 1 and image.shape[1] == latent_channels
+            ):
+                image = image_processor.resize(image, height, width)
+                prompt_image = image
+                image = image_processor.preprocess(image, height, width)
+                image = image.unsqueeze(2)
+
+                # Store preprocessed image and prompt image in request
+                req.preprocessed_image = image
+                req.prompt_image = prompt_image
+
+        return requests
+
+    return pre_process_func
+
+
+def get_qwen_image_edit_post_process_func(
+    od_config: OmniDiffusionConfig,
+):
+    """Post-processing function for QwenImageEditPipeline."""
+    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["temporal_downsample"]) if "temporal_downsample" 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
+
+
+def calculate_dimensions(target_area: float, ratio: float):
+    """Calculate width and height from target area and aspect ratio."""
+    width = math.sqrt(target_area * ratio)
+    height = width / ratio
+
+    width = round(width / 32) * 32
+    height = round(height / 32) * 32
+
+    return width, height
+
+
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[list[int]] = None,
+    sigmas: Optional[list[float]] = 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`.
+    """
+    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
+
+
+def retrieve_latents(
+    encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "argmax"
+):
+    """Retrieve latents from VAE encoder output."""
+    if hasattr(encoder_output, "latent_dist"):
+        return (
+            encoder_output.latent_dist.mode()
+            if sample_mode == "argmax"
+            else encoder_output.latent_dist.sample(generator)
+        )
+    elif hasattr(encoder_output, "latents"):
+        return encoder_output.latents
+    else:
+        raise AttributeError("Could not access latents of provided encoder_output")
+
+
+class QwenImageEditPipeline(
+    nn.Module,
+):
+    def __init__(
+        self,
+        *,
+        od_config: OmniDiffusionConfig,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.od_config = od_config
+        self.weights_sources = [
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=od_config.model,
+                subfolder="transformer",
+                revision=None,
+                prefix="transformer.",
+                fall_back_to_pt=True,
+            )
+        ]
+        self.device = get_local_device()
+        model = od_config.model
+
+        # 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.processor = Qwen2VLProcessor.from_pretrained(
+            model, subfolder="processor", local_files_only=local_files_only
+        )
+
+        self.stage = None
+
+        self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8
+        self.latent_channels = self.vae.config.z_dim if getattr(self, "vae", None) else 16
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)
+        self.tokenizer_max_length = 1024
+        # Edit prompt template - different from generation template
+        self.prompt_template_encode = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n<|im_start|>assistant\n"  # noqa: E501
+        self.prompt_template_encode_start_idx = 64
+        self.default_sample_size = 128
+
+    def check_inputs(
+        self,
+        prompt,
+        height,
+        width,
+        image=None,
+        negative_prompt=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        prompt_embeds_mask=None,
+        negative_prompt_embeds_mask=None,
+        callback_on_step_end_tensor_inputs=None,
+        max_sequence_length=None,
+    ):
+        if height % (self.vae_scale_factor * 2) != 0 or width % (self.vae_scale_factor * 2) != 0:
+            logger.warning(
+                f"`height` and `width` have to be divisible by {self.vae_scale_factor * 2} "
+                f"but are {height} and {width}. Dimensions will be resized accordingly"
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+        if negative_prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Make sure to only forward one of the two."
+            )
+
+        if prompt_embeds is not None and prompt_embeds_mask is None:
+            raise ValueError(
+                "If `prompt_embeds` are provided, `prompt_embeds_mask` also have to be passed. "
+                "Make sure to generate `prompt_embeds_mask` from the same text encoder "
+                "that was used to generate `prompt_embeds`."
+            )
+        if negative_prompt_embeds is not None and negative_prompt_embeds_mask is None:
+            raise ValueError(
+                "If `negative_prompt_embeds` are provided, `negative_prompt_embeds_mask` also have to be passed. "
+                "Make sure to generate `negative_prompt_embeds_mask` from the same text encoder "
+                "that was used to generate `negative_prompt_embeds`."
+            )
+
+        if max_sequence_length is not None and max_sequence_length > 1024:
+            raise ValueError(f"`max_sequence_length` cannot be greater than 1024 but is {max_sequence_length}")
+
+    def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor):
+        bool_mask = mask.bool()
+        valid_lengths = bool_mask.sum(dim=1)
+        selected = hidden_states[bool_mask]
+        split_result = torch.split(selected, valid_lengths.tolist(), dim=0)
+
+        return split_result
+
+    def _get_qwen_prompt_embeds(
+        self,
+        prompt: Union[str, list[str]] = None,
+        image: Optional[torch.Tensor] = None,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        dtype = dtype or self.text_encoder.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+
+        template = self.prompt_template_encode
+        drop_idx = self.prompt_template_encode_start_idx
+        txt = [template.format(e) for e in prompt]
+
+        model_inputs = self.processor(
+            text=txt,
+            images=image,
+            padding=True,
+            return_tensors="pt",
+        ).to(self.device)
+
+        outputs = self.text_encoder(
+            input_ids=model_inputs.input_ids,
+            attention_mask=model_inputs.attention_mask,
+            pixel_values=model_inputs.pixel_values,
+            image_grid_thw=model_inputs.image_grid_thw,
+            output_hidden_states=True,
+        )
+
+        hidden_states = outputs.hidden_states[-1]
+        split_hidden_states = self._extract_masked_hidden(hidden_states, model_inputs.attention_mask)
+        split_hidden_states = [e[drop_idx:] for e in split_hidden_states]
+        attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states]
+        max_seq_len = max([e.size(0) for e in split_hidden_states])
+        prompt_embeds = torch.stack(
+            [torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states]
+        )
+        encoder_attention_mask = torch.stack(
+            [torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list]
+        )
+
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=self.device)
+
+        return prompt_embeds, encoder_attention_mask
+
+    def _get_qwen_prompt_embeds(
+        self,
+        prompt: Union[str, list[str]] = None,
+        image: Optional[Union[PIL.Image.Image, torch.Tensor]] = None,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        """Get prompt embeddings with image support for editing."""
+        dtype = dtype or self.text_encoder.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+
+        template = self.prompt_template_encode
+        drop_idx = self.prompt_template_encode_start_idx
+        txt = [template.format(e) for e in prompt]
+
+        # Use processor to handle both text and image inputs
+        model_inputs = self.processor(
+            text=txt,
+            images=image,
+            padding=True,
+            return_tensors="pt",
+        ).to(self.device)
+
+        outputs = self.text_encoder(
+            input_ids=model_inputs.input_ids,
+            attention_mask=model_inputs.attention_mask,
+            pixel_values=model_inputs.pixel_values,
+            image_grid_thw=model_inputs.image_grid_thw,
+            output_hidden_states=True,
+        )
+
+        hidden_states = outputs.hidden_states[-1]
+        split_hidden_states = self._extract_masked_hidden(hidden_states, model_inputs.attention_mask)
+        split_hidden_states = [e[drop_idx:] for e in split_hidden_states]
+        attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states]
+        max_seq_len = max([e.size(0) for e in split_hidden_states])
+        prompt_embeds = torch.stack(
+            [torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states]
+        )
+        encoder_attention_mask = torch.stack(
+            [torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list]
+        )
+
+        prompt_embeds = prompt_embeds.to(dtype=dtype)
+
+        return prompt_embeds, encoder_attention_mask
+
+    def encode_prompt(
+        self,
+        prompt: Union[str, list[str]],
+        image: Optional[torch.Tensor] = None,
+        num_images_per_prompt: int = 1,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        prompt_embeds_mask: Optional[torch.Tensor] = None,
+        max_sequence_length: int = 1024,
+    ):
+        r"""
+
+        Args:
+            prompt (`str` or `list[str]`, *optional*):
+                prompt to be encoded
+            image (`torch.Tensor`, *optional*):
+                image to be encoded
+            num_images_per_prompt (`int`):
+                number of images that should be generated per prompt
+            prompt_embeds (`torch.Tensor`, *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.
+        """
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0]
+
+        if prompt_embeds is None:
+            prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt, image)
+
+        _, seq_len, _ = prompt_embeds.shape
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+        prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len)
+
+        return prompt_embeds, prompt_embeds_mask
+
+    @staticmethod
+    def _pack_latents(latents, batch_size, num_channels_latents, height, width):
+        latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
+        latents = latents.permute(0, 2, 4, 1, 3, 5)
+        latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4)
+
+        return latents
+
+    @staticmethod
+    def _unpack_latents(latents, height, width, vae_scale_factor):
+        batch_size, num_patches, channels = latents.shape
+
+        # VAE applies 8x compression on images but we must also account for packing which requires
+        # latent height and width to be divisible by 2.
+        height = 2 * (int(height) // (vae_scale_factor * 2))
+        width = 2 * (int(width) // (vae_scale_factor * 2))
+
+        latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2)
+        latents = latents.permute(0, 3, 1, 4, 2, 5)
+
+        latents = latents.reshape(batch_size, channels // (2 * 2), 1, height, width)
+
+        return latents
+
+    def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator):
+        if isinstance(generator, list):
+            image_latents = [
+                retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i], sample_mode="argmax")
+                for i in range(image.shape[0])
+            ]
+            image_latents = torch.cat(image_latents, dim=0)
+        else:
+            image_latents = retrieve_latents(self.vae.encode(image), generator=generator, sample_mode="argmax")
+        latents_mean = (
+            torch.tensor(self.vae.config.latents_mean)
+            .view(1, self.latent_channels, 1, 1, 1)
+            .to(image_latents.device, image_latents.dtype)
+        )
+        latents_std = (
+            torch.tensor(self.vae.config.latents_std)
+            .view(1, self.latent_channels, 1, 1, 1)
+            .to(image_latents.device, image_latents.dtype)
+        )
+        image_latents = (image_latents - latents_mean) / latents_std
+
+        return image_latents
+
+    def prepare_latents(
+        self,
+        image,
+        batch_size,
+        num_channels_latents,
+        height,
+        width,
+        dtype,
+        device,
+        generator,
+        latents=None,
+    ):
+        # VAE applies 8x compression on images but we must also account for packing which requires
+        # latent height and width to be divisible by 2.
+        height = 2 * (int(height) // (self.vae_scale_factor * 2))
+        width = 2 * (int(width) // (self.vae_scale_factor * 2))
+
+        shape = (batch_size, 1, num_channels_latents, height, width)
+
+        image_latents = None
+        if image is not None:
+            image = image.to(device=device, dtype=dtype)
+            if image.shape[1] != self.latent_channels:
+                image_latents = self._encode_vae_image(image=image, generator=generator)
+            else:
+                image_latents = image
+            if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0:
+                # expand init_latents for batch_size
+                additional_image_per_prompt = batch_size // image_latents.shape[0]
+                image_latents = torch.cat([image_latents] * additional_image_per_prompt, dim=0)
+            elif batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] != 0:
+                raise ValueError(
+                    f"Cannot duplicate `image` of batch size {image_latents.shape[0]} to {batch_size} text prompts."
+                )
+            else:
+                image_latents = torch.cat([image_latents], dim=0)
+
+            image_latent_height, image_latent_width = image_latents.shape[3:]
+            image_latents = self._pack_latents(
+                image_latents, batch_size, num_channels_latents, image_latent_height, image_latent_width
+            )
+
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+            latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width)
+        else:
+            latents = latents.to(device=device, dtype=dtype)
+
+        return latents, image_latents
+
+    def prepare_timesteps(self, num_inference_steps, sigmas, image_seq_len):
+        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
+        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),
+        )
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            sigmas=sigmas,
+            mu=mu,
+        )
+        return timesteps, num_inference_steps
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def attention_kwargs(self):
+        return self._attention_kwargs
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def current_timestep(self):
+        return self._current_timestep
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    def diffuse(
+        self,
+        prompt_embeds,
+        prompt_embeds_mask,
+        negative_prompt_embeds,
+        negative_prompt_embeds_mask,
+        latents,
+        image_latents,
+        img_shapes,
+        txt_seq_lens,
+        negative_txt_seq_lens,
+        timesteps,
+        do_true_cfg,
+        guidance,
+        true_cfg_scale,
+    ):
+        """Diffusion loop with optional image conditioning."""
+        self.scheduler.set_begin_index(0)
+        for i, t in enumerate(timesteps):
+            if self.interrupt:
+                continue
+            self._current_timestep = t
+            # broadcast to batch dimension and place on same device/dtype as latents
+            timestep = t.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype)
+
+            # Concatenate image latents with noise latents if available
+            latent_model_input = latents
+            if image_latents is not None:
+                latent_model_input = torch.cat([latents, image_latents], dim=1)
+
+            noise_pred = self.transformer(
+                hidden_states=latent_model_input,
+                timestep=timestep / 1000,
+                guidance=guidance,
+                encoder_hidden_states_mask=prompt_embeds_mask,
+                encoder_hidden_states=prompt_embeds,
+                img_shapes=img_shapes,
+                txt_seq_lens=txt_seq_lens,
+                attention_kwargs=self.attention_kwargs,
+                return_dict=False,
+            )[0]
+            noise_pred = noise_pred[:, : latents.size(1)]
+
+            if do_true_cfg:
+                neg_noise_pred = self.transformer(
+                    hidden_states=latent_model_input,
+                    timestep=timestep / 1000,
+                    guidance=guidance,
+                    encoder_hidden_states_mask=negative_prompt_embeds_mask,
+                    encoder_hidden_states=negative_prompt_embeds,
+                    img_shapes=img_shapes,
+                    txt_seq_lens=negative_txt_seq_lens,
+                    attention_kwargs=self.attention_kwargs,
+                    return_dict=False,
+                )[0]
+                neg_noise_pred = neg_noise_pred[:, : latents.size(1)]
+                comb_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
+                cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True)
+                noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True)
+                noise_pred = comb_pred * (cond_norm / noise_norm)
+            # compute the previous noisy sample x_t -> x_t-1
+            latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+        return latents
+
+    def forward(
+        self,
+        req: OmniDiffusionRequest,
+        prompt: Union[str, list[str]] = "",
+        negative_prompt: Union[str, list[str]] = "",
+        image: Optional[Union[PIL.Image.Image, torch.Tensor]] = None,
+        true_cfg_scale: float = 4.0,
+        height: int | None = None,
+        width: int | None = None,
+        num_inference_steps: int = 50,
+        sigmas: Optional[list[float]] = None,
+        guidance_scale: float = 1.0,
+        num_images_per_prompt: int = 1,
+        generator: Optional[Union[torch.Generator, list[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        prompt_embeds_mask: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        negative_prompt_embeds_mask: Optional[torch.Tensor] = None,
+        output_type: Optional[str] = "pil",
+        attention_kwargs: Optional[dict[str, Any]] = None,
+        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
+        max_sequence_length: int = 512,
+    ) -> DiffusionOutput:
+        """Forward pass for image editing."""
+        prompt = req.prompt if req.prompt is not None else prompt
+        negative_prompt = req.negative_prompt if req.negative_prompt is not None else negative_prompt
+
+        # Get preprocessed image from request (pre-processing is done in DiffusionEngine)
+        if hasattr(req, "preprocessed_image"):
+            prompt_image = req.prompt_image
+            image = req.preprocessed_image
+            calculated_height = req.calculated_height
+            calculated_width = req.calculated_width
+            height = req.height
+            width = req.width
+        else:
+            # fallback to run pre-processing in pipeline (debug only)
+            image_size = image[0].size if isinstance(image, list) else image.size
+            calculated_width, calculated_height, _ = calculate_dimensions(1024 * 1024, image_size[0] / image_size[1])
+            height = height or calculated_height
+            width = width or calculated_width
+
+            multiple_of = self.vae_scale_factor * 2
+            width = width // multiple_of * multiple_of
+            height = height // multiple_of * multiple_of
+
+            if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels):
+                image = self.image_processor.resize(image, calculated_height, calculated_width)
+                prompt_image = image
+                image = self.image_processor.preprocess(image, calculated_height, calculated_width)
+                image = image.unsqueeze(2)
+
+        num_inference_steps = req.num_inference_steps or num_inference_steps
+        generator = req.generator or generator
+        true_cfg_scale = req.true_cfg_scale or true_cfg_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
+
+        # 1. check inputs
+        # 2. encode prompts
+        # 3. prepare latents and timesteps
+        # 4. diffusion process
+        # 5. decode latents
+        # 6. post-process outputs
+        self.check_inputs(
+            prompt,
+            height,
+            width,
+            image,
+            negative_prompt,
+            prompt_embeds,
+            negative_prompt_embeds,
+            prompt_embeds_mask,
+            negative_prompt_embeds_mask,
+            callback_on_step_end_tensor_inputs,
+            max_sequence_length,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._attention_kwargs = attention_kwargs
+        self._current_timestep = None
+        self._interrupt = False
+
+        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 = prompt_embeds.shape[0]
+
+        has_neg_prompt = negative_prompt is not None or (
+            negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None
+        )
+
+        do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
+        prompt_embeds, prompt_embeds_mask = self.encode_prompt(
+            prompt=prompt,
+            image=prompt_image,  # Use resized image for prompt encoding
+            prompt_embeds=prompt_embeds,
+            prompt_embeds_mask=prompt_embeds_mask,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+        )
+
+        if do_true_cfg:
+            negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt(
+                prompt=negative_prompt,
+                image=prompt_image,  # Use same resized image for negative prompt encoding
+                prompt_embeds=negative_prompt_embeds,
+                prompt_embeds_mask=negative_prompt_embeds_mask,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=max_sequence_length,
+            )
+
+        num_channels_latents = self.transformer.in_channels // 4
+        # random noise latents, and image latents encoded by vae
+        latents, image_latents = self.prepare_latents(
+            image,
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            self.device,
+            generator,
+            latents,
+        )
+        img_shapes = [
+            [
+                (1, height // self.vae_scale_factor // 2, width // self.vae_scale_factor // 2),
+                (1, calculated_height // self.vae_scale_factor // 2, calculated_width // self.vae_scale_factor // 2),
+            ]
+        ] * batch_size
+
+        timesteps, num_inference_steps = self.prepare_timesteps(num_inference_steps, sigmas, latents.shape[1])
+        self._num_timesteps = len(timesteps)
+
+        # handle guidance
+        if self.transformer.guidance_embeds:
+            guidance = torch.full([1], guidance_scale, dtype=torch.float32)
+            guidance = guidance.expand(latents.shape[0])
+        else:
+            guidance = None
+
+        if self.attention_kwargs is None:
+            self._attention_kwargs = {}
+
+        txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None
+        negative_txt_seq_lens = (
+            negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None
+        )
+
+        latents = self.diffuse(
+            prompt_embeds,
+            prompt_embeds_mask,
+            negative_prompt_embeds,
+            negative_prompt_embeds_mask,
+            latents,
+            image_latents,
+            img_shapes,
+            txt_seq_lens,
+            negative_txt_seq_lens,
+            timesteps,
+            do_true_cfg,
+            guidance,
+            true_cfg_scale,
+        )
+
+        self._current_timestep = None
+        if output_type == "latent":
+            image = latents
+        else:
+            latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+            latents = latents.to(self.vae.dtype)
+            latents_mean = (
+                torch.tensor(self.vae.config.latents_mean)
+                .view(1, self.vae.config.z_dim, 1, 1, 1)
+                .to(latents.device, latents.dtype)
+            )
+            latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
+                latents.device, latents.dtype
+            )
+            latents = latents / latents_std + latents_mean
+            image = self.vae.decode(latents, return_dict=False)[0][:, :, 0]
+
+        return DiffusionOutput(output=image)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights)
@@ -8,7 +8,6 @@ from typing import Any, Optional, Union
 import torch
 import torch.nn as nn
 from diffusers.models.attention import FeedForward
-from diffusers.models.attention_dispatch import dispatch_attention_fn

 # TODO replace this with vLLM implementation
 from diffusers.models.embeddings import TimestepEmbedding, Timesteps
@@ -17,6 +16,10 @@ from diffusers.models.normalization import AdaLayerNormContinuous
 from vllm.logger import init_logger
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import QKVParallelLinear, ReplicatedLinear
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.data import OmniDiffusionConfig

 logger = init_logger(__name__)

@@ -262,12 +265,12 @@ class QwenImageCrossAttention(nn.Module):
        self.norm_added_q = RMSNorm(head_dim, eps=eps)
        self.norm_added_k = RMSNorm(head_dim, eps=eps)

-        # self.attn = Attention(
-        #     num_heads=num_heads,
-        #     head_size=self.head_dim,
-        #     scale=1.0,  # softmax_scale
-        #     attn_type=AttentionType.ENCODER,
-        # )
+        self.attn = Attention(
+            num_heads=num_heads,
+            head_size=self.head_dim,
+            softmax_scale=1.0 / (self.head_dim**0.5),
+            causal=False,
+        )

    def forward(
        self,
@@ -316,20 +319,11 @@ class QwenImageCrossAttention(nn.Module):
        joint_value = torch.cat([txt_value, img_value], dim=1)

        # Compute joint attention
-        # joint_hidden_states = self.attn(
-        #     joint_query,
-        #     joint_key,
-        #     joint_value,
-        # )
-        joint_hidden_states = dispatch_attention_fn(
+        joint_hidden_states = self.attn(
            joint_query,
            joint_key,
            joint_value,
-            dropout_p=0.0,
-            is_causal=False,
        )
-
-        # Reshape back
        joint_hidden_states = joint_hidden_states.flatten(2, 3)
        joint_hidden_states = joint_hidden_states.to(joint_query.dtype)

@@ -368,19 +362,6 @@ class QwenImageTransformerBlock(nn.Module):
            nn.Linear(dim, 6 * dim, bias=True),  # For scale, shift, gate for norm1 and norm2
        )
        self.img_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
-        # self.attn = Attention(
-        #     query_dim=dim,
-        #     cross_attention_dim=None,  # Enable cross attention for joint computation
-        #     added_kv_proj_dim=dim,  # Enable added KV projections for text stream
-        #     dim_head=attention_head_dim,
-        #     heads=num_attention_heads,
-        #     out_dim=dim,
-        #     context_pre_only=False,
-        #     bias=True,
-        #     processor=QwenDoubleStreamAttnProcessor2_0(),
-        #     qk_norm=qk_norm,
-        #     eps=eps,
-        # )
        self.attn = QwenImageCrossAttention(
            dim=dim,
            num_heads=num_attention_heads,
@@ -506,6 +487,7 @@ class QwenImageTransformer2DModel(nn.Module):
    # _repeated_blocks = ["QwenImageTransformerBlock"]
    def __init__(
        self,
+        od_config: OmniDiffusionConfig,
        patch_size: int = 2,
        in_channels: int = 64,
        out_channels: Optional[int] = 16,
@@ -517,7 +499,8 @@ class QwenImageTransformer2DModel(nn.Module):
        axes_dims_rope: tuple[int, int, int] = (16, 56, 56),
    ):
        super().__init__()
-        self.config = None
+        model_config = od_config.tf_model_config
+        num_layers = model_config.num_layers
        self.in_channels = in_channels
        self.out_channels = out_channels or in_channels
        self.inner_dim = num_attention_heads * attention_head_dim
@@ -624,79 +607,39 @@ class QwenImageTransformer2DModel(nn.Module):

        return Transformer2DModelOutput(sample=output)

-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
-        # Collect incoming weights into a dict for flexible mapping and grouping
-        weights_map: dict[str, torch.Tensor] = {name: tensor for name, tensor in weights}
-
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        stacked_params_mapping = [
-            (
-                "to_qkv",
-                [
-                    {"aliases": ["q_proj", "to_q", "q", "query"], "order": 0},
-                    {"aliases": ["k_proj", "to_k", "k", "key"], "order": 1},
-                    {"aliases": ["v_proj", "to_v", "v", "value"], "order": 2},
-                ],
-            ),
-            (
-                "add_kv_proj",
-                [
-                    {"aliases": ["add_q_proj", "add_q"], "order": 0},
-                    {"aliases": ["add_k_proj", "add_k"], "order": 1},
-                    {"aliases": ["add_v_proj", "add_v"], "order": 2},
-                ],
-            ),
+            # (param_name, shard_name, shard_id)
+            # self-attn
+            (".to_qkv", ".to_q", "q"),
+            (".to_qkv", ".to_k", "k"),
+            (".to_qkv", ".to_v", "v"),
+            # cross-attn
+            (".add_kv_proj", ".add_q_proj", "q"),
+            (".add_kv_proj", ".add_k_proj", "k"),
+            (".add_kv_proj", ".add_v_proj", "v"),
        ]

-        # Helper that attempts to gather shards per base key using aliases and merges them
-        for target, shards in stacked_params_mapping:
-            base_candidates: dict[str, dict[str, dict[int, tuple[str, torch.Tensor]]]] = {}
-            for key in list(weights_map.keys()):
-                for shard in shards:
-                    for alias in shard["aliases"]:
-                        weight_suffix = f".{alias}.weight"
-                        bias_suffix = f".{alias}.bias"
-                        if key.endswith(weight_suffix):
-                            base = key[: -len(weight_suffix)]
-                            base_entry = base_candidates.setdefault(base, {"weight": {}, "bias": {}})
-                            base_entry["weight"][shard["order"]] = (
-                                key,
-                                weights_map[key],
-                            )
-                            break
-                        if key.endswith(bias_suffix):
-                            base = key[: -len(bias_suffix)]
-                            base_entry = base_candidates.setdefault(base, {"weight": {}, "bias": {}})
-                            base_entry["bias"][shard["order"]] = (key, weights_map[key])
-                            break
-
-            for base, entry in base_candidates.items():
-                if len(entry["weight"]) == len(shards):
-                    ordered = [entry["weight"][order] for order in sorted(entry["weight"])]
-                    combined = torch.cat([tensor for _, tensor in ordered], dim=0)
-                    for key, _ in ordered:
-                        weights_map.pop(key, None)
-                    weights_map[base + f".{target}.weight"] = combined
-                if len(entry["bias"]) == len(shards):
-                    ordered_bias = [entry["bias"][order] for order in sorted(entry["bias"])]
-                    bcombined = torch.cat([tensor for _, tensor in ordered_bias], dim=0)
-                    for key, _ in ordered_bias:
-                        weights_map.pop(key, None)
-                    weights_map[base + f".{target}.bias"] = bcombined
-
-        # Now perform the parameter copy using the possibly-updated weights_map
        params_dict = dict(self.named_parameters())
-        for name, loaded_weight in weights_map.items():
-            # Handle potential name mismatches for FeedForward layers
-            if name.endswith(".proj.weight") and name not in params_dict:
-                mapped_name = name.replace(".proj.weight", ".weight")
-            elif name.endswith(".proj.bias") and name not in params_dict:
-                mapped_name = name.replace(".proj.bias", ".bias")
-            else:
-                mapped_name = name

-            if mapped_name in params_dict:
-                param = params_dict[mapped_name]
-                with torch.no_grad():
-                    param.copy_(loaded_weight)
+        # we need to load the buffers for beta and eps (XIELU)
+        for name, buffer in self.named_buffers():
+            if name.endswith(".beta") or name.endswith(".eps"):
+                params_dict[name] = buffer
+
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
            else:
-                logger.warning(f"Weight {name} (mapped to {mapped_name}) not found in model params.")
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params
Author	SHA1	Message	Date
Kevin H. Luu	61e2c3e458	push Signed-off-by: Kevin H. Luu <khluu000@gmail.com>	2025-12-08 16:44:19 -08:00
Kevin H. Luu	560ecc0af2	push Signed-off-by: Kevin H. Luu <khluu000@gmail.com>	2025-12-08 16:41:04 -08:00
Kevin H. Luu	c9647d4688	push Signed-off-by: Kevin H. Luu <khluu000@gmail.com>	2025-12-08 16:39:58 -08:00
Kevin H. Luu	a8b5f0b9ae	push Signed-off-by: Kevin H. Luu <khluu000@gmail.com>	2025-12-08 16:35:25 -08:00
Kevin H. Luu	84c0a26bc9	push Signed-off-by: Kevin H. Luu <khluu000@gmail.com>	2025-12-08 16:04:10 -08:00
Kevin H. Luu	6940d7ce91	push Signed-off-by: Kevin H. Luu <khluu000@gmail.com>	2025-12-08 15:34:04 -08:00
Kevin H. Luu	c0ca921bd7	push Signed-off-by: Kevin H. Luu <khluu000@gmail.com>	2025-12-08 15:32:47 -08:00
Kevin H. Luu	07282be582	push Signed-off-by: Kevin H. Luu <khluu000@gmail.com>	2025-12-08 15:29:32 -08:00
Kevin H. Luu	6b7982afb0	push Signed-off-by: Kevin H. Luu <khluu000@gmail.com>	2025-12-08 15:20:47 -08:00
Isotr0py	83c2723113	[Model] Fuse Z-Image's `qkv_proj` and `gate_up_proj` (#226 ) Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>	2025-12-08 23:21:25 +08:00
Hongsheng Liu	3bf21396c6	[typo] fix doc readme (#242 ) Signed-off-by: hsliu <liuhongsheng4@huawei.com>	2025-12-08 22:23:20 +08:00
rein yang	baf1086d55	[CI] Add Qwen3-omni offline UT (#216 ) Signed-off-by: Rein Yang <ruiruyang2@gmail.com>	2025-12-08 18:11:39 +08:00
Hongsheng Liu	892f837c5a	update docs/readme.md and design folder (#234 ) Signed-off-by: hsliu <liuhongsheng4@huawei.com>	2025-12-08 08:24:16 +00:00
Samit	e8ea4784c6	[Model] Add Qwen-Image-Edit (#196 ) Signed-off-by: samithuang <285365963@qq.com> Signed-off-by: Samit <285365963@qq.com>	2025-12-08 08:19:27 +00:00
SamitandHongsheng Liu	4ac3e6e7d0	[Docs] Add basic usage examples for diffusion (#222 ) Signed-off-by: samithuang <285365963@qq.com> Signed-off-by: Samit <285365963@qq.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2025-12-08 00:00:25 -08:00
Gao Han	45dc189d8d	[Bugfix] Support passing vllm cli args to online serving in vLLM-Omni (#206 ) Signed-off-by: Gaohan123 <gaohan19@huawei.com> Signed-off-by: Gaohan123 <hgaoaf@connect.ust.hk>	2025-12-08 14:59:53 +08:00
Peiqi Yin	bbb5a72871	[BugFix] Fix qwen3omni thinker batching. (#207 ) Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>	2025-12-08 04:44:13 +00:00
Jiangyun Zhu	aa689e6af0	[Core] add clean up method for diffusion engine (#219 ) Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>	2025-12-08 03:52:07 +00:00
Gao Han	02b1b36c6d	[Bugfix] Fix ci bug for qwen2.5-omni (#230 ) Signed-off-by: Gaohan123 <hgaoaf@connect.ust.hk>	2025-12-08 02:43:48 +08:00
Baoyuan Qi	cfe7ad47c1	Fix: Safe handling for multimodal_config to avoid 'NoneType' object h… (#227 ) Signed-off-by: Baoyuan Qi <qibaoyuan@126.com>	2025-12-07 11:40:17 +08:00
Isotr0py	059f45e48d	[Diffusion] Refactor diffusion models weights loading (#157 ) Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>	2025-12-06 11:12:13 +00:00
汪志鹏	08fb978c46	[Docs]Fix doc wrong link (#223 ) Signed-off-by: princepride <wangzhipeng628@gmail.com>	2025-12-06 22:41:03 +13:00
Canlin Guo Isotr0py Gao Han	f75c21bcd0	[E2E] Add Qwen2.5-Omni model test with OmniRunner (#168 ) Signed-off-by: gcanlin <canlinguosdu@gmail.com> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Gao Han <gaohan19@huawei.com>	2025-12-06 06:20:26 +00:00
Jiangyun ZhuandSamit	98fe15762d	[Diffusion] Init Attention Backends and Selector for Diffusion (#115 ) Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Signed-off-by: samithuang <285365963@qq.com> Co-authored-by: Samit <285365963@qq.com>	2025-12-05 10:46:20 +00:00
Canlin Guo	6079b14040	[NPU][Refactor] Rename Diffusion* to Generation* (#211 ) Signed-off-by: gcanlin <canlinguosdu@gmail.com>	2025-12-05 08:23:37 +00:00
Chenguang Zheng	a3f5cd55ee	[bugfix] qwen2.5 omni does not support chunked prefill now (#193 ) Signed-off-by: ZHENG Chenguang <645327136@qq.com>	2025-12-05 04:12:17 +00:00
Zhou TaichangandGao Han	c70b167f70	[Engine] Remove Diffusion_XX which duplicates with Generation_XX (#163 ) Signed-off-by: tzhouam <tzhouam@connect.ust.hk> Co-authored-by: Gao Han <gaohan19@huawei.com>	2025-12-05 03:04:23 +00:00
汪志鹏	f3c69df419	[FixBug]online serving fails for high-resolution videos (#198 ) Signed-off-by: princepride <wangzhipeng628@gmail.com> Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>	2025-12-05 02:29:01 +00:00
Isotr0py	128248650c	[Chore]: Refactor diffusion model registry to reuse vLLM's ModelRegistry (#200 ) Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>	2025-12-05 11:53:39 +13:00
汪志鹏	522ce30a85	[Doc]Fixed some spelling errors (#199 ) Signed-off-by: princepride <wangzhipeng628@gmail.com>	2025-12-04 08:43:28 -08:00
Yueqian Lin	1406c6ec71	[Misc] Add stage config for Qwen3-Omni-30B-A3B-Thinking (#172 ) Signed-off-by: linyueqian <linyueqian@outlook.com>	2025-12-04 08:41:58 -08:00
Jiangyun ZhuandKevin H. Luu	630cbe1ee6	[CI] add diffusion ci (#174 ) Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Signed-off-by: Kevin H. Luu <khluu000@gmail.com> Co-authored-by: Kevin H. Luu <khluu000@gmail.com>	2025-12-04 15:23:22 +00: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
				`@@ -0,0 +1 @@`
				`## OmniConnector for fully disaggregated serving (update soon)`
				`@@ -0,0 +1 @@`
				`Architecture design of the AR module (update soon)`
				`@@ -0,0 +1 @@`
				`Architecture design of the DiT module (update soon)`
				`@@ -0,0 +1 @@`
				`Architecture design of the entrypoint (update soon)`