push

Signed-off-by: Kevin H. Luu <khluu000@gmail.com>

.buildkite/ci_config.yaml

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

.buildkite/jobs/build.yaml

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

.buildkite/jobs/tests.yaml

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

.buildkite/pipeline.yml

@@ -1,31 +0,0 @@
-steps:
-  - label: ":docker: Build image"
-    key: image-build
-    commands:
-      - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-      - "docker build --file docker/Dockerfile.ci -t vllm-omni-ci ."
-      - "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
-      - "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
-    agents:
-      queue: "cpu_queue_premerge_us_east_1"
-
-  - label: "Simple Unit Test"
-    commands:
-      - ".buildkite/scripts/simple_test.sh"
-    agents:
-      queue: "cpu_queue_premerge"
-  - label: "Diffusion Model Test"
-    depends_on: image-build
-    commands:
-      - pytest -s -v tests/single_stage/test_diffusion_model.py
-    agents:
-      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"

docs/.nav.yml

@@ -32,7 +32,13 @@ nav:
     - design/index.md
     - design/architecture_overview.md
     - design/vllm_omni_design.md
-    - design/mrs_design.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

docs/README.md

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

docs/api/README.md

@@ -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
 
@@ -108,13 +107,11 @@ Worker classes and model runners for distributed inference.
 - [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][]

docs/configuration/stage_configs/qwen2_5_omni.yaml

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

docs/design/feature/omni_connector_design.md

@@ -0,0 +1 @@
+## OmniConnector for fully disaggregated serving (update soon)

docs/design/index.md

@@ -6,4 +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)
+
+## Feature Design Documents
+
+- [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)

docs/design/module/ar_module_design.md

@@ -0,0 +1 @@
+Architecture design of the AR module (update soon)

docs/design/module/dit_module_design.md

@@ -0,0 +1 @@
+Architecture design of the DiT module (update soon)

docs/design/module/entrypoint_module_design.md

@@ -0,0 +1 @@
+Architecture design of the entrypoint (update soon)

docs/getting_started/quickstart.md

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

docs/models/supported_models.md

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

docs/user_guide/examples/online_serving/qwen3_omni.md

@@ -114,6 +114,10 @@ The gradio script supports the following arguments:
     ``````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<-- "examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh"

examples/offline_inference/qwen_image/README.md

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

examples/offline_inference/qwen_image/gradio_demo.py

@@ -105,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]]

examples/offline_inference/qwen_image/image_edit.py

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

examples/offline_inference/qwen_image/text_to_image.py

@@ -65,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,
     )
 

examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml

@@ -14,7 +14,7 @@ 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
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.9
       enforce_eager: true
       trust_remote_code: true

pyproject.toml

@@ -156,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]

tests/conftest.py

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

tests/multi_stages/conftest.py

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

tests/multi_stages/stage_configs/qwen2_5_omni_ci.yaml

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

tests/multi_stages/test_qwen2_5_omni.py

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

tests/multi_stages/utils.py

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

tests/multi_stages_h100/conftest.py

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

tests/multi_stages_h100/stage_configs/qwen3_omni_ci.yaml

@@ -0,0 +1,98 @@
+# Stage config for running Qwen3-Omni-MoE with 3-stage architecture
+# Stage 0: Thinker (multimodal understanding + text generation)
+# Stage 1: Talker (text embeddings → 8-layer RVQ codec codes)
+# Stage 2: Code2Wav (8-layer RVQ codes → audio waveform)
+
+# 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.6
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: latent  # Output hidden states for talker
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      hf_config_name: thinker_config
+      tensor_parallel_size: 2
+      load_format: dummy
+    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: 100
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
+
+  - stage_id: 1
+    runtime:
+       devices: "1"
+       max_batch_size: 1
+    engine_args:
+       model_stage: talker
+       model_arch: Qwen3OmniMoeForConditionalGeneration
+       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+       gpu_memory_utilization: 0.5
+       enforce_eager: true
+       trust_remote_code: true
+       engine_output_type: latent  # Output codec codes for code2wav
+      #  tensor_parallel_size: 2
+       enable_prefix_caching: false
+       distributed_executor_backend: "mp"
+       hf_config_name: talker_config
+       load_format: dummy
+    engine_input_source: [0]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
+    # final_output: true
+    # final_output_type: text
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 100
+      seed: 42
+      detokenize: False
+      repetition_penalty: 1.05
+      stop_token_ids: [2150]
+
+  - stage_id: 2
+    runtime:
+      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.omni_generation_scheduler.OmniGenerationScheduler
+      enforce_eager: true
+      trust_remote_code: true
+      enable_prefix_caching: false
+      engine_output_type: audio  # Final output: audio waveform
+      gpu_memory_utilization: 0.1
+      distributed_executor_backend: "mp"
+      max_num_batched_tokens: 1000000
+      hf_config_name: thinker_config
+      load_format: dummy
+    engine_input_source: [1]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
+    final_output: true
+    final_output_type: audio
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 200
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.1

tests/multi_stages_h100/test_qwen3_omni.py

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

tests/single_stage/test_diffusion_model.py

@@ -1,9 +1,20 @@
+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
 
-models = ["Tongyi-MAI/Z-Image-Turbo"]
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+
+models = ["Tongyi-MAI/Z-Image-Turbo", "riverclouds/qwen_image_random"]
 
 
 @pytest.mark.parametrize("model_name", models)
@@ -25,4 +36,4 @@ def test_diffusion_model(model_name: str):
     # check image size
     assert images[0].width == width
     assert images[0].height == height
-    images[0].save("z_image_output.png")
+    images[0].save("image_output.png")

tests/utils.py

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

vllm_omni/diffusion/data.py

@@ -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)
@@ -23,6 +49,8 @@ class OmniDiffusionConfig:
 
     dtype: torch.dtype = torch.bfloat16
 
+    tf_model_config: TransformerConfig = field(default_factory=TransformerConfig)
+
     # Attention
     # attention_backend: str = None
 
@@ -214,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"}

vllm_omni/diffusion/diffusion_engine.py

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

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -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}"
+        #         )

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py

@@ -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,9 +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.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,
 )
@@ -240,9 +243,18 @@ class QwenImagePipeline(
     ):
         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)
 
@@ -255,7 +267,7 @@ class QwenImagePipeline(
         self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
             self.device
         )
-        self.transformer = QwenImageTransformer2DModel()
+        self.transformer = QwenImageTransformer2DModel(od_config=od_config)
         self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
 
         self.stage = None
@@ -731,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)

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py

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

vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py

@@ -16,8 +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__)
 
@@ -485,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,
@@ -496,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
@@ -603,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

vllm_omni/diffusion/models/z_image/pipeline_z_image.py

@@ -18,6 +18,7 @@
 import inspect
 import json
 import os
+from collections.abc import Iterable
 from typing import Any, Callable
 
 import torch
@@ -28,9 +29,11 @@ from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
 from diffusers.utils import logging
 from diffusers.utils.torch_utils import randn_tensor
 from transformers import AutoModel, AutoTokenizer
+from vllm.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.z_image.z_image_transformer import (
     ZImageTransformer2DModel,
 )
@@ -143,11 +146,20 @@ class ZImagePipeline(nn.Module):
     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._execution_device = get_local_device()
         model = od_config.model
         local_files_only = os.path.exists(model)
@@ -598,51 +610,6 @@ class ZImagePipeline(nn.Module):
 
         return DiffusionOutput(output=image)
 
-    def load_weights(self):
-        self.load_transformer()
-
-    def load_transformer(self):
-        import glob
-        import os
-
-        # Define the weight iterator
-        def weight_iterator(transformer_path):
-            if not os.path.exists(transformer_path):
-                logger.warning(f"Path {transformer_path} does not exist.")
-                return
-
-            # Look for safetensors first
-            safetensors_files = glob.glob(os.path.join(transformer_path, "*.safetensors"))
-            if safetensors_files:
-                try:
-                    from safetensors.torch import load_file
-                except ImportError:
-                    logger.warning("safetensors not installed, cannot load .safetensors files.")
-                    return
-
-                for file_path in safetensors_files:
-                    state_dict = load_file(file_path)
-                    for name, tensor in state_dict.items():
-                        yield name, tensor
-            else:
-                # Fallback to bin
-                bin_files = glob.glob(os.path.join(transformer_path, "*.bin"))
-                for file_path in bin_files:
-                    state_dict = torch.load(file_path)
-                    for name, tensor in state_dict.items():
-                        yield name, tensor
-
-        try:
-            # Get model path from config or download from HF
-            model_name = self.od_config.model
-            if os.path.exists(model_name):
-                model_path = model_name
-            else:
-                model_path = download_weights_from_hf_specific(model_name, None, ["*"])
-
-            transformer_path = os.path.join(model_path, "transformer")
-            self.transformer.load_weights(weight_iterator(transformer_path))
-
-        except Exception as e:
-            print(f"An error occurred: {e}")
-            raise e
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights)

vllm_omni/diffusion/models/z_image/z_image_transformer.py

@@ -21,9 +21,15 @@ from typing import Optional
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 from torch.nn.utils.rnn import pad_sequence
+from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.layer import Attention
 
@@ -89,16 +95,13 @@ class ZImageAttention(nn.Module):
         self.head_dim = dim // num_heads
         self.qk_norm = qk_norm
 
-        self.to_q = nn.Linear(dim, dim, bias=False)
-        self.to_k = nn.Linear(dim, self.head_dim * num_kv_heads, bias=False)
-        self.to_v = nn.Linear(dim, self.head_dim * num_kv_heads, bias=False)
-        # TODO enable this when we refactor weight loader
-        # self.to_qkv = QKVParallelLinear(
-        #     hidden_size=dim,
-        #     head_size=self.head_dim,
-        #     total_num_heads=num_heads,
-        #     disable_tp=True,
-        # )
+        self.to_qkv = QKVParallelLinear(
+            hidden_size=dim,
+            head_size=self.head_dim,
+            total_num_heads=num_heads,
+            disable_tp=True,
+            bias=False,
+        )
 
         assert qk_norm is True
         self.norm_q = RMSNorm(self.head_dim, eps=eps)
@@ -119,11 +122,8 @@ class ZImageAttention(nn.Module):
         attention_mask: torch.Tensor,
         freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None,
     ):
-        query = self.to_q(hidden_states)
-        key = self.to_k(hidden_states)
-        value = self.to_v(hidden_states)
-        # qkv, _ = self.to_qkv(hidden_states)
-        # query, key, value = qkv.chunk(3, dim=-1)
+        qkv, _ = self.to_qkv(hidden_states)
+        query, key, value = qkv.chunk(3, dim=-1)
 
         query = query.unflatten(-1, (self.num_heads, -1))
         key = key.unflatten(-1, (self.num_heads, -1))
@@ -170,15 +170,24 @@ class ZImageAttention(nn.Module):
 class FeedForward(nn.Module):
     def __init__(self, dim: int, hidden_dim: int):
         super().__init__()
-        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
-        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
-        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
-
-    def _forward_silu_gating(self, x1, x3):
-        return F.silu(x1) * x3
+        self.w13 = MergedColumnParallelLinear(
+            dim,
+            [hidden_dim] * 2,
+            bias=False,
+            disable_tp=True,
+            return_bias=False,
+        )
+        self.act = SiluAndMul()
+        self.w2 = RowParallelLinear(
+            hidden_dim,
+            dim,
+            bias=False,
+            disable_tp=True,
+            return_bias=False,
+        )
 
     def forward(self, x):
-        return self.w2(self._forward_silu_gating(self.w1(x), self.w3(x)))
+        return self.w2(self.act(self.w13(x)))
 
 
 class ZImageTransformerBlock(nn.Module):
@@ -642,33 +651,33 @@ class ZImageTransformer2DModel(nn.Module):
 
         return x, {}
 
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
-        weights_map: dict[str, torch.Tensor] = {name: tensor for name, tensor in weights}
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            # self-attn
+            (".to_qkv", ".to_q", "q"),
+            (".to_qkv", ".to_k", "k"),
+            (".to_qkv", ".to_v", "v"),
+            # ffn
+            (".w13", ".w1", 0),
+            (".w13", ".w3", 1),
+        ]
+
         params_dict = dict(self.named_parameters())
 
-        for name, loaded_weight in weights_map.items():
-            # In ZImageAttention, to_q, to_k, to_v are separate.
-            # The weight names are like '...attention.to_q.weight'.
-            # The logic from qwen_image_transformer for stacking is not needed here.
-            # We can do a direct mapping.
-
-            # The names in the state dict seem to match the model structure.
-            # Let's try a direct load.
-            if name in params_dict:
+        loaded_params = set[str]()
+        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]
-                with torch.no_grad():
-                    param.copy_(loaded_weight)
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
             else:
-                # The provided weight list might use a different naming convention.
-                # For example, `layers.0.attention.to_out.0.weight` vs `...`
-                # Let's try to handle simple cases.
-
-                # Replace dots with hyphens for ModuleDict keys
-                # e.g. all_x_embedder.2-1.weight -> all_x_embedder.2-1.weight
-                # The names seem to match already.
-
-                # If there are any specific mismatches, they can be handled here.
-                # For now, we assume a direct mapping and warn if a weight is not found.
-                import logging
-
-                logging.warning(f"Weight {name} not found in model parameters.")
+                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

vllm_omni/diffusion/omni_diffusion.py

@@ -7,7 +7,7 @@ from dataclasses import fields
 from vllm.logger import init_logger
 from vllm.transformers_utils.config import get_hf_file_to_dict
 
-from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 
@@ -53,6 +53,11 @@ class OmniDiffusion:
             od_config.model,
         )
         od_config.model_class_name = config_dict.get("_class_name", None)
+        tf_config_dict = get_hf_file_to_dict(
+            "transformer/config.json",
+            od_config.model,
+        )
+        od_config.tf_model_config = TransformerConfig.from_dict(tf_config_dict)
 
         self.engine: DiffusionEngine = DiffusionEngine.make_engine(od_config)
 
@@ -82,3 +87,12 @@ class OmniDiffusion:
 
     def _run_engine(self, requests: list[OmniDiffusionRequest]):
         return self.engine.step(requests)
+
+    def close(self) -> None:
+        self.engine.close()
+
+    def __del__(self):  # pragma: no cover - best effort cleanup
+        try:
+            self.close()
+        except Exception:
+            pass

vllm_omni/diffusion/registry.py

@@ -14,6 +14,11 @@ _DIFFUSION_MODELS = {
         "pipeline_qwen_image",
         "QwenImagePipeline",
     ),
+    "QwenImageEditPipeline": (
+        "qwen_image",
+        "pipeline_qwen_image_edit",
+        "QwenImageEditPipeline",
+    ),
     "ZImagePipeline": (
         "z_image",
         "pipeline_z_image",
@@ -54,19 +59,38 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     # `post_process_func` function must be placed in {mod_folder}/{mod_relname}.py,
     # where mod_folder and mod_relname are  defined and mapped using `_DIFFUSION_MODELS` via the `arch` key
     "QwenImagePipeline": "get_qwen_image_post_process_func",
+    "QwenImageEditPipeline": "get_qwen_image_edit_post_process_func",
     "ZImagePipeline": "get_post_process_func",
 }
 
+_DIFFUSION_PRE_PROCESS_FUNCS = {
+    # arch: pre_process_func
+    # `pre_process_func` function must be placed in {mod_folder}/{mod_relname}.py,
+    # where mod_folder and mod_relname are  defined and mapped using `_DIFFUSION_MODELS` via the `arch` key
+    "QwenImageEditPipeline": "get_qwen_image_edit_pre_process_func",
+}
+
+
+def _load_process_func(od_config: OmniDiffusionConfig, func_name: str):
+    """Load and return a process function from the appropriate module."""
+    mod_folder, mod_relname, _ = _DIFFUSION_MODELS[od_config.model_class_name]
+    module_name = f"vllm_omni.diffusion.models.{mod_folder}.{mod_relname}"
+    module = importlib.import_module(module_name)
+    func = getattr(module, func_name)
+    return func(od_config)
+
 
 def get_diffusion_post_process_func(od_config: OmniDiffusionConfig):
-    if od_config.model_class_name in _DIFFUSION_POST_PROCESS_FUNCS:
-        mod_folder, mod_relname, _ = _DIFFUSION_MODELS[od_config.model_class_name]
-        func_name = _DIFFUSION_POST_PROCESS_FUNCS[od_config.model_class_name]
-        module_name = f"vllm_omni.diffusion.models.{mod_folder}.{mod_relname}"
-        module = importlib.import_module(module_name)
-        post_process_func = getattr(module, func_name)
-        return post_process_func(od_config)
-    else:
+    if od_config.model_class_name not in _DIFFUSION_POST_PROCESS_FUNCS:
         raise ValueError(
             f"Post process function for model class {od_config.model_class_name} not found in diffusion model registry."
         )
+    func_name = _DIFFUSION_POST_PROCESS_FUNCS[od_config.model_class_name]
+    return _load_process_func(od_config, func_name)
+
+
+def get_diffusion_pre_process_func(od_config: OmniDiffusionConfig):
+    if od_config.model_class_name not in _DIFFUSION_PRE_PROCESS_FUNCS:
+        return None  # Return None if no pre-processing function is registered (for backward compatibility)
+    func_name = _DIFFUSION_PRE_PROCESS_FUNCS[od_config.model_class_name]
+    return _load_process_func(od_config, func_name)

vllm_omni/diffusion/scheduler.py

@@ -20,7 +20,8 @@ class Scheduler:
         return cls._instance
 
     def initialize(self, od_config: OmniDiffusionConfig):
-        if hasattr(self, "context") and not self.context.closed:
+        existing_context = getattr(self, "context", None)
+        if existing_context is not None and not existing_context.closed:
             logger.warning("SyncSchedulerClient is already initialized. Re-initializing.")
             self.close()
 
@@ -65,6 +66,9 @@ class Scheduler:
         """Closes the socket and terminates the context."""
         if hasattr(self, "context"):
             self.context.term()
+        self.context = None
+        self.mq = None
+        self.result_mq = None
 
 
 # Singleton instance for easy access

vllm_omni/diffusion/worker/gpu_worker.py

@@ -2,20 +2,25 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import multiprocessing as mp
 import os
+import time
 
 import torch
 import zmq
-from vllm.config import VllmConfig, set_current_vllm_config
+from vllm.config import LoadConfig, VllmConfig, set_current_vllm_config
 from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
 from vllm.distributed.parallel_state import (
     init_distributed_environment,
     initialize_model_parallel,
 )
 from vllm.logger import init_logger
-from vllm.model_executor.model_loader.utils import set_default_torch_dtype
+from vllm.utils import DeviceMemoryProfiler, GiB_bytes
 
-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
-from vllm_omni.diffusion.registry import initialize_model
+from vllm_omni.diffusion.data import (
+    SHUTDOWN_MESSAGE,
+    DiffusionOutput,
+    OmniDiffusionConfig,
+)
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 
 logger = init_logger(__name__)
@@ -60,13 +65,23 @@ class GPUWorker:
 
         init_distributed_environment(world_size=world_size, rank=rank)
         initialize_model_parallel(tensor_model_parallel_size=world_size)
+        logger.info(f"Worker {self.rank}: Initialized device and distributed environment.")
 
-        with device:
-            with set_default_torch_dtype(self.od_config.dtype):
-                self.pipeline = initialize_model(self.od_config)
-                self.pipeline.load_weights()
-                self.pipeline.eval()
-        logger.info(f"Worker {self.rank}: Initialized device, model, and distributed environment.")
+        load_config = LoadConfig()
+        model_loader = DiffusersPipelineLoader(load_config)
+        time_before_load = time.perf_counter()
+        with DeviceMemoryProfiler() as m:
+            self.pipeline = model_loader.load_model(
+                od_config=self.od_config,
+                load_device=f"cuda:{rank}",
+            )
+        time_after_load = time.perf_counter()
+
+        logger.info(
+            "Model loading took %.4f GiB and %.6f seconds",
+            m.consumed_memory / GiB_bytes,
+            time_after_load - time_before_load,
+        )
         logger.info(f"Worker {self.rank}: Model loaded successfully.")
 
     @torch.inference_mode()
@@ -80,6 +95,14 @@ class GPUWorker:
         output = self.pipeline.forward(req)
         return output
 
+    def shutdown(self) -> None:
+        if torch.distributed.is_initialized():
+            try:
+                torch.distributed.destroy_process_group()
+                logger.info("Worker %s: Destroyed process group", self.rank)
+            except Exception as exc:  # pragma: no cover - best effort cleanup
+                logger.warning("Worker %s: Failed to destroy process group: %s", self.rank, exc)
+
 
 class WorkerProc:
     """Wrapper that runs one Worker in a separate process."""
@@ -151,6 +174,14 @@ class WorkerProc:
                 )
                 continue
 
+            if reqs == SHUTDOWN_MESSAGE:
+                logger.info("Worker %s: Received shutdown message", self.gpu_id)
+                self._running = False
+                continue
+            if reqs is None:
+                logger.warning("Worker %s: Received empty payload, ignoring", self.gpu_id)
+                continue
+
             # 2: execute, make sure a reply is always sent
             try:
                 output = self.worker.execute_model(reqs, self.od_config)
@@ -169,6 +200,10 @@ class WorkerProc:
                 continue
 
         logger.info("event loop terminated.")
+        try:
+            self.worker.shutdown()
+        except Exception as exc:  # pragma: no cover - best effort cleanup
+            logger.warning("Worker %s: Shutdown encountered an error: %s", self.gpu_id, exc)
         # if self.result_sender is not None:
         #     self.result_sender.close()
         self.context.term()

vllm_omni/diffusion/worker/npu/npu_worker.py

@@ -2,21 +2,26 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import multiprocessing as mp
 import os
+import time
 
 import torch
 import zmq
 from transformers import PretrainedConfig
-from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
+from vllm.config import LoadConfig, ModelConfig, VllmConfig, set_current_vllm_config
 from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
 from vllm.distributed.parallel_state import (
     init_distributed_environment,
     initialize_model_parallel,
 )
 from vllm.logger import init_logger
-from vllm.model_executor.model_loader.utils import set_default_torch_dtype
+from vllm.utils import DeviceMemoryProfiler, GiB_bytes
 
-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
-from vllm_omni.diffusion.registry import initialize_model
+from vllm_omni.diffusion.data import (
+    SHUTDOWN_MESSAGE,
+    DiffusionOutput,
+    OmniDiffusionConfig,
+)
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 
 logger = init_logger(__name__)
@@ -60,19 +65,27 @@ class NPUWorker:
             def __getattr__(self, name):
                 return "fake"
 
-        vllm_config = VllmConfig(model_config=ModelConfig(hf_config=_FakePretrainedConfig()))
+        vllm_config = VllmConfig(model_config=ModelConfig(hf_config=_FakePretrainedConfig()), load_config=LoadConfig())
         vllm_config.parallel_config.tensor_parallel_size = self.od_config.num_gpus
         set_current_vllm_config(vllm_config)
 
         init_distributed_environment(world_size=world_size, rank=rank)
         initialize_model_parallel(tensor_model_parallel_size=world_size)
 
-        with device:
-            with set_default_torch_dtype(self.od_config.dtype):
-                self.pipeline = initialize_model(self.od_config)
-                self.pipeline.load_weights()
-                self.pipeline.eval()
-        logger.info(f"Worker {self.rank}: Initialized device, model, and distributed environment.")
+        model_loader = DiffusersPipelineLoader(vllm_config.load_config)
+        time_before_load = time.perf_counter()
+        with DeviceMemoryProfiler() as m:
+            self.pipeline = model_loader.load_model(
+                od_config=self.od_config,
+                load_device=f"npu:{rank}",
+            )
+        time_after_load = time.perf_counter()
+
+        logger.info(
+            "Model loading took %.4f GiB and %.6f seconds",
+            m.consumed_memory / GiB_bytes,
+            time_after_load - time_before_load,
+        )
         logger.info(f"Worker {self.rank}: Model loaded successfully.")
 
     @torch.inference_mode()
@@ -86,6 +99,14 @@ class NPUWorker:
         output = self.pipeline.forward(req)
         return output
 
+    def shutdown(self) -> None:
+        if torch.distributed.is_initialized():
+            try:
+                torch.distributed.destroy_process_group()
+                logger.info("Worker %s: Destroyed process group", self.rank)
+            except Exception as exc:  # pragma: no cover - best effort cleanup
+                logger.warning("Worker %s: Failed to destroy process group: %s", self.rank, exc)
+
 
 class NPUWorkerProc:
     """Wrapper that runs one Worker in a separate process."""
@@ -157,6 +178,14 @@ class NPUWorkerProc:
                 )
                 continue
 
+            if reqs == SHUTDOWN_MESSAGE:
+                logger.info("Worker %s: Received shutdown message", self.gpu_id)
+                self._running = False
+                continue
+            if reqs is None:
+                logger.warning("Worker %s: Received empty payload, ignoring", self.gpu_id)
+                continue
+
             # 2: execute, make sure a reply is always sent
             try:
                 output = self.worker.execute_model(reqs, self.od_config)
@@ -175,6 +204,10 @@ class NPUWorkerProc:
                 continue
 
         logger.info("event loop terminated.")
+        try:
+            self.worker.shutdown()
+        except Exception as exc:  # pragma: no cover - best effort cleanup
+            logger.warning("Worker %s: Shutdown encountered an error: %s", self.gpu_id, exc)
         # if self.result_sender is not None:
         #     self.result_sender.close()
         self.context.term()

vllm_omni/engine/arg_utils.py

@@ -51,6 +51,11 @@ class OmniEngineArgs(EngineArgs):
         # Create OmniModelConfig by copying all base config attributes
         # and adding the new omni-specific fields
         config_dict = base_config.__dict__.copy()
+        # FIXME(Isotr0py): This is a temporary workaround for multimodal_config
+        config_dict = {
+            **(getattr(mm := config_dict.pop("multimodal_config", None), "__dict__", mm or {})),
+            **config_dict,
+        }
 
         # Add the new omni-specific fields
         config_dict["stage_id"] = self.stage_id

vllm_omni/entrypoints/async_omni.py

@@ -3,6 +3,7 @@ import multiprocessing as mp
 import os
 import socket
 import time
+from argparse import Namespace
 from collections.abc import AsyncGenerator, Iterable, Mapping, Sequence
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import Any, Optional, Union
@@ -61,21 +62,23 @@ class AsyncOmni(EngineClient):
 
     Args:
         model: Model name or path to load
-        stage_configs_path: Optional path to YAML file containing stage
-            configurations. If None, configurations are loaded from the model.
-        log_stats: Whether to enable statistics logging
-        log_file: Optional path prefix for log files. If provided, logs will
-            be written to files with stage-specific suffixes.
-        init_sleep_seconds: Number of seconds to sleep between starting
-            each stage process during initialization
-        shm_threshold_bytes: Threshold in bytes for using shared memory
-            for IPC. Objects larger than this threshold will use shared memory.
-        batch_timeout: Timeout in seconds for batching requests within a stage
-        init_timeout: Timeout in seconds for waiting for all stages to initialize
+        cli_args: Namespace object containing command-line arguments.
+            Expected attributes include:
+            - stage_configs_path: Optional path to YAML file containing stage
+              configurations. If None, configurations are loaded from the model.
+            - log_stats: Whether to enable statistics logging
+            - log_file: Optional path prefix for log files. If provided, logs will
+              be written to files with stage-specific suffixes.
+            - init_sleep_seconds: Number of seconds to sleep between starting
+              each stage process during initialization
+            - shm_threshold_bytes: Threshold in bytes for using shared memory
+              for IPC. Objects larger than this threshold will use shared memory.
+            - batch_timeout: Timeout in seconds for batching requests within a stage
+            - init_timeout: Timeout in seconds for waiting for all stages to initialize
         **kwargs: Additional keyword arguments passed to stage engines
 
     Example:
-        >>> async_llm = AsyncOmni(model="Qwen/Qwen2.5-Omni-7B")
+        >>> async_llm = AsyncOmni(model="Qwen/Qwen2.5-Omni-7B", cli_args=args)
         >>> async for output in async_llm.generate(
         ...     prompt="Hello",
         ...     request_id="req-1",
@@ -87,33 +90,29 @@ class AsyncOmni(EngineClient):
     def __init__(
         self,
         model: str,
-        stage_configs_path: Optional[str] = None,
-        log_stats: bool = False,
-        log_file: Optional[str] = None,
-        init_sleep_seconds: int = 30,
-        shm_threshold_bytes: int = 65536,
-        batch_timeout: int = 10,
-        init_timeout: int = 60000,
+        cli_args: Namespace,
         **kwargs: Any,
     ):
-        self.batch_timeout = batch_timeout
-        self._enable_stats: bool = bool(log_stats)
+        self.batch_timeout = cli_args.batch_timeout
+        self._enable_stats: bool = bool(cli_args.log_stats)
 
-        if stage_configs_path is None:
-            self.stage_configs = load_stage_configs_from_model(model)
+        base_engine_args = AsyncOmniEngineArgs.from_cli_args(cli_args).__dict__.copy()
+
+        if cli_args.stage_configs_path is None:
+            self.stage_configs = load_stage_configs_from_model(model, base_engine_args)
         else:
-            self.stage_configs = load_stage_configs_from_yaml(stage_configs_path)
+            self.stage_configs = load_stage_configs_from_yaml(cli_args.stage_configs_path, base_engine_args)
 
         self.stage_list: list[OmniStage] = []
         self.default_sampling_params_list: list[SamplingParams] = []
         # Optional file handler for orchestrator
-        self._log_file = log_file
+        self._log_file = cli_args.log_file
         if self._log_file:
             remove_old_logs(self._log_file, len(self.stage_configs))
             configure_orchestrator_logger(logger, self._log_file)
 
         self._stats_file, self._overall_stats_file = init_stats_paths(self._enable_stats, self._log_file)
-        self._initialize_stages(model, init_sleep_seconds, shm_threshold_bytes, init_timeout)
+        self._initialize_stages(model, cli_args.init_sleep_seconds, cli_args.shm_threshold_bytes, cli_args.init_timeout)
 
     def _initialize_stages(
         self,

vllm_omni/entrypoints/cli/serve.py

@@ -60,6 +60,41 @@ class OmniServeCommand(CLISubcommand):
             default=None,
             help="Path to the stage configs file. If not specified, the stage configs will be loaded from the model.",
         )
+        serve_parser.add_argument(
+            "--init-sleep-seconds",
+            type=int,
+            default=30,
+            help="The number of seconds to sleep before initializing the stages.",
+        )
+        serve_parser.add_argument(
+            "--init-timeout",
+            type=int,
+            default=60000,
+            help="The timeout for initializing the stages.",
+        )
+        serve_parser.add_argument(
+            "--shm-threshold-bytes",
+            type=int,
+            default=65536,
+            help="The threshold for the shared memory size.",
+        )
+        serve_parser.add_argument(
+            "--log-stats",
+            action="store_true",
+            help="Enable logging the stats.",
+        )
+        serve_parser.add_argument(
+            "--log-file",
+            type=str,
+            default=None,
+            help="The path to the log file.",
+        )
+        serve_parser.add_argument(
+            "--batch-timeout",
+            type=int,
+            default=10,
+            help="The timeout for the batch.",
+        )
         return serve_parser
 
 

vllm_omni/entrypoints/omni.py

@@ -47,3 +47,14 @@ class Omni:
         if hasattr(self.instance, "generate"):
             return getattr(self.instance, "generate")(*args, **kwargs)
         raise AttributeError(f"'{self.instance.__class__.__name__}' has no attribute 'generate'")
+
+    def close(self) -> None:
+        close_method = getattr(self.instance, "close", None)
+        if callable(close_method):
+            close_method()
+
+    def __del__(self):  # pragma: no cover - best effort cleanup
+        try:
+            self.close()
+        except Exception:
+            pass

vllm_omni/entrypoints/omni_stage.py

@@ -735,7 +735,7 @@ async def _stage_worker_async(
         stage_id,
         list(engine_args.keys()),
     )
-    omni_engine_args = AsyncOmniEngineArgs(model=model, **engine_args)
+    omni_engine_args = AsyncOmniEngineArgs(**engine_args)
     usage_context = UsageContext.OPENAI_API_SERVER
     vllm_config = omni_engine_args.create_engine_config(usage_context=usage_context)
     stage_engine = AsyncOmniStageLLM.from_vllm_config(

vllm_omni/entrypoints/openai/api_server.py

@@ -185,10 +185,7 @@ async def build_async_omni_from_stage_config(
     async_omni: Optional[EngineClient] = None
 
     try:
-        if getattr(args, "stage_configs_path", None):
-            async_omni = AsyncOmni(model=args.model, stage_configs_path=args.stage_configs_path)
-        else:
-            async_omni = AsyncOmni(model=args.model)
+        async_omni = AsyncOmni(model=args.model, cli_args=args)
 
         # # Don't keep the dummy data in memory
         # await async_llm.reset_mm_cache()

vllm_omni/entrypoints/utils.py

@@ -1,5 +1,8 @@
 import os
+from collections import Counter
+from dataclasses import asdict, is_dataclass
 from pathlib import Path
+from typing import Any
 
 from omegaconf import OmegaConf
 from vllm.transformers_utils.config import get_config
@@ -10,7 +13,56 @@ from vllm_omni.utils import detect_device_type
 PROJECT_ROOT = Path(__file__).parent.parent.parent
 
 
-def load_stage_configs_from_model(model: str) -> list:
+def _convert_dataclasses_to_dict(obj: Any) -> Any:
+    """Recursively convert non-serializable objects to OmegaConf-compatible types.
+
+    This is needed because OmegaConf cannot handle:
+    - Dataclass objects with Literal type annotations (e.g., StructuredOutputsConfig)
+    - Counter objects (from collections or vllm.utils)
+    - Set objects
+    - Other non-primitive types
+    """
+    # IMPORTANT: Check Counter BEFORE dict, since Counter is a subclass of dict
+    # Handle Counter objects (convert to dict)
+    # Check by class name first to catch both collections.Counter and vllm.utils.Counter
+    if hasattr(obj, "__class__") and obj.__class__.__name__ == "Counter":
+        try:
+            return dict(obj)
+        except (TypeError, ValueError):
+            # If Counter can't be converted to dict, return empty dict
+            return {}
+    # Also check isinstance for collections.Counter (must be before dict check)
+    if isinstance(obj, Counter):
+        return dict(obj)
+    # Handle set objects (convert to list)
+    if isinstance(obj, set):
+        return list(obj)
+    # Handle dataclass objects
+    # Note: asdict() recursively converts nested dataclasses but not Counter objects,
+    # so we need to recursively process the result
+    if is_dataclass(obj):
+        result = asdict(obj)
+        # Recursively process the result to convert any Counter objects
+        return _convert_dataclasses_to_dict(result)
+    # Handle dictionaries (recurse into values)
+    # Note: This must come AFTER Counter check since Counter is a dict subclass
+    if isinstance(obj, dict):
+        return {k: _convert_dataclasses_to_dict(v) for k, v in obj.items()}
+    # Handle lists and tuples (recurse into items)
+    if isinstance(obj, (list, tuple)):
+        return type(obj)(_convert_dataclasses_to_dict(item) for item in obj)
+    # Try to convert any dict-like object (has keys/values methods) to dict
+    if hasattr(obj, "keys") and hasattr(obj, "values") and not isinstance(obj, (str, bytes)):
+        try:
+            return {k: _convert_dataclasses_to_dict(v) for k, v in obj.items()}
+        except (TypeError, ValueError, AttributeError):
+            # If conversion fails, return as-is
+            return obj
+    # Primitive types and other objects that OmegaConf can handle
+    return obj
+
+
+def load_stage_configs_from_model(model: str, base_engine_args: dict = {}) -> list:
     """Load stage configurations from model's default config file.
 
     Loads stage configurations based on the model type and device type.
@@ -35,7 +87,9 @@ def load_stage_configs_from_model(model: str) -> list:
         device_config_file = f"vllm_omni/model_executor/stage_configs/{device_type}/{model_type}.yaml"
         device_config_path = PROJECT_ROOT / device_config_file
         if os.path.exists(device_config_path):
-            stage_configs = load_stage_configs_from_yaml(config_path=str(device_config_path))
+            stage_configs = load_stage_configs_from_yaml(
+                config_path=str(device_config_path), base_engine_args=base_engine_args
+            )
             return stage_configs
 
     # Fall back to default config
@@ -43,11 +97,11 @@ def load_stage_configs_from_model(model: str) -> list:
     stage_config_path = PROJECT_ROOT / stage_config_file
     if not os.path.exists(stage_config_path):
         raise FileNotFoundError(f"Stage config file {stage_config_path} not found")
-    stage_configs = load_stage_configs_from_yaml(config_path=str(stage_config_path))
+    stage_configs = load_stage_configs_from_yaml(config_path=str(stage_config_path), base_engine_args=base_engine_args)
     return stage_configs
 
 
-def load_stage_configs_from_yaml(config_path: str) -> list:
+def load_stage_configs_from_yaml(config_path: str, base_engine_args: dict = {}) -> list:
     """Load stage configurations from a YAML file.
 
     Args:
@@ -57,4 +111,14 @@ def load_stage_configs_from_yaml(config_path: str) -> list:
         List of stage configuration dictionaries from the file's stage_args
     """
     config_data = OmegaConf.load(config_path)
-    return config_data.stage_args
+    stage_args = config_data.stage_args
+    # Convert any nested dataclass objects to dicts before creating OmegaConf
+    base_engine_args = _convert_dataclasses_to_dict(base_engine_args)
+    base_engine_args = OmegaConf.create(base_engine_args)
+    for stage_arg in stage_args:
+        base_engine_args_tmp = base_engine_args.copy()
+        # Update base_engine_args with stage-specific engine_args if they exist
+        if hasattr(stage_arg, "engine_args") and stage_arg.engine_args is not None:
+            base_engine_args_tmp = OmegaConf.merge(base_engine_args_tmp, stage_arg.engine_args)
+        stage_arg.engine_args = base_engine_args_tmp
+    return stage_args

vllm_omni/model_executor/model_loader/weight_utils.py

@@ -3,10 +3,15 @@ from pathlib import Path
 from typing import Optional, Union
 
 import huggingface_hub
-from huggingface_hub import snapshot_download
+import vllm.envs as envs
 from vllm.logger import init_logger
 from vllm.model_executor.model_loader.weight_utils import DisabledTqdm, get_lock
 
+if envs.VLLM_USE_MODELSCOPE:
+    from modelscope.hub.snapshot_download import snapshot_download
+else:
+    from huggingface_hub import snapshot_download
+
 logger = init_logger(__name__)
 
 
@@ -37,6 +42,7 @@ def download_weights_from_hf_specific(
     """
     assert len(allow_patterns) > 0
     local_only = huggingface_hub.constants.HF_HUB_OFFLINE
+    download_kwargs = {"tqdm_class": DisabledTqdm} if not envs.VLLM_USE_MODELSCOPE else {}
 
     logger.info("Using model weights format %s", allow_patterns)
     # Use file lock to prevent multiple processes from
@@ -49,9 +55,9 @@ def download_weights_from_hf_specific(
                 allow_patterns=allow_pattern,
                 ignore_patterns=ignore_patterns,
                 cache_dir=cache_dir,
-                tqdm_class=DisabledTqdm,
                 revision=revision,
                 local_files_only=local_only,
+                **download_kwargs,
             )
             # If we have downloaded weights for this allow_pattern,
             # we don't need to check the rest.

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -156,6 +156,9 @@ class Qwen3OmniMoeForConditionalGeneration(
                 ]
             )
 
+            # for CI: Initialize special tokens embeddings early to avoid AttributeError when loading dummy weights
+            self._init_special_tokens_embeddings()
+
         elif self.model_stage == "code2wav":
             self.thinker = None
             self.talker = None
@@ -342,9 +345,9 @@ class Qwen3OmniMoeForConditionalGeneration(
                     and thinker_tts_embeds.shape[1] == 3
                 ):
                     bos_eos_pad = thinker_tts_embeds.to(text_hidden_states.device).chunk(3, dim=1)  # 3 * [1,1,H]
-                    multimodal_outputs["tts_bos_embed"] = bos_eos_pad[0]
-                    multimodal_outputs["tts_eos_embed"] = bos_eos_pad[1]
-                    multimodal_outputs["tts_pad_embed"] = bos_eos_pad[2]
+                    multimodal_outputs["tts_bos_embed"] = [bos_eos_pad[0]]
+                    multimodal_outputs["tts_eos_embed"] = [bos_eos_pad[1]]
+                    multimodal_outputs["tts_pad_embed"] = [bos_eos_pad[2]]
             except Exception:
                 # Best-effort; absence will be handled by talker with fallbacks
                 pass

vllm_omni/worker/gpu_ar_model_runner.py

@@ -291,6 +291,11 @@ class GPUARModelRunner(OmniGPUModelRunner):
                         # Case 1: tensor aligned on token dimension
                         if isinstance(v, torch.Tensor) and v.shape[0] == hidden_states_cpu.shape[0]:
                             mm_payload[k] = v.detach().to("cpu")[prev_logits_index : logits_index + 1].contiguous()
+                        elif isinstance(v, torch.Tensor) and v.shape[0] != hidden_states_cpu.shape[0]:
+                            logger.error(
+                                f"Error in merge multimodal outputs: Tensor dimension mismatch, \
+                                          {v.shape} != {hidden_states_cpu.shape} for {k}"
+                            )
                         # Case 2: nested dict of tensors aligned on token dimension (e.g., selected_hidden_layers)
                         elif isinstance(v, dict):
                             sub_dict: dict[str, torch.Tensor] = {}
@@ -302,7 +307,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
                             if sub_dict:
                                 mm_payload[k] = sub_dict
                         elif isinstance(v, list):
-                            element: torch.Tensor = v[0]
+                            element = v[0]
+                            if isinstance(element, torch.Tensor):
+                                element = element.detach().to("cpu").contiguous()
                             multimodal_outputs[k] = v[1:] if len(v) > 1 else v
                             mm_payload[k] = element
                     except Exception as e: