[BugFix] Adapt to pin_memory in vLLM v0.24.0 (#4860)

Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>

.buildkite/pipeline-intel.yaml

@@ -10,7 +10,7 @@ steps:
           DOCKER_BUILDKIT: "1"
           # Buildkite will automatically replace this with the actual commit hash
           VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
-          VLLM_VERSION: "v0.23.0"
+          VLLM_VERSION: "v0.24.0"
         priority: 100
         timeout_in_minutes: 60
         soft_fail: false

.buildkite/test-amd-merge.yml

@@ -7,17 +7,29 @@ steps:
 
 - group: ":card_index_dividers: Simple Test"
   steps:
-    - label: "Simple · Diffusion & Model Executor Test"
+    - label: "Simple · Diffusion & Model Executor Test (1/2)"
       agent_pool: mi325_1
       depends_on: amd-build
-      timeout_in_minutes: 50
+      timeout_in_minutes: 70
       mirror_hardwares: [amdproduction]
       grade: Blocking
       commands:
         - export VLLM_ROCM_USE_AITER=0
         # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
         # that is not supported on CPU
-        - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+        # CPU tests on ROCm gpu machines are very slow, so we give them more time.
+        - "timeout 60m pytest -sv tests/diffusion -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+
+    - label: "Simple · Diffusion & Model Executor Test (2/2)"
+      agent_pool: mi325_1
+      depends_on: amd-build
+      timeout_in_minutes: 70
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      commands:
+        - export VLLM_ROCM_USE_AITER=0
+        # CPU tests on ROCm gpu machines are very slow, so we give them more time.
+        - "timeout 60m pytest -sv tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
 
     - label: "Simple · Other Test"
       agent_pool: mi325_1
@@ -158,7 +170,7 @@ steps:
       grade: Blocking
       commands:
         - |
-          timeout 20m bash -c '
+          timeout 50m bash -c '
             export VLLM_LOGGING_LEVEL=DEBUG
             export VLLM_WORKER_MULTIPROC_METHOD=spawn
             export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
@@ -170,13 +182,14 @@ steps:
       depends_on: amd-build
       mirror_hardwares: [amdproduction]
       grade: Blocking
+      timeout_in_minutes: 50
       commands:
         - |
-          timeout 30m bash -c '
+          timeout 40m bash -c '
             export VLLM_LOGGING_LEVEL=DEBUG
             export VLLM_WORKER_MULTIPROC_METHOD=spawn
             export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-            pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model"
+            pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m 'advanced_model and cuda' --run-level 'advanced_model'
           '
 
 # TODO: Bagel test on ROCm is very unstable. @tjtanaa

.github/CODEOWNERS

@@ -22,30 +22,30 @@
 /.buildkite/ @yenuo26 @congw729
 
 # Engine and model_executor
-/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan
-/vllm_omni/engine/ @tzhouam @yinpeiqi @fake0fan
-/vllm_omni/entrypoints/ @tzhouam @yinpeiqi @fake0fan
+/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan @Sy0307
+/vllm_omni/engine/ @tzhouam @yinpeiqi @fake0fan @Sy0307
+/vllm_omni/entrypoints/ @tzhouam @yinpeiqi @fake0fan @alex-jw-brooks
 /vllm_omni/model_executor/ @tzhouam @gcanlin
 /vllm_omni/worker/ @tzhouam @gcanlin
 
 # Diffusion
-/vllm_omni/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666
+/vllm_omni/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666 @xuechendi
 /vllm_omni/lora/ @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao
 
 # Diffusion CPU offloading
-/vllm_omni/diffusion/offloader/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao
+/vllm_omni/diffusion/offloader/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao @xuechendi
 
 # Omni and TTS support
 /vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao
-/vllm_omni/model_executor/models/ @linyueqian @ZeldaHuang @princepride @yuanheng-zhao @amy-why-3459 @Bounty-hunter
+/vllm_omni/model_executor/models/ @linyueqian @ZeldaHuang @princepride @yuanheng-zhao @amy-why-3459 @Bounty-hunter @Sy0307 @gcanlin
 
 # Benchmarks
 /benchmarks/accuracy/ @david6666666
 /benchmarks/build_dataset/ @linyueqian @yuanheng-zhao
-/benchmarks/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa
+/benchmarks/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @xuechendi
 /benchmarks/distributed/ @princepride
 /benchmarks/glm_image/ @princepride @JaredforReal
-/benchmarks/tts/ @linyueqian @ZeldaHuang @princepride
+/benchmarks/tts/ @linyueqian @ZeldaHuang @princepride @Sy0307 @gcanlin
 
 # Usability and Observability
 /vllm_omni/deploy/ @tzhouam @lishunyang12
@@ -53,13 +53,13 @@
 /vllm_omni/profiler/ @gcanlin
 
 # ComfyUI
-/apps/ @fhfuih
+/apps/ @fhfuih @alex-jw-brooks
 
 # Large Scale Serving
 /vllm_omni/distributed/ @princepride @yuanheng-zhao @xuechendi @natureofnature @fake0fan
 
 # Quantization and configuration
-/vllm_omni/config/ @lishunyang12
+/vllm_omni/config/ @lishunyang12 @alex-jw-brooks
 /vllm_omni/quantization/ @david6666666 @Isotr0py @lishunyang12
 
 # Hardware plugins and accelerators

docker/Dockerfile.npu

@@ -1,5 +1,5 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.23.0-4885d8e
+ARG VLLM_ASCEND_TAG=v0.24.0
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 # WORKDIR /vllm-workspace/vllm

docker/Dockerfile.npu.a3

@@ -1,5 +1,5 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.23.0-4885d8e-a3
+ARG VLLM_ASCEND_TAG=v0.24.0-a3
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 # WORKDIR /vllm-workspace/vllm

docker/Dockerfile.npu.ci

@@ -1,5 +1,5 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.23.0-4885d8e
+ARG VLLM_ASCEND_TAG=v0.24.0
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 ARG APP_DIR=/vllm-workspace/vllm-omni

docker/Dockerfile.npu.ci.a3

@@ -1,5 +1,5 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.23.0-4885d8e-a3
+ARG VLLM_ASCEND_TAG=v0.24.0-a3
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 ARG APP_DIR=/vllm-workspace/vllm-omni

docker/Dockerfile.xpu

@@ -107,7 +107,9 @@ ENV VLLM_TARGET_DEVICE=xpu
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
 
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --no-build-isolation .
+    uv pip uninstall triton triton-xpu && \
+    uv pip install triton-xpu==3.7.1 && \
+    uv pip install --no-build-isolation --no-deps .
 
 CMD ["/bin/bash"]
 
@@ -143,12 +145,11 @@ COPY . .
 
 ENV VLLM_OMNI_TARGET_DEVICE=xpu
 RUN --mount=type=cache,target=/root/.cache/uv \
-    echo "triton @ https://example.com/blocked" > /tmp/vllm_omni_constraints.txt && \
-    PIP_CONSTRAINT=/tmp/vllm_omni_constraints.txt \
     uv pip install --no-build-isolation \
     --extra-index-url https://mirrors.ustc.edu.cn/pypi/simple/ \
     ".[dev]" && \
-    uv pip uninstall oneccl oneccl-devel
+    uv pip uninstall triton oneccl oneccl-devel && \
+    uv pip install triton-xpu==3.7.1 --reinstall
 
 FROM vllm-omni AS vllm-omni-openai
 

docs/community/governance.md

@@ -33,16 +33,20 @@ Lead maintainers are responsible for the overall direction and strategy of the p
 
 Committers have write access and merge rights. They typically have deep expertise in specific areas of this project and shepherd the community contributions:
 
+- [@alex-jw-brooks](https://github.com/alex-jw-brooks): Frontend, Configuration, and Entrypoint Support
+- [@amy-why-3459](https://github.com/amy-why-3459): Omni Modality Model and World Model Support
 - [@david6666666](https://github.com/david6666666): Quantization and Community Relationship
-- [@gcanlin](https://github.com/gcanlin): Hardware plugin and NPU integration
+- [@gcanlin](https://github.com/gcanlin): Hardware plugin, NPU integration, and TTS Model Support
 - [@Isotr0py](https://github.com/Isotr0py): Diffusion and Quantization
 - [@linyueqian](https://github.com/linyueqian): TTS and Omni Support
 - [@lishunyang12](https://github.com/lishunyang12): Quantization and Configuration
 - [@princepride](https://github.com/princepride): Diffusion and Omni Support
 - [@RuixiangMa](https://github.com/RuixiangMa): Diffusion models, parallel, cache, and docs
 - [@SamitHuang](https://github.com/SamitHuang): RL, Diffusion, and cache
+- [@Sy0307](https://github.com/Sy0307): TTS Model and AR Engine Support
 - [@tzhouam](https://github.com/tzhouam): Engine and New Model Support
 - [@wtomin](https://github.com/wtomin): Diffusion models, parallel, and docs
+- [@xuechendi](https://github.com/xuechendi): Diffusion and Intel GPU Support
 - [@ZeldaHuang](https://github.com/ZeldaHuang): Omni Support
 - [@ZJY0516](https://github.com/ZJY0516): Diffusion attention backend, kernel fusion, and CustomOp
 - [@yuanheng-zhao](https://github.com/yuanheng-zhao): Diffusion cache, offload, and Omni Support
@@ -53,7 +57,7 @@ This section breaks down diffusion-related responsibilities. If you have PRs tou
 
 ### Runtime-related
 
-- Diffusion models: @RuixiangMa, @wtomin
+- Diffusion models: @RuixiangMa, @wtomin, @xuechendi
 
 ### Optimization
 
@@ -61,6 +65,7 @@ This section breaks down diffusion-related responsibilities. If you have PRs tou
 - Attention backend: @ZJY0516
 - Cache: @yuanheng-zhao, @RuixiangMa, @SamitHuang
 - Offload: @yuanheng-zhao
+- Intel GPU support: @xuechendi
 - Quantization: @lishunyang12, @david6666666
 - Kernel fusion / communication-computation: @ZJY0516
 

docs/contributing/model/adding_diffusion_model.md

@@ -382,7 +382,7 @@ def forward(
         for p in req.prompts
     ]
 
-    # Extract sampling parameters
+    # Extract common sampling parameters
     sampling_params = req.sampling_params
     num_inference_steps = sampling_params.num_inference_steps or 50
     guidance_scale = sampling_params.guidance_scale or 7.5

docs/design/feature/cfg_parallel.md

@@ -126,16 +126,16 @@ Call `self.diffuse` in your pipeline's forward function:
 
 ```python
 import torch.nn as nn
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+
 class YourModelPipeline(nn.Module, CFGParallelMixin):
     def forward(
         self,
-        prompt: str,
-        negative_prompt: str | None = None,
-        guidance_scale: float = 3.5,
-        num_inference_steps: int = 50,
-        **kwargs,
-    ):
-        # Encode prompts, Initialize latents, Get timesteps
+        req: DiffusionRequestBatch,
+    ) -> list[DiffusionOutput]:
+        # Read prompts and sampling params from the request
+        ...
+        # Encode prompts, initialize latents, get timesteps
         ...
         # Run diffusion loop (calls the mixin's diffuse method)
         latents = self.diffuse(

docs/features/custom_pipeline.md

@@ -54,20 +54,24 @@ Create a custom pipeline class that extends an existing pipeline. In this exampl
 
 ```python
 # custom_pipeline.py
-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit import QwenImageEditPipeline
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 import torch
 
 class CustomPipeline(QwenImageEditPipeline):
     def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
         super().__init__(od_config=od_config, prefix=prefix)
 
-    def forward(self, req, prompt=None, negative_prompt=None, **kwargs):
+    def forward(self, req: DiffusionRequestBatch):
+        # Optionally customize sampling parameters
+        actual_num_steps = req.sampling_params.num_inference_steps or 50
+        req.sampling_params.num_inference_steps = actual_num_steps
+
         # Call parent's forward to get normal output
-        output = super().forward(req=req, prompt=prompt, negative_prompt=negative_prompt, **kwargs)
+        output = super().forward(req=req)
 
         # Add custom trajectory data
-        actual_num_steps = req.sampling_params.num_inference_steps or kwargs.get('num_inference_steps', 50)
         output.trajectory_timesteps = torch.linspace(1000, 0, actual_num_steps, dtype=torch.float32)
         output.trajectory_latents = torch.randn(actual_num_steps, 1, 16, 64, 64, dtype=torch.float32)
 

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

@@ -1,19 +1,16 @@
 # --8<-- [start:requirements]
 
-For detailed hardware and software requirements, please refer to the [vllm-ascend installation documentation](https://docs.vllm.ai/projects/ascend/en/latest/installation.html).
+For detailed hardware and software requirements, please refer to the [vLLM-Ascend installation documentation](https://docs.vllm.ai/projects/ascend/en/latest/installation.html).
 
 # --8<-- [end:requirements]
 # --8<-- [start:installation-release]
 
-The recommended way to use vLLM-Omni on NPU is through the vllm-ascend pre-built Docker images:
+The recommended way to use vLLM-Omni on NPU is through the vLLM-Ascend pre-built Docker images:
 
 ```bash
-# Update the vllm-ascend image
-# Atlas A2:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.18.0rc1
-# Atlas A3:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.18.0rc1-a3
-export IMAGE=quay.io/ascend/vllm-ascend:v0.18.0rc1
+# vLLM-Ascend image
+# Atlas A2: quay.io/atlas-ci/vllm-ascend:v0.24.0
+# Atlas A3: quay.io/atlas-ci/vllm-ascend:v0.24.0-a3
 docker run --rm \
     --name vllm-omni-npu \
     --shm-size=1g \
@@ -31,11 +28,11 @@ docker run --rm \
     -v /etc/ascend_install.info:/etc/ascend_install.info \
     -v /root/.cache:/root/.cache \
     -p 8000:8000 \
-    -it $IMAGE bash
+    -it quay.io/atlas-ci/vllm-ascend:0.24.0-a3 bash
 
 # Inside the container, install vLLM-Omni from source
 cd /vllm-workspace
-git clone -b v0.18.0 https://github.com/vllm-project/vllm-omni.git
+git clone -b v0.24.0 https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
 pip install -v -e . --no-build-isolation
 # or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
@@ -87,14 +84,11 @@ Supported images as following.
 Here's an example deployment command that has been verified on 4 x NPUs:
 
 ```bash
-# Atlas A2:
-# export IMAGE=quay.io/ascend/vllm-omni:v0.18.0
-# Atlas A3:
-# export IMAGE=quay.io/ascend/vllm-omni:v0.18.0-a3
-export IMAGE=quay.io/ascend/vllm-omni:v0.18.0
+# Atlas A2: quay.io/ascend/vllm-omni:v0.24.0
+# Atlas A3: quay.io/ascend/vllm-omni:v0.24.0-a3
 docker run --rm \
-    --name vllm-omni-npu \
-    --shm-size=1g \
+    --name vllm-omni-a3 \
+    --shm-size=4g \
     --device /dev/davinci0 \
     --device /dev/davinci1 \
     --device /dev/davinci2 \
@@ -106,10 +100,17 @@ docker run --rm \
     -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
     -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
     -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
+    --device /dev/davinci_manager \
+    --device /dev/devmm_svm \
+    --device /dev/hisi_hdc \
+    -v /usr/local/dcmi:/usr/local/dcmi \
+    -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
+    -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
+    -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
     -v /etc/ascend_install.info:/etc/ascend_install.info \
-    -v /root/.cache:/root/.cache \
-    -p 8000:8000 \
-    -it $IMAGE bash
+    -v ~/.cache:/root/.cache \
+    -p 8091:8091 \
+    -it quay.io/ascend/vllm-omni:v0.24.0-a3 bash
 ```
 
 !!! tip

docs/user_guide/examples/offline_inference/image_to_video.md

@@ -3,7 +3,9 @@
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video>.
 
 
-This example demonstrates how to generate videos from images using Wan2.2 Image-to-Video models with vLLM-Omni's offline inference API.
+This example demonstrates how to generate videos from images using Wan2.2,
+LTX-2/LTX-2.3, and HunyuanVideo-1.5 Image-to-Video models with vLLM-Omni's
+offline inference API.
 
 ## Local CLI Usage
 
@@ -51,13 +53,41 @@ python image_to_video.py \
   --output i2v_output.mp4
 ```
 
+### LTX-2.3
+
+```bash
+python image_to_video.py \
+  --model dg845/LTX-2.3-Diffusers \
+  --model-class-name LTX23ImageToVideoPipeline \
+  --image cherry_blossom.jpg \
+  --prompt "Cherry blossoms swaying gently in the breeze with synchronized ambient sound" \
+  --negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \
+  --height 384 \
+  --width 512 \
+  --num-frames 25 \
+  --guidance-scale 4.0 \
+  --num-inference-steps 20 \
+  --frame-rate 24 \
+  --fps 24 \
+  --output ltx23_i2v_output.mp4
+```
+
+Use a Diffusers-format checkpoint such as `dg845/LTX-2.3-Diffusers`; the
+upstream `Lightricks/LTX-2.3` raw safetensors repo is not directly loadable by
+this pipeline. Pass `--model-class-name LTX23ImageToVideoPipeline` to select
+the LTX-2.3 image-to-video pipeline.
+
 Key arguments:
 
-- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V).
+- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V, or
+  LTX-2/LTX-2.3).
 - `--image`: Path to input image (required).
 - `--prompt`: Text description of desired motion/animation.
-- `--height/--width`: Output resolution (auto-calculated from image if not set). Dimensions should be multiples of 16.
-- `--num-frames`: Number of frames (default 81).
+- `--height/--width`: Output resolution (auto-calculated from image if not set).
+  Wan dimensions should be multiples of 16; LTX dimensions should be multiples
+  of 32.
+- `--num-frames`: Number of frames (model-specific default; LTX-style models
+  work best with `8k + 1`).
 - `--guidance-scale` and `--guidance-scale-high`: CFG scale (applied to low/high-noise stages for MoE).
 - `--negative-prompt`: Optional list of artifacts to suppress.
 - `--boundary-ratio`: Boundary split ratio for two-stage MoE models.
@@ -65,6 +95,7 @@ Key arguments:
 - `--sample-solver`: Wan2.2 sampling solver. Use `unipc` for the default multistep solver, or `euler` for Lightning/Distill checkpoints.
 - `--num-inference-steps`: Number of denoising steps (default 50).
 - `--fps`: Frames per second for the saved MP4 (requires `diffusers` export_to_video).
+- `--audio-sample-rate`: fallback audio sample rate for embedded audio.
 - `--output`: Path to save the generated video.
 - `--vae-use-slicing`: Enable VAE slicing for memory optimization.
 - `--vae-use-tiling`: Enable VAE tiling for memory optimization.

docs/user_guide/examples/offline_inference/text_to_video.md

@@ -12,6 +12,7 @@ For backend selection and SageAttention usage, see the [Diffusion Attention Back
 | Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM (BF16) |
 |---|---|---|---|---|---|
 | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
+| `dg845/LTX-2.3-Diffusers` | 384x512 | 25 | 20 | 4.0 | 96GB-class GPU |
 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB |
 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
 
@@ -50,6 +51,30 @@ python text_to_video.py \
   --output ltx2_out.mp4
 ```
 
+### LTX-2.3
+
+```bash
+python text_to_video.py \
+  --model dg845/LTX-2.3-Diffusers \
+  --model-class-name LTX23Pipeline \
+  --prompt "Cherry blossoms swaying gently in the breeze with synchronized ambient sound" \
+  --negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \
+  --height 384 \
+  --width 512 \
+  --num-frames 25 \
+  --num-inference-steps 20 \
+  --guidance-scale 4.0 \
+  --frame-rate 24 \
+  --fps 24 \
+  --audio-sample-rate 48000 \
+  --output ltx23_t2v_output.mp4
+```
+
+Use a Diffusers-format checkpoint such as `dg845/LTX-2.3-Diffusers`; the
+upstream `Lightricks/LTX-2.3` raw safetensors repo is not directly loadable by
+this pipeline. Pass `--model-class-name LTX23Pipeline` to select the LTX-2.3
+text-to-video pipeline explicitly.
+
 ### HunyuanVideo-1.5 (480p)
 
 ```bash
@@ -110,6 +135,8 @@ python text_to_video.py \
 ### Common
 
 - `--model`: Diffusers model ID or local path.
+- `--model-class-name`: Optional explicit pipeline class. Use `LTX23Pipeline`
+  for LTX-2.3 text-to-video.
 - `--prompt`: text description (string).
 - `--height/--width`: output resolution. Default depends on model.
 - `--num-frames`: number of frames. Default depends on model.
@@ -124,7 +151,8 @@ python text_to_video.py \
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
 - `--enable-layerwise-offload`: enable layerwise offloading on DiT modules.
 - `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2).
-- `--audio-sample-rate`: audio sample rate for embedded audio (when the pipeline returns audio).
+- `--audio-sample-rate`: audio sample rate for embedded audio (when the
+  pipeline returns audio; LTX-2.3 outputs 48kHz audio).
 - `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF).
 - `--flow-shift`: scheduler flow_shift parameter.
 

docs/user_guide/examples/online_serving/diffusers_pipeline_adapter.md

@@ -9,7 +9,8 @@ vLLM-Omni supports running diffusion models with the diffusers backend, directly
 
 The diffusers backend is a black-box adapter. Its primary focus is to serve diffusion models online.
 Currently, the following features are NOT yet supported.
-It is not guaranteed whether they will be supported in the future.
+Community contributions are welcome to add these features when they can be
+implemented with clear Diffusers or vLLM-Omni behavior.
 
 - CFG parallel execution
 - Sequence parallel execution
@@ -66,6 +67,35 @@ This is suitable for sampling parameters not available through the vLLM-Omni int
 When a parameter is available in the vLLM-Omni interface, it will be adapted here.
 But if that parameter is simultaneously set in both the vLLM-Omni interface and `diffusers_call_kwargs`, the **former** will take precedence (because it is set at request time).
 
+### Quantization
+
+Use `diffusers_load_kwargs` for Diffusers-native quantization options. If
+`diffusers_load_kwargs["quantization_config"]` is provided as a dictionary, the
+diffusers backend builds a Diffusers quantization config and lets Diffusers
+validate it before calling `DiffusionPipeline.from_pretrained()`.
+
+The diffusers backend also provides a small compatibility shortcut for
+vLLM-Omni quantization configs when `diffusers_load_kwargs` does not already
+contain `quantization_config`. Currently this maps online/dynamic `fp8` and
+online/dynamic `int8` to Diffusers/TorchAO dynamic quantization for transformer
+components such as `transformer` or `transformer_2`. This path requires
+`torchao` to be installed.
+
+For example, the CLI can request dynamic FP8 through the vLLM-Omni interface:
+
+```bash
+vllm serve "Qwen/Qwen-Image" \
+    --omni \
+    --diffusion-load-format diffusers \
+    --quantization-config '{"method": "fp8"}'
+```
+
+Other vLLM-Omni quantization methods, such as `gguf`, `modelopt`, `mxfp4`,
+`mxfp8`, serialized checkpoints, static FP8 configs, and layer-name skip lists
+such as `ignored_layers`, are intentionally not translated by this compatibility
+shortcut. Use Diffusers-native configuration through `diffusers_load_kwargs` or
+a native vLLM-Omni pipeline for those cases.
+
 ### Attention Backends
 
 The diffusers backend converts

docs/user_guide/quantization/fp8.md

@@ -44,10 +44,16 @@ pip install vllm-omni[quack]
 pip install 'quack-kernels[cu13]' --extra-index-url https://download.pytorch.org/whl/cu130
 ```
 
-It is enabled automatically once installed (no flag needed) and is **Blackwell-only**:
-on Hopper/Ada the CUTLASS FP8 kernel already fuses bias, so quack is not used there.
-Set `VLLM_OMNI_USE_QUACK_FP8=0` to force the FlashInfer path. If `quack-kernels` is
-not installed, FP8 still works — it just keeps the unfused FlashInfer path.
+It is enabled automatically once installed (no flag needed) and is **datacenter
+Blackwell only** (`sm_100` / `sm_101` / `sm_103`, compute capability `10.x`, e.g.
+B200): quack's CuteDSL GEMM uses the 5th-gen `tcgen05` tensor-core MMA, which exists
+only on those parts. On Hopper/Ada the CUTLASS FP8 kernel already fuses bias, and on
+workstation/consumer Blackwell (`sm_120` / `sm_121`, compute capability `12.x`, e.g.
+RTX PRO 6000 / RTX 50-series) `tcgen05` is absent — so quack is **not** auto-enabled
+there and FlashInfer's native FP8 path is used instead. Set
+`VLLM_OMNI_USE_QUACK_FP8=1` to force quack on, or `VLLM_OMNI_USE_QUACK_FP8=0` to force
+the FlashInfer path. If `quack-kernels` is not installed, FP8 still works — it just
+keeps the unfused FlashInfer path.
 
 #### Compile cache and warmup
 

docs/user_guide/quantization/mxfp4.md

@@ -39,7 +39,7 @@ vLLM-Omni provides two quantization methods with different scale structures:
 | NVIDIA Blackwell GPU (SM 100+) | ⭕ |
 | NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ |
 | NVIDIA Ampere GPU (SM 80+) | ⭕ |
-| AMD ROCm | ⭕ |
+| AMD ROCm (gfx950 / MI355X) | ✅ |
 | Intel XPU | ⭕ |
 | Ascend NPU (Atlas 950 A5) | ✅ |
 

docs/user_guide/quantization/overview.md

@@ -20,7 +20,7 @@ type has a different quantization scope.
 | NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
 | NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
 | NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
-| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ❌ |
+| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ✅ | ⭕ | ⭕ | ❌ |
 | Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
 | Ascend NPU | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ |
 

examples/offline_inference/custom_pipeline/image_to_image/custom_pipeline.py

@@ -2,14 +2,12 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import logging
-from typing import Any
 
-import PIL.Image
 import torch
 
-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit import QwenImageEditPipeline
-from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 
 logger = logging.getLogger(__name__)
 
@@ -18,58 +16,14 @@ class CustomPipeline(QwenImageEditPipeline):
     def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
         super().__init__(od_config=od_config, prefix=prefix)
 
-    def forward(
-        self,
-        req: OmniDiffusionRequest,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        image: PIL.Image.Image | torch.Tensor | None = None,
-        true_cfg_scale: float = 4.0,
-        height: int | None = None,
-        width: int | None = None,
-        num_inference_steps: int = 50,
-        sigmas: list[float] | None = None,
-        guidance_scale: float = 1.0,
-        num_images_per_prompt: int = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.Tensor | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        prompt_embeds_mask: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds_mask: torch.Tensor | None = None,
-        output_type: str | None = "pil",
-        attention_kwargs: dict[str, Any] | None = None,
-        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
-        max_sequence_length: int = 512,
-    ) -> DiffusionOutput:
+    def forward(self, req: DiffusionRequestBatch):
         """Forward pass for image editing with dummy trajectory data."""
-        # Call parent's forward to get the normal output
-        output = super().forward(
-            req=req,
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            image=image,
-            true_cfg_scale=true_cfg_scale,
-            height=height,
-            width=width,
-            num_inference_steps=num_inference_steps,
-            sigmas=sigmas,
-            guidance_scale=guidance_scale,
-            num_images_per_prompt=num_images_per_prompt,
-            generator=generator,
-            latents=latents,
-            prompt_embeds=prompt_embeds,
-            prompt_embeds_mask=prompt_embeds_mask,
-            negative_prompt_embeds=negative_prompt_embeds,
-            negative_prompt_embeds_mask=negative_prompt_embeds_mask,
-            output_type=output_type,
-            attention_kwargs=attention_kwargs,
-            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
-            max_sequence_length=max_sequence_length,
-        )
+        # Customize num_inference_steps that best suits your model
+        actual_num_steps = req.sampling_params.num_inference_steps or 50
+        req.sampling_params.num_inference_steps = actual_num_steps
 
-        # Get actual num_inference_steps used
-        actual_num_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        # Call parent's forward to get the normal output
+        output = super().forward(req=req)
 
         # Create dummy trajectory data
         dummy_trajectory_latents = torch.randn(actual_num_steps, 1, 16, 64, 64, dtype=torch.float32)

examples/offline_inference/image_to_video/README.md

@@ -1,7 +1,7 @@
 # Image-To-Video
 
-This shared example generates videos from images with VACE, Wan2.2, LTX2,
-HunyuanVideo-1.5, Cosmos3, and other compatible pipelines.
+This shared example generates videos from images with VACE, Wan2.2,
+LTX-2/LTX-2.3, HunyuanVideo-1.5, Cosmos3, and other compatible pipelines.
 
 ## Local CLI Usage
 
@@ -49,6 +49,30 @@ python image_to_video.py \
   --output i2v_output.mp4
 ```
 
+### LTX-2.3
+
+```bash
+python image_to_video.py \
+  --model dg845/LTX-2.3-Diffusers \
+  --model-class-name LTX23ImageToVideoPipeline \
+  --image cherry_blossom.jpg \
+  --prompt "Cherry blossoms swaying gently in the breeze with synchronized ambient sound" \
+  --negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \
+  --height 384 \
+  --width 512 \
+  --num-frames 25 \
+  --guidance-scale 4.0 \
+  --num-inference-steps 20 \
+  --frame-rate 24 \
+  --fps 24 \
+  --output ltx23_i2v_output.mp4
+```
+
+Use a Diffusers-format checkpoint such as `dg845/LTX-2.3-Diffusers`; the
+upstream `Lightricks/LTX-2.3` raw safetensors repo is not directly loadable by
+this pipeline. Pass `--model-class-name LTX23ImageToVideoPipeline` to select
+the LTX-2.3 image-to-video pipeline.
+
 ### Cosmos3
 
 ```bash
@@ -66,15 +90,19 @@ python image_to_video.py \
 
 Key arguments:
 
-- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V).
+- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V,
+  LTX-2/LTX-2.3, Cosmos3, or VACE).
 - `--image`: Path to the first-frame or source image.
 - `--last-image`: Optional last-frame condition for models such as VACE.
 - `--mask-image`: Optional inpainting mask. White pixels are regenerated and black pixels are preserved.
 - `--reference-image`: Optional reference image; repeat it to provide multiple references.
 - `--extra-body`: JSON object of model-specific generation params, filtered against the model's declared `extra_body_params` (see [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)). Used by Cosmos3.
 - `--prompt`: Text description of desired motion/animation.
-- `--height/--width`: Output resolution (auto-calculated from image if not set). Dimensions should be multiples of 16.
-- `--num-frames`: Number of frames (default: model-specific — Wan 81, LTX2 121, Cosmos3 189).
+- `--height/--width`: Output resolution (auto-calculated from image if not set).
+  Wan dimensions should be multiples of 16; LTX dimensions should be multiples
+  of 32.
+- `--num-frames`: Number of frames (model-specific default; LTX-style models
+  work best with `8k + 1`; Cosmos3 defaults to 189).
 - `--guidance-scale` and `--guidance-scale-high`: CFG scale (applied to low/high-noise stages for MoE).
 - `--negative-prompt`: Optional list of artifacts to suppress.
 - `--boundary-ratio`: Boundary split ratio for two-stage MoE models.
@@ -82,6 +110,7 @@ Key arguments:
 - `--sample-solver`: Wan2.2 sampling solver. Use `unipc` for the default multistep solver, or `euler` for Lightning/Distill checkpoints.
 - `--num-inference-steps`: Number of denoising steps (default: model-specific — Wan 50, LTX2 40, Cosmos3 35).
 - `--fps`: Frames per second for the saved MP4 (requires `diffusers` export_to_video).
+- `--audio-sample-rate`: fallback audio sample rate for embedded audio.
 - `--output`: Path to save the generated video.
 - `--vae-use-slicing`: Enable VAE slicing for memory optimization.
 - `--vae-use-tiling`: Enable VAE tiling for memory optimization.

examples/offline_inference/image_to_video/image_to_video.py

@@ -2,7 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 """
-Image-to-Video generation example using Wan2.2 I2V/TI2V models, LTX2, HunyuanVideo-1.5, or Wan2.1 VACE.
+Image-to-Video generation example using Wan2.2 I2V/TI2V models, LTX2/LTX-2.3,
+HunyuanVideo-1.5, Cosmos3, or Wan2.1 VACE.
 
 Supports:
 - Wan2.2-I2V-A14B-Diffusers: MoE model with CLIP image encoder
@@ -27,6 +28,13 @@ Usage:
         --num-frames 121 --num-inference-steps 40 --guidance-scale 4.0 \
         --frame-rate 24 --fps 24 --output ltx2_i2v.mp4
 
+    # LTX-2.3 image-to-video
+    python image_to_video.py --model dg845/LTX-2.3-Diffusers \
+        --model-class-name LTX23ImageToVideoPipeline \
+        --image input.jpg --prompt "A cinematic dolly shot of a boat" \
+        --height 384 --width 512 --num-frames 25 --num-inference-steps 20 \
+        --frame-rate 24 --fps 24 --output ltx23_i2v.mp4
+
     # HunyuanVideo-1.5 I2V (480p)
     python image_to_video.py --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v \
         --image input.jpg --prompt "A cat playing with yarn" \
@@ -79,17 +87,20 @@ parse_profiler_config = functools.partial(parse_json_object, flag_name="--profil
 
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(
-        description="Generate a video from one or more images (Wan2.2, LTX2, HunyuanVideo-1.5, or Wan2.1 VACE)."
+        description=(
+            "Generate a video from one or more images "
+            "(Wan2.2, LTX2/LTX-2.3, HunyuanVideo-1.5, Cosmos3, or Wan2.1 VACE)."
+        )
     )
     parser.add_argument(
         "--model",
         default="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
-        help="Diffusers I2V model ID or local path (Wan2.2, HunyuanVideo-1.5, or Wan2.1 VACE).",
+        help="Diffusers I2V model ID or local path (Wan2.2, LTX2/LTX-2.3, HunyuanVideo-1.5, Cosmos3, or Wan2.1 VACE).",
     )
     parser.add_argument(
         "--model-class-name",
         default=None,
-        help="Override model class name (e.g., LTX2ImageToVideoPipeline).",
+        help="Override model class name (e.g., LTX2ImageToVideoPipeline or LTX23ImageToVideoPipeline).",
     )
     parser.add_argument("--image", help="Path to the first-frame or source image.")
     parser.add_argument("--last-image", help="Path to a last-frame condition (used by models such as VACE).")
@@ -113,7 +124,10 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--width", type=int, default=None, help="Video width (auto-calculated from image if not set).")
     parser.add_argument("--num-frames", type=int, default=None, help="Number of frames. Default: model-specific.")
     parser.add_argument(
-        "--num-inference-steps", type=int, default=None, help="Sampling steps. Default: model-specific."
+        "--num-inference-steps",
+        type=int,
+        default=None,
+        help="Sampling steps. Default: model-specific.",
     )
     parser.add_argument("--boundary-ratio", type=float, default=0.875, help="Boundary split ratio for MoE models.")
     parser.add_argument(
@@ -312,10 +326,8 @@ def main():
     generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
     model_name = str(args.model).lower() if args.model is not None else ""
     model_class_name = args.model_class_name
-    is_ltx2 = "ltx2" in model_name or (model_class_name and "ltx2" in model_class_name.lower())
+    is_ltx2 = model_class_name in {"LTX2ImageToVideoPipeline", "LTX23ImageToVideoPipeline"}
     is_cosmos = "cosmos" in model_name or (model_class_name is not None and "cosmos" in model_class_name.lower())
-    if model_class_name is None and is_ltx2:
-        model_class_name = "LTX2ImageToVideoPipeline"
 
     image = PIL.Image.open(args.image).convert("RGB") if args.image else None
     last_image = PIL.Image.open(args.last_image).convert("RGB") if args.last_image else None
@@ -340,7 +352,15 @@ def main():
             16,
         )
     elif is_ltx2:
-        d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = 24, 4.0, 121, 40, 5.0, 512 * 768, 32
+        d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = (
+            24,
+            4.0,
+            121,
+            40,
+            5.0,
+            512 * 768,
+            32,
+        )
     else:  # Wan2.2 / HunyuanVideo-1.5
         d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = 16, 5.0, 81, 50, 5.0, 480 * 832, 16
 
@@ -462,7 +482,7 @@ def main():
         f" tensor_parallel_size={args.tensor_parallel_size}, vae_patch_parallel_size={args.vae_patch_parallel_size},"
         f" pipeline_parallel_size={args.pipeline_parallel_size}"
     )
-    print(f"  Video size: {args.width}x{args.height}")
+    print(f"  Video size: {width}x{height}")
     print(f"{'=' * 60}\n")
 
     prompt_dict = build_image_to_video_prompt(
@@ -509,6 +529,7 @@ def main():
     print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
 
     audio = None
+    audio_sample_rate = args.audio_sample_rate
     if isinstance(frames, list):
         frames = frames[0] if frames else None
 
@@ -519,11 +540,13 @@ def main():
             )
         if frames.multimodal_output and "audio" in frames.multimodal_output:
             audio = frames.multimodal_output["audio"]
+            audio_sample_rate = frames.multimodal_output.get("audio_sample_rate", audio_sample_rate)
         if frames.is_pipeline_output and frames.request_output is not None:
             inner_output = frames.request_output
             if isinstance(inner_output, OmniRequestOutput):
                 if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
                     audio = inner_output.multimodal_output["audio"]
+                    audio_sample_rate = inner_output.multimodal_output.get("audio_sample_rate", audio_sample_rate)
                 frames = inner_output
         if isinstance(frames, OmniRequestOutput):
             if frames.images:
@@ -531,6 +554,7 @@ def main():
                     frames, audio = frames.images[0]
                 elif len(frames.images) == 1 and isinstance(frames.images[0], dict):
                     audio = frames.images[0].get("audio")
+                    audio_sample_rate = frames.images[0].get("audio_sample_rate", audio_sample_rate)
                     frames = frames.images[0].get("frames") or frames.images[0].get("video")
                 else:
                     frames = frames.images
@@ -543,6 +567,7 @@ def main():
             frames, audio = first_item
         elif isinstance(first_item, dict):
             audio = first_item.get("audio")
+            audio_sample_rate = first_item.get("audio_sample_rate", audio_sample_rate)
             frames = first_item.get("frames") or first_item.get("video")
         elif isinstance(first_item, list):
             frames = first_item
@@ -551,6 +576,7 @@ def main():
         frames, audio = frames
     elif isinstance(frames, dict):
         audio = frames.get("audio")
+        audio_sample_rate = frames.get("audio_sample_rate", audio_sample_rate)
         frames = frames.get("frames") or frames.get("video")
 
     if frames is None:
@@ -669,7 +695,7 @@ def main():
             frames_u8,
             audio_np,
             fps=float(fps),
-            audio_sample_rate=args.audio_sample_rate,
+            audio_sample_rate=audio_sample_rate,
         )
         with open(str(output_path), "wb") as f:
             f.write(video_bytes)

examples/offline_inference/text_to_video/text_to_video.md

@@ -8,6 +8,7 @@ A unified script for text-to-video generation. Supports multiple models with mod
 |---|---|---|---|---|---|
 | `Wan-AI/Wan2.1-VACE-1.3B-diffusers` | 480x832 | 81 | 30 | 5.0 | ~20 GiB (RTX 5090, VAE tiling) |
 | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
+| `dg845/LTX-2.3-Diffusers` | 384x512 | 25 | 20 | 4.0 | 96GB-class GPU |
 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB |
 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
 | `nvidia/Cosmos3-Nano` | 720x1280 | 189 | 35 | 6.0 | ~46 GiB (peak, 720p) |
@@ -70,6 +71,30 @@ python text_to_video.py \
   --output ltx2_out.mp4
 ```
 
+### LTX-2.3
+
+```bash
+python text_to_video.py \
+  --model dg845/LTX-2.3-Diffusers \
+  --model-class-name LTX23Pipeline \
+  --prompt "Cherry blossoms swaying gently in the breeze with synchronized ambient sound" \
+  --negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \
+  --height 384 \
+  --width 512 \
+  --num-frames 25 \
+  --num-inference-steps 20 \
+  --guidance-scale 4.0 \
+  --frame-rate 24 \
+  --fps 24 \
+  --audio-sample-rate 48000 \
+  --output ltx23_t2v_output.mp4
+```
+
+Use a Diffusers-format checkpoint such as `dg845/LTX-2.3-Diffusers`; the
+upstream `Lightricks/LTX-2.3` raw safetensors repo is not directly loadable by
+this pipeline. Pass `--model-class-name LTX23Pipeline` to select the LTX-2.3
+text-to-video pipeline explicitly.
+
 ### HunyuanVideo-1.5 (480p)
 
 ```bash
@@ -187,6 +212,8 @@ python text_to_video.py \
 ### Common
 
 - `--model`: Diffusers model ID or local path.
+- `--model-class-name`: Optional explicit pipeline class. Use `LTX23Pipeline`
+  for LTX-2.3 text-to-video.
 - `--prompt`: text description (string).
 - `--height/--width`: output resolution. Default depends on model.
 - `--num-frames`: number of frames. Default depends on model.
@@ -202,7 +229,8 @@ python text_to_video.py \
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
 - `--enable-layerwise-offload`: enable layerwise offloading on DiT modules.
 - `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2).
-- `--audio-sample-rate`: audio sample rate for embedded audio (when the pipeline returns audio).
+- `--audio-sample-rate`: audio sample rate for embedded audio (when the
+  pipeline returns audio; LTX-2.3 outputs 48kHz audio).
 - `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF).
 - `--flow-shift`: scheduler flow_shift parameter.
 - `--extra-body`: JSON object of model-specific generation params, filtered against the model's declared `extra_body_params` (see [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)). Used by Cosmos3 (see above).

examples/offline_inference/text_to_video/text_to_video.py

@@ -455,6 +455,7 @@ def main():
         print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
 
     audio = None
+    audio_sample_rate = args.audio_sample_rate
     if isinstance(frames, list):
         frames = frames[0] if frames else None
 
@@ -465,11 +466,13 @@ def main():
             )
         if frames.multimodal_output and "audio" in frames.multimodal_output:
             audio = frames.multimodal_output["audio"]
+            audio_sample_rate = frames.multimodal_output.get("audio_sample_rate", audio_sample_rate)
         if frames.is_pipeline_output and frames.request_output is not None:
             inner_output = frames.request_output
             if isinstance(inner_output, OmniRequestOutput):
                 if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
                     audio = inner_output.multimodal_output["audio"]
+                    audio_sample_rate = inner_output.multimodal_output.get("audio_sample_rate", audio_sample_rate)
                 frames = inner_output
         if isinstance(frames, OmniRequestOutput):
             if frames.images:
@@ -477,6 +480,7 @@ def main():
                     frames, audio = frames.images[0]
                 elif len(frames.images) == 1 and isinstance(frames.images[0], dict):
                     audio = frames.images[0].get("audio")
+                    audio_sample_rate = frames.images[0].get("audio_sample_rate", audio_sample_rate)
                     frames = frames.images[0].get("frames") or frames.images[0].get("video")
                 else:
                     frames = frames.images
@@ -489,6 +493,7 @@ def main():
             frames, audio = first_item
         elif isinstance(first_item, dict):
             audio = first_item.get("audio")
+            audio_sample_rate = first_item.get("audio_sample_rate", audio_sample_rate)
             frames = first_item.get("frames") or first_item.get("video")
         elif isinstance(first_item, list):
             frames = first_item
@@ -497,6 +502,7 @@ def main():
         frames, audio = frames
     elif isinstance(frames, dict):
         audio = frames.get("audio")
+        audio_sample_rate = frames.get("audio_sample_rate", audio_sample_rate)
         frames = frames.get("frames") or frames.get("video")
 
     if frames is None:
@@ -611,7 +617,7 @@ def main():
             frames_u8,
             audio_np,
             fps=float(args.fps),
-            audio_sample_rate=args.audio_sample_rate,
+            audio_sample_rate=audio_sample_rate,
         )
         with open(str(output_path), "wb") as f:
             f.write(video_bytes)

examples/online_serving/diffusers_pipeline_adapter/README.md

@@ -6,7 +6,8 @@ vLLM-Omni supports running diffusion models with the diffusers backend, directly
 
 The diffusers backend is a black-box adapter. Its primary focus is to serve diffusion models online.
 Currently, the following features are NOT yet supported.
-It is not guaranteed whether they will be supported in the future.
+Community contributions are welcome to add these features when they can be
+implemented with clear Diffusers or vLLM-Omni behavior.
 
 - CFG parallel execution
 - Sequence parallel execution
@@ -68,6 +69,35 @@ But if that parameter is simultaneously set in both the vLLM-Omni interface and
     Consider referring to [`OmniDiffusionSamplingParams`](../../../vllm_omni/inputs/data.py) for the default sampling parameters in the vLLM-Omni interface,
     and the corresponding [diffusers pipeline](https://huggingface.co/docs/diffusers/main/en/api/pipelines/overview)'s `__call__` function documentation.
 
+### Quantization
+
+Use `diffusers_load_kwargs` for Diffusers-native quantization options. If
+`diffusers_load_kwargs["quantization_config"]` is provided as a dictionary, the
+diffusers backend builds a Diffusers quantization config and lets Diffusers
+validate it before calling `DiffusionPipeline.from_pretrained()`.
+
+The diffusers backend also provides a small compatibility shortcut for
+vLLM-Omni quantization configs when `diffusers_load_kwargs` does not already
+contain `quantization_config`. Currently this maps online/dynamic `fp8` and
+online/dynamic `int8` to Diffusers/TorchAO dynamic quantization for transformer
+components such as `transformer` or `transformer_2`. This path requires
+`torchao` to be installed.
+
+For example, the CLI can request dynamic FP8 through the vLLM-Omni interface:
+
+```bash
+vllm serve "Qwen/Qwen-Image" \
+    --omni \
+    --diffusion-load-format diffusers \
+    --quantization-config '{"method": "fp8"}'
+```
+
+Other vLLM-Omni quantization methods, such as `gguf`, `modelopt`, `mxfp4`,
+`mxfp8`, serialized checkpoints, static FP8 configs, and layer-name skip lists
+such as `ignored_layers`, are intentionally not translated by this compatibility
+shortcut. Use Diffusers-native configuration through `diffusers_load_kwargs` or
+a native vLLM-Omni pipeline for those cases.
+
 ### Attention Backends
 
 The diffusers backend converts

recipes/LTX/LTX-2.3.md

@@ -6,21 +6,23 @@
 
 - Vendor: Lightricks
 - Model: `dg845/LTX-2.3-Diffusers`
-- Task: Text-to-video with synchronized audio generation
+- Task: Text-to-video and image-to-video with synchronized audio generation
 - Mode: Online serving (pure diffusion)
 - Maintainer: @oglok
 
 ## When to use this recipe
 
-Use this recipe when you want to serve LTX-2.3 for text-to-video generation
-with audio. The model generates videos up to 20+ seconds at 768x512 resolution
-with 48kHz audio, all from a single text prompt. Requires a GPU with at least
-96GB VRAM due to the 22B parameter transformer (~44GB weights) plus text
+Use this recipe when you want to serve LTX-2.3 for text-to-video or
+image-to-video generation with audio. The model generates videos up to 20+
+seconds at 768x512 resolution with 48kHz audio from text prompts, and can
+condition video generation on an initial image. Start with a 96GB-class GPU for
+validation because the model combines a 22B parameter transformer with text
 encoder, VAE, and vocoder components.
 
 ## References
 
-- Model: <https://huggingface.co/dg845/LTX-2.3-Diffusers>
+- Upstream raw checkpoints: <https://huggingface.co/Lightricks/LTX-2.3>
+- Diffusers-format checkpoint: <https://huggingface.co/dg845/LTX-2.3-Diffusers>
 - Requires `diffusers >= 0.38.0` (install from git: `pip install git+https://github.com/huggingface/diffusers.git`)
 
 ## Serving
@@ -34,7 +36,16 @@ vllm serve dg845/LTX-2.3-Diffusers \
   --stage-init-timeout 600
 ```
 
-### Verification
+For image-to-video:
+
+```bash
+vllm serve dg845/LTX-2.3-Diffusers \
+  --omni \
+  --model-class-name LTX23ImageToVideoPipeline \
+  --stage-init-timeout 600
+```
+
+### T2V Verification
 
 ```bash
 # Health check
@@ -73,19 +84,48 @@ curl -X POST http://localhost:8000/v1/videos \
   -F "guidance_scale=4.0"
 ```
 
+### I2V Verification
+
+Start the server with `--model-class-name LTX23ImageToVideoPipeline`, then pass
+an image reference. The I2V pipeline requires `input_reference` or
+`image_reference` because the first frame is the conditioning input.
+
+```bash
+curl http://localhost:8000/health
+
+curl -X POST http://localhost:8000/v1/videos/sync \
+  -F "prompt=A plush toy astronaut gently waving while the camera slowly pushes in." \
+  -F "input_reference=@/absolute/path/to/reference.png" \
+  -F "model=dg845/LTX-2.3-Diffusers" \
+  -F "num_frames=81" \
+  -F "fps=24" \
+  -F "size=768x512" \
+  -F "num_inference_steps=30" \
+  -F "guidance_scale=4.0" \
+  -F "seed=42" \
+  -o ltx23_i2v.mp4
+```
+
 ### Notes
 
-- Memory usage: Model loads at ~36 GiB, peaks at ~62 GiB during inference
+- Memory usage: run a latest-head serving smoke or DFX benchmark on the target
+  GPU before publishing load/peak VRAM numbers.
+- Checkpoint format: use a Diffusers-format checkpoint such as
+  `dg845/LTX-2.3-Diffusers`; the upstream `Lightricks/LTX-2.3` repository ships
+  raw safetensors and does not contain the subfolder configs required by this
+  pipeline loader.
 - Key flags:
-  - `--stage-init-timeout 600`: Required for the initial `torch.compile` warmup (~90-140 seconds on first request)
-  - `--model-class-name LTX23Pipeline`: Selects the LTX-2.3 pipeline (not LTX-2)
+  - `--stage-init-timeout 600`: Gives the initial `torch.compile` warmup enough time on large checkpoints
+  - `--model-class-name LTX23Pipeline`: Selects the LTX-2.3 text-to-video pipeline (not LTX-2)
+  - `--model-class-name LTX23ImageToVideoPipeline`: Selects the LTX-2.3 image-to-video pipeline
+  - `input_reference=@...`: Uploads the I2V reference image directly
 - Audio: 48kHz AAC via BWE vocoder, automatically synced with video
 - CPU offloading: Text encoder (Gemma-3-12B), connectors, VAE, audio VAE, and vocoder stay on CPU and are moved to GPU only when needed
 - Supported resolutions: 768x512, 512x384 (must be divisible by 32)
 - Frame rate: 24 fps
 - Duration: Controlled by `num_frames` (frames = duration_seconds * 24 + 1)
 - Known limitations:
-  - No image-to-video support yet (LTX23ImageToVideoPipeline is a placeholder)
+  - LTX-2.3 I2V currently supports first-frame image conditioning, matching the LTX image-conditioning contract.
   - Requires `diffusers >= 0.38.0` (not yet on PyPI, install from git)
 
 ## Hardware Support
@@ -99,13 +139,16 @@ curl -X POST http://localhost:8000/v1/videos \
 - OS: Ubuntu 22.04
 - Python: 3.10+
 - Driver / runtime: CUDA 13.0, Driver 580.126.09
-- vLLM version: 0.19.x
-- vLLM-Omni version: 0.19.x
+- vLLM version: Match the PR or release checkout being validated
+- vLLM-Omni version or commit: Use the commit you are deploying from
 
-### Validated configurations
+### Benchmarking
 
-| Duration | Frames | Resolution | Steps | Guidance | Inference Time | Peak VRAM |
-|----------|--------|------------|-------|----------|----------------|-----------|
-| 3s | 81 | 768x512 | 30 | 4.0 | ~110s | ~62 GB |
-| 10s | 241 | 768x512 | 30 | 4.0 | ~130s | ~62 GB |
-| 20s | 481 | 768x512 | 30 | 4.0 | ~420s | ~62 GB |
+Do not copy latency or VRAM numbers between commits or machines. Record them
+only after running a latest-head sweep on the target hardware.
+
+For formal PR benchmarking, reuse
+`tests/dfx/perf/tests/test_ltx2_3_vllm_omni.json` with
+`tests/dfx/perf/scripts/run_diffusion_benchmark.py`. That config captures the
+single-device eager baseline and CFG-parallel=2 case for a small 384x512,
+25-frame, 20-step workload.

recipes/mistralai/Voxtral-TTS.md

@@ -27,6 +27,79 @@ examples in this repository.
 
 ## GPU
 
+### 1x NVIDIA GeForce RTX 5090 32GB
+
+#### Environment
+
+- OS: Ubuntu 24.04.4 LTS
+- Python: 3.12.3
+- PyTorch: 2.11.0+cu130
+- Driver / runtime: NVIDIA driver 580.159.03, CUDA 13.0
+- GPU: NVIDIA GeForce RTX 5090, 32 GB (`32607 MiB`)
+- vLLM version: 0.23.0
+- vLLM-Omni version or commit: 0.23.0rc2.dev140+gbe60d7c7c (`be60d7c7`)
+- Tested with: `mistral-common==1.11.5`
+
+#### Command
+
+```bash
+vllm serve mistralai/Voxtral-4B-TTS-2603 \
+    --omni \
+    --port 8091 \
+    --trust-remote-code
+```
+
+The deploy config is auto-loaded from `vllm_omni/deploy/voxtral_tts.yaml`.
+
+#### Verification
+
+Health check and voices:
+
+```bash
+curl -s -o /tmp/voxtral_health.txt -w '%{http_code}' http://127.0.0.1:8091/health
+curl -s http://127.0.0.1:8091/v1/audio/voices
+```
+
+Non-streaming WAV:
+
+```bash
+curl -X POST http://127.0.0.1:8091/v1/audio/speech \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "mistralai/Voxtral-4B-TTS-2603",
+    "input": "Hello, this is Voxtral TTS running on an NVIDIA GeForce RTX 5090 with vLLM-Omni.",
+    "voice": "casual_female",
+    "language": "English",
+    "response_format": "wav"
+  }' \
+  --output /tmp/voxtral_5090.wav
+```
+
+Streaming PCM:
+
+```bash
+curl -X POST http://127.0.0.1:8091/v1/audio/speech \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "mistralai/Voxtral-4B-TTS-2603",
+    "input": "Hello, this is Voxtral TTS streaming PCM on an NVIDIA GeForce RTX 5090.",
+    "voice": "casual_female",
+    "language": "English",
+    "stream": true,
+    "response_format": "pcm"
+  }' \
+  --output /tmp/voxtral_5090_stream.pcm
+
+ffmpeg -f s16le -ar 24000 -ac 1 -i /tmp/voxtral_5090_stream.pcm /tmp/voxtral_5090_stream.wav -y
+ffprobe -hide_banner /tmp/voxtral_5090_stream.wav
+```
+
+#### Notes
+
+- Output: 24 kHz mono WAV/PCM; both non-streaming and streaming requests passed.
+- Memory usage: ~27.3 GiB after warmup (Stage 0 ~25.5 GiB, Stage 1 ~1.7 GiB).
+- Key flags: `--omni` is required; `--trust-remote-code` is recommended.
+
 ### 1x RTX 4090 24GB
 
 #### Environment

tests/dfx/perf/scripts/run_benchmark.py

@@ -94,7 +94,7 @@ def omni_server(request):
 benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping)
 
 
-@pytest.fixture(params=benchmark_indices)
+@pytest.fixture
 def benchmark_params(request, omni_server):
     """Benchmark parameters fixture with proper parametrization"""
     test_name, param_index = request.param

tests/diffusion/cache/test_cache_dit.py

@@ -13,7 +13,7 @@ import torch
 from cache_dit.caching.cache_blocks.pattern_0_1_2 import CachedBlocks_Pattern_0_1_2
 
 import vllm_omni.diffusion.cache.cache_dit_backend as cd_backend
-from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig, CacheDiTBackend
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig, CacheDiTBackend, cache_summary
 from vllm_omni.diffusion.data import DiffusionCacheConfig
 from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
 from vllm_omni.diffusion.models.helios.helios_transformer import HeliosTransformer3DModel
@@ -193,3 +193,12 @@ def test_ltx2_cache_dit_receives_audio_as_encoder(init_fake_tp_group):
     # Pattern_0 maps (hidden_states, encoder_hidden_states) to (video, audio)
     assert torch.equal(captured["hidden_states"], video_in)
     assert torch.equal(captured["encoder_hidden_states"], audio_in)
+
+
+def test_summary_with_no_transformer_is_nonfatal():
+    """Regression test for https://github.com/vllm-project/vllm-omni/issues/4325."""
+
+    class FakePipeline:
+        pass
+
+    cache_summary(pipeline=FakePipeline())

tests/diffusion/diffusion_backend/test_diffusers_backend.py

@@ -67,6 +67,50 @@ def _make_request(**overrides) -> OmniDiffusionRequest:
     return OmniDiffusionRequest(**defaults)
 
 
+def _patch_fake_diffusers_quantization_backends(mocker):
+    class FakePipelineQuantizationConfig:
+        def __init__(self, *, quant_mapping=None, quant_backend=None, quant_kwargs=None, components_to_quantize=None):
+            self.quant_mapping = quant_mapping
+            self.quant_backend = quant_backend
+            self.quant_kwargs = quant_kwargs
+            self.components_to_quantize = components_to_quantize
+
+    class FakeTorchAoConfig:
+        def __init__(self, quant_type, modules_to_not_convert=None):
+            self.quant_type = quant_type
+            self.modules_to_not_convert = modules_to_not_convert
+
+        @classmethod
+        def from_dict(cls, config):
+            return cls(quant_type=config["quant_type"])
+
+    class FakeInt8DynamicActivationInt8WeightConfig:
+        pass
+
+    mocker.patch(
+        "vllm_omni.diffusion.models.diffusers_adapter.quantization_utils.PipelineQuantizationConfig",
+        FakePipelineQuantizationConfig,
+    )
+    mocker.patch(
+        "vllm_omni.diffusion.models.diffusers_adapter.quantization_utils.TorchAoConfig",
+        FakeTorchAoConfig,
+    )
+    mocker.patch(
+        "vllm_omni.diffusion.models.diffusers_adapter.quantization_utils._get_torchao_quant_type_cls",
+        return_value=FakeInt8DynamicActivationInt8WeightConfig,
+    )
+    mocker.patch(
+        "vllm_omni.diffusion.models.diffusers_adapter.quantization_utils._get_diffusers_quantization_config_cls",
+        return_value=FakeTorchAoConfig,
+    )
+    return FakePipelineQuantizationConfig, FakeTorchAoConfig, FakeInt8DynamicActivationInt8WeightConfig
+
+
+class _TransformerBackedPipeline(DiffusionPipeline):
+    def __init__(self, transformer):
+        pass
+
+
 def _make_batch(**overrides) -> DiffusionRequestBatch:
     """Wrap a single request in a DiffusionRequestBatch, matching the forward contract."""
     return DiffusionRequestBatch(requests=[_make_request(**overrides)])
@@ -98,7 +142,7 @@ class TestPipelineArgumentsHandling:
 
     @pytest.mark.parametrize(
         "feature_id",
-        ["cfg_parallel", "ulysses", "ring", "teacache", "cache_dit", "enforce_eager", "quantization"],
+        ["cfg_parallel", "ulysses", "ring", "teacache", "cache_dit", "enforce_eager", "unsupported_quantization"],
     )
     def test_adapter_guard_unsupported_feature(self, feature_id):
         if feature_id == "cfg_parallel":
@@ -128,14 +172,188 @@ class TestPipelineArgumentsHandling:
             )
         elif feature_id == "enforce_eager":
             od_config = _make_od_config(enforce_eager=True)
-        elif feature_id == "quantization":
-            od_config = _make_od_config(quantization_config=SimpleNamespace(quant_method="fp8"))
+        elif feature_id == "unsupported_quantization":
+            od_config = _make_od_config(quantization_config=SimpleNamespace(quant_method="gguf"))
         else:
             raise ValueError(f"Unknown feature ID: {feature_id}")
 
         with pytest.raises(NotImplementedError):
             DiffusersAdapterPipeline(od_config=od_config)
 
+    def test_adapter_load_weights_injects_supported_quantization_config(self, mocker):
+        (
+            FakePipelineQuantizationConfig,
+            FakeTorchAoConfig,
+            FakeInt8DynamicActivationInt8WeightConfig,
+        ) = _patch_fake_diffusers_quantization_backends(mocker)
+
+        class MockPipeline:
+            def __call__(self, prompt=None):
+                return None
+
+            def to(self, device):
+                return self
+
+        mock_from_pretrained = mocker.patch(
+            "vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter.DiffusionPipeline.from_pretrained",
+            return_value=MockPipeline(),
+        )
+        mocker.patch(
+            "vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter.DiffusionPipeline.load_config",
+            return_value={
+                "_class_name": "TransformerBackedPipeline",
+                "scheduler": ["diffusers", "SchedulerMixin"],
+                "transformer": ["diffusers", "Transformer2DModel"],
+                "transformer_2": ["diffusers", "Transformer2DModel"],
+                "vae": ["diffusers", "AutoencoderKL"],
+            },
+        )
+
+        adapter = DiffusersAdapterPipeline(
+            od_config=_make_od_config(
+                quantization_config=SimpleNamespace(quant_method="int8"),
+                diffusers_pipeline_cls=_TransformerBackedPipeline,
+            )
+        )
+        adapter.load_weights()
+
+        pipeline_quant_config = mock_from_pretrained.call_args.kwargs["quantization_config"]
+        torchao_config = pipeline_quant_config.quant_mapping["transformer"]
+        assert isinstance(pipeline_quant_config, FakePipelineQuantizationConfig)
+        assert isinstance(torchao_config, FakeTorchAoConfig)
+        assert isinstance(torchao_config.quant_type, FakeInt8DynamicActivationInt8WeightConfig)
+        assert isinstance(
+            pipeline_quant_config.quant_mapping["transformer_2"].quant_type,
+            FakeInt8DynamicActivationInt8WeightConfig,
+        )
+
+    @pytest.mark.parametrize(
+        ("pipeline_config", "diffusers_load_kwargs", "diffusers_pipeline_cls"),
+        [
+            (
+                {
+                    "_class_name": "NoTransformerPipeline",
+                    "scheduler": ["diffusers", "SchedulerMixin"],
+                    "unet": ["diffusers", "UNet2DConditionModel"],
+                    "vae": ["diffusers", "AutoencoderKL"],
+                },
+                {},
+                _TransformerBackedPipeline,
+            ),
+            (
+                {
+                    "_class_name": "TransformerPipeline",
+                    "scheduler": ["diffusers", "SchedulerMixin"],
+                    "transformer": ["diffusers", "Transformer2DModel"],
+                    "vae": ["diffusers", "AutoencoderKL"],
+                },
+                {"transformer": None},
+                None,
+            ),
+        ],
+    )
+    def test_adapter_load_weights_rejects_converted_quantization_without_transformer_component(
+        self,
+        pipeline_config,
+        diffusers_load_kwargs,
+        diffusers_pipeline_cls,
+        mocker,
+    ):
+        _patch_fake_diffusers_quantization_backends(mocker)
+        mocker.patch(
+            "vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter.DiffusionPipeline.load_config",
+            return_value=pipeline_config,
+        )
+        mock_from_pretrained = mocker.patch(
+            "vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter.DiffusionPipeline.from_pretrained"
+        )
+
+        adapter = DiffusersAdapterPipeline(
+            od_config=_make_od_config(
+                quantization_config=SimpleNamespace(quant_method="int8"),
+                diffusers_load_kwargs=diffusers_load_kwargs,
+                diffusers_pipeline_cls=diffusers_pipeline_cls,
+            )
+        )
+
+        with pytest.raises(NotImplementedError, match="transformer"):
+            adapter.load_weights()
+        mock_from_pretrained.assert_not_called()
+
+    def test_adapter_rejects_dduf_file_with_converted_quantization(self):
+        with pytest.raises(NotImplementedError, match="dduf_file"):
+            DiffusersAdapterPipeline(
+                od_config=_make_od_config(
+                    quantization_config=SimpleNamespace(quant_method="int8"),
+                    diffusers_load_kwargs={"dduf_file": "model.dduf"},
+                )
+            )
+
+    def test_adapter_load_weights_preserves_diffusers_native_quantization_config(self, mocker):
+        class MockPipeline:
+            def __call__(self, prompt=None):
+                return None
+
+            def to(self, device):
+                return self
+
+        diffusers_native_quantization_config = object()
+        mock_from_pretrained = mocker.patch(
+            "vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter.DiffusionPipeline.from_pretrained",
+            return_value=MockPipeline(),
+        )
+
+        adapter = DiffusersAdapterPipeline(
+            od_config=_make_od_config(
+                quantization_config=SimpleNamespace(quant_method="gguf"),
+                diffusers_load_kwargs={
+                    "dduf_file": "model.dduf",
+                    "quantization_config": diffusers_native_quantization_config,
+                },
+            )
+        )
+        adapter.load_weights()
+
+        kwargs = mock_from_pretrained.call_args.kwargs
+        assert kwargs["quantization_config"] is diffusers_native_quantization_config
+        assert kwargs["dduf_file"] == "model.dduf"
+
+    def test_adapter_load_weights_builds_diffusers_native_quantization_config_from_dict(self, mocker):
+        FakePipelineQuantizationConfig, FakeTorchAoConfig, _ = _patch_fake_diffusers_quantization_backends(mocker)
+
+        class MockPipeline:
+            def __call__(self, prompt=None):
+                return None
+
+            def to(self, device):
+                return self
+
+        mock_from_pretrained = mocker.patch(
+            "vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter.DiffusionPipeline.from_pretrained",
+            return_value=MockPipeline(),
+        )
+
+        adapter = DiffusersAdapterPipeline(
+            od_config=_make_od_config(
+                diffusers_load_kwargs={
+                    "quantization_config": {
+                        "quant_mapping": {
+                            "transformer": {
+                                "quant_method": "torchao",
+                                "quant_type": "fake-int8-config",
+                            },
+                        },
+                    },
+                },
+            )
+        )
+        adapter.load_weights()
+
+        quant_config = mock_from_pretrained.call_args.kwargs["quantization_config"]
+        assert isinstance(quant_config, FakePipelineQuantizationConfig)
+        assert isinstance(quant_config.quant_mapping["transformer"], FakeTorchAoConfig)
+        assert quant_config.quant_mapping["transformer"].quant_type == "fake-int8-config"
+
     def test_adapter_guard_unknown_output_type(self, mocker):
         """Test that the adapter wraps an unknown output type as-is.
         This is useful when `return_dict=True` and the diffusers pipeline returns an OrderedDict subclass."""

tests/diffusion/diffusion_backend/test_diffusers_quantization_real.py

@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Optional real-backend checks for Diffusers quantization config conversion."""
+
+from __future__ import annotations
+
+import importlib.util
+from pathlib import Path
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _load_quantization_utils():
+    module_path = (
+        Path(__file__).resolve().parents[3] / "vllm_omni/diffusion/models/diffusers_adapter/quantization_utils.py"
+    )
+    spec = importlib.util.spec_from_file_location("diffusers_quantization_utils_real_test", module_path)
+    assert spec is not None
+    assert spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_convert_diffusers_quantization_config_from_real_torchao_dict():
+    diffusers = pytest.importorskip("diffusers")
+    torchao_quantization = pytest.importorskip("torchao.quantization")
+    quant_type_cls = getattr(torchao_quantization, "Int8DynamicActivationInt8WeightConfig", None)
+    if quant_type_cls is None:
+        pytest.skip("torchao does not expose Int8DynamicActivationInt8WeightConfig")
+
+    quantization_utils = _load_quantization_utils()
+
+    real_torchao_config = diffusers.TorchAoConfig(quant_type=quant_type_cls())
+    serialized_torchao_config = real_torchao_config.to_dict()
+    assert serialized_torchao_config["quant_type"].keys() == {"default"}
+
+    load_kwargs = {
+        "quantization_config": {
+            "quant_mapping": {
+                "transformer": serialized_torchao_config,
+            },
+        },
+    }
+
+    quantization_utils.convert_diffusers_quantization_config(load_kwargs)
+
+    quant_config = load_kwargs["quantization_config"]
+    component_config = quant_config.quant_mapping["transformer"]
+    assert isinstance(quant_config, diffusers.PipelineQuantizationConfig)
+    assert isinstance(component_config, diffusers.TorchAoConfig)
+    assert isinstance(
+        component_config.quant_type,
+        quant_type_cls,
+    )

tests/diffusion/diffusion_backend/test_diffusers_quantization_utils.py

@@ -0,0 +1,152 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Focused tests for Diffusers backend quantization conversion helpers."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.diffusion.models.diffusers_adapter import quantization_utils
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+class _FakePipelineQuantizationConfig:
+    def __init__(self, *, quant_mapping=None, quant_backend=None, quant_kwargs=None, components_to_quantize=None):
+        self.quant_mapping = quant_mapping
+        self.quant_backend = quant_backend
+        self.quant_kwargs = quant_kwargs
+        self.components_to_quantize = components_to_quantize
+
+
+class _FakeTorchAoConfig:
+    def __init__(self, quant_type):
+        self.quant_type = quant_type
+
+    @classmethod
+    def from_dict(cls, config):
+        return cls(quant_type=config["quant_type"])
+
+
+class _FakeFloat8DynamicActivationFloat8WeightConfig:
+    pass
+
+
+class _FakeInt8DynamicActivationInt8WeightConfig:
+    pass
+
+
+@pytest.fixture(autouse=True)
+def patch_quantization_backends(monkeypatch):
+    monkeypatch.setattr(
+        quantization_utils,
+        "PipelineQuantizationConfig",
+        _FakePipelineQuantizationConfig,
+    )
+    monkeypatch.setattr(quantization_utils, "TorchAoConfig", _FakeTorchAoConfig)
+
+    def fake_get_torchao_quant_type_cls(class_name: str):
+        return {
+            "Float8DynamicActivationFloat8WeightConfig": _FakeFloat8DynamicActivationFloat8WeightConfig,
+            "Int8DynamicActivationInt8WeightConfig": _FakeInt8DynamicActivationInt8WeightConfig,
+        }[class_name]
+
+    monkeypatch.setattr(quantization_utils, "_get_torchao_quant_type_cls", fake_get_torchao_quant_type_cls)
+    monkeypatch.setattr(quantization_utils, "_get_diffusers_quantization_config_cls", lambda _: _FakeTorchAoConfig)
+
+
+def _quant_config(method: str, **kwargs):
+    return SimpleNamespace(get_name=lambda: method, **kwargs)
+
+
+@pytest.mark.parametrize(
+    ("method", "expected_type"),
+    [
+        ("fp8", _FakeFloat8DynamicActivationFloat8WeightConfig),
+        ("int8", _FakeInt8DynamicActivationInt8WeightConfig),
+    ],
+)
+def test_apply_injects_converted_quantization_config(method, expected_type):
+    od_config = SimpleNamespace(quantization_config=_quant_config(method))
+    load_kwargs = {}
+
+    injected = quantization_utils.apply_diffusers_quantization_config(
+        od_config,
+        load_kwargs,
+        {"transformer": ["diffusers", "Transformer2DModel"]},
+    )
+
+    assert injected is True
+    torchao_config = load_kwargs["quantization_config"].quant_mapping["transformer"]
+    assert isinstance(torchao_config.quant_type, expected_type)
+
+
+def test_apply_injects_converted_quantization_config_for_transformer_2():
+    od_config = SimpleNamespace(quantization_config=_quant_config("int8"))
+    load_kwargs = {}
+
+    quantization_utils.apply_diffusers_quantization_config(
+        od_config,
+        load_kwargs,
+        {
+            "transformer": ["diffusers", "Transformer2DModel"],
+            "transformer_2": ["diffusers", "Transformer2DModel"],
+        },
+    )
+
+    assert sorted(load_kwargs["quantization_config"].quant_mapping) == ["transformer", "transformer_2"]
+
+
+@pytest.mark.parametrize("method", ["gguf", "modelopt", "mxfp4", "mxfp8", "inc"])
+def test_unsupported_methods_fail_explicitly(method):
+    with pytest.raises(NotImplementedError, match=method):
+        quantization_utils.ensure_supported_diffusers_quantization(_quant_config(method))
+
+
+@pytest.mark.parametrize(
+    "quant_config",
+    [
+        _quant_config("fp8", activation_scheme="static"),
+        _quant_config("int8", is_checkpoint_int8_serialized=True),
+        _quant_config("fp8", weight_block_size=[128, 128]),
+        _quant_config("int8", ignored_layers=["transformer.proj_out"]),
+    ],
+)
+def test_ambiguous_mappings_fail_explicitly(quant_config):
+    with pytest.raises(NotImplementedError):
+        quantization_utils.ensure_supported_diffusers_quantization(quant_config)
+
+
+def test_apply_preserves_diffusers_load_kwargs_quantization_config(mocker):
+    od_config = SimpleNamespace(quantization_config=_quant_config("fp8"))
+    existing = object()
+    load_kwargs = {"quantization_config": existing}
+    mock_warning = mocker.patch.object(quantization_utils.logger, "warning")
+
+    injected = quantization_utils.apply_diffusers_quantization_config(od_config, load_kwargs, {})
+
+    assert injected is False
+    assert load_kwargs["quantization_config"] is existing
+    assert "Using the Diffusers-native quantization_config" in mock_warning.call_args.args[0]
+
+
+def test_convert_diffusers_quantization_config_from_dict():
+    load_kwargs = {
+        "quantization_config": {
+            "quant_mapping": {
+                "transformer": {
+                    "quant_method": "torchao",
+                    "quant_type": "fake-int8-config",
+                }
+            }
+        }
+    }
+
+    quantization_utils.convert_diffusers_quantization_config(load_kwargs)
+
+    quant_config = load_kwargs["quantization_config"]
+    assert isinstance(quant_config, _FakePipelineQuantizationConfig)
+    assert isinstance(quant_config.quant_mapping["transformer"], _FakeTorchAoConfig)
+    assert quant_config.quant_mapping["transformer"].quant_type == "fake-int8-config"

tests/diffusion/distributed/test_expert_parallel_layout.py

@@ -0,0 +1,282 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+
+import vllm_omni.diffusion.distributed.parallel_state as parallel_state
+
+
+class _FakeGroup:
+    def __init__(
+        self,
+        group_ranks: list[list[int]],
+        local_rank: int,
+        parallel_mode: str,
+        **kwargs,
+    ) -> None:
+        self.group_ranks = group_ranks
+        self.parallel_mode = parallel_mode
+        self.device_group = object()
+        self.ulysses_group = kwargs.get("ulysses_group")
+        self.ring_group = kwargs.get("ring_group")
+        self.local_group = next(group for group in group_ranks if local_rank in group)
+        self.world_size = len(self.local_group)
+        self.rank_in_group = self.local_group.index(local_rank)
+
+    def destroy(self) -> None:
+        pass
+
+
+@pytest.mark.cpu
+@pytest.mark.core_model
+def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
+    """MoE+EP rank layout should map SP->PCP, CFG*DP->DP, and TP*SP*CFG*DP->EP."""
+    local_rank = 0
+    world_size = 32
+    created_groups: list[_FakeGroup] = []
+
+    def fake_init_model_parallel_group(
+        group_ranks,
+        local_rank,
+        backend,
+        parallel_mode=None,
+        group_name=None,
+        **kwargs,
+    ):
+        del backend, group_name
+        group = _FakeGroup(
+            [list(ranks) for ranks in group_ranks],
+            local_rank,
+            parallel_mode or "",
+            **kwargs,
+        )
+        created_groups.append(group)
+        return group
+
+    fake_world_group = SimpleNamespace(
+        rank_in_group=local_rank,
+        local_rank=local_rank,
+        device_group=object(),
+    )
+    fake_forward_context = SimpleNamespace(omni_diffusion_config=SimpleNamespace(is_moe=True))
+    monkeypatch.setattr(parallel_state.torch.distributed, "is_initialized", lambda: True)
+    monkeypatch.setattr(parallel_state.torch.distributed, "get_world_size", lambda: world_size)
+    monkeypatch.setattr(parallel_state.torch.distributed, "get_backend", lambda *_args, **_kwargs: "gloo")
+    monkeypatch.setattr(parallel_state.torch.distributed, "new_group", lambda ranks: tuple(ranks))
+    monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
+    monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context)
+    monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
+    monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
+
+    for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
+        monkeypatch.setattr(parallel_state, name, None)
+    for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
+        monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
+
+    parallel_state.initialize_model_parallel(
+        tensor_parallel_size=2,
+        sequence_parallel_size=2,
+        ulysses_degree=2,
+        ring_degree=1,
+        pipeline_parallel_size=2,
+        cfg_parallel_size=2,
+        data_parallel_size=2,
+        enable_expert_parallel=True,
+        backend="gloo",
+    )
+
+    assert parallel_state.vllm_parallel_state._PCP is parallel_state._SP
+    assert parallel_state.vllm_parallel_state._PCP.world_size == 2
+    assert parallel_state._DP.world_size == 2
+    assert parallel_state.vllm_parallel_state._DP is not parallel_state._DP
+    assert parallel_state.vllm_parallel_state._DP.world_size == 4
+    assert parallel_state.vllm_parallel_state._EP.world_size == 16
+    assert parallel_state.vllm_parallel_state._TP.world_size == 2
+    assert parallel_state._PP.world_size == 2
+
+    assert parallel_state.vllm_parallel_state._PCP.local_group == [0, 2]
+    assert parallel_state.vllm_parallel_state._DP.local_group == [0, 8, 16, 24]
+    assert parallel_state.vllm_parallel_state._EP.local_group == [
+        0,
+        1,
+        2,
+        3,
+        8,
+        9,
+        10,
+        11,
+        16,
+        17,
+        18,
+        19,
+        24,
+        25,
+        26,
+        27,
+    ]
+    assert parallel_state.get_expert_parallel_group_ranks() == [
+        [
+            0,
+            1,
+            2,
+            3,
+            8,
+            9,
+            10,
+            11,
+            16,
+            17,
+            18,
+            19,
+            24,
+            25,
+            26,
+            27,
+        ],
+        [
+            4,
+            5,
+            6,
+            7,
+            12,
+            13,
+            14,
+            15,
+            20,
+            21,
+            22,
+            23,
+            28,
+            29,
+            30,
+            31,
+        ],
+    ]
+
+    ep_groups = [group.local_group for group in created_groups if group.parallel_mode == "expert"]
+    assert ep_groups == [parallel_state.vllm_parallel_state._EP.local_group]
+
+
+@pytest.mark.cpu
+@pytest.mark.core_model
+def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch):
+    """vLLM DP should keep diffusion DP when expert parallelism is not enabled."""
+    local_rank = 0
+    world_size = 8
+
+    def fake_init_model_parallel_group(
+        group_ranks,
+        local_rank,
+        backend,
+        parallel_mode=None,
+        group_name=None,
+        **kwargs,
+    ):
+        del backend, group_name
+        return _FakeGroup(
+            [list(ranks) for ranks in group_ranks],
+            local_rank,
+            parallel_mode or "",
+            **kwargs,
+        )
+
+    fake_world_group = SimpleNamespace(
+        rank_in_group=local_rank,
+        local_rank=local_rank,
+        device_group=object(),
+    )
+    monkeypatch.setattr(parallel_state.torch.distributed, "is_initialized", lambda: True)
+    monkeypatch.setattr(parallel_state.torch.distributed, "get_world_size", lambda: world_size)
+    monkeypatch.setattr(parallel_state.torch.distributed, "get_backend", lambda *_args, **_kwargs: "gloo")
+    monkeypatch.setattr(parallel_state.torch.distributed, "new_group", lambda ranks: tuple(ranks))
+    monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
+    monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
+    monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
+
+    for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
+        monkeypatch.setattr(parallel_state, name, None)
+    for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
+        monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
+
+    parallel_state.initialize_model_parallel(
+        tensor_parallel_size=2,
+        sequence_parallel_size=1,
+        ulysses_degree=1,
+        ring_degree=1,
+        pipeline_parallel_size=1,
+        cfg_parallel_size=2,
+        data_parallel_size=2,
+        enable_expert_parallel=False,
+        backend="gloo",
+    )
+
+    assert parallel_state._DP.world_size == 2
+    assert parallel_state.vllm_parallel_state._DP is parallel_state._DP
+    assert parallel_state.vllm_parallel_state._DP.world_size == 2
+    assert parallel_state.vllm_parallel_state._DP.local_group == [0, 4]
+    assert parallel_state.vllm_parallel_state._PCP is None
+    assert parallel_state.vllm_parallel_state._EP is None
+    assert parallel_state._EXPERT_PARALLEL_GROUP_RANKS is None
+
+
+@pytest.mark.cpu
+@pytest.mark.core_model
+def test_non_moe_ep_fails_before_vllm_ep_remap(monkeypatch):
+    """Non-MoE diffusion configs should not create vLLM PCP/DP/EP remap state."""
+    local_rank = 0
+    world_size = 8
+
+    def fake_init_model_parallel_group(
+        group_ranks,
+        local_rank,
+        backend,
+        parallel_mode=None,
+        group_name=None,
+        **kwargs,
+    ):
+        del backend, group_name
+        return _FakeGroup(
+            [list(ranks) for ranks in group_ranks],
+            local_rank,
+            parallel_mode or "",
+            **kwargs,
+        )
+
+    fake_world_group = SimpleNamespace(
+        rank_in_group=local_rank,
+        local_rank=local_rank,
+        device_group=object(),
+    )
+    fake_forward_context = SimpleNamespace(omni_diffusion_config=SimpleNamespace(is_moe=False))
+    monkeypatch.setattr(parallel_state.torch.distributed, "is_initialized", lambda: True)
+    monkeypatch.setattr(parallel_state.torch.distributed, "get_world_size", lambda: world_size)
+    monkeypatch.setattr(parallel_state.torch.distributed, "get_backend", lambda *_args, **_kwargs: "gloo")
+    monkeypatch.setattr(parallel_state.torch.distributed, "new_group", lambda ranks: tuple(ranks))
+    monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
+    monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context)
+    monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
+    monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
+
+    for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
+        monkeypatch.setattr(parallel_state, name, None)
+    for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
+        monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
+
+    with pytest.raises(RuntimeError, match="Expert parallelism enabled for a non-MoE model"):
+        parallel_state.initialize_model_parallel(
+            tensor_parallel_size=2,
+            sequence_parallel_size=2,
+            ulysses_degree=2,
+            ring_degree=1,
+            pipeline_parallel_size=1,
+            cfg_parallel_size=2,
+            data_parallel_size=1,
+            enable_expert_parallel=True,
+            backend="gloo",
+        )
+
+    assert parallel_state.vllm_parallel_state._PCP is None
+    assert parallel_state.vllm_parallel_state._EP is None
+    assert parallel_state._EXPERT_PARALLEL_GROUP_RANKS is None

tests/diffusion/lora/test_lora_manager.py

@@ -575,3 +575,45 @@ def test_lora_manager_discovers_bagel_component(monkeypatch):
     assert "bagel.language_model.qkv_proj" in manager._lora_modules
     # Verify the module was actually replaced in the tree (not just recorded)
     assert isinstance(pipeline.bagel.language_model.qkv_proj, _DummyBaseLayerWithLoRA)
+
+
+def test_lora_manager_discovers_unet_component(monkeypatch):
+    """Verify that _replace_layers_with_lora finds layers under 'unet'."""
+    import vllm_omni.diffusion.lora.manager as manager_mod
+
+    monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA)
+
+    def _fake_from_layer_diffusion(*, layer: torch.nn.Module, **_kwargs):
+        if isinstance(layer, _FakeLinearBase):
+            return _DummyBaseLayerWithLoRA(layer)
+        return layer
+
+    replace_calls: list[str] = []
+
+    monkeypatch.setattr(manager_mod, "from_layer_diffusion", _fake_from_layer_diffusion)
+    monkeypatch.setattr(
+        manager_mod,
+        "replace_submodule",
+        lambda root, name, sub: fake_replace_submodule(root, name, sub, replace_calls),
+    )
+
+    # Pipeline with a 'unet' component (no 'transformer')
+    pipeline = torch.nn.Module()
+    pipeline.unet = torch.nn.Module()
+    pipeline.unet.down_block = torch.nn.Module()
+    pipeline.unet.down_block.proj = _FakeLinearBase()
+
+    manager = DiffusionLoRAManager(
+        pipeline=pipeline,
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        max_cached_adapters=1,
+    )
+
+    peft_helper = type("_PH", (), {"r": 1})()
+    manager._replace_layers_with_lora(peft_helper)
+
+    assert "down_block.proj" in replace_calls
+    assert "unet.down_block.proj" in manager._lora_modules
+    # Verify the module was actually replaced in the tree (not just recorded)
+    assert isinstance(pipeline.unet.down_block.proj, _DummyBaseLayerWithLoRA)

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py

@@ -916,6 +916,49 @@ def test_prepare_latents_i2v_encodes_only_conditioning_frame(make_cosmos3_pipeli
     assert velocity_mask.tolist() == [[[[[0.0]], [[1.0]], [[1.0]]]]]
 
 
+@pytest.mark.parametrize("mode", ["policy", "forward_dynamics"])
+def test_prepare_action_video_latents_encodes_only_conditioning_frame(make_cosmos3_pipeline, mode: str) -> None:
+    pipeline = make_cosmos3_pipeline()
+    video = torch.zeros(1, 3, 9, 24, 32)
+
+    latents, velocity_mask, condition_latents = pipeline._prepare_latents_action_video(
+        video,
+        mode,
+        24,
+        32,
+        9,
+        torch.Generator(device="cpu").manual_seed(0),
+        image_size=torch.tensor([1, 3, 16, 24]),
+    )
+
+    assert pipeline.vae.encode_input_shapes == [(1, 3, 1, 24, 32)]
+    assert latents.shape == (1, 2, 3, 2, 3)
+    assert condition_latents.shape == latents.shape
+    torch.testing.assert_close(condition_latents[:, :, 0], torch.ones(1, 2, 2, 3))
+    assert torch.count_nonzero(condition_latents[:, :, 1:]) == 0
+    torch.testing.assert_close(latents[:, :, 0:1], condition_latents[:, :, 0:1])
+    assert velocity_mask.tolist() == [[[[[0.0]], [[1.0]], [[1.0]]]]]
+
+
+def test_prepare_inverse_dynamics_latents_encodes_full_video(make_cosmos3_pipeline) -> None:
+    pipeline = make_cosmos3_pipeline()
+    video = torch.zeros(1, 3, 9, 16, 24)
+
+    latents, velocity_mask, condition_latents = pipeline._prepare_latents_action_video(
+        video,
+        "inverse_dynamics",
+        16,
+        24,
+        9,
+        torch.Generator(device="cpu").manual_seed(0),
+    )
+
+    assert pipeline.vae.encode_input_shapes == [(1, 3, 9, 16, 24)]
+    assert condition_latents.shape == (1, 2, 3, 2, 3)
+    torch.testing.assert_close(latents, condition_latents)
+    assert torch.count_nonzero(velocity_mask) == 0
+
+
 def test_diffuse_covers_cfg_i2v_and_multimodal_steps(make_cosmos3_pipeline) -> None:
     pipeline = make_cosmos3_pipeline()
     latents = torch.zeros(1, 2, 1, 1, 1)

tests/diffusion/models/flux/test_flux_pipeline.py

@@ -21,6 +21,8 @@ def _make_flux_sampling(**overrides):
         "true_cfg_scale": 4.0,
         "num_outputs_per_prompt": 0,
         "latents": None,
+        "output_type": None,
+        "max_sequence_length": None,
     }
     values.update(overrides)
     return SimpleNamespace(**values)
@@ -131,7 +133,7 @@ def test_forward_collates_request_prompt_tensors_for_flux(monkeypatch):
                     "negative_prompt_embeds": negative_prompt_embeds_a,
                     "negative_pooled_prompt_embeds": negative_pooled_prompt_embeds_a,
                 },
-                sampling_params=_make_flux_sampling(generator=gen_a, latents=latents_a),
+                sampling_params=_make_flux_sampling(generator=gen_a, latents=latents_a, output_type="latent"),
             ),
             SimpleNamespace(
                 request_id="flux-prompt-b",
@@ -145,12 +147,12 @@ def test_forward_collates_request_prompt_tensors_for_flux(monkeypatch):
                         "negative_pooled_prompt_embeds": [negative_pooled_prompt_embeds_b],
                     },
                 },
-                sampling_params=_make_flux_sampling(generator=gen_b, latents=latents_b),
+                sampling_params=_make_flux_sampling(generator=gen_b, latents=latents_b, output_type="latent"),
             ),
         ]
     )
 
-    outputs = pipeline.forward(batch, output_type="latent")
+    outputs = pipeline.forward(batch)
 
     assert encode_calls[0]["prompt"] is None
     assert encode_calls[0]["prompt_2"] is None

tests/diffusion/models/hunyuan_image3/test_multi_resolution.py

@@ -524,11 +524,11 @@ class TestAr2DiffusionIntegration:
         self,
         token_ids: list[int],
         prompt_kwargs: dict[str, Any] | None = None,
-    ) -> tuple[Any, Any]:
+    ) -> dict[str, Any] | None:
         """Invoke ``ar2diffusion`` with a single fake AR output.
 
-        Returns ``(diffusion_inputs, ar_predicted_sizes)`` so callers can
-        assert on the resolved height/width.
+        Returns the diffusion request payload so callers can assert on the
+        resolved height/width.
         """
         from vllm_omni.model_executor.stage_input_processors.hunyuan_image3 import (
             ar2diffusion,
@@ -601,8 +601,8 @@ class TestAr2DiffusionIntegration:
             token_ids=[150000, 150001, token_id],
             prompt_kwargs={"height": 1024, "width": 1024, "image_base_size": 1024},
         )
-        assert len(result) == 1
-        diffusion_input = result[0]
+        assert result is not None
+        diffusion_input = result
         assert diffusion_input["height"] == expected_h, (
             f"ratio_idx={ratio_idx}: expected height {expected_h}, got {diffusion_input['height']}"
         )
@@ -617,8 +617,8 @@ class TestAr2DiffusionIntegration:
             token_ids=[150000, 150001],  # no ratio tokens
             prompt_kwargs={"height": 768, "width": 1024, "image_base_size": 1024},
         )
-        assert len(result) == 1
-        diffusion_input = result[0]
+        assert result is not None
+        diffusion_input = result
         assert diffusion_input["height"] == 768
         assert diffusion_input["width"] == 1024
 
@@ -633,9 +633,9 @@ class TestAr2DiffusionIntegration:
                 token_ids=[150000],
                 prompt_kwargs={"height": 512, "width": 512, "image_base_size": 1024},
             )
-        assert len(result) == 1
-        assert result[0]["height"] == 512
-        assert result[0]["width"] == 512
+        assert result is not None
+        assert result["height"] == 512
+        assert result["width"] == 512
 
     # -- cache behaviour -----------------------------------------------------
 

tests/diffusion/models/ltx2/test_ltx2_3_forward_parallel.py

@@ -0,0 +1,436 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Unit tests for LTX-2.3 forward-parallel behavior."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestCFGParallelIntegration:
+    def test_ltx23_pipeline_has_cfg_parallel_mixin(self):
+        """LTX23Pipeline must use the shared CFG parallel implementation."""
+        from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+
+        assert issubclass(LTX23Pipeline, CFGParallelMixin)
+
+    def test_ltx2_transformer_has_dit_cache_config(self):
+        """Ensure LTX2 has a Cache DiT adapter config and that it uses separate CFG."""
+        from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
+        from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
+
+        adapter_config = getattr(LTX2VideoTransformer3DModel, "_cache_dit_adapter_config")
+        assert isinstance(adapter_config, CacheDiTAdapterConfig)
+        assert adapter_config.has_separate_cfg
+
+
+class TestCFGParallelHelpers:
+    """Test LTX-2.3 CFG helper math without loading model weights."""
+
+    def test_combine_cfg_noise_matches_x0_space_formula(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+
+        pipe = object.__new__(LTX23Pipeline)
+        video_sample = torch.tensor([[[1.0, -2.0]]])
+        audio_sample = torch.tensor([[[0.5, 3.0]]])
+        video_pos = torch.tensor([[[0.2, -0.3]]])
+        video_neg = torch.tensor([[[-0.4, 0.1]]])
+        audio_pos = torch.tensor([[[0.7, -0.2]]])
+        audio_neg = torch.tensor([[[0.1, 0.4]]])
+        video_sigma = torch.tensor(0.25)
+        audio_sigma = torch.tensor(0.5)
+        scale = 4.0
+
+        video_combined, audio_combined = pipe.combine_cfg_noise(
+            (video_pos, audio_pos),
+            (video_neg, audio_neg),
+            scale,
+            video_latents=video_sample,
+            audio_latents=audio_sample,
+            video_sigma=video_sigma,
+            audio_sigma=audio_sigma,
+        )
+
+        x0_video_cond = video_sample - video_pos * video_sigma
+        x0_video_uncond = video_sample - video_neg * video_sigma
+        x0_video_guided = x0_video_cond + (scale - 1) * (x0_video_cond - x0_video_uncond)
+        expected_video = (video_sample - x0_video_guided) / video_sigma
+
+        x0_audio_cond = audio_sample - audio_pos * audio_sigma
+        x0_audio_uncond = audio_sample - audio_neg * audio_sigma
+        x0_audio_guided = x0_audio_cond + (scale - 1) * (x0_audio_cond - x0_audio_uncond)
+        expected_audio = (audio_sample - x0_audio_guided) / audio_sigma
+        assert torch.allclose(video_combined, expected_video)
+        assert torch.allclose(audio_combined, expected_audio)
+
+    def test_two_rank_cfg_parallel_smoke_uses_rank_local_branch_and_x0_formula(self, monkeypatch):
+        from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23
+
+        pipe = object.__new__(ltx23.LTX23Pipeline)
+        video_sample = torch.tensor([[[1.0, -2.0]]])
+        audio_sample = torch.tensor([[[0.5, 3.0, -1.0]]])
+        video_pos = torch.tensor([[[0.2, -0.3]]])
+        video_neg = torch.tensor([[[-0.4, 0.1]]])
+        audio_pos = torch.tensor([[[0.7, -0.2, 0.3]]])
+        audio_neg = torch.tensor([[[0.1, 0.4, -0.5]]])
+        video_sigma = torch.tensor(0.25)
+        audio_sigma = torch.tensor(0.5)
+        scale = 4.0
+
+        class FakeCfgGroup:
+            def all_gather(self, tensor, separate_tensors=True):
+                assert separate_tensors
+                if tensor.shape == video_pos.shape:
+                    return [video_pos, video_neg]
+                return [audio_pos, audio_neg]
+
+        monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 2)
+        monkeypatch.setattr(ltx23, "get_cfg_group", lambda: FakeCfgGroup())
+
+        expected_video = ltx23.LTX23Pipeline._combine_x0_space_cfg(
+            video_sample,
+            video_pos,
+            video_neg,
+            video_sigma,
+            scale,
+        )
+        expected_audio = ltx23.LTX23Pipeline._combine_x0_space_cfg(
+            audio_sample,
+            audio_pos,
+            audio_neg,
+            audio_sigma,
+            scale,
+        )
+
+        for rank, expected_branch in ((0, "positive"), (1, "negative")):
+            calls = []
+            monkeypatch.setattr(ltx23, "get_classifier_free_guidance_rank", lambda rank=rank: rank)
+
+            def fake_predict_noise(**kwargs):
+                calls.append(kwargs["branch"])
+                if kwargs["branch"] == "positive":
+                    return video_pos, audio_pos
+                return video_neg, audio_neg
+
+            object.__setattr__(pipe, "predict_noise", fake_predict_noise)
+            video_combined, audio_combined = pipe.predict_noise_with_parallel_cfg(
+                true_cfg_scale=scale,
+                positive_kwargs={"branch": "positive"},
+                negative_kwargs={"branch": "negative"},
+                cfg_normalize=False,
+                video_latents=video_sample,
+                audio_latents=audio_sample,
+                video_sigma=video_sigma,
+                audio_sigma=audio_sigma,
+            )
+
+            assert calls == [expected_branch]
+            torch.testing.assert_close(video_combined, expected_video)
+            torch.testing.assert_close(audio_combined, expected_audio)
+
+        assert "_cfg_video_latents" not in pipe.__dict__
+        assert "_cfg_audio_latents" not in pipe.__dict__
+
+
+class TestCFGParallelForwardPath:
+    """Test the LTX-2.3 CFG-parallel denoising path without loading model weights."""
+
+    def test_forward_collates_request_prompt_embeds_and_mask_aliases(self, monkeypatch):
+        from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23
+        from vllm_omni.diffusion.request import OmniDiffusionRequest
+        from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+        pipe = object.__new__(ltx23.LTX23Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.device = torch.device("cpu")
+        pipe.tokenizer_max_length = 4
+        monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 1)
+
+        class StopAtEncodePromptError(Exception):
+            pass
+
+        captured = {}
+
+        def fake_encode_prompt(**kwargs):
+            captured.update(kwargs)
+            raise StopAtEncodePromptError
+
+        object.__setattr__(pipe, "encode_prompt", fake_encode_prompt)
+
+        prompt_embeds_a = torch.zeros(2, 3)
+        prompt_embeds_b = torch.ones(2, 3)
+        negative_prompt_embeds_a = torch.full((2, 3), 2.0)
+        negative_prompt_embeds_b = torch.full((2, 3), 3.0)
+        prompt_attention_mask_a = torch.tensor([True, True])
+        prompt_attention_mask_b = torch.tensor([True, False])
+        negative_attention_mask_a = torch.tensor([False, True])
+        negative_attention_mask_b = torch.tensor([False, False])
+
+        requests = [
+            OmniDiffusionRequest(
+                prompt={
+                    "prompt": "prompt-a",
+                    "negative_prompt": "negative-a",
+                    "prompt_embeds": prompt_embeds_a,
+                    "negative_prompt_embeds": negative_prompt_embeds_a,
+                    "prompt_attention_mask": prompt_attention_mask_a,
+                    "negative_prompt_attention_mask": negative_attention_mask_a,
+                },
+                sampling_params=OmniDiffusionSamplingParams(
+                    height=32,
+                    width=32,
+                    num_frames=1,
+                    frame_rate=1.0,
+                    num_inference_steps=2,
+                ),
+                request_id="ltx23-prompt-local-a",
+            ),
+            OmniDiffusionRequest(
+                prompt={
+                    "prompt": "prompt-b",
+                    "negative_prompt": "negative-b",
+                    "prompt_embeds": prompt_embeds_b,
+                    "negative_prompt_embeds": negative_prompt_embeds_b,
+                    "attention_mask": prompt_attention_mask_b,
+                    "negative_attention_mask": negative_attention_mask_b,
+                },
+                sampling_params=OmniDiffusionSamplingParams(
+                    height=32,
+                    width=32,
+                    num_frames=1,
+                    frame_rate=1.0,
+                    num_inference_steps=2,
+                ),
+                request_id="ltx23-prompt-local-b",
+            ),
+        ]
+
+        with pytest.raises(StopAtEncodePromptError):
+            pipe.forward(DiffusionRequestBatch(requests=requests))
+
+        assert captured["prompt"] is None
+        assert captured["negative_prompt"] is None
+        torch.testing.assert_close(
+            captured["prompt_embeds"],
+            torch.stack([prompt_embeds_a, prompt_embeds_b], dim=0),
+        )
+        torch.testing.assert_close(
+            captured["negative_prompt_embeds"],
+            torch.stack([negative_prompt_embeds_a, negative_prompt_embeds_b], dim=0),
+        )
+        torch.testing.assert_close(
+            captured["prompt_attention_mask"],
+            torch.stack([prompt_attention_mask_a, prompt_attention_mask_b], dim=0),
+        )
+        torch.testing.assert_close(
+            captured["negative_prompt_attention_mask"],
+            torch.stack([negative_attention_mask_a, negative_attention_mask_b], dim=0),
+        )
+
+    @pytest.mark.parametrize(("cfg_rank", "expected_prompt_value"), [(0, 1.0), (1, 0.0)])
+    @pytest.mark.parametrize(
+        ("frame_rate_input", "audio_sampling_rate", "expected_frame_rate"),
+        [(1.0, 1, 1.0), (None, 24, 24.0)],
+    )
+    def test_forward_cfg_parallel_steps_video_and_audio_scheduler(
+        self,
+        monkeypatch,
+        cfg_rank,
+        expected_prompt_value,
+        frame_rate_input,
+        audio_sampling_rate,
+        expected_frame_rate,
+    ):
+        from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23
+        from vllm_omni.diffusion.request import OmniDiffusionRequest
+        from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+        pipe = object.__new__(ltx23.LTX23Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.device = torch.device("cpu")
+        pipe.tokenizer_max_length = 1
+        pipe.vae_spatial_compression_ratio = 32
+        pipe.vae_temporal_compression_ratio = 1
+        pipe.transformer_spatial_patch_size = 1
+        pipe.transformer_temporal_patch_size = 1
+        pipe.audio_sampling_rate = audio_sampling_rate
+        pipe.audio_hop_length = 1
+        pipe.audio_vae_temporal_compression_ratio = 1
+        pipe.audio_vae_mel_compression_ratio = 1
+        pipe.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=1))
+        pipe.tokenizer = SimpleNamespace(padding_side="left")
+        pipe.vae = SimpleNamespace(
+            latents_mean=torch.zeros(2),
+            latents_std=torch.ones(2),
+            config=SimpleNamespace(scaling_factor=1.0),
+        )
+        pipe.audio_vae = SimpleNamespace(
+            latents_mean=torch.zeros(2),
+            latents_std=torch.ones(2),
+            config=SimpleNamespace(mel_bins=2, latent_channels=1),
+        )
+
+        video_pos = torch.tensor([[[0.2, -0.3]]])
+        video_neg = torch.tensor([[[-0.4, 0.1]]])
+        audio_pos = torch.tensor([[[0.7, -0.2]]])
+        audio_neg = torch.tensor([[[0.1, 0.4]]])
+
+        class FakeCfgGroup:
+            def all_gather(self, tensor, separate_tensors=True):
+                assert separate_tensors
+                if torch.equal(tensor, video_pos) or torch.equal(tensor, video_neg):
+                    return [video_pos, video_neg]
+                if torch.equal(tensor, audio_pos) or torch.equal(tensor, audio_neg):
+                    return [audio_pos, audio_neg]
+                raise AssertionError(f"Unexpected gathered tensor: {tensor}")
+
+        monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 2)
+        monkeypatch.setattr(ltx23, "get_classifier_free_guidance_rank", lambda: cfg_rank)
+        monkeypatch.setattr(ltx23, "get_cfg_group", lambda: FakeCfgGroup())
+
+        def fake_retrieve_timesteps(scheduler, num_inference_steps, device, timesteps, sigmas=None, mu=None):
+            scheduler.sigmas = torch.tensor([0.25, 0.25], device=device)
+            return torch.tensor([1.0, 0.5], device=device), 2
+
+        monkeypatch.setattr(ltx23, "retrieve_timesteps", fake_retrieve_timesteps)
+
+        class FakeScheduler:
+            def __init__(self, name="video", calls=None):
+                self.name = name
+                self.calls = [] if calls is None else calls
+                self.config = {
+                    "max_image_seq_len": 4096,
+                    "base_image_seq_len": 1024,
+                    "base_shift": 0.95,
+                    "max_shift": 2.05,
+                }
+                self.sigmas = torch.tensor([0.25, 0.25])
+
+            def __deepcopy__(self, memo):
+                return FakeScheduler("audio", self.calls)
+
+            def step(self, noise_pred, t, latents, return_dict=False, generator=None):
+                self.calls.append((self.name, noise_pred.clone(), t.clone(), latents.clone()))
+                return (latents - noise_pred,)
+
+        class FakeConnectors:
+            def to(self, device):
+                return self
+
+            def __call__(self, prompt_embeds, prompt_attention_mask, padding_side):
+                assert padding_side == "left"
+                assert prompt_embeds.shape[0] == 2
+                return prompt_embeds, prompt_embeds, prompt_attention_mask
+
+        rope_video_fps: list[float] = []
+
+        class FakeRope:
+            def prepare_video_coords(self, batch_size, num_frames, height, width, device, fps):
+                rope_video_fps.append(fps)
+                return torch.zeros(batch_size, num_frames * height * width, 3, device=device)
+
+            def prepare_audio_coords(self, batch_size, num_frames, device):
+                return torch.zeros(batch_size, num_frames, 1, device=device)
+
+        class FakeTransformer:
+            def __init__(self):
+                self.config = SimpleNamespace(in_channels=2)
+                self.rope = FakeRope()
+                self.audio_rope = FakeRope()
+                self.calls = []
+
+            def __call__(self, **kwargs):
+                self.calls.append(kwargs)
+                expected_prompt = torch.full((1, 1, 1), expected_prompt_value)
+                torch.testing.assert_close(kwargs["encoder_hidden_states"], expected_prompt)
+                torch.testing.assert_close(kwargs["audio_encoder_hidden_states"], expected_prompt)
+                assert kwargs["hidden_states"].shape == (1, 1, 2)
+                assert kwargs["audio_hidden_states"].shape == (1, 1, 2)
+                if cfg_rank == 0:
+                    return video_pos, audio_pos
+                return video_neg, audio_neg
+
+        class DummyProgress:
+            def __enter__(self):
+                return self
+
+            def __exit__(self, exc_type, exc, tb):
+                return False
+
+            def update(self):
+                pass
+
+        pipe.scheduler = FakeScheduler()
+        pipe.connectors = FakeConnectors()
+        pipe.transformer = FakeTransformer()
+        object.__setattr__(pipe, "progress_bar", lambda total: DummyProgress())
+
+        def fake_encode_prompt(**kwargs):
+            return (
+                torch.ones(1, 1, 1),
+                torch.ones(1, 1, dtype=torch.bool),
+                torch.zeros(1, 1, 1),
+                torch.ones(1, 1, dtype=torch.bool),
+            )
+
+        object.__setattr__(pipe, "encode_prompt", fake_encode_prompt)
+
+        video_latents = torch.tensor([[[1.0, -2.0]]])
+        audio_latents = torch.tensor([[[0.5, 3.0]]])
+        req = OmniDiffusionRequest(
+            prompt={"prompt": "prompt", "negative_prompt": "negative"},
+            sampling_params=OmniDiffusionSamplingParams(
+                height=32,
+                width=32,
+                num_frames=1,
+                frame_rate=frame_rate_input,
+                num_inference_steps=2,
+                guidance_scale=4.0,
+                latents=video_latents,
+                audio_latents=audio_latents,
+                output_type="latent",
+            ),
+            request_id="ltx23-cfg-parallel-forward-test",
+        )
+
+        output = pipe.forward(DiffusionRequestBatch(requests=[req]))[0]
+
+        expected_video_noise = ltx23.LTX23Pipeline._combine_x0_space_cfg(
+            video_latents,
+            video_pos,
+            video_neg,
+            pipe.scheduler.sigmas[0],
+            4.0,
+        )
+        expected_audio_noise = ltx23.LTX23Pipeline._combine_x0_space_cfg(
+            audio_latents,
+            audio_pos,
+            audio_neg,
+            pipe.scheduler.sigmas[0],
+            4.0,
+        )
+        scheduler_call_names = [call[0] for call in pipe.scheduler.calls]
+        assert scheduler_call_names == ["video", "audio", "video", "audio"]
+        assert len(pipe.transformer.calls) == 2
+        torch.testing.assert_close(pipe.scheduler.calls[0][1], expected_video_noise)
+        torch.testing.assert_close(pipe.scheduler.calls[1][1], expected_audio_noise)
+        torch.testing.assert_close(pipe.scheduler.calls[2][1], expected_video_noise)
+        torch.testing.assert_close(pipe.scheduler.calls[3][1], expected_audio_noise)
+        torch.testing.assert_close(pipe.scheduler.calls[2][3], video_latents - expected_video_noise)
+        torch.testing.assert_close(pipe.scheduler.calls[3][3], audio_latents - expected_audio_noise)
+
+        video_out, audio_out = output.output
+        torch.testing.assert_close(video_out, (video_latents - 2 * expected_video_noise).reshape(1, 2, 1, 1, 1))
+        torch.testing.assert_close(audio_out, (audio_latents - 2 * expected_audio_noise).reshape(1, 1, 1, 2))
+
+        # fps regression guard: an omitted request fps (frame_rate_input=None) must resolve
+        # to the model's own 24.0 default, not crash on None; a provided rate is passed through.
+        assert rope_video_fps
+        assert all(fps == expected_frame_rate for fps in rope_video_fps)

tests/diffusion/models/ltx2/test_ltx2_3_image2video_pipeline.py

@@ -0,0 +1,244 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Unit tests for the LTX-2.3 image-to-video pipeline."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_ltx23_request_pipe(cls):
+    pipe = object.__new__(cls)
+    torch.nn.Module.__init__(pipe)
+    pipe.device = torch.device("cpu")
+    pipe.tokenizer_max_length = 99
+    return pipe
+
+
+class TestLTX23ImageToVideoForwardStages:
+    def test_forward_resolves_request_image_and_delegates_to_shared_forward_impl(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+        from vllm_omni.diffusion.request import OmniDiffusionRequest
+        from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+        pipe = _make_ltx23_request_pipe(LTX23ImageToVideoPipeline)
+        image = torch.zeros(3, 8, 8)
+        req = DiffusionRequestBatch(
+            [
+                OmniDiffusionRequest(
+                    prompt={
+                        "prompt": "make the image move",
+                        "negative_prompt": "jitter",
+                        "multi_modal_data": {"image": image},
+                    },
+                    sampling_params=OmniDiffusionSamplingParams(
+                        height=384,
+                        width=512,
+                        num_frames=25,
+                        num_inference_steps=2,
+                    ),
+                    request_id="ltx23-i2v-forward-stage-delegation",
+                )
+            ]
+        )
+        seen = {}
+
+        def fake_forward_impl(req_arg, request_inputs, **kwargs):
+            seen["req"] = req_arg
+            seen["request_inputs"] = request_inputs
+            seen["kwargs"] = kwargs
+            return ["i2v-delegated"]
+
+        object.__setattr__(pipe, "_forward_impl", fake_forward_impl)
+
+        output = pipe.forward(req, noise_scale=0.5)
+
+        assert output == ["i2v-delegated"]
+        assert seen["req"] is req
+        assert seen["request_inputs"].prompt == ["make the image move"]
+        assert seen["request_inputs"].negative_prompt == ["jitter"]
+        assert seen["kwargs"]["image"] is image
+        assert seen["kwargs"]["noise_scale"] == 0.5
+
+    def test_denoise_timestep_kwargs_masks_video_only(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+
+        pipe = object.__new__(LTX23ImageToVideoPipeline)
+        ts = torch.tensor([2.0, 4.0])
+        denoise_ctx = SimpleNamespace(
+            conditioning_mask=torch.tensor([[1.0, 0.0]]),
+            conditioning_mask_for_model=torch.tensor([[1.0, 0.0], [0.0, 1.0]]),
+        )
+
+        kwargs = pipe._denoise_timestep_kwargs(
+            ts,
+            SimpleNamespace(cfg_parallel_ready=False),
+            denoise_ctx,
+        )
+
+        torch.testing.assert_close(kwargs["timestep"], torch.tensor([[0.0, 2.0], [4.0, 0.0]]))
+        torch.testing.assert_close(kwargs["audio_timestep"], ts)
+        torch.testing.assert_close(kwargs["sigma"], ts)
+
+
+class TestLTX23ImageToVideoPipeline:
+    def test_ltx23_i2v_pipeline_reuses_ltx23_semantics(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+
+        assert issubclass(LTX23ImageToVideoPipeline, LTX23Pipeline)
+        assert LTX23ImageToVideoPipeline.support_image_input is True
+
+    def test_ltx23_i2v_rejects_multi_image_prompt_list(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+
+        image = object()
+
+        assert LTX23ImageToVideoPipeline._resolve_single_prompt_image([image]) is image
+        with pytest.raises(ValueError, match="exactly one image per prompt"):
+            LTX23ImageToVideoPipeline._resolve_single_prompt_image([object(), object()])
+
+    def test_ltx23_i2v_additional_image_resolution_is_tensor_safe(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+
+        image = torch.zeros(1, 3, 4, 4)
+        additional = {
+            "preprocessed_image": None,
+            "pixel_values": image,
+            "image": torch.ones_like(image),
+        }
+
+        assert LTX23ImageToVideoPipeline._resolve_additional_image(additional) is image
+
+    def test_ltx23_i2v_packed_latents_are_not_noised(self, monkeypatch):
+        import vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video as ltx23_i2v
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+
+        pipe = object.__new__(LTX23ImageToVideoPipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.vae_spatial_compression_ratio = 1
+        pipe.vae_temporal_compression_ratio = 1
+        pipe.transformer_spatial_patch_size = 1
+        pipe.transformer_temporal_patch_size = 1
+
+        def fake_randn_tensor(shape, generator=None, device=None, dtype=None):
+            raise AssertionError("packed I2V latents should not be noised")
+
+        monkeypatch.setattr(ltx23_i2v, "randn_tensor", fake_randn_tensor)
+
+        latents = torch.tensor([[[10.0, 11.0], [20.0, 21.0], [30.0, 31.0]]])
+
+        out, conditioning_mask = pipe.prepare_latents(
+            image=None,
+            batch_size=1,
+            num_channels_latents=2,
+            height=1,
+            width=1,
+            num_frames=3,
+            noise_scale=1.0,
+            dtype=torch.float32,
+            device=torch.device("cpu"),
+            latents=latents,
+        )
+
+        torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
+        torch.testing.assert_close(out, latents)
+
+    def test_ltx23_i2v_5d_latents_noise_preserves_conditioning_frame(self, monkeypatch):
+        import vllm_omni.diffusion.models.ltx2.pipeline_ltx2 as ltx2
+        import vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video as ltx23_i2v
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+
+        pipe = object.__new__(LTX23ImageToVideoPipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.vae_spatial_compression_ratio = 1
+        pipe.vae_temporal_compression_ratio = 1
+        pipe.transformer_spatial_patch_size = 1
+        pipe.transformer_temporal_patch_size = 1
+        pipe.vae = SimpleNamespace(
+            latents_mean=torch.zeros(2),
+            latents_std=torch.ones(2),
+            config=SimpleNamespace(scaling_factor=1.0),
+        )
+
+        def fake_randn_tensor(shape, generator=None, device=None, dtype=None):
+            return torch.ones(shape, device=device, dtype=dtype)
+
+        monkeypatch.setattr(ltx23_i2v, "randn_tensor", fake_randn_tensor)
+        monkeypatch.setattr(ltx2, "randn_tensor", fake_randn_tensor)
+
+        latents = torch.tensor([[[[[10.0]], [[20.0]], [[30.0]]], [[[11.0]], [[21.0]], [[31.0]]]]])
+
+        out, conditioning_mask = pipe.prepare_latents(
+            image=None,
+            batch_size=1,
+            num_channels_latents=2,
+            height=1,
+            width=1,
+            num_frames=3,
+            noise_scale=1.0,
+            dtype=torch.float32,
+            device=torch.device("cpu"),
+            latents=latents,
+        )
+
+        torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
+        torch.testing.assert_close(out, torch.tensor([[[10.0, 11.0], [1.0, 1.0], [1.0, 1.0]]]))
+
+    def test_ltx23_i2v_video_step_preserves_conditioning_frame(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+
+        pipe = object.__new__(LTX23ImageToVideoPipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.transformer_spatial_patch_size = 1
+        pipe.transformer_temporal_patch_size = 1
+
+        class FakeScheduler:
+            def step(self, noise_pred, t, latents, return_dict=False):
+                return (latents + noise_pred + t,)
+
+        pipe.scheduler = FakeScheduler()
+        latents = torch.tensor([[[1.0], [2.0], [3.0]]])
+        noise_pred = torch.full_like(latents, 10.0)
+
+        out = pipe._step_video_latents_i2v(
+            noise_pred,
+            latents,
+            torch.tensor(0.5),
+            latent_num_frames=3,
+            latent_height=1,
+            latent_width=1,
+        )
+
+        torch.testing.assert_close(out[:, :1], latents[:, :1])
+        torch.testing.assert_close(out[:, 1:], latents[:, 1:] + noise_pred[:, 1:] + 0.5)
+
+
+class TestLTX23ImageToVideoModule:
+    """Test that pipeline_ltx2_3_image2video.py exposes I2V entry points."""
+
+    def test_i2v_classes_importable(self):
+        """I2V classes must be importable from the implementation module."""
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+
+        assert LTX23ImageToVideoPipeline is not None
+
+    def test_post_process_func_importable(self):
+        """get_ltx2_post_process_func must be importable for registry lookup."""
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import get_ltx2_post_process_func
+
+        assert callable(get_ltx2_post_process_func)
+
+    def test_i2v_class_matches_package_export(self):
+        """Package export must point at the I2V implementation module."""
+        from vllm_omni.diffusion.models.ltx2 import LTX23ImageToVideoPipeline as PackageExported
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import (
+            LTX23ImageToVideoPipeline as ModuleExported,
+        )
+
+        assert PackageExported is ModuleExported

tests/diffusion/models/ltx2/test_ltx2_3_vae_tiling.py

@@ -0,0 +1,282 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Unit tests for LTX-2.3 VAE tiling and distributed decode behavior."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestLTX23OutputRank:
+    def test_single_process_rank_is_output_rank(self, monkeypatch):
+        import vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 as ltx23
+
+        monkeypatch.setattr(ltx23.torch.distributed, "is_initialized", lambda: False)
+
+        assert ltx23._is_output_rank() is True
+        assert ltx23._should_decode_video_on_rank(SimpleNamespace(is_distributed_enabled=lambda: False)) is True
+
+    def test_non_output_rank_skips_decode_unless_vae_decode_is_distributed(self, monkeypatch):
+        import vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 as ltx23
+
+        monkeypatch.setattr(ltx23.torch.distributed, "is_initialized", lambda: True)
+        monkeypatch.setattr(ltx23.torch.distributed, "get_rank", lambda: 1)
+
+        assert ltx23._is_output_rank() is False
+        assert ltx23._should_decode_video_on_rank(SimpleNamespace(is_distributed_enabled=lambda: False)) is False
+        assert ltx23._should_decode_video_on_rank(SimpleNamespace(is_distributed_enabled=lambda: True)) is True
+
+
+class TestLTX23VaeDistributedDecode:
+    """Test LTX-2.3 distributed VAE helpers without loading weights."""
+
+    def test_ltx23_video_vae_is_distributed_tile_only_class(self):
+        from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
+            DistributedAutoencoderKLLTX2Video,
+        )
+        from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import DistributedVaeMixin
+
+        assert issubclass(DistributedAutoencoderKLLTX2Video, DistributedVaeMixin)
+        assert not hasattr(DistributedAutoencoderKLLTX2Video, "patch_split")
+
+    def test_ltx23_vae_executor_gathers_known_tile_shapes_and_returns_empty_on_non_rank0(self):
+        from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import LTX2VaeExecutor
+        from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+            DistributedOperator,
+            GridSpec,
+            TileTask,
+        )
+
+        z = torch.zeros(1, 1, 1, 1, 1)
+        tile_output_shapes = {
+            0: (1, 1, 1, 2, 2),
+            1: (1, 1, 1, 2, 1),
+            2: (1, 1, 1, 1, 2),
+            3: (1, 1, 1, 1, 1),
+        }
+        tasks = [
+            TileTask(0, (0, 0), z, workload=4),
+            TileTask(1, (0, 1), z, workload=2),
+            TileTask(2, (1, 0), z, workload=2),
+            TileTask(3, (1, 1), z, workload=1),
+        ]
+        grid_spec = GridSpec(
+            split_dims=(3, 4),
+            grid_shape=(2, 2),
+            tile_spec={
+                "max_tile_output_shape": (1, 1, 1, 2, 2),
+                "tile_output_shapes": tile_output_shapes,
+            },
+            output_dtype=torch.float32,
+        )
+        seen = {}
+
+        def exec_tile(task):
+            return torch.full(tile_output_shapes[task.tile_id], float(task.tile_id + 1))
+
+        def merge_tiles(coord_tensor_map, passed_grid_spec):
+            seen["merged_shapes"] = {coord: tuple(tile.shape) for coord, tile in coord_tensor_map.items()}
+            assert passed_grid_spec is grid_spec
+            return torch.stack(
+                [
+                    coord_tensor_map[(0, 0)].flatten()[0],
+                    coord_tensor_map[(0, 1)].flatten()[0],
+                    coord_tensor_map[(1, 0)].flatten()[0],
+                    coord_tensor_map[(1, 1)].flatten()[0],
+                ]
+            )
+
+        operator = DistributedOperator(split=lambda _z: (tasks, grid_spec), exec=exec_tile, merge=merge_tiles)
+
+        rank0_executor = object.__new__(LTX2VaeExecutor)
+        rank0_executor.parallel_size = 2
+        rank0_executor.world_size = 2
+        rank0_executor.rank = 0
+
+        def gather_rank0(local_tile_tensor):
+            assigned = rank0_executor._balance_tasks(tasks, 2)
+            rank1_results = [(task.tile_id, exec_tile(task)) for task in assigned[1]]
+            rank1_tile_tensor = rank0_executor._pack_local_tiles_without_meta(
+                rank1_results,
+                list(local_tile_tensor.shape),
+                z.device,
+                torch.float32,
+            )
+            seen["rank0_gather_shape"] = tuple(local_tile_tensor.shape)
+            return [local_tile_tensor, rank1_tile_tensor]
+
+        def fail_final_sync(*_args, **_kwargs):
+            raise AssertionError("broadcast_result=False should not sync the final result")
+
+        rank0_executor.gather_tensors = gather_rank0
+        rank0_executor._sync_final_result = fail_final_sync
+
+        rank0_result = rank0_executor.execute(z, operator, broadcast_result=False)
+
+        torch.testing.assert_close(rank0_result, torch.tensor([1.0, 2.0, 3.0, 4.0]))
+        assert seen["rank0_gather_shape"] == (2, 1, 1, 1, 2, 2)
+        assert seen["merged_shapes"] == {
+            (0, 0): (1, 1, 1, 2, 2),
+            (0, 1): (1, 1, 1, 2, 1),
+            (1, 0): (1, 1, 1, 1, 2),
+            (1, 1): (1, 1, 1, 1, 1),
+        }
+
+        non_rank0_executor = object.__new__(LTX2VaeExecutor)
+        non_rank0_executor.parallel_size = 2
+        non_rank0_executor.world_size = 2
+        non_rank0_executor.rank = 1
+
+        def gather_rank1(local_tile_tensor):
+            seen["rank1_gather_shape"] = tuple(local_tile_tensor.shape)
+            return None
+
+        def fail_non_rank0_merge(*_args, **_kwargs):
+            raise AssertionError("non-rank0 should not merge gathered tiles")
+
+        non_rank0_executor.gather_tensors = gather_rank1
+        non_rank0_executor._sync_final_result = fail_final_sync
+
+        empty_result = non_rank0_executor.execute(
+            z,
+            DistributedOperator(
+                split=lambda _z: (tasks, grid_spec),
+                exec=exec_tile,
+                merge=fail_non_rank0_merge,
+            ),
+            broadcast_result=False,
+        )
+
+        assert tuple(empty_result.shape) == (0,)
+        assert seen["rank1_gather_shape"] == (2, 1, 1, 1, 2, 2)
+
+
+class TestLTX23VaeTiling:
+    """Test LTX-2.3 video VAE tile helpers without loading weights."""
+
+    def test_ltx23_video_vae_tile_split_uses_native_ltx23_tile_geometry(self):
+        from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
+            DistributedAutoencoderKLLTX2Video,
+        )
+
+        vae = SimpleNamespace(
+            spatial_compression_ratio=32,
+            tile_sample_min_height=512,
+            tile_sample_min_width=512,
+            tile_sample_stride_height=448,
+            tile_sample_stride_width=448,
+            temporal_compression_ratio=8,
+            dtype=torch.float32,
+        )
+
+        z = torch.zeros(1, 2, 5, 16, 24)
+        tasks, grid_spec = DistributedAutoencoderKLLTX2Video.tile_split(vae, z)
+
+        assert grid_spec.grid_shape == (2, 2)
+        assert grid_spec.split_dims == (3, 4)
+        assert grid_spec.tile_spec["sample_height"] == 512
+        assert grid_spec.tile_spec["sample_width"] == 768
+        assert grid_spec.tile_spec["blend_height"] == 64
+        assert grid_spec.tile_spec["blend_width"] == 64
+        assert grid_spec.tile_spec["max_tile_output_shape"] == (1, 3, 33, 512, 512)
+        assert grid_spec.tile_spec["tile_output_shapes"] == {
+            0: (1, 3, 33, 512, 512),
+            1: (1, 3, 33, 512, 320),
+            2: (1, 3, 33, 64, 512),
+            3: (1, 3, 33, 64, 320),
+        }
+        assert [task.grid_coord for task in tasks] == [(0, 0), (0, 1), (1, 0), (1, 1)]
+        assert [tuple(task.tensor.shape) for task in tasks] == [
+            (1, 2, 5, 16, 16),
+            (1, 2, 5, 16, 10),
+            (1, 2, 5, 2, 16),
+            (1, 2, 5, 2, 10),
+        ]
+        assert [task.workload for task in tasks] == [5 * 16 * 16, 5 * 16 * 10, 5 * 2 * 16, 5 * 2 * 10]
+
+    def test_ltx23_video_vae_tile_merge_blends_and_crops_like_tiled_decode(self):
+        from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
+            DistributedAutoencoderKLLTX2Video,
+        )
+        from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import GridSpec
+
+        class FakeVae:
+            def __init__(self):
+                self.blend_calls = []
+
+            def clear_cache(self):
+                pass
+
+            def blend_v(self, _previous, current, blend_height):
+                self.blend_calls.append(("v", blend_height))
+                return current
+
+            def blend_h(self, _previous, current, blend_width):
+                self.blend_calls.append(("h", blend_width))
+                return current
+
+        fake_vae = FakeVae()
+        grid_spec = GridSpec(
+            split_dims=(3, 4),
+            grid_shape=(2, 2),
+            tile_spec={
+                "sample_height": 10,
+                "sample_width": 10,
+                "blend_height": 1,
+                "blend_width": 2,
+                "tile_sample_stride_height": 5,
+                "tile_sample_stride_width": 5,
+            },
+        )
+        tiles = {
+            (0, 0): torch.full((1, 3, 2, 6, 6), 1.0),
+            (0, 1): torch.full((1, 3, 2, 6, 6), 2.0),
+            (1, 0): torch.full((1, 3, 2, 6, 6), 3.0),
+            (1, 1): torch.full((1, 3, 2, 6, 6), 4.0),
+        }
+
+        merged = DistributedAutoencoderKLLTX2Video.tile_merge(fake_vae, tiles, grid_spec)
+
+        assert merged.shape == (1, 3, 2, 10, 10)
+        assert fake_vae.blend_calls == [("h", 2), ("v", 1), ("v", 1), ("h", 2)]
+        torch.testing.assert_close(merged[:, :, :, :5, :5], torch.ones(1, 3, 2, 5, 5))
+        torch.testing.assert_close(merged[:, :, :, :5, 5:], torch.full((1, 3, 2, 5, 5), 2.0))
+        torch.testing.assert_close(merged[:, :, :, 5:, :5], torch.full((1, 3, 2, 5, 5), 3.0))
+        torch.testing.assert_close(merged[:, :, :, 5:, 5:], torch.full((1, 3, 2, 5, 5), 4.0))
+
+    def test_ltx23_video_vae_tiled_decode_dispatches_to_tile_operator(self):
+        from vllm_omni.diffusion.distributed.autoencoders import autoencoder_kl_ltx2
+
+        z = torch.zeros(1, 2, 1, 16, 24)
+        expected = torch.ones(1, 3, 1, 512, 768)
+        seen = {}
+
+        class FakeExecutor:
+            def execute(self, tensor, operator, broadcast_result=True):
+                seen["tensor"] = tensor
+                seen["operator"] = operator
+                seen["broadcast_result"] = broadcast_result
+                return expected
+
+        vae = SimpleNamespace(distributed_executor=FakeExecutor(), is_distributed_enabled=lambda: True)
+        vae.tile_split = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_split.__get__(vae)
+        vae.tile_exec = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_exec.__get__(vae)
+        vae.tile_merge = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_merge.__get__(vae)
+
+        output = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tiled_decode(
+            vae,
+            z,
+            temb=torch.tensor(0.5),
+            return_dict=False,
+        )
+
+        assert len(output) == 1
+        assert output[0] is expected
+        assert seen["tensor"] is z
+        assert seen["broadcast_result"] is False
+        assert seen["operator"].split.__name__ == "tile_split"
+        assert seen["operator"].merge.__name__ == "tile_merge"

tests/diffusion/models/qwen_image/test_qwen_image_max_sequence_length.py

@@ -1,4 +1,3 @@
-import inspect
 from types import SimpleNamespace
 
 import pytest
@@ -193,6 +192,7 @@ def _make_request_batch_prompt_sampling(**overrides):
         "true_cfg_scale": None,
         "guidance_scale_provided": False,
         "guidance_scale": 1.0,
+        "output_type": None,
     }
     values.update(overrides)
     return SimpleNamespace(**values)
@@ -344,16 +344,3 @@ def test_qwen_edit_validator_excludes_image_placeholders_from_budget(pipeline_cl
 
     with pytest.raises(AssertionError, match="text encoder should not run"):
         pipeline.encode_prompt(prompt="short prompt")
-
-
-@pytest.mark.parametrize(
-    "pipeline_class",
-    [
-        QwenImagePipeline,
-        QwenImageLayeredPipeline,
-        QwenImageEditPipeline,
-        QwenImageEditPlusPipeline,
-    ],
-)
-def test_forward_max_sequence_length_default_is_1024(pipeline_class: type):
-    assert inspect.signature(pipeline_class.forward).parameters["max_sequence_length"].default == 1024

tests/diffusion/models/sd3/test_sd3_pipeline.py

@@ -21,6 +21,7 @@ def _make_sd3_sampling(**overrides):
         "generator": None,
         "latents": None,
         "guidance_scale": 4.0,
+        "output_type": None,
     }
     values.update(overrides)
     return SimpleNamespace(**values)

tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py

@@ -20,6 +20,12 @@ class _StubTransformer(nn.Module):
         return torch.float32
 
 
+class _StubTextEncoder(nn.Module):
+    @property
+    def dtype(self) -> torch.dtype:
+        return torch.float32
+
+
 class _StubScheduler:
     def __init__(self, timesteps: list[int]) -> None:
         self.timesteps = torch.tensor(timesteps, dtype=torch.int64)
@@ -41,12 +47,30 @@ def _noop_progress_bar(*args, **kwargs):
     yield _Bar()
 
 
+def _stub_encode_prompt(
+    prompt,
+    negative_prompt=None,
+    do_classifier_free_guidance=True,
+    num_videos_per_prompt=1,
+    max_sequence_length=512,
+    device=None,
+    dtype=None,
+):
+    del negative_prompt, do_classifier_free_guidance, device, dtype
+    batch_size = 1 if isinstance(prompt, str) else len(prompt)
+    n = batch_size * num_videos_per_prompt
+    hidden_size = 8
+    prompt_embeds = torch.zeros(n, max_sequence_length, hidden_size)
+    return prompt_embeds, None
+
+
 def _make_pipeline() -> Wan22Pipeline:
     pipeline = object.__new__(Wan22Pipeline)
     nn.Module.__init__(pipeline)
     pipeline.device = torch.device("cpu")
     pipeline.transformer = _StubTransformer()
     pipeline.transformer_2 = None
+    pipeline.text_encoder = _StubTextEncoder()
     pipeline.transformer_config = SimpleNamespace(patch_size=(1, 2, 2), in_channels=4, out_channels=4)
     pipeline.scheduler = _StubScheduler([9, 5])
     pipeline.od_config = SimpleNamespace(flow_shift=5.0)
@@ -61,15 +85,14 @@ def _make_pipeline() -> Wan22Pipeline:
     pipeline._num_timesteps = None
     pipeline._current_timestep = None
     pipeline.check_inputs = lambda **kwargs: None
+    pipeline.encode_prompt = _stub_encode_prompt  # type: ignore[method-assign]
     pipeline.prepare_latents = lambda **kwargs: torch.zeros((1, 4, 1, 8, 8), dtype=torch.float32)
     pipeline.progress_bar = _noop_progress_bar
     return pipeline
 
 
-def test_forward_delegates_denoising_to_diffuse(monkeypatch) -> None:
+def test_forward_delegates_denoising_to_diffuse() -> None:
     pipeline = _make_pipeline()
-
-    prompt_embeds = torch.randn(1, 8)
     captured: dict[str, object] = {}
 
     def _fake_diffuse(**kwargs):
@@ -86,8 +109,8 @@ def test_forward_delegates_denoising_to_diffuse(monkeypatch) -> None:
             width=None,
             num_frames=1,
             num_inference_steps=2,
-            guidance_scale_provided=False,
-            guidance_scale=None,
+            guidance_scale_provided=True,
+            guidance_scale=1.0,
             guidance_scale_2=None,
             boundary_ratio=None,
             generator=None,
@@ -95,14 +118,16 @@ def test_forward_delegates_denoising_to_diffuse(monkeypatch) -> None:
             num_outputs_per_prompt=1,
             max_sequence_length=32,
             latents=None,
+            output_type="latent",
             extra_args={},
         ),
     )
     batch = DiffusionRequestBatch(requests=[mock_req])
 
-    output = pipeline.forward(batch, prompt_embeds=prompt_embeds, output_type="latent", guidance_scale=1.0)
+    output = pipeline.forward(batch)
 
     assert torch.equal(output.output, torch.ones((1, 4, 1, 8, 8)))
+    assert torch.equal(captured["prompt_embeds"], torch.zeros(1, 32, 8))
     assert torch.equal(captured["timesteps"], pipeline.scheduler.timesteps)
     assert captured["guidance_low"] == 1.0
     assert captured["guidance_high"] == 1.0

tests/diffusion/quantization/test_mxfp4_config.py

@@ -537,3 +537,118 @@ def test_single_scale_row_parallel_tp2_loader_simulation():
 
     for rank in range(_TP2):
         assert _shard(ckpt_weight_scale, layer.weight_scale, rank, _TP2, "input_dim").shape == layer.weight_scale.shape
+
+
+# ---------------------------------------------------------------------------
+# ROCm MXFP4 (gfx950) — get_quant_method dispatch
+#
+# These run on CPU: the platform check, gcnArchName probe and aiter op
+# registration are all mocked (stdlib unittest.mock + monkeypatch, so no
+# pytest-mock dependency).  They cover only the dispatch + weight-allocation
+# logic added by the ROCm PR — the AITER GEMM / quant kernels require real
+# gfx950 hardware and are intentionally NOT exercised here.
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def _rocm_platform(monkeypatch: pytest.MonkeyPatch):
+    """Make current_omni_platform report ROCm, and stub the aiter
+    custom-op registration so ROCmMxfp4*Method can be constructed without aiter."""
+    from vllm_omni.platforms import current_omni_platform
+    from vllm_omni.quantization import mxfp4_config
+
+    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: False)
+    monkeypatch.setattr(current_omni_platform, "is_rocm", lambda: True)
+    monkeypatch.setattr(mxfp4_config, "_register_rocm_mxfp4_op", lambda: None)
+
+
+def _patch_gcn_arch(monkeypatch: pytest.MonkeyPatch, arch: str) -> None:
+    """Patch torch.cuda.get_device_properties(...).gcnArchName to return `arch`."""
+    from types import SimpleNamespace
+
+    monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 0)
+    monkeypatch.setattr(
+        torch.cuda,
+        "get_device_properties",
+        lambda *a, **k: SimpleNamespace(gcnArchName=arch),
+    )
+
+
+def _fake_linear_layer():
+    """A stand-in that passes isinstance(layer, LinearBase) without a real layer."""
+    from unittest.mock import MagicMock
+
+    from vllm.model_executor.layers.linear import LinearBase
+
+    return MagicMock(spec=LinearBase)
+
+
+def test_rocm_online_dispatch_returns_rocm_method(_rocm_platform, monkeypatch):
+    """ROCm + gfx950 + online checkpoint must return ROCmMxfp4OnlineLinearMethod."""
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, ROCmMxfp4OnlineLinearMethod
+
+    _patch_gcn_arch(monkeypatch, "gfx950:sramecc+:xnack-")
+    cfg = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=False)
+
+    method = cfg.get_quant_method(_fake_linear_layer(), "blocks.0.attn1.to_q")
+    assert isinstance(method, ROCmMxfp4OnlineLinearMethod)
+
+
+def test_rocm_ignored_layer_returns_unquantized(_rocm_platform, monkeypatch):
+    """A prefix in ignored_layers must return UnquantizedLinearMethod before the gfx950 probe."""
+    from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
+
+    _patch_gcn_arch(monkeypatch, "gfx950:sramecc+:xnack-")
+    cfg = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=False, ignored_layers=["proj_out"])
+
+    assert isinstance(cfg.get_quant_method(_fake_linear_layer(), "proj_out"), UnquantizedLinearMethod)
+
+
+def test_rocm_non_gfx950_raises(_rocm_platform, monkeypatch):
+    """MXFP4 on ROCm requires gfx950; any other arch must raise."""
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
+
+    _patch_gcn_arch(monkeypatch, "gfx942:sramecc+:xnack-")
+    cfg = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=False)
+
+    with pytest.raises(NotImplementedError, match="gfx950"):
+        cfg.get_quant_method(_fake_linear_layer(), "blocks.0.attn1.to_q")
+
+
+# ---------------------------------------------------------------------------
+# ROCm MXFP4 — create_weights (online lazy meta-device placeholder)
+#
+# The base ROCmMxfp4LinearMethod is abstract (no create_weights); the online
+# subclass gets create_weights from _LazyWeightMixin, which registers a BF16
+# weight on the meta device to be materialised at load time.
+# ---------------------------------------------------------------------------
+
+
+def test_rocm_create_weights_column_parallel_tp2(_rocm_platform):
+    """Column-parallel TP=2: meta BF16 weight has the output halved, input full."""
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, ROCmMxfp4OnlineLinearMethod
+
+    method = ROCmMxfp4OnlineLinearMethod(DiffusionMXFP4Config())
+    layer = _create_weights(method, input_size_per_partition=_TP2_K, output_partition_sizes=[_TP2_N // _TP2])
+
+    assert layer.weight.shape == (_TP2_N // _TP2, _TP2_K)
+    assert layer.weight.dtype == torch.bfloat16
+    assert layer.weight.device.type == "meta"
+    assert layer.weight.input_dim == 1
+    assert layer.weight.output_dim == 0
+    assert layer.logical_widths == [_TP2_N // _TP2]
+    assert layer.weight_block_size is None
+
+
+def test_rocm_create_weights_row_parallel_tp2(_rocm_platform):
+    """Row-parallel TP=2: meta BF16 weight has the input halved, output full."""
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, ROCmMxfp4OnlineLinearMethod
+
+    method = ROCmMxfp4OnlineLinearMethod(DiffusionMXFP4Config())
+    layer = _create_weights(method, input_size_per_partition=_TP2_K // _TP2, output_partition_sizes=[_TP2_N])
+
+    assert layer.weight.shape == (_TP2_N, _TP2_K // _TP2)
+    assert layer.input_size_per_partition == _TP2_K // _TP2
+    assert layer.output_size_per_partition == _TP2_N

tests/diffusion/test_diffusion_step_pipeline.py

@@ -30,6 +30,9 @@ from vllm_omni.diffusion.ipc import (
     pack_diffusion_output_shm,
     unpack_diffusion_output_shm,
 )
+from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import (
+    DiffusionPipelineProfilerMixin,
+)
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched import StepScheduler
 from vllm_omni.diffusion.sched.interface import (
@@ -97,6 +100,51 @@ class _StepPipeline:
         return DiffusionOutput(output=torch.tensor([state.step_index], dtype=torch.float32))
 
 
+class _ProfilingStepPipeline(_StepPipeline):
+    enable_diffusion_pipeline_profiler = True
+
+    def __init__(self):
+        super().__init__()
+        self._stage_durations: dict[str, float] = {}
+
+    @property
+    def stage_durations(self) -> dict[str, float]:
+        return dict(self._stage_durations)
+
+    def clear_profiler_records(self) -> None:
+        self._stage_durations.clear()
+
+    def prepare_encode(self, state, **kwargs):
+        result = super().prepare_encode(state, **kwargs)
+        self._stage_durations["QwenImagePipeline.text_encoder.forward"] = 1.0
+        return result
+
+    def denoise_step(self, input_batch, **kwargs):
+        result = super().denoise_step(input_batch, **kwargs)
+        self._stage_durations["QwenImagePipeline.diffuse"] = 2.0
+        return result
+
+    def post_decode(self, state, **kwargs):
+        result = super().post_decode(state, **kwargs)
+        self._stage_durations["QwenImagePipeline.vae.decode"] = 3.0
+        return result
+
+
+class _AutoDenoiseProfilerPipeline(DiffusionPipelineProfilerMixin):
+    _PROFILER_TARGETS: list[str] = []
+
+    def __init__(self):
+        self.setup_diffusion_pipeline_profiler(
+            enable_diffusion_pipeline_profiler=True,
+        )
+
+    def forward(self):
+        return None
+
+    def denoise_step(self):
+        return "ok"
+
+
 class _InterruptingStepPipeline(_StepPipeline):
     interrupt = True
 
@@ -130,6 +178,9 @@ class _FakePeakMemoryPlatform:
         index = min(self.reset_calls - 1, len(self._reserved_mb) - 1)
         return int((self._reserved_mb[index] - 100) * 1024**2)
 
+    def is_available(self) -> bool:
+        return True
+
 
 class _IdentityNoiseTransformer(torch.nn.Module):
     def forward(self, x: torch.Tensor, **kwargs):
@@ -437,6 +488,16 @@ def test_input_batch_cached_repack_keeps_static_prompt_fields_for_same_compositi
     torch.testing.assert_close(repacked.prompt_embeds, torch.ones(1, 2, 3))
 
 
+@pytest.mark.cpu
+def test_step_profiler_reports_denoise_step_as_diffuse():
+    pipeline = _AutoDenoiseProfilerPipeline()
+
+    assert pipeline.denoise_step() == "ok"
+
+    assert any(key.endswith(".diffuse") for key in pipeline.stage_durations)
+    assert not any(key.endswith(".denoise_step") for key in pipeline.stage_durations)
+
+
 @pytest.mark.cpu
 class TestRunner:
     """DiffusionModelRunner.execute_stepwise"""
@@ -469,6 +530,53 @@ class TestRunner:
         assert runner.pipeline.scheduler_calls == 2
         assert runner.pipeline.decode_calls == 1
 
+    def test_stepwise_output_includes_stage_and_peak_metrics(self, monkeypatch):
+        runner = _make_runner()
+        runner.pipeline = _ProfilingStepPipeline()
+        req = _make_step_request()
+        reset_calls = []
+        monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+        monkeypatch.setattr(
+            model_runner_module.current_omni_platform,
+            "is_available",
+            lambda: True,
+        )
+        monkeypatch.setattr(
+            model_runner_module.current_omni_platform,
+            "reset_peak_memory_stats",
+            lambda: reset_calls.append(True),
+        )
+        monkeypatch.setattr(
+            model_runner_module.current_omni_platform,
+            "max_memory_reserved",
+            lambda: 2 * 1024**2,
+        )
+        monkeypatch.setattr(
+            model_runner_module.current_omni_platform,
+            "max_memory_allocated",
+            lambda: 1024**2,
+        )
+
+        DiffusionModelRunner.execute_stepwise(
+            runner,
+            _make_scheduler_output(req, step_id=0),
+        )
+        result = DiffusionModelRunner.execute_stepwise(
+            runner,
+            _make_cached_scheduler_output(step_id=1),
+        )
+
+        output = result.get_request_output("req-1")
+        assert output.finished is True
+        assert output.result is not None
+        assert output.result.peak_memory_mb == 2
+        assert output.result.stage_durations == {
+            "QwenImagePipeline.text_encoder.forward": 1.0,
+            "QwenImagePipeline.diffuse": 4.0,
+            "QwenImagePipeline.vae.decode": 3.0,
+        }
+        assert reset_calls == [True]
+
     def test_carries_peak_memory_across_stepwise_request_lifecycle(self, monkeypatch):
         runner = _make_runner()
         req = _make_step_request()

tests/e2e/accuracy/test_qwen_image.py

@@ -29,7 +29,7 @@ HEIGHT = 512
 NUM_INFERENCE_STEPS = 20
 TRUE_CFG_SCALE = 4.0
 SEED = 42
-SSIM_THRESHOLD = 0.97
+SSIM_THRESHOLD = 0.94
 PSNR_THRESHOLD = 30.0
 
 MODEL_2512_ID = "Qwen/Qwen-Image-2512"

tests/e2e/online_serving/test_voxcpm2_tts_expansion.py

@@ -74,7 +74,6 @@ def test_voice_clone_streaming_001(omni_server, openai_client) -> None:
 
 @hardware_test(res={"cuda": "L4"}, num_cards=1)
 @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-@pytest.mark.skip(reason="#4335")
 def test_response_format_001(omni_server, openai_client) -> None:
     """
     Test voice-clone non-stream PCM output via OpenAI API.
@@ -93,5 +92,6 @@ def test_response_format_001(omni_server, openai_client) -> None:
         "voice": "default",
         "ref_audio": REF_AUDIO_URL,
         "min_audio_bytes": _MIN_AUDIO_BYTES,
+        "min_hnr_db": -2.0,
     }
     openai_client.send_audio_speech_request(request_config)

tests/entrypoints/openai_api/test_serving_speech.py

@@ -2530,7 +2530,7 @@ def test_api_server_delete_voice_exception_returns_500(mocker: MockerFixture):
 def test_api_server_create_speech_without_handler_returns_404(mocker: MockerFixture):
     fake_base = _patch_api_server_base(mocker)
     raw_request = _make_api_server_request(None, path="/v1/audio/speech")
-    raw_request.app.state.openai_serving_tokenization = fake_base
+    raw_request.app.state.serving_tokenization = fake_base
     request = OpenAICreateSpeechRequest(input="Hello")
 
     response = asyncio.run(api_server_module.create_speech(request, raw_request))
@@ -2543,7 +2543,7 @@ def test_api_server_create_speech_without_handler_returns_404(mocker: MockerFixt
 def test_api_server_create_speech_batch_without_handler_returns_404(mocker: MockerFixture):
     fake_base = _patch_api_server_base(mocker)
     raw_request = _make_api_server_request(None, path="/v1/audio/speech/batch")
-    raw_request.app.state.openai_serving_tokenization = fake_base
+    raw_request.app.state.serving_tokenization = fake_base
     request = BatchSpeechRequest(items=[SpeechBatchItem(input="hi")])
 
     response = asyncio.run(api_server_module.create_speech_batch(request, raw_request))
@@ -2612,7 +2612,7 @@ def test_api_server_create_audio_generate_without_handler_returns_404(mocker: Mo
     fake_base = _patch_api_server_base(mocker)
     raw_request = _make_api_server_request(None, path="/v1/audio/generate")
     raw_request.app.state.openai_serving_audio_generate = None
-    raw_request.app.state.openai_serving_tokenization = fake_base
+    raw_request.app.state.serving_tokenization = fake_base
     request = OpenAICreateAudioGenerateRequest(input="a bird singing")
 
     response = asyncio.run(api_server_module.create_audio_generate(request, raw_request))

tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_streaming.py

@@ -0,0 +1,235 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for ``Qwen3OmniMoeCode2Wav.chunked_decode_streaming``.
+
+Background
+----------
+The Code2Wav decoder is a causal conv / transposed-conv stack whose forward
+emits a fixed number ``C`` *fewer* samples than ``frames * total_upsample`` (the
+right-edge trim, ``C == 555`` for Qwen3-Omni-30B-A3B). ``chunked_decode_streaming``
+used to slice ``[left_context * up : code_seq_len * up]``, assuming the full
+length. Because the real output is ``C`` short, the upper bound clamped and every
+streaming chunk silently dropped its last ``C`` samples -> a gap/click at every
+chunk boundary plus ~1.2% cumulative time compression.
+
+The fix measures ``tail`` from the actual decoder output and shifts each chunk's
+start back by ``tail`` so the chunk refills the gap the previous chunk left.
+
+Why this is path-agnostic (eager *and* CUDA graph)
+--------------------------------------------------
+Originally the CUDA-graph wrapper returned a *surplus* (nominal ``actual * up``,
+with stale padded-tail samples), so ``tail`` evaluated to 0 and the graph path
+looked immune. Since #4466 (``CUDAGraphDecoderWrapper._trim_replay_output``) the
+graph wrapper returns the eager-equivalent ``actual * up - C`` instead, so the
+same per-chunk tail drop happens on the graph branch too and the same fix
+repairs it. These tests cover both branches; ``_FakeWrapper`` models the
+post-#4466 contract (``decode`` returns ``actual * up - C``).
+
+These are pure length/slice-arithmetic tests: CPU-only, no weights, no GPU.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav import (
+    Qwen3OmniMoeCode2Wav,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_UP = 8  # total_upsample
+_TRIM = 5  # causal right-edge trim C, 0 < C < up (mirrors 555 < 1920)
+_Q = 4  # num_quantizers
+
+
+class _ShortTrimDecoder:
+    """Minimal stand-in exposing only what ``chunked_decode_streaming`` touches.
+
+    Models a *shift-equivariant* causal decoder: the waveform value at an
+    absolute output sample ``q`` is ``floor(q / up)`` == the index of the frame
+    that sample belongs to. With codes set so frame ``i`` carries value ``i``,
+    any window covering a frame (with enough right context) reproduces the same
+    global waveform -- which is exactly the property that lets the next chunk
+    refill the tail the previous chunk dropped. The forward emits ``trim`` fewer
+    samples than ``frames * up`` (the right-edge trim that lacks right context).
+    """
+
+    def __init__(self, *, up=_UP, trim=_TRIM, num_quantizers=_Q, cudagraph=False):
+        self.total_upsample = up
+        self.trim = trim
+        self.config = SimpleNamespace(num_quantizers=num_quantizers)
+        self._cudagraph_enabled = cudagraph
+        self._cudagraph_wrapper = _FakeWrapper(self) if cudagraph else None
+
+    def _decode(self, codes: torch.Tensor) -> torch.Tensor:
+        # codes: [B, Q, W]; frame value (row 0) == absolute frame index.
+        batch, _, width = codes.shape
+        out_len = max(0, width * self.total_upsample - self.trim)
+        frame_of = torch.arange(out_len) // self.total_upsample  # [out_len]
+        vals = codes[:, 0, :].to(torch.float32)  # [B, W]
+        out = vals.gather(1, frame_of.unsqueeze(0).expand(batch, -1))  # [B, out_len]
+        return out.unsqueeze(1)  # [B, 1, out_len]
+
+    # ``chunked_decode_streaming`` calls ``self(codes)`` on the eager branch.
+    def __call__(self, codes: torch.Tensor) -> torch.Tensor:
+        return self._decode(codes)
+
+
+class _FakeWrapper:
+    """Models ``CUDAGraphDecoderWrapper.decode`` after #4466: returns the
+    eager-equivalent short length (``actual * up - C``), values correct."""
+
+    def __init__(self, decoder: _ShortTrimDecoder):
+        self.decoder = decoder
+
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        return self.decoder._decode(codes)
+
+
+def _make_codes(frames: int, q: int = _Q) -> torch.Tensor:
+    """codes[0, :, i] == i so each frame carries its own absolute index."""
+    idx = torch.arange(frames, dtype=torch.long)
+    return idx.view(1, 1, frames).expand(1, q, frames).contiguous()
+
+
+def _decode_chunk(model, window, lc):
+    """Call the real production method for one streaming window."""
+    win_frames = window.shape[-1]
+    out = Qwen3OmniMoeCode2Wav.chunked_decode_streaming(
+        model,
+        window,
+        left_context_size=[lc],
+        seq_token_counts=[win_frames * model.config.num_quantizers],
+    )
+    return out[0].reshape(-1)
+
+
+def _drive_async(model, codes, *, initial=4, chunk=25, left_context=25):
+    """Replicate qwen3_omni's ``talker2code2wav_async_chunk`` 4/25/25 schedule."""
+    total = codes.shape[-1]
+    chunks = []
+    emitted = 0
+    first = True
+    while emitted < total:
+        target = min(initial, total) if first else min(emitted + chunk, total)
+        lc = min(emitted, left_context)
+        window = codes[:, :, emitted - lc : target]
+        chunks.append(_decode_chunk(model, window, lc))
+        emitted = target
+        first = False
+    return torch.cat(chunks, dim=-1)
+
+
+def _single_shot(model, codes):
+    """Ground-truth: one full-utterance forward (length T*up - trim)."""
+    return model(codes)[0].reshape(-1)
+
+
+def test_streaming_reconstructs_single_shot_eager():
+    """Eager: concatenated streaming chunks == the single-shot full decode."""
+    model = _ShortTrimDecoder()
+    codes = _make_codes(116)
+
+    gt = _single_shot(model, codes)
+    streamed = _drive_async(model, codes)
+
+    assert streamed.shape[-1] == gt.shape[-1]
+    torch.testing.assert_close(streamed, gt)
+
+
+def test_streaming_reconstructs_single_shot_cudagraph():
+    """CUDA-graph branch (post-#4466 contract): identical exact reconstruction.
+
+    Answers the first-principles question: the non-eager path is NOT immune on
+    current main -- it drops the same tail without the fix, and the fix repairs
+    it the same way.
+    """
+    model = _ShortTrimDecoder(cudagraph=True)
+    codes = _make_codes(116)
+
+    gt = _single_shot(model, codes)  # forward() is eager regardless of the branch
+    streamed = _drive_async(model, codes)
+
+    assert streamed.shape[-1] == gt.shape[-1]
+    torch.testing.assert_close(streamed, gt)
+
+
+def test_eager_and_cudagraph_streaming_match():
+    """Both decode branches must yield bit-identical streamed audio."""
+    codes = _make_codes(116)
+    eager = _drive_async(_ShortTrimDecoder(cudagraph=False), codes)
+    graph = _drive_async(_ShortTrimDecoder(cudagraph=True), codes)
+    torch.testing.assert_close(eager, graph)
+
+
+def test_per_chunk_tail_drop_is_repaired_not_present():
+    """The defect is exactly ``C`` lost per internal boundary; the fix removes it.
+
+    Without the fix the streamed length would be short by
+    ``(num_internal_boundaries) * trim``; assert the boundary count and that the
+    repaired stream has the full single-shot length (no deficit).
+    """
+    model = _ShortTrimDecoder()
+    codes = _make_codes(116)
+    gt = _single_shot(model, codes)
+    streamed = _drive_async(model, codes)
+
+    # 116 frames @ initial=4, chunk=25 -> targets 4,29,54,79,104,116 => 5 boundaries
+    num_internal_boundaries = 5
+    buggy_len = gt.shape[-1] - num_internal_boundaries * _TRIM
+    assert streamed.shape[-1] == gt.shape[-1]
+    assert streamed.shape[-1] != buggy_len  # the pre-fix slice produced buggy_len
+
+
+def test_first_chunk_start_is_not_shifted_below_zero():
+    """First chunk has left_context==0 -> start = max(0, 0 - tail) = 0 (unchanged)."""
+    model = _ShortTrimDecoder()
+    codes = _make_codes(4)  # single initial chunk, lc=0
+    out = _decode_chunk(model, codes, lc=0)
+    # length == frames*up - trim, values == frame indices (GT prefix)
+    assert out.shape[-1] == 4 * _UP - _TRIM
+    torch.testing.assert_close(out, _single_shot(model, codes))
+
+
+def test_no_shift_when_decoder_returns_nominal_length():
+    """C==0 decoder (e.g. Qwen3-TTS): tail==0 -> genuine no-op.
+
+    With no trim, the fixed slice must equal the original
+    ``[left_context*up : code_seq_len*up]`` -- every chunk emits exactly
+    ``ctx_len*up`` samples and streaming reconstructs the full (untrimmed) audio.
+    """
+    model = _ShortTrimDecoder(trim=0)
+    codes = _make_codes(116)
+    gt = _single_shot(model, codes)
+    streamed = _drive_async(model, codes)
+
+    assert gt.shape[-1] == 116 * _UP  # no trim -> nominal full length
+    assert streamed.shape[-1] == gt.shape[-1]
+    torch.testing.assert_close(streamed, gt)
+
+
+def test_batched_streaming_only_shifts_the_boundary_row():
+    """Per-row tail: in a shared window the shorter row had right context from
+    the longer row's frames, so it lost nothing (tail==0, no shift); only the
+    row at the decoded boundary (code_seq_len == window length) is shifted."""
+    model = _ShortTrimDecoder()
+    # One window of W=30 frames, two rows. codes[:, 0, :] are absolute indices.
+    idx = torch.arange(30, dtype=torch.long)
+    window = idx.view(1, 1, 30).expand(2, _Q, 30).contiguous()
+
+    long_frames, short_frames = 30, 20  # row0 spans the full window, row1 is shorter
+    lc = 5
+    out = Qwen3OmniMoeCode2Wav.chunked_decode_streaming(
+        model,
+        window,
+        left_context_size=[lc, lc],
+        seq_token_counts=[long_frames * _Q, short_frames * _Q],
+    )
+
+    batch_len = 30 * _UP - _TRIM  # actual decoder output width
+    # Row 0 (boundary row): tail == TRIM -> start = lc*up - TRIM
+    assert out[0].shape[-1] == batch_len - (lc * _UP - _TRIM)
+    # Row 1 (shorter): code_seq_len*up < batch_len -> tail == 0 -> start = lc*up
+    assert out[1].shape[-1] == short_frames * _UP - lc * _UP

tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py

@@ -34,6 +34,7 @@ def _make_minimal_talker(
     """
     model = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration)
     model.talker_config = SimpleNamespace(codec_pad_id=7, num_code_groups=16)
+    model._embedding_dtype = torch.bfloat16
     if tts_pad_embed is None:
         tts_pad_embed = torch.zeros((1, 4), dtype=torch.bfloat16)
     model._tts_pad_embed = tts_pad_embed
@@ -82,6 +83,7 @@ def _make_minimal_builder(
     builder._codec_embed = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device)
     builder._residual_code_embeddings = lambda: []
     builder._speaker_encoder = None
+    builder._embedding_dtype = torch.bfloat16
     builder._tts_pad_embed_buffer = (
         tts_pad_embed if tts_pad_embed is not None else torch.zeros((1, 4), dtype=torch.bfloat16)
     )
@@ -254,6 +256,36 @@ def test_decode_compacts_long_trailing_text_after_large_offset():
     assert torch.equal(update["hidden_states"]["trailing_text"], trailing_text[65:])
 
 
+def test_decode_replay_span_embeds_all_tokens_without_mutating_decode_state():
+    model = _make_minimal_talker()
+
+    def fake_embed_input_ids(input_ids):
+        return input_ids.to(torch.float32).reshape(-1, 1, 1).expand(-1, 1, 4)
+
+    model.embed_input_ids = fake_embed_input_ids
+    trailing_text = torch.arange(12, dtype=torch.float32).reshape(3, 4)
+    last_hidden = torch.full((4,), 2.0, dtype=torch.float32)
+
+    out_ids, out_embeds, update = model.preprocess(
+        input_ids=torch.tensor([101, 202, 303], dtype=torch.long),
+        input_embeds=None,
+        text=["hello"],
+        task_type=["Base"],
+        hidden_states={"trailing_text": trailing_text, "last": last_hidden},
+        meta={"talker_text_offset": 1, "codec_streaming": True},
+        _omni_is_prefill=False,
+        _omni_num_computed_tokens=2400,
+        _omni_prompt_len=213,
+    )
+
+    assert out_ids.tolist() == [101, 202, 303]
+    assert torch.equal(
+        out_embeds.cpu(),
+        torch.tensor([[101.0] * 4, [202.0] * 4, [303.0] * 4], dtype=torch.bfloat16),
+    )
+    assert update == {"meta": {"codec_streaming": True}}
+
+
 def test_decode_batch_preprocess_matches_decode_state_updates():
     tts_pad = torch.full((1, 4), -1.0, dtype=torch.bfloat16)
     model = _make_minimal_talker(tts_pad_embed=tts_pad)

tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_ref_codes.py

@@ -74,3 +74,18 @@ def test_make_omni_output_omits_ref_key_when_no_request_has_one() -> None:
     )
 
     assert "ref" not in out.multimodal_outputs["codes"]
+
+
+def test_make_omni_output_keeps_hidden_states_for_replay_spans_without_audio_codes() -> None:
+    hidden = torch.arange(6 * 8, dtype=torch.float32).reshape(6, 8)
+
+    out = _make_talker().make_omni_output(
+        hidden,
+        model_intermediate_buffer=[
+            {"meta": {"codec_streaming": True}},
+            _info(1, None),
+        ],
+    )
+
+    assert torch.equal(out.text_hidden_states, hidden)
+    assert out.multimodal_outputs["codes"]["audio"].shape == (1, _NUM_CODE_GROUPS)

tests/model_executor/models/test_kv_cache_scale_mapper.py

@@ -0,0 +1,148 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for vLLM's KV-cache scale mapper load path."""
+
+from collections.abc import Iterable
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_REPO_ROOT = Path(__file__).resolve().parents[3]
+_STALE_API_FILES = [
+    "vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py",
+    "vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py",
+    "vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py",
+    "vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py",
+]
+
+_SOURCE_SCALE_NAME = "layers.0.self_attn.k_proj.output_scale"
+_MAPPED_SCALE_NAME = "layers.0.self_attn.attn.k_scale"
+
+
+class _RecordingCacheScaleMapper:
+    def __init__(self) -> None:
+        self.called = False
+
+    def apply(
+        self,
+        weights: Iterable[tuple[str, torch.Tensor]],
+    ) -> Iterable[tuple[str, torch.Tensor]]:
+        self.called = True
+        for name, weight in weights:
+            yield name.replace(".k_proj.output_scale", ".attn.k_scale"), weight
+
+
+class _QuantConfigWithoutOldCacheScale:
+    def __init__(self) -> None:
+        self.mapper = _RecordingCacheScaleMapper()
+
+    def get_cache_scale_mapper(self) -> _RecordingCacheScaleMapper:
+        return self.mapper
+
+
+class _FakeModel:
+    def __init__(self, quant_config: _QuantConfigWithoutOldCacheScale) -> None:
+        self.quant_config = quant_config
+        self.param = torch.nn.Parameter(torch.zeros(()), requires_grad=False)
+
+    def named_parameters(
+        self,
+        remove_duplicate: bool = True,
+    ) -> Iterable[tuple[str, torch.nn.Parameter]]:
+        del remove_duplicate
+        return [(_MAPPED_SCALE_NAME, self.param)]
+
+
+class _FakeHunyuanModel(_FakeModel):
+    def __init__(self, quant_config: _QuantConfigWithoutOldCacheScale) -> None:
+        super().__init__(quant_config)
+        self.config = SimpleNamespace(
+            tie_word_embeddings=False,
+            num_attention_heads=1,
+            num_key_value_heads=1,
+            use_cla=False,
+        )
+
+    def get_expert_mapping(self) -> tuple[list[object], dict[str, object]]:
+        return [], {}
+
+    def _split_qkv_weight(self, weight: torch.Tensor) -> torch.Tensor:
+        return weight
+
+
+@pytest.mark.parametrize("rel_path", _STALE_API_FILES)
+def test_model_loaders_do_not_call_removed_get_cache_scale(rel_path: str) -> None:
+    """Affected loaders must not call the removed vLLM get_cache_scale API."""
+    source = (_REPO_ROOT / rel_path).read_text(encoding="utf-8")
+    assert ".get_cache_scale(" not in source
+
+
+@pytest.mark.parametrize(
+    ("module_name", "class_name", "fake_model_cls"),
+    [
+        (
+            "vllm_omni.model_executor.models.hunyuan_image3.hunyuan_image3",
+            "HunyuanModel",
+            _FakeHunyuanModel,
+        ),
+        (
+            "vllm_omni.model_executor.models.qwen2_5_omni.qwen2_old",
+            "Qwen2Model",
+            _FakeModel,
+        ),
+    ],
+)
+def test_auto_weights_loader_delegated_loaders_accept_mapped_cache_scale(
+    module_name: str,
+    class_name: str,
+    fake_model_cls: type[_FakeModel],
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Delegated loaders must not expect the removed get_cache_scale method."""
+    module = __import__(module_name, fromlist=[class_name])
+    monkeypatch.setattr(module, "is_pp_missing_parameter", lambda *_: False)
+
+    fake_model = fake_model_cls(_QuantConfigWithoutOldCacheScale())
+    load_weights = getattr(module, class_name).load_weights
+
+    loaded = load_weights(fake_model, [(_MAPPED_SCALE_NAME, torch.ones(()))])
+
+    assert loaded == {_MAPPED_SCALE_NAME}
+    torch.testing.assert_close(fake_model.param.data, torch.ones(()))
+
+
+@pytest.mark.parametrize(
+    ("module_name", "class_name"),
+    [
+        (
+            "vllm_omni.model_executor.models.mammoth_moda2.mammoth_moda2",
+            "MammothModa2Qwen2ForCausalLM",
+        ),
+        (
+            "vllm_omni.model_executor.models.mimo_audio.mimo_audio_llm",
+            "MiMoAudioLLMForConditionalGeneration",
+        ),
+    ],
+)
+def test_direct_custom_loaders_apply_cache_scale_mapper(
+    module_name: str,
+    class_name: str,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Direct custom loaders must apply the new mapper before fall-through load."""
+    module = __import__(module_name, fromlist=[class_name])
+    monkeypatch.setattr(module, "is_pp_missing_parameter", lambda *_: False)
+
+    quant_config = _QuantConfigWithoutOldCacheScale()
+    fake_model = _FakeModel(quant_config)
+    load_weights = getattr(module, class_name).load_weights
+
+    loaded = load_weights(fake_model, [(_SOURCE_SCALE_NAME, torch.ones(()))])
+
+    assert quant_config.mapper.called
+    assert loaded == {_MAPPED_SCALE_NAME}
+    torch.testing.assert_close(fake_model.param.data, torch.ones(()))

tests/model_executor/stage_input_processors/test_hunyuan_image3_bridge.py

@@ -45,6 +45,33 @@ def test_truncate_at_cot_end_strips_tail_after_recaption_marker():
     assert text == "body text</recaption>"
 
 
+def test_ar2diffusion_returns_one_request_payload_for_request_level_batching():
+    result = ar2diffusion(
+        [_source_output([100], text="thought")],
+        prompt={"prompt": "edit"},
+    )
+
+    assert isinstance(result, dict)
+    assert result["prompt"] == "edit"
+
+
+def test_ar2diffusion_uses_parent_output_when_companions_are_present():
+    result = ar2diffusion(
+        [
+            _source_output([100], text="parent thought"),
+            _source_output([200], text="companion thought"),
+        ],
+        prompt={"prompt": "edit"},
+    )
+
+    assert result is not None
+    assert result["extra"]["ar_generated_text"] == "parent thought"
+
+
+def test_ar2diffusion_returns_none_without_parent_output():
+    assert ar2diffusion([], prompt={"prompt": "edit"}) is None
+
+
 def test_ar2diffusion_applies_ratio_and_truncates_tail_without_tokenizer(monkeypatch: pytest.MonkeyPatch):
     real_import = builtins.__import__
 
@@ -67,10 +94,10 @@ def test_ar2diffusion_applies_ratio_and_truncates_tail_without_tokenizer(monkeyp
         prompt=[{"prompt": "edit", "height": 64, "width": 64}],
     )
 
-    assert len(result) == 1
-    assert (result[0]["height"], result[0]["width"]) == (512, 2048)
-    assert result[0]["extra"]["ar_generated_text"] == "decoded without special tokens"
-    assert "ar_token_ids" not in result[0]["extra"]
+    assert result is not None
+    assert (result["height"], result["width"]) == (512, 2048)
+    assert result["extra"]["ar_generated_text"] == "decoded without special tokens"
+    assert "ar_token_ids" not in result["extra"]
 
 
 def test_ar2diffusion_forwards_custom_system_prompt_body():
@@ -88,5 +115,6 @@ def test_ar2diffusion_forwards_custom_system_prompt_body():
         ],
     )
 
-    assert result[0]["use_system_prompt"] == "custom"
-    assert result[0]["system_prompt"] == marker
+    assert result is not None
+    assert result["use_system_prompt"] == "custom"
+    assert result["system_prompt"] == marker

tests/quantization/test_quack_fp8_gate.py

@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the quack FP8 auto-enable gate.
+
+quack's CuteDSL FP8 GEMM uses the 5th-gen ``tcgen05`` tensor-core MMA, which
+exists only on datacenter Blackwell (``sm_100a`` / ``sm_101a`` / ``sm_103a``,
+compute capability ``10.x``). It must NOT auto-enable on workstation/consumer
+Blackwell (``sm_120`` / ``sm_121``, cc ``12.x``, e.g. RTX PRO 6000 / RTX
+50-series), where ``tcgen05`` is absent and every FP8 GEMM would fall back to
+FlashInfer one call at a time (catastrophically slow). See
+``vllm_omni/quantization/quack_fp8.py``.
+
+Regression guard for the ``>= 10`` gate that matched cc ``12.x`` too. These
+tests are hardware-free: the ``torch.cuda`` probes are monkeypatched, so no GPU
+is required.
+"""
+
+import pytest
+import torch
+
+from vllm_omni.quantization import quack_fp8
+
+_ENV = "VLLM_OMNI_USE_QUACK_FP8"
+
+
+def _fake_cuda(monkeypatch: pytest.MonkeyPatch, capability: tuple[int, int]) -> None:
+    """Pretend a CUDA device with the given (major, minor) capability is present."""
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: capability)
+
+
+@pytest.mark.parametrize(
+    "capability, expected",
+    [
+        ((10, 0), True),  # sm_100 datacenter Blackwell — tcgen05 present
+        ((10, 3), True),  # sm_103 datacenter Blackwell
+        ((12, 0), False),  # sm_120 workstation Blackwell — the regression
+        ((12, 1), False),  # sm_121 consumer Blackwell
+        ((9, 0), False),  # Hopper — CUTLASS already fuses bias, quack unused
+        ((8, 9), False),  # Ada
+    ],
+)
+def test_auto_enable_only_on_datacenter_blackwell(
+    monkeypatch: pytest.MonkeyPatch,
+    capability: tuple[int, int],
+    expected: bool,
+) -> None:
+    monkeypatch.delenv(_ENV, raising=False)
+    _fake_cuda(monkeypatch, capability)
+    assert quack_fp8._is_quack_capable() is expected
+    assert quack_fp8.quack_enabled() is expected
+
+
+def test_no_cuda_disables_quack(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.delenv(_ENV, raising=False)
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+    assert quack_fp8._is_quack_capable() is False
+    assert quack_fp8.quack_enabled() is False
+
+
+@pytest.mark.parametrize("value", ["1", "true", "yes", "on", "TRUE", "On"])
+def test_env_override_forces_on_even_on_sm120(monkeypatch: pytest.MonkeyPatch, value: str) -> None:
+    # sm_120 auto-disables, but an explicit truthy override forces quack on
+    # (e.g. once CuteDSL ships sm_120a support).
+    _fake_cuda(monkeypatch, (12, 0))
+    monkeypatch.setenv(_ENV, value)
+    assert quack_fp8.quack_enabled() is True
+
+
+@pytest.mark.parametrize("value", ["0", "false", "no", "off", ""])
+def test_env_override_forces_off_even_on_datacenter_blackwell(monkeypatch: pytest.MonkeyPatch, value: str) -> None:
+    # sm_100 auto-enables, but an explicit non-truthy override forces quack off.
+    _fake_cuda(monkeypatch, (10, 0))
+    monkeypatch.setenv(_ENV, value)
+    assert quack_fp8.quack_enabled() is False

tests/utils/test_tracking_parser.py

@@ -16,7 +16,12 @@ from vllm_omni.config.stage_config import (
     StageDeployConfig,
     StagePipelineConfig,
 )
-from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace
+from vllm_omni.utils.tracking_parser import (
+    UNSET,
+    TrackingArgumentParser,
+    TrackingNamespace,
+    build_shadow_kwargs,
+)
 
 ### Fake pipeline/deploy config for integration tests
 
@@ -326,6 +331,132 @@ def test_omitted_nargs():
     assert ns.items is None
 
 
+### Tests for in-place mutating actions (append / extend / count).
+# These actions mutate their default in place, which would crash on the bare
+# ``UNSET`` sentinel. The shadow parser remaps them to a non-mutating store-style
+# action (see build_shadow_kwargs in tracking_parser), so the real namespace
+# still accumulates while explicit-arg tracking keeps working.
+def test_append_action_omitted():
+    """Omitted append args parse without error and aren't marked explicit."""
+    p = TrackingArgumentParser()
+    p.add_argument("--tag", action="append")
+    ns = p.parse_args([])
+    assert ns.explicit_keys == set()
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.tag is None
+
+
+def test_append_action_explicit():
+    """Repeated append args are collected and marked explicit."""
+    p = TrackingArgumentParser()
+    p.add_argument("--tag", action="append")
+    ns = p.parse_args(["--tag", "a", "--tag", "b"])
+    assert ns.explicit_keys == {"tag"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.tag == ["a", "b"]
+
+
+def test_append_const_action_omitted():
+    """Omitted append_const args aren't marked explicit."""
+    p = TrackingArgumentParser()
+    p.add_argument("--dbg", action="append_const", const="x", dest="flags")
+    ns = p.parse_args([])
+    assert ns.explicit_keys == set()
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.flags is None
+
+
+def test_append_const_action_explicit():
+    """Repeated append_const args accumulate the const and are explicit."""
+    p = TrackingArgumentParser()
+    p.add_argument("--dbg", action="append_const", const="x", dest="flags")
+    ns = p.parse_args(["--dbg", "--dbg"])
+    assert ns.explicit_keys == {"flags"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.flags == ["x", "x"]
+
+
+def test_extend_action_omitted():
+    """Omitted extend args parse without error and aren't marked explicit."""
+    p = TrackingArgumentParser()
+    p.add_argument("--items", action="extend", nargs="+")
+    ns = p.parse_args([])
+    assert ns.explicit_keys == set()
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.items is None
+
+
+def test_extend_action_explicit():
+    """Repeated extend args are flattened into one list and marked explicit."""
+    p = TrackingArgumentParser()
+    p.add_argument("--items", action="extend", nargs="+")
+    ns = p.parse_args(["--items", "a", "b", "--items", "c"])
+    assert ns.explicit_keys == {"items"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.items == ["a", "b", "c"]
+
+
+def test_count_action_omitted():
+    """Omitted count args keep their default and aren't marked explicit."""
+    p = TrackingArgumentParser()
+    p.add_argument("-v", "--verbose", action="count", default=0)
+    ns = p.parse_args([])
+    assert ns.explicit_keys == set()
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.verbose == 0
+
+
+def test_count_action_explicit():
+    """Repeated count flags increment and are marked explicit."""
+    p = TrackingArgumentParser()
+    p.add_argument("-v", "--verbose", action="count", default=0)
+    ns = p.parse_args(["-vv"])
+    assert ns.explicit_keys == {"verbose"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.verbose == 2
+
+
+def test_group_append_action_explicit():
+    """append args added via a group go through the same shadow-default path."""
+    p = TrackingArgumentParser()
+    g = p.add_argument_group("TestGroup")
+    g.add_argument("--tag", action="append")
+    ns = p.parse_args(["--tag", "a", "--tag", "b"])
+    assert ns.explicit_keys == {"tag"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.tag == ["a", "b"]
+
+
+@pytest.mark.parametrize(
+    ("action", "expected"),
+    [
+        ("append", "store"),
+        ("extend", "store"),
+        ("append_const", "store_const"),
+        ("count", "store_const"),
+    ],
+)
+def test_build_shadow_kwargs_remaps_mutating_actions(action, expected):
+    """Mutating actions are remapped to a store-style action with an UNSET default."""
+    shadow = build_shadow_kwargs({"action": action, "const": "c"})
+    assert shadow["action"] == expected
+    assert shadow["default"] is UNSET
+
+
+def test_build_shadow_kwargs_count_gets_marker_const():
+    """count carries no const, so the remapped store_const gets a non-UNSET marker."""
+    shadow = build_shadow_kwargs({"action": "count"})
+    assert shadow["action"] == "store_const"
+    assert shadow["const"] is not UNSET
+
+
+def test_build_shadow_kwargs_leaves_other_actions_untouched():
+    """Non-mutating actions keep their action and only get the UNSET default."""
+    shadow = build_shadow_kwargs({"action": "store_true"})
+    assert shadow["action"] == "store_true"
+    assert shadow["default"] is UNSET
+
+
 def test_parse_known_args_tracking():
     """Ensure parse_known_args is also trackable"""
     p = TrackingArgumentParser()

tests/worker/test_omni_gpu_model_runner.py

@@ -275,7 +275,7 @@ def test_talker_mtp_forward_passes_qwen3_tts_subtalker_sampling_params_to_talker
     runner = _make_runner(req_ids=("r1",), hidden_size=4)
     runner.requests["r1"].sampling_params = SimpleNamespace(
         seed=42,
-        extra_args={"qwen3_tts_request_seed": 42},
+        extra_args={"tts_local_seed": 42},
     )
     runner.talker_mtp = CaptureTalkerMTP()
     runner.vllm_config = SimpleNamespace(
@@ -319,11 +319,11 @@ def test_talker_mtp_forward_keeps_explicit_seeded_requests_scalar(monkeypatch):
     runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4)
     runner.requests["r1"].sampling_params = SimpleNamespace(
         seed=11,
-        extra_args={"qwen3_tts_request_seed": 11},
+        extra_args={"tts_local_seed": 11},
     )
     runner.requests["r2"].sampling_params = SimpleNamespace(
         seed=22,
-        extra_args={"qwen3_tts_request_seed": 22},
+        extra_args={"tts_local_seed": 22},
     )
     runner.talker_mtp = CaptureTalkerMTP()
     runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(subtalker_sampling_params={}))

vllm_omni/benchmarks/data_modules/seed_tts_eval.py

@@ -36,7 +36,11 @@ Enable with ``SEED_TTS_WER_EVAL=1`` or ``--seed-tts-wer-eval``. Install optional
     pip install 'vllm-omni[dev]'
 
 Env: ``SEED_TTS_EVAL_DEVICE`` (e.g. ``cuda:0``, ``cpu``); ``SEED_TTS_HF_WHISPER_MODEL``
-defaults to ``openai/whisper-large-v3`` (override for debugging only).
+defaults to ``openai/whisper-large-v3`` (override for debugging only). Set
+``SEED_TTS_WER_SAVE_AUDIO_DIR`` to save the captured 24 kHz mono WAV used by
+WER evaluation for each synthesized utterance.
+Streaming PCM is decoded using ``VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE`` /
+``VLLM_OMNI_BENCH_AUDIO_CHANNELS`` (default: 24 kHz mono).
 """
 
 from __future__ import annotations
@@ -50,12 +54,14 @@ import string
 import tempfile
 import threading
 import wave
+from pathlib import Path
 from typing import Any
 
 import numpy as np
 from vllm.benchmarks.datasets import SampleRequest
 
 from vllm_omni.benchmarks.data_modules.seed_tts_dataset import SeedTTSSampleRequest
+from vllm_omni.metrics.definitions import stream_pcm_format_from_env
 
 logger = logging.getLogger(__name__)
 
@@ -87,6 +93,30 @@ def pcm_s16le_mono_to_wav_bytes(pcm: bytes, *, sample_rate: int = 24000) -> byte
     return buf.getvalue()
 
 
+def _safe_filename_part(value: Any, *, default: str = "item", max_len: int = 96) -> str:
+    text = str(value or "").strip() or default
+    safe = "".join(ch if ch.isalnum() or ch in ("-", "_", ".") else "_" for ch in text)
+    safe = safe.strip("._") or default
+    return safe[:max_len]
+
+
+def _save_seed_tts_eval_audio(
+    pcm: bytes,
+    *,
+    output_dir: Path | None,
+    index: int,
+    utterance_id: Any,
+    locale: str,
+) -> str | None:
+    if output_dir is None:
+        return None
+    output_dir.mkdir(parents=True, exist_ok=True)
+    stem = f"{index:05d}_{_safe_filename_part(utterance_id)}_{_safe_filename_part(locale)}"
+    path = output_dir / f"{stem}.wav"
+    path.write_bytes(pcm_s16le_mono_to_wav_bytes(pcm, sample_rate=24000))
+    return str(path)
+
+
 def _get_eval_device() -> str:
     explicit = os.environ.get("SEED_TTS_EVAL_DEVICE", "").strip()
     if explicit:
@@ -146,12 +176,27 @@ def process_one_official(hypo: str, truth: str, lang: str) -> tuple[float, str,
     return wer, raw_truth, raw_hypo
 
 
-def _pcm_s16le_to_f32_16k(pcm: bytes, pcm_sample_rate: int = 24000) -> np.ndarray:
+def _pcm_s16le_to_f32_16k(
+    pcm: bytes,
+    pcm_sample_rate: int | None = None,
+    channels: int | None = None,
+) -> np.ndarray:
     import scipy.signal
 
     if not pcm:
         return np.zeros(0, dtype=np.float32)
+    if pcm_sample_rate is None or channels is None:
+        env_sample_rate, env_channels = stream_pcm_format_from_env()
+        if pcm_sample_rate is None:
+            pcm_sample_rate = env_sample_rate
+        if channels is None:
+            channels = env_channels
     raw = np.frombuffer(pcm, dtype=np.int16).astype(np.float32) / 32768.0
+    if channels > 1:
+        usable = (len(raw) // channels) * channels
+        if usable <= 0:
+            return np.zeros(0, dtype=np.float32)
+        raw = raw[:usable].reshape(-1, channels).mean(axis=1)
     target_len = int(len(raw) * 16000 / pcm_sample_rate)
     if target_len <= 0:
         return np.zeros(0, dtype=np.float32)
@@ -512,13 +557,18 @@ def compute_seed_tts_wer_metrics(
     sim_skipped_no_ref = 0
     utmos_failed = 0
     utmos_on = _eval_submetric_enabled("SEED_TTS_UTMOS_EVAL", default=False)
+    save_audio_raw = os.environ.get("SEED_TTS_WER_SAVE_AUDIO_DIR", "").strip()
+    save_audio_dir = Path(save_audio_raw).expanduser() if save_audio_raw else None
+    saved_audio = 0
+    save_audio_failed = 0
 
-    for req, out in zip(input_requests, outputs, strict=True):
+    for index, (req, out) in enumerate(zip(input_requests, outputs, strict=True)):
         assert isinstance(req, SeedTTSSampleRequest)
         ref = req.prompt
         locale = req.seed_tts_locale or "en"
         row_lang = "zh" if locale.lower().startswith("zh") else "en"
         utmos_v: float | None = None
+        audio_path: str | None = None
 
         if not out.success:
             request_failed += 1
@@ -545,6 +595,23 @@ def compute_seed_tts_wer_metrics(
                     }
                 )
             continue
+        try:
+            audio_path = _save_seed_tts_eval_audio(
+                pcm,
+                output_dir=save_audio_dir,
+                index=index,
+                utterance_id=req.seed_tts_utterance_id,
+                locale=locale,
+            )
+            if audio_path:
+                saved_audio += 1
+        except OSError as e:
+            save_audio_failed += 1
+            logger.warning(
+                "Seed-TTS WER audio save failed for utterance=%s: %s",
+                req.seed_tts_utterance_id,
+                e,
+            )
 
         wav_16k = _pcm_s16le_to_f32_16k(pcm)
         if len(wav_16k) == 0:
@@ -555,6 +622,7 @@ def compute_seed_tts_wer_metrics(
                         "utterance_id": req.seed_tts_utterance_id,
                         "locale": locale,
                         "error": "empty_audio",
+                        "audio_path": audio_path,
                     }
                 )
             continue
@@ -609,6 +677,7 @@ def compute_seed_tts_wer_metrics(
                         "locale": locale,
                         "error": "asr_exception",
                         "detail": str(e)[:500],
+                        "audio_path": audio_path,
                     }
                 )
             continue
@@ -621,6 +690,7 @@ def compute_seed_tts_wer_metrics(
                         "utterance_id": req.seed_tts_utterance_id,
                         "locale": locale,
                         "error": "empty_asr",
+                        "audio_path": audio_path,
                     }
                 )
             continue
@@ -637,6 +707,7 @@ def compute_seed_tts_wer_metrics(
                         "locale": locale,
                         "error": "wer_compute_failed",
                         "detail": str(e)[:500],
+                        "audio_path": audio_path,
                     }
                 )
             continue
@@ -673,6 +744,8 @@ def compute_seed_tts_wer_metrics(
                 "reference_raw": raw_truth,
                 "asr_raw": raw_hypo,
             }
+            if audio_path:
+                row["audio_path"] = audio_path
             if sim_v is not None:
                 row["sim"] = sim_v
             if utmos_v is not None:
@@ -697,7 +770,11 @@ def compute_seed_tts_wer_metrics(
         "seed_tts_utmos_mean": statistics.fmean(utmos_values) if utmos_values else None,
         "seed_tts_utmos_median": statistics.median(utmos_values) if utmos_values else None,
         "seed_tts_utmos_failed": utmos_failed,
+        "seed_tts_saved_audio": saved_audio,
+        "seed_tts_save_audio_failed": save_audio_failed,
     }
+    if save_audio_dir is not None:
+        result["seed_tts_save_audio_dir"] = str(save_audio_dir)
     if include_per_item:
         result["seed_tts_wer_eval_items"] = items
     return result
@@ -729,6 +806,13 @@ def print_seed_tts_wer_summary(metrics: dict[str, Any]) -> None:
     print("{:<40} {:<10}".format("Request failed:", metrics.get("seed_tts_request_failed", 0)))
     print("{:<40} {:<10}".format("No PCM captured:", metrics.get("seed_tts_no_pcm", 0)))
     print("{:<40} {:<10}".format("ASR / WER failed:", metrics.get("seed_tts_asr_failed", 0)))
+    save_dir = metrics.get("seed_tts_save_audio_dir")
+    if save_dir:
+        print("{:<40} {:<10}".format("Saved eval WAVs:", metrics.get("seed_tts_saved_audio", 0)))
+        print("{:<40} {}".format("Saved eval WAV dir:", save_dir))
+        failed = int(metrics.get("seed_tts_save_audio_failed", 0) or 0)
+        if failed:
+            print("{:<40} {:<10}".format("Save eval WAV errors:", failed))
     if sim_ev or metrics.get("seed_tts_sim_skipped_no_ref") or metrics.get("seed_tts_sim_failed"):
         print("{:<40} {:<10}".format("SIM evaluated (higher ~ closer):", sim_ev))
         sm = metrics.get("seed_tts_sim_mean")

vllm_omni/benchmarks/metrics/metrics.py

@@ -63,11 +63,9 @@ _STAGE_BENCHMARK_FIELDS = [
     (defs.VLLM_TPOTS, _FLOAT_LIST_TYPE, field(default_factory=list)),
     (defs.VLLM_ITLS, _FLOAT_LIST_TYPE, field(default_factory=list)),
     (defs.AUDIO_TTFPS, _FLOAT_LIST_TYPE, field(default_factory=list)),
-    (defs.AUDIO_RTFS, _FLOAT_LIST_TYPE, field(default_factory=list)),
     (defs.AUDIO_DURATIONS, _FLOAT_LIST_TYPE, field(default_factory=list)),
     (defs.AUDIO_FRAMES, _INT_LIST_TYPE, field(default_factory=list)),
     (defs.MISSING_AUDIO_DURATION_COUNT, int, field(default=0)),
-    (defs.MISSING_AUDIO_RTF_COUNT, int, field(default=0)),
     (defs.STAGE_GEN_TIMES_MS, _FLOAT_LIST_TYPE, field(default_factory=list)),
     (defs.POSTPROCESS_TIMES_MS, _FLOAT_LIST_TYPE, field(default_factory=list)),
     ("output_unit_type", str, field(default="")),
@@ -443,12 +441,6 @@ def _print_audio_stage_metrics(
             getattr(sm, defs.AUDIO_TTFPS),
             selected_percentiles,
         )
-    if defs.AUDIO_RTF in selected_percentile_metrics:
-        audio_rtfs = getattr(sm, defs.AUDIO_RTFS)
-        if audio_rtfs:
-            _print_percentile_metric("Real Time Factor", "AUDIO_RTF", audio_rtfs, selected_percentiles, to_ms=False)
-        elif getattr(sm, defs.MISSING_AUDIO_RTF_COUNT) > 0:
-            print("{:<40} {:<10}".format("AUDIO_RTF skipped:", "missing stage-local audio duration/sample rate"))
     if defs.AUDIO_DURATION in selected_percentile_metrics:
         audio_durations = getattr(sm, defs.AUDIO_DURATIONS)
         if audio_durations:
@@ -643,11 +635,9 @@ def _build_stage_metrics_from_outputs(
             total_output = sum(int((info or {}).get(defs.NUM_TOKENS_OUT) or 0) for _, info in rows)
 
         audio_ttfps: list[float] = []
-        audio_rtfs: list[float] = []
         audio_durations: list[float] = []
         audio_frames: list[int] = []
         missing_audio_duration_count = 0
-        missing_audio_rtf_count = 0
         if is_audio_stage:
             for _, info in rows:
                 audio_ttfp_ms = float((info or {}).get(defs.SERVING_TIME_TO_FIRST_OUTPUT_MS) or 0.0)
@@ -665,12 +655,6 @@ def _build_stage_metrics_from_outputs(
                 else:
                     missing_audio_duration_count += 1
 
-                rtf = float((info or {}).get(defs.AUDIO_RTF) or 0.0)
-                if rtf > 0:
-                    audio_rtfs.append(rtf)
-                else:
-                    missing_audio_rtf_count += 1
-
         result.append(
             StageBenchmarkMetrics(
                 **{
@@ -685,11 +669,9 @@ def _build_stage_metrics_from_outputs(
                     defs.VLLM_TPOTS: vllm_tpots,
                     defs.VLLM_ITLS: vllm_itls,
                     defs.AUDIO_TTFPS: audio_ttfps,
-                    defs.AUDIO_RTFS: audio_rtfs,
                     defs.AUDIO_DURATIONS: audio_durations,
                     defs.AUDIO_FRAMES: audio_frames,
                     defs.MISSING_AUDIO_DURATION_COUNT: missing_audio_duration_count,
-                    defs.MISSING_AUDIO_RTF_COUNT: missing_audio_rtf_count,
                     defs.STAGE_GEN_TIMES_MS: stage_gen_times_ms,
                     defs.POSTPROCESS_TIMES_MS: postprocess_times_ms,
                     "output_unit_type": output_unit_type,

vllm_omni/benchmarks/patch/patch.py

@@ -35,8 +35,6 @@ from vllm.benchmarks.lib.endpoint_request_func import (
 from vllm.logger import init_logger
 from vllm.tokenizers import TokenizerLike
 
-logger = init_logger(__name__)
-
 from vllm_omni.benchmarks.audio_continuity import compute_continuity_stats
 from vllm_omni.benchmarks.data_modules.daily_omni_dataset import DailyOmniDataset, DailyOmniSampleRequest
 from vllm_omni.benchmarks.data_modules.random_multi_modal_dataset import OmniRandomMultiModalDataset
@@ -51,9 +49,9 @@ from vllm_omni.benchmarks.data_modules.sound_effect_dataset import SoundEffectDa
 from vllm_omni.benchmarks.data_modules.ttsd_dataset import TTSDDataset
 from vllm_omni.metrics import definitions as defs
 
+logger = init_logger(__name__)
+
 _AUDIO_CONTINUITY_THRESHOLD_ENV = "VLLM_OMNI_BENCH_AUDIO_CONTINUITY_THRESHOLD_S"
-_AUDIO_SAMPLE_RATE_ENV = "VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE"
-_AUDIO_CHANNELS_ENV = "VLLM_OMNI_BENCH_AUDIO_CHANNELS"
 RETURN_STAGE_METRICS_FIELD = "return_stage_metrics"
 _IMAGE_STAGE_METRICS_BACKENDS = frozenset({"openai-image-edits-omni"})
 _PRINT_STAGE = False
@@ -111,30 +109,28 @@ def _audio_continuity_threshold_s() -> float:
     return max(value, 0.0)
 
 
-def _audio_pcm_format() -> tuple[int, int]:
-    """Return ``(sample_rate, channels)`` of the streamed PCM response.
+def _pcm_s16le_to_seed_tts_wer_bytes(
+    pcm_bytes: bytes,
+    *,
+    sample_rate: int,
+    channels: int,
+) -> bytes:
+    """Normalize streamed raw PCM to the 24 kHz mono PCM used by Seed-TTS WER."""
+    if not pcm_bytes:
+        return b""
+    channels = max(1, int(channels))
+    pcm = np.frombuffer(pcm_bytes, dtype=np.int16)
+    if channels > 1:
+        usable = (pcm.size // channels) * channels
+        pcm = pcm[:usable].reshape(-1, channels).mean(axis=1)
+    pcm_f32 = pcm.astype(np.float32) / 32767.0
+    if int(sample_rate) != 24000 and pcm_f32.size:
+        from vllm.multimodal.audio import AudioResampler
 
-    Defaults to 24 kHz mono (Qwen3-TTS / VoxCPM2 / most MOSS-TTS variants).
-    Override for models with a different codec output format (e.g.
-    MOSS-TTS-Local-Transformer-v1.5, which streams 48 kHz stereo PCM) via
-    ``VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE`` / ``VLLM_OMNI_BENCH_AUDIO_CHANNELS`` -
-    getting this wrong silently skews audio_duration/audio_rtf/underrun stats.
-    """
-    sample_rate = 24000
-    channels = 1
-    raw_sr = os.environ.get(_AUDIO_SAMPLE_RATE_ENV)
-    if raw_sr:
-        try:
-            sample_rate = int(raw_sr)
-        except ValueError:
-            logger.warning("Invalid %s=%r; using default %d", _AUDIO_SAMPLE_RATE_ENV, raw_sr, sample_rate)
-    raw_ch = os.environ.get(_AUDIO_CHANNELS_ENV)
-    if raw_ch:
-        try:
-            channels = int(raw_ch)
-        except ValueError:
-            logger.warning("Invalid %s=%r; using default %d", _AUDIO_CHANNELS_ENV, raw_ch, channels)
-    return sample_rate, channels
+        resampler = AudioResampler(target_sr=24000)
+        pcm_f32 = resampler.resample(pcm_f32, orig_sr=int(sample_rate))
+    pcm_f32 = np.clip(pcm_f32, -1.0, 1.0)
+    return (pcm_f32 * 32767).astype(np.int16).tobytes()
 
 
 get_samples_old = datasets.get_samples
@@ -1163,9 +1159,8 @@ async def async_request_openai_audio_speech(
     output = MixRequestFuncOutput()
     output.prompt_len = request_func_input.prompt_len
 
-    # PCM format: 16-bit signed; sample_rate/channels are model-dependent
-    # (see _audio_pcm_format docstring).
-    sample_rate, channels = _audio_pcm_format()
+    # PCM format: 16-bit signed; sample_rate/channels are model-dependent.
+    sample_rate, channels = defs.stream_pcm_format_from_env()
     sample_width = 2  # 16-bit = 2 bytes
 
     st = time.perf_counter()
@@ -1216,7 +1211,15 @@ async def async_request_openai_audio_speech(
                 output.audio_continuity_ok = continuity.is_continuous
                 output.audio_underrun_event_count = continuity.underrun_event_count
                 if pcm_capture is not None and pcm_capture:
-                    output.tts_output_pcm_bytes = bytes(pcm_capture)
+                    try:
+                        output.tts_output_pcm_bytes = _pcm_s16le_to_seed_tts_wer_bytes(
+                            bytes(pcm_capture),
+                            sample_rate=sample_rate,
+                            channels=channels,
+                        )
+                    except Exception as ex:
+                        logger.warning("Seed-TTS WER PCM normalization failed: %s", ex)
+                        output.tts_output_pcm_bytes = bytes(pcm_capture)
                 elif capture_wer_pcm:
                     ct = response.headers.get("Content-Type", "")
                     logger.warning(

vllm_omni/deploy/moss_tts_local.yaml

@@ -14,8 +14,10 @@ connectors:
       connector_get_sleep_s: 0.005
       connector_get_max_wait_first_chunk: 1000
       connector_get_max_wait: 300
-      codec_chunk_frames: 15          # 15 frames * 80 ms/frame = 1.2 s chunks
+      initial_codec_chunk_frames: 1   # 1 frame * 80 ms/frame = 80 ms first audio chunk
+      codec_chunk_frames: 15          # 15 frames * 80 ms/frame = 1.2 s steady chunks
       codec_left_context_frames: 0
+      codec_stream_slots: 8
 
 stages:
   - stage_id: 0
@@ -37,7 +39,7 @@ stages:
       repetition_penalty: 1.1
 
   - stage_id: 1
-    max_num_seqs: 1
+    max_num_seqs: 8
     gpu_memory_utilization: 0.15
     enforce_eager: true
     enable_prefix_caching: false

vllm_omni/diffusion/cache/cache_dit_backend.py

@@ -59,9 +59,12 @@ CUSTOM_DIT_ENABLERS: dict[str, Callable] = {}
 
 # Small helper to centralize cache-dit summaries.
 def cache_summary(pipeline: Any, details: bool = True) -> None:
-    cache_dit.summary(pipeline.transformer, details=details)
+    if hasattr(pipeline, "transformer"):
+        cache_dit.summary(pipeline.transformer, details=details)
     if hasattr(pipeline, "transformer_2"):
         cache_dit.summary(pipeline.transformer_2, details=details)
+    if not hasattr(pipeline, "transformer") and not hasattr(pipeline, "transformer_2"):
+        logger.warning("CacheDiT summary failed; this pipeline has no defined transformer attribute")
 
 
 def default_get_pipeline_transformer(pipeline: Any) -> Any:

vllm_omni/diffusion/distributed/parallel_state.py

@@ -35,7 +35,6 @@ from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion import envs
-from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.forward_context import get_forward_context
 from vllm_omni.platforms import current_omni_platform
 
@@ -61,6 +60,11 @@ _DP: GroupCoordinator | None = None
 _FS: GroupCoordinator | None = None  # Fully Sharded (HSDP shard dimension)
 _DIT: GroupCoordinator | None = None
 
+# Rank-layout metadata for expert parallelism. This is not a process group;
+# it is reused by platform-specific runtimes that must build companion groups
+# with the same rank layout as vLLM EP.
+_EXPERT_PARALLEL_GROUP_RANKS: list[list[int]] | None = None
+
 
 def generate_masked_orthogonal_rank_groups(
     world_size: int, parallel_size: list[int], mask: list[bool]
@@ -190,7 +194,6 @@ class RankGenerator:
         self.fs = fs
         self.rank_offset = rank_offset
         self.world_size = tp * sp * pp * cfg * dp
-        self.ep = tp * sp * cfg * dp  # EP level exclude PP
 
         self.name_to_size = {
             "tp": self.tp,
@@ -253,15 +256,6 @@ class RankGenerator:
                 ranks.append(group)
             return ranks
 
-        if token == "ep":
-            ranks = []
-            num_pp_stages = self.pp
-            for i in range(num_pp_stages):
-                start = i * self.ep + self.rank_offset
-                end = start + self.ep
-                ranks.append(list(range(start, end)))
-            return ranks
-
         mask = self.get_mask(self.order, token)
         ranks = generate_masked_orthogonal_rank_groups(self.world_size, self.ordered_size, mask)
         if self.rank_offset > 0:
@@ -309,6 +303,12 @@ def get_ring_parallel_rank():
     return get_sp_group().ring_rank
 
 
+def get_expert_parallel_group_ranks() -> list[list[int]]:
+    assert vllm_parallel_state._EP is not None, "expert parallel group is not initialized"
+    assert _EXPERT_PARALLEL_GROUP_RANKS is not None, "expert parallel group ranks are not initialized"
+    return _EXPERT_PARALLEL_GROUP_RANKS
+
+
 # PP
 def get_pp_group() -> PipelineGroupCoordinator:
     assert _PP is not None, "pipeline model parallel group is not initialized"
@@ -782,6 +782,13 @@ def initialize_model_parallel(
         fs=fully_shard_degree,
         order="tp-sp-pp-cfg-dp",
     )
+    use_moe_parallel_mapping = False
+    if enable_expert_parallel:
+        od_config = get_forward_context().omni_diffusion_config
+        use_moe_parallel_mapping = bool(od_config and od_config.is_moe)
+        if not use_moe_parallel_mapping:
+            raise RuntimeError("Expert parallelism enabled for a non-MoE model")
+
     sp_group_ranks = rank_generator.get_ranks("sp")
     global _DP
     assert _DP is None, "data parallel group is already initialized"
@@ -828,6 +835,10 @@ def initialize_model_parallel(
         ulysses_group=ulysses_pg,
         ring_group=ring_pg,
     )
+    if use_moe_parallel_mapping:
+        # Diffusion normally uses its own SP group. Map it to vLLM PCP only for
+        # expert-parallel runtimes that rely on vLLM FusedMoE group semantics.
+        vllm_parallel_state._PCP = _SP
 
     assert vllm_parallel_state._TP is None, "Tensor parallel group is already initialized"
     vllm_parallel_state._TP = init_model_parallel_group(
@@ -836,8 +847,17 @@ def initialize_model_parallel(
         backend=backend,
         parallel_mode="tensor",
     )
+    if use_moe_parallel_mapping and cfg_parallel_size > 1:
+        # CFG is a diffusion-specific replica dimension. Fold it into vLLM DP
+        # only when constructing the vLLM EP layout for expert-parallel paths.
+        vllm_parallel_state._DP = init_model_parallel_group(
+            group_ranks=rank_generator.get_ranks("cfg-dp"),
+            local_rank=get_world_group().local_rank,
+            backend=backend,
+            parallel_mode="data",
+        )
 
-    global _FS
+    global _FS, _EXPERT_PARALLEL_GROUP_RANKS
     assert _FS is None, "fully shard group is already initialized"
     _FS = init_model_parallel_group(
         group_ranks=rank_generator.get_ranks("fs", independent_ranks=True),
@@ -846,34 +866,40 @@ def initialize_model_parallel(
         parallel_mode="fully_shard",
     )
 
-    if enable_expert_parallel:
-        od_config: OmniDiffusionConfig | None = get_forward_context().omni_diffusion_config
-        if od_config and od_config.is_moe:
-            vllm_parallel_state._EP = init_model_parallel_group(
-                group_ranks=rank_generator.get_ranks("ep"),
-                local_rank=get_world_group().local_rank,
-                backend=backend,
-                parallel_mode="expert",
-            )
-        else:
-            raise RuntimeError("Expert parallelism enabled for a non-MoE model ")
+    _EXPERT_PARALLEL_GROUP_RANKS = None
+    if use_moe_parallel_mapping:
+        ep_group_ranks = rank_generator.get_ranks("tp-sp-cfg-dp")
+        vllm_parallel_state._EP = init_model_parallel_group(
+            group_ranks=ep_group_ranks,
+            local_rank=get_world_group().local_rank,
+            backend=backend,
+            parallel_mode="expert",
+        )
+        _EXPERT_PARALLEL_GROUP_RANKS = ep_group_ranks
 
     init_dit_group(dit_parallel_size, backend)
 
 
 def destroy_model_parallel():
     """Set the groups to none and destroy them."""
-    global _DP
+    global _DP, _CFG, _SP, _PP, _FS, _EXPERT_PARALLEL_GROUP_RANKS
+
+    if vllm_parallel_state._DP and vllm_parallel_state._DP is not _DP:
+        vllm_parallel_state._DP.destroy()
+    vllm_parallel_state._DP = None
+
     if _DP:
         _DP.destroy()
     _DP = None
 
-    global _CFG
     if _CFG:
         _CFG.destroy()
     _CFG = None
 
-    global _SP
+    if vllm_parallel_state._PCP and vllm_parallel_state._PCP is not _SP:
+        vllm_parallel_state._PCP.destroy()
+    vllm_parallel_state._PCP = None
+
     if _SP:
         _SP.destroy()
     _SP = None
@@ -885,13 +911,12 @@ def destroy_model_parallel():
     if vllm_parallel_state._EP:
         vllm_parallel_state._EP.destroy()
     vllm_parallel_state._EP = None
+    _EXPERT_PARALLEL_GROUP_RANKS = None
 
-    global _PP
     if _PP:
         _PP.destroy()
     _PP = None
 
-    global _FS
     if _FS:
         _FS.destroy()
     _FS = None

vllm_omni/diffusion/lora/manager.py

@@ -368,7 +368,17 @@ class DiffusionLoRAManager:
 
         # Default denoising components plus any a pipeline opts into via
         # ``_lora_components`` (opt-in; other pipelines unchanged).
-        default_components = ("transformer", "transformer_2", "dit", "bagel")
+        #
+        # NOTE: SDXL-style pipelines expose the denoiser as ``unet``.
+        # Without scanning this component, adapters can load/activate while
+        # effectively applying to zero layers, producing base-identical output.
+        default_components = (
+            "transformer",
+            "transformer_2",
+            "dit",
+            "bagel",
+            "unet",
+        )
         extra_components = tuple(getattr(self.pipeline, "_lora_components", ()) or ())
         for component_name in (*default_components, *extra_components):
             if not hasattr(self.pipeline, component_name):

vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py

@@ -77,6 +77,7 @@ from .action import (
     normalize_action_mode,
     pad_action_to_dim,
     resolve_domain_id,
+    vision_condition_indexes,
 )
 from .transfer import (
     Cosmos3TransferConfig,
@@ -2009,8 +2010,21 @@ class Cosmos3OmniDiffusersPipeline(
         generator: torch.Generator,
         image_size: Any | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """Prepare video latents for action modes with mode-specific conditioning."""
+        """Prepare video latents for action modes with mode-specific conditioning.
+
+        Policy and forward-dynamics modes condition only latent frame zero.
+        The Wan VAE is temporally causal, so encoding only the first pixel
+        frame preserves that latent while avoiding unused future-frame work.
+        Inverse dynamics conditions every latent and keeps the full encode.
+        """
         del height, width
+        if video_tensor.ndim == 4:
+            video_tensor = video_tensor.unsqueeze(0)
+        if video_tensor.ndim != 5 or video_tensor.shape[0] != 1 or video_tensor.shape[1] != 3:
+            raise ValueError(f"Cosmos3 video tensor must have shape [1, 3, T, H, W], got {tuple(video_tensor.shape)}.")
+        if video_tensor.shape[2] < 1:
+            raise ValueError("Cosmos3 action video tensor must contain at least one frame.")
+
         C = self.transformer.latent_channel_size
         T_lat = (num_frames - 1) // self.vae_scale_factor_temporal + 1
         latent_hw = self._latent_hw_from_image_size(image_size)
@@ -2026,12 +2040,24 @@ class Cosmos3OmniDiffusersPipeline(
             device=self.device,
             dtype=self.dtype,
         )
-        cond_latent = self._encode_video_tensor(video_tensor, image_size=image_size)
-        if cond_latent.shape[2:] != noise.shape[2:]:
+        condition_indexes = vision_condition_indexes(mode, num_frames, self.vae_scale_factor_temporal)
+        condition_video = video_tensor[:, :, :1] if condition_indexes == [0] else video_tensor
+        cond_prefix_latent = self._encode_video_tensor(condition_video, image_size=image_size)
+        expected_prefix = (1, C, max(condition_indexes) + 1, H_lat, W_lat)
+        if (
+            cond_prefix_latent.shape[0] != expected_prefix[0]
+            or cond_prefix_latent.shape[1] != expected_prefix[1]
+            or cond_prefix_latent.shape[2] < expected_prefix[2]
+            or cond_prefix_latent.shape[3:] != expected_prefix[3:]
+        ):
             raise ValueError(
                 "Cosmos3 action video latent shape mismatch: "
-                f"encoded={tuple(cond_latent.shape)}, expected={tuple(noise.shape)}."
+                f"encoded={tuple(cond_prefix_latent.shape)}, expected at least {expected_prefix}."
             )
+
+        condition_latents = torch.zeros_like(noise)
+        for index in condition_indexes:
+            condition_latents[:, :, index : index + 1] = cond_prefix_latent[:, :, index : index + 1]
         condition_mask = build_vision_condition_mask(
             mode,
             num_frames,
@@ -2039,9 +2065,9 @@ class Cosmos3OmniDiffusersPipeline(
             device=self.device,
             dtype=self.dtype,
         )
-        latents = condition_mask * cond_latent + (1.0 - condition_mask) * noise
+        latents = condition_mask * condition_latents + (1.0 - condition_mask) * noise
         velocity_mask = 1.0 - condition_mask
-        return latents, velocity_mask, cond_latent
+        return latents, velocity_mask, condition_latents
 
     def _prepare_action_latents(
         self,

vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py

@@ -24,6 +24,11 @@ from torch import nn
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import BasePipelineUtils, get_pipeline_utils
+from vllm_omni.diffusion.models.diffusers_adapter.quantization_utils import (
+    apply_diffusers_quantization_config,
+    convert_diffusers_quantization_config,
+    ensure_supported_diffusers_quantization,
+)
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.inputs.data import OmniPromptType, OmniTextPrompt
@@ -31,6 +36,20 @@ from vllm_omni.platforms import current_omni_platform
 
 logger = logging.getLogger(__name__)
 
+_DIFFUSERS_CONFIG_LOAD_KWARGS = {
+    "cache_dir",
+    "dduf_entries",
+    "force_download",
+    "local_dir",
+    "local_dir_use_symlinks",
+    "local_files_only",
+    "proxies",
+    "revision",
+    "subfolder",
+    "token",
+    "user_agent",
+}
+
 
 class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
     """Black-box adapter that delegates full pipeline execution to a diffusers pipeline.
@@ -80,12 +99,19 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
             "torch_dtype": dtype,
             **self.od_config.diffusers_load_kwargs,
         }
-        logger.debug(f"Loading diffusers pipeline with kwargs: {load_kwargs}")
+        convert_diffusers_quantization_config(load_kwargs)
 
         pipeline_class = self.od_config.diffusers_pipeline_cls
         pipeline_class_name = pipeline_class.__name__ if pipeline_class is not None else None
         self._pipeline_utils = get_pipeline_utils(pipeline_class_name)
         self._pipeline_utils.update_load_kwargs(self.od_config, load_kwargs)
+        component_names = (
+            self._load_diffusers_component_names(model_id, load_kwargs)
+            if self.od_config.quantization_config is not None and "quantization_config" not in load_kwargs
+            else {}
+        )
+        apply_diffusers_quantization_config(self.od_config, load_kwargs, component_names)
+        logger.debug(f"Loading diffusers pipeline with kwargs: {load_kwargs}")
 
         self._pipeline = DiffusionPipeline.from_pretrained(model_id, **load_kwargs)
         self._pipeline_utils.apply_post_load_updates(self._pipeline, self.od_config)
@@ -190,10 +216,38 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
                 "Eager execution is not supported with the diffusers backend. "
                 "Use a native pipeline for continuous batching mode."
             )
-        if self.od_config.quantization_config is not None:
-            raise NotImplementedError(
-                "Quantization is not supported with the diffusers backend. Use a native pipeline for quantization."
+        if (
+            self.od_config.quantization_config is not None
+            and "quantization_config" not in self.od_config.diffusers_load_kwargs
+        ):
+            ensure_supported_diffusers_quantization(self.od_config.quantization_config)
+            if self.od_config.diffusers_load_kwargs.get("dduf_file"):
+                raise NotImplementedError(
+                    "Diffusers backend quantization conversion does not support "
+                    "diffusers_load_kwargs.dduf_file yet. The preflight component "
+                    "discovery would need to mirror Diffusers' DDUF config loading. "
+                    "Use diffusers_load_kwargs.quantization_config for a native "
+                    "Diffusers quantization config, or omit dduf_file for vLLM-Omni "
+                    "quantization conversion."
+                )
+
+    def _load_diffusers_component_names(
+        self,
+        model_id: str,
+        load_kwargs: dict[str, Any],
+    ) -> dict[str, Any]:
+        config_load_kwargs = {k: load_kwargs[k] for k in _DIFFUSERS_CONFIG_LOAD_KWARGS if k in load_kwargs}
+        pipeline_config = DiffusionPipeline.load_config(model_id, **config_load_kwargs)
+        return {
+            name: value
+            for name, value in pipeline_config.items()
+            if (
+                isinstance(value, list)
+                and len(value) > 0
+                and value[0] is not None
+                and (name not in load_kwargs or load_kwargs[name] is not None)
             )
+        }
 
     # ------------------------------------------------------------------
     # Wrap settings, inputs, and outputs

vllm_omni/diffusion/models/diffusers_adapter/quantization_utils.py

@@ -0,0 +1,266 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Quantization helpers for the Diffusers backend."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from typing import Any
+
+from diffusers import PipelineQuantizationConfig, TorchAoConfig
+
+logger = logging.getLogger(__name__)
+
+_DIFFUSERS_TRANSFORMER_COMPONENTS = ("transformer", "transformer_2")
+_QuantizationConfigSpec = tuple[Callable[[Any], None], str]
+
+
+def ensure_supported_diffusers_quantization(quant_config: Any) -> None:
+    """Validate that a vLLM-Omni quantization config has a Diffusers mapping."""
+
+    method = _get_quant_method_name(quant_config)
+    validator, _ = _get_quantization_config_spec(method)
+    validator(quant_config)
+
+
+def convert_diffusers_quantization_config(load_kwargs: dict[str, Any]) -> None:
+    """Build Diffusers-native quantization config objects from dict kwargs."""
+
+    quant_config = load_kwargs.get("quantization_config")
+    if isinstance(quant_config, dict):
+        load_kwargs["quantization_config"] = _build_diffusers_native_quantization_config(quant_config)
+
+
+def apply_diffusers_quantization_config(
+    od_config: Any,
+    load_kwargs: dict[str, Any],
+    component_names: dict[str, Any],
+) -> bool:
+    """Inject a courtesy-converted quantization_config into load kwargs.
+
+    ``diffusers_load_kwargs`` is the canonical Diffusers backend configuration
+    path, so an explicit ``quantization_config`` already present in ``load_kwargs``
+    is never replaced. Returns whether a config was injected by this helper.
+    """
+
+    quant_config = getattr(od_config, "quantization_config", None)
+    if quant_config is None:
+        return False
+
+    if "quantization_config" in load_kwargs:
+        logger.warning(
+            "Both vLLM-Omni quantization_config and diffusers_load_kwargs.quantization_config "
+            "were provided for the diffusers backend. Using the Diffusers-native "
+            "quantization_config from diffusers_load_kwargs."
+        )
+        return False
+
+    transformer_components = _get_diffusers_transformer_components(component_names)
+    load_kwargs["quantization_config"] = _build_diffusers_quantization_config(
+        quant_config,
+        transformer_components,
+    )
+    return True
+
+
+def _normalize_method_name(method: Any) -> str:
+    method = getattr(method, "value", method)
+    return str(method).lower().replace("-", "_")
+
+
+def _get_quant_method_name(quant_config: Any) -> str:
+    get_name = getattr(quant_config, "get_name", None)
+    if get_name is not None:
+        return _normalize_method_name(get_name() if callable(get_name) else get_name)
+
+    method = getattr(quant_config, "quant_method", None)
+    if method is None:
+        method = getattr(quant_config, "method", None)
+    if method is None:
+        raise NotImplementedError(
+            "Diffusers backend quantization conversion requires a quantization "
+            "config with get_name(), quant_method, or method."
+        )
+    return _normalize_method_name(method)
+
+
+def _ensure_no_ignored_layers(quant_config: Any) -> None:
+    ignored_layers = getattr(quant_config, "ignored_layers", None)
+    if not ignored_layers:
+        ignored_layers = getattr(quant_config, "modules_to_not_convert", None)
+    if ignored_layers:
+        raise NotImplementedError(
+            "Diffusers backend quantization conversion does not map vLLM "
+            "ignored_layers/modules_to_not_convert names to Diffusers module "
+            "names. Use diffusers_load_kwargs for a native Diffusers config."
+        )
+
+
+def _get_torchao_quant_type_cls(class_name: str) -> type[Any]:
+    try:
+        import torchao.quantization as torchao_quantization
+    except ImportError as exc:
+        raise ImportError(
+            "Diffusers backend quantization conversion for fp8/int8 requires "
+            "torchao. Install torchao or pass a Diffusers-native "
+            "quantization_config through diffusers_load_kwargs."
+        ) from exc
+
+    try:
+        return getattr(torchao_quantization, class_name)
+    except AttributeError as exc:
+        raise ImportError(f"torchao.quantization.{class_name} is required for this quantization mapping.") from exc
+
+
+def _build_diffusers_native_quantization_config(config: dict[str, Any]) -> Any:
+    config = config.copy()
+    quant_mapping = config.get("quant_mapping")
+    if isinstance(quant_mapping, dict):
+        config["quant_mapping"] = {
+            component_name: _build_diffusers_component_quant_config(component_config)
+            for component_name, component_config in quant_mapping.items()
+        }
+    return PipelineQuantizationConfig(**config)
+
+
+def _build_diffusers_component_quant_config(config: Any) -> Any:
+    if not isinstance(config, dict):
+        return config
+
+    quant_method = config.get("quant_method")
+    if quant_method is None:
+        raise ValueError("Diffusers quant_mapping entries must provide quant_method.")
+
+    quant_method = _normalize_method_name(quant_method)
+    config_cls = _get_diffusers_quantization_config_cls(quant_method)
+    from_dict = getattr(config_cls, "from_dict", None)
+    if from_dict is not None:
+        return from_dict(config)
+
+    init_kwargs = config.copy()
+    init_kwargs.pop("quant_method", None)
+    return config_cls(**init_kwargs)
+
+
+def _get_diffusers_quantization_config_cls(quant_method: str) -> type[Any]:
+    from diffusers.quantizers.auto import AUTO_QUANTIZATION_CONFIG_MAPPING as DIFFUSERS_QUANT_CONFIG_MAPPING
+
+    config_cls = DIFFUSERS_QUANT_CONFIG_MAPPING.get(quant_method)
+    if config_cls is not None:
+        return config_cls
+
+    try:
+        from transformers.quantizers.auto import (
+            AUTO_QUANTIZATION_CONFIG_MAPPING as TRANSFORMERS_QUANT_CONFIG_MAPPING,
+        )
+    except ImportError:
+        TRANSFORMERS_QUANT_CONFIG_MAPPING = {}
+
+    config_cls = TRANSFORMERS_QUANT_CONFIG_MAPPING.get(quant_method)
+    if config_cls is None:
+        raise ValueError(f"Diffusers quant_mapping entry uses unknown quant_method={quant_method!r}.")
+    return config_cls
+
+
+def _build_torchao_pipeline_quant_config(
+    torchao_quant_type_name: str,
+    transformer_components: list[str],
+) -> Any:
+    """Build a Diffusers config after the caller has validated the vLLM config."""
+
+    quant_type_cls = _get_torchao_quant_type_cls(torchao_quant_type_name)
+    return PipelineQuantizationConfig(
+        quant_mapping={
+            component_name: TorchAoConfig(
+                quant_type=quant_type_cls(),
+            )
+            for component_name in transformer_components
+        }
+    )
+
+
+def _get_diffusers_transformer_components(component_names: dict[str, Any]) -> list[str]:
+    transformer_components = [
+        component_name for component_name in component_names if component_name in _DIFFUSERS_TRANSFORMER_COMPONENTS
+    ]
+    if transformer_components:
+        return transformer_components
+
+    available_components = ", ".join(sorted(component_names.keys())) or "<none>"
+    raise NotImplementedError(
+        "Diffusers backend quantization conversion currently only supports "
+        "pipelines with transformer or transformer_2 components. "
+        f"Found pipeline components: {available_components}. Use "
+        "diffusers_load_kwargs for a native Diffusers quantization config."
+    )
+
+
+def _validate_fp8_quant_config(quant_config: Any) -> None:
+    if getattr(quant_config, "is_checkpoint_fp8_serialized", False):
+        raise NotImplementedError(
+            "Diffusers backend fp8 conversion only supports online/dynamic "
+            "TorchAO quantization; serialized vLLM fp8 checkpoints are not mapped."
+        )
+
+    activation_scheme = getattr(quant_config, "activation_scheme", "dynamic")
+    if activation_scheme != "dynamic":
+        raise NotImplementedError(
+            f"Diffusers backend fp8 conversion only supports activation_scheme='dynamic'. Got {activation_scheme!r}."
+        )
+
+    weight_block_size = getattr(quant_config, "weight_block_size", None)
+    if weight_block_size is not None:
+        raise NotImplementedError(
+            "Diffusers backend fp8 conversion does not map vLLM weight_block_size "
+            "to TorchAO. Use diffusers_load_kwargs for a native Diffusers config."
+        )
+
+    _ensure_no_ignored_layers(quant_config)
+
+
+def _validate_int8_quant_config(quant_config: Any) -> None:
+    if getattr(quant_config, "is_checkpoint_int8_serialized", False):
+        raise NotImplementedError(
+            "Diffusers backend int8 conversion only supports online/dynamic "
+            "TorchAO quantization; serialized vLLM int8 checkpoints are not mapped."
+        )
+
+    activation_scheme = getattr(quant_config, "activation_scheme", "dynamic")
+    if activation_scheme != "dynamic":
+        raise NotImplementedError(
+            f"Diffusers backend int8 conversion only supports activation_scheme='dynamic'. Got {activation_scheme!r}."
+        )
+
+    _ensure_no_ignored_layers(quant_config)
+
+
+_QUANTIZATION_CONFIG_SPECS: dict[str, _QuantizationConfigSpec] = {
+    "fp8": (_validate_fp8_quant_config, "Float8DynamicActivationFloat8WeightConfig"),
+    "int8": (_validate_int8_quant_config, "Int8DynamicActivationInt8WeightConfig"),
+}
+
+
+def _get_quantization_config_spec(method: str) -> _QuantizationConfigSpec:
+    spec = _QUANTIZATION_CONFIG_SPECS.get(method)
+    if spec is None:
+        raise NotImplementedError(
+            f"Diffusers backend quantization conversion does not support {method!r}. "
+            "Use diffusers_load_kwargs for a native Diffusers quantization config, "
+            "or use a native vLLM-Omni pipeline for this quantization method."
+        )
+    return spec
+
+
+def _build_diffusers_quantization_config(
+    quant_config: Any,
+    transformer_components: list[str],
+) -> Any:
+    """Build a Diffusers PipelineQuantizationConfig from a supported config."""
+
+    method = _get_quant_method_name(quant_config)
+    _, torchao_quant_type_name = _get_quantization_config_spec(method)
+    return _build_torchao_pipeline_quant_config(
+        torchao_quant_type_name,
+        transformer_components,
+    )

vllm_omni/diffusion/models/flux/pipeline_flux.py

@@ -502,38 +502,14 @@ class FluxPipeline(
             return False
         return True
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | list[str] | None = None,
-        prompt_2: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        negative_prompt_2: str | list[str] | None = None,
-        true_cfg_scale: float = 1.0,
-        height: int | None = None,
-        width: int | None = None,
-        num_inference_steps: int = 28,
-        sigmas: list[float] | None = None,
-        guidance_scale: float = 3.5,
-        num_images_per_prompt: int = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.FloatTensor | None = None,
-        prompt_embeds: torch.FloatTensor | None = None,
-        pooled_prompt_embeds: torch.FloatTensor | None = None,
-        negative_prompt_embeds: torch.FloatTensor | None = None,
-        negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
-        output_type: str | None = "pil",
-        return_dict: bool = True,
-        joint_attention_kwargs: dict[str, Any] | None = None,
-        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
-        max_sequence_length: int = 512,
-    ) -> list[DiffusionOutput]:
+    def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]:
         """Forward pass for flux."""
         # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
         # TODO: May be some data formatting operations on the API side. Hack for now.
         sampling_params_list = req.sampling_params_list
         common_sampling_params = sampling_params_list[0]
-        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts]
+
         if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
             negative_prompt = None
         elif req.prompts:
@@ -541,40 +517,39 @@ class FluxPipeline(
 
         height = common_sampling_params.height or self.default_sample_size * self.vae_scale_factor
         width = common_sampling_params.width or self.default_sample_size * self.vae_scale_factor
-        num_inference_steps = common_sampling_params.num_inference_steps or num_inference_steps
-        sigmas = common_sampling_params.sigmas or sigmas
-        guidance_scale = (
-            common_sampling_params.guidance_scale
-            if common_sampling_params.guidance_scale is not None
-            else guidance_scale
-        )
-        true_cfg_scale = common_sampling_params.true_cfg_scale or true_cfg_scale
+        num_inference_steps = common_sampling_params.num_inference_steps or 28
+        sigmas = common_sampling_params.sigmas
+        guidance_scale = common_sampling_params.guidance_scale
+        true_cfg_scale = common_sampling_params.true_cfg_scale or 1.0
         num_images_per_prompt = (
-            common_sampling_params.num_outputs_per_prompt
-            if common_sampling_params.num_outputs_per_prompt > 0
-            else num_images_per_prompt
+            common_sampling_params.num_outputs_per_prompt if common_sampling_params.num_outputs_per_prompt > 0 else 1
         )
-        generator = req.collate_request_generators(num_images_per_prompt, generator)
-        latents = req.collate_request_tensors("latents", latents)
+        max_sequence_length = common_sampling_params.max_sequence_length or 512
+        generator = req.collate_request_generators(num_images_per_prompt, None)
+        latents = req.collate_request_tensors("latents", None)
         prompt_fields = DiffusionRequestBatch.collate_prompt_field_map(
             req.prompts,
             {
-                "prompt_embeds": prompt_embeds,
-                "negative_prompt_embeds": negative_prompt_embeds,
-                "pooled_prompt_embeds": pooled_prompt_embeds,
-                "negative_pooled_prompt_embeds": negative_pooled_prompt_embeds,
+                "prompt_embeds": None,
+                "negative_prompt_embeds": None,
+                "pooled_prompt_embeds": None,
+                "negative_pooled_prompt_embeds": None,
             },
         )
         prompt_embeds = prompt_fields["prompt_embeds"]
         negative_prompt_embeds = prompt_fields["negative_prompt_embeds"]
         pooled_prompt_embeds = prompt_fields["pooled_prompt_embeds"]
         negative_pooled_prompt_embeds = prompt_fields["negative_pooled_prompt_embeds"]
+        prompt_2 = None
+        negative_prompt_2 = None
         if prompt_embeds is not None:
             prompt = None
-            prompt_2 = None
         if negative_prompt_embeds is not None:
             negative_prompt = None
-            negative_prompt_2 = None
+
+        output_type = common_sampling_params.output_type or "pil"
+        joint_attention_kwargs: dict[str, Any] | None = None
+        callback_on_step_end_tensor_inputs = ["latents"]
 
         # 1. Check inputs. Raise error if not correct
         self.check_inputs(
@@ -598,9 +573,7 @@ class FluxPipeline(
         self._interrupt = False
 
         # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
+        if prompt is not None:
             batch_size = len(prompt)
         else:
             batch_size = prompt_embeds.shape[0]

vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py

@@ -750,108 +750,7 @@ class Flux2KleinPipeline(
     def interrupt(self):
         return self._interrupt
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        image: PIL.Image.Image | list[PIL.Image.Image] | None = None,
-        reference_image: PIL.Image.Image | list[PIL.Image.Image] | None = None,
-        mask_image: PIL.Image.Image | list[PIL.Image.Image] | None = None,
-        prompt: str | list[str] | None = None,
-        height: int | None = None,
-        width: int | None = None,
-        num_inference_steps: int = 50,
-        sigmas: list[float] | None = None,
-        strength: float = 1.0,
-        guidance_scale: float | None = 4.0,
-        num_images_per_prompt: int = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.Tensor | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        output_type: str | None = "pil",
-        return_dict: bool = True,
-        attention_kwargs: dict[str, Any] | None = None,
-        callback_on_step_end: Callable[[int, int, dict], None] | None = None,
-        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
-        max_sequence_length: int = 512,
-        text_encoder_out_layers: tuple[int, ...] = (9, 18, 27),
-        padding_mask_crop: int | None = None,
-    ) -> DiffusionOutput:
-        r"""
-        Function invoked when calling the pipeline for generation.
-
-        Args:
-            image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or list of these):
-                `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
-                numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
-                or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
-                list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image
-                latents as `image`, but if passing latents directly it is not encoded again.
-            prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
-                instead.
-            guidance_scale (`float`, *optional*, defaults to 4.0):
-                Guidance scale as defined in [Classifier-Free Diffusion
-                Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2.
-                of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
-                `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
-                the text `prompt`, usually at the expense of lower image quality. For step-wise distilled models,
-                `guidance_scale` is ignored.
-            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
-                The height in pixels of the generated image. This is set to 1024 by default for the best results.
-            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
-                The width in pixels of the generated image. This is set to 1024 by default for the best results.
-            num_inference_steps (`int`, *optional*, defaults to 50):
-                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
-                expense of slower inference.
-            sigmas (`List[float]`, *optional*):
-                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
-                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
-                will be used.
-            num_images_per_prompt (`int`, *optional*, defaults to 1):
-                The number of images to generate per prompt.
-            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
-                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
-                to make generation deterministic.
-            latents (`torch.Tensor`, *optional*):
-                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
-                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
-                tensor will be generated by sampling using the supplied random `generator`.
-            prompt_embeds (`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.
-            negative_prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated negative text embeddings. Note that "" is used as the negative prompt in this pipeline.
-                If not provided, will be generated from "".
-            output_type (`str`, *optional*, defaults to `"pil"`):
-                The output format of the generate image. Choose between
-                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`~pipelines.qwenimage.QwenImagePipelineOutput`] instead of a plain tuple.
-            attention_kwargs (`dict`, *optional*):
-                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
-                `self.processor` in
-                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
-            callback_on_step_end (`Callable`, *optional*):
-                A function that calls at the end of each denoising steps during the inference. The function is called
-                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
-                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
-                `callback_on_step_end_tensor_inputs`.
-            callback_on_step_end_tensor_inputs (`List`, *optional*):
-                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
-                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
-                `._callback_tensor_inputs` attribute of your pipeline class.
-            max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`.
-            text_encoder_out_layers (`Tuple[int]`):
-                Layer indices to use in the `text_encoder` to derive the final prompt embeddings.
-
-        Examples:
-
-        Returns:
-            [`~pipelines.flux2.Flux2PipelineOutput`] or `tuple`: [`~pipelines.flux2.Flux2PipelineOutput`] if
-            `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the
-            generated images.
-        """
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
         if len(req.prompts) > 1:
             logger.warning(
                 """This model only supports a single prompt, not a batched request.""",
@@ -878,39 +777,32 @@ class Flux2KleinPipeline(
         else:
             image = PIL.Image.open(raw_image) if isinstance(raw_image, str) else cast(PIL.Image.Image, raw_image)
 
-        height = req.sampling_params.height or height
-        width = req.sampling_params.width or width
+        prompt_embeds = None
+        negative_prompt_embeds = None
+
+        height = req.sampling_params.height
+        width = req.sampling_params.width
         num_inference_steps = (
-            req.sampling_params.num_inference_steps
-            if req.sampling_params.num_inference_steps is not None
-            else num_inference_steps
+            50 if req.sampling_params.num_inference_steps is None else req.sampling_params.num_inference_steps
         )
-        sigmas = req.sampling_params.sigmas or sigmas
-        guidance_scale = (
-            req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale
-        )
-        generator = req.sampling_params.generator or generator
+        sigmas = req.sampling_params.sigmas
+        guidance_scale = req.sampling_params.guidance_scale
+        generator = req.sampling_params.generator
         num_images_per_prompt = (
-            req.sampling_params.num_outputs_per_prompt
-            if req.sampling_params.num_outputs_per_prompt > 0
-            else num_images_per_prompt
+            req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1
         )
-        max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length
-        text_encoder_out_layers = req.sampling_params.extra_args.get("text_encoder_out_layers", text_encoder_out_layers)
+        max_sequence_length = req.sampling_params.max_sequence_length or 512
+        latents = req.sampling_params.latents
+        strength = req.sampling_params.strength if req.sampling_params.strength is not None else 1.0
 
-        req_prompt_embeds = [p.get("prompt_embeds") if not isinstance(p, str) else None for p in req.prompts]
-        if any(p is not None for p in req_prompt_embeds):
-            # If at list one prompt is provided as an embedding,
-            # Then assume that the user wants to provide embeddings for all prompts, and enter this if block
-            # If the user in fact provides mixed input format, req_prompt_embeds will have some None's
-            # And `torch.stack` automatically raises an exception for us
-            prompt_embeds = torch.stack(req_prompt_embeds)  # type: ignore # intentionally expect TypeError
-
-        req_negative_prompt_embeds = [
-            p.get("negative_prompt_embeds") if not isinstance(p, str) else None for p in req.prompts
-        ]
-        if any(p is not None for p in req_negative_prompt_embeds):
-            negative_prompt_embeds = torch.stack(req_negative_prompt_embeds)  # type: ignore # intentionally expect TypeError
+        text_encoder_out_layers: tuple[int, ...] = req.sampling_params.extra_args.get(
+            "text_encoder_out_layers", (9, 18, 27)
+        )
+        padding_mask_crop: int | None = req.sampling_params.extra_args.get("padding_mask_crop")
+        attention_kwargs: dict[str, Any] | None = None
+        callback_on_step_end: Callable[[int, int, dict], None] | None = None
+        callback_on_step_end_tensor_inputs = ["latents"]
+        output_type = req.sampling_params.output_type or "pil"
 
         # 1. Check inputs. Raise error if not correct
         self.check_inputs(
@@ -930,12 +822,7 @@ class Flux2KleinPipeline(
         self._interrupt = False
 
         # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
+        batch_size = 1
 
         device = self._execution_device
 

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -29,7 +29,7 @@ from transformers.modeling_outputs import (
     CausalLMOutputWithPast,
 )
 from transformers.modeling_utils import PreTrainedModel
-from vllm.config import CacheConfig
+from vllm.config import CacheConfig, get_current_vllm_config
 from vllm.distributed import (
     get_tensor_model_parallel_world_size,
 )
@@ -1624,6 +1624,7 @@ class HunYuanSparseMoeBlock(nn.Module):
         else:
             self.shared_mlp = None
 
+        enable_expert_parallel = get_current_vllm_config().parallel_config.enable_expert_parallel
         self.experts = HunyuanFusedMoE(
             shared_experts=self.shared_mlp,
             num_experts=self.n_routed_experts,
@@ -1635,7 +1636,7 @@ class HunYuanSparseMoeBlock(nn.Module):
             prefix=f"{prefix}.experts",
             enable_eplb=self.enable_eplb,
             num_redundant_experts=self.n_redundant_experts,
-            pcp_size=1,
+            pcp_size=None if enable_expert_parallel else 1,
         )
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py

@@ -495,66 +495,35 @@ class LongCatImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfile
         noise_pred = comb_pred * scale
         return noise_pred
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        height: int | None = None,
-        width: int | None = None,
-        num_inference_steps: int = 50,
-        sigmas: list[float] | None = None,
-        guidance_scale: float = 4.5,
-        num_images_per_prompt: int = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.FloatTensor | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        output_type: str | None = "pil",
-        return_dict: bool = True,
-        joint_attention_kwargs: dict[str, Any] | None = None,
-        enable_cfg_renorm: bool | None = True,
-        cfg_renorm_min: float | None = 0.0,
-        enable_prompt_rewrite: bool | None = True,
-    ) -> DiffusionOutput:
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
         # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
         # TODO: May be some data formatting operations on the API side. Hack for now.
-        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts]
         if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
             negative_prompt = None
         elif req.prompts:
             negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
 
-        height = req.sampling_params.height or height or self.default_sample_size * self.vae_scale_factor
-        width = req.sampling_params.width or width or self.default_sample_size * self.vae_scale_factor
-        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
-        sigmas = req.sampling_params.sigmas or sigmas
-        generator = req.sampling_params.generator or generator
-        guidance_scale = (
-            req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale
-        )
+        prompt_embeds = None
+        negative_prompt_embeds = None
+
+        height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor
+        width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor
+        num_inference_steps = req.sampling_params.num_inference_steps or 50
+        sigmas = req.sampling_params.sigmas
+        generator = req.sampling_params.generator
+        guidance_scale = req.sampling_params.guidance_scale
         num_images_per_prompt = (
-            req.sampling_params.num_outputs_per_prompt
-            if req.sampling_params.num_outputs_per_prompt is not None
-            else num_images_per_prompt
+            req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1
         )
-        enable_prompt_rewrite = req.sampling_params.extra_args.get("enable_prompt_rewrite", enable_prompt_rewrite)
-        enable_cfg_renorm = req.sampling_params.extra_args.get("enable_cfg_renorm", enable_cfg_renorm)
-        cfg_renorm_min = req.sampling_params.extra_args.get("cfg_renorm_min", cfg_renorm_min)
+        latents = req.sampling_params.latents
 
-        req_prompt_embeds = [p.get("prompt_embeds") if not isinstance(p, str) else None for p in req.prompts]
-        if any(p is not None for p in req_prompt_embeds):
-            # If at list one prompt is provided as an embedding,
-            # Then assume that the user wants to provide embeddings for all prompts, and enter this if block
-            # If the user in fact provides mixed input format, req_prompt_embeds will have some None's
-            # And `torch.stack` automatically raises an exception for us
-            prompt_embeds = torch.stack(req_prompt_embeds)  # type: ignore # intentionally expect TypeError
-
-        req_negative_prompt_embeds = [
-            p.get("negative_prompt_embeds") if not isinstance(p, str) else None for p in req.prompts
-        ]
-        if any(p is not None for p in req_negative_prompt_embeds):
-            negative_prompt_embeds = torch.stack(req_negative_prompt_embeds)  # type: ignore # intentionally expect TypeError
+        extra_args = getattr(req.sampling_params, "extra_args", {}) or {}
+        enable_prompt_rewrite: bool = extra_args.get("enable_prompt_rewrite", True)
+        enable_cfg_renorm: bool = extra_args.get("enable_cfg_renorm", True)
+        cfg_renorm_min: float = extra_args.get("cfg_renorm_min", 0.0)
+        joint_attention_kwargs: dict[str, Any] | None = None
+        output_type = req.sampling_params.output_type or "pil"
 
         self.check_inputs(
             prompt,
@@ -571,12 +540,7 @@ class LongCatImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfile
         self._interrupt = False
 
         # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
+        batch_size = len(prompt)
 
         device = self.device
         if enable_prompt_rewrite and prompt is not None:

vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py

@@ -6,7 +6,7 @@ import math
 import os
 import re
 from collections.abc import Iterable
-from typing import Any, ClassVar, cast
+from typing import ClassVar, cast
 
 import numpy as np
 import PIL.Image
@@ -542,24 +542,7 @@ class LongCatImageEditPipeline(
                 f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
             )
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        image: PIL.Image.Image | torch.Tensor | None = None,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        num_inference_steps: int = 50,
-        sigmas: list[float] | None = None,
-        guidance_scale: float = 3.5,
-        num_images_per_prompt: int | None = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.FloatTensor | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        output_type: str | None = "pil",
-        return_dict: bool = True,
-        joint_attention_kwargs: dict[str, Any] | None = None,
-    ) -> DiffusionOutput:
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
         # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
         # TODO: May be some data formatting operations on the API side. Hack for now.
         if len(req.prompts) > 1:
@@ -570,27 +553,22 @@ class LongCatImageEditPipeline(
         first_prompt = req.prompts[0]
         prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
         negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt")
-        prompt_embeds = None if isinstance(first_prompt, str) else first_prompt.get("prompt_embeds")
-        negative_prompt_embeds = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt_embeds")  # type: ignore # Why it is list[torch.Tensor] in OmniTokenInputs or OmniEmbedsPrompt? Doesn't make sense
 
-        sigmas = req.sampling_params.sigmas or sigmas
-        guidance_scale = (
-            req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale
-        )
-        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        sigmas = req.sampling_params.sigmas
+        guidance_scale = req.sampling_params.guidance_scale
+        num_inference_steps = req.sampling_params.num_inference_steps or 50
         num_images_per_prompt = (
-            req.sampling_params.num_outputs_per_prompt
-            if req.sampling_params.num_outputs_per_prompt is not None
-            else num_images_per_prompt
+            req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1
         )
-        generator = req.sampling_params.generator or generator
+        generator = req.sampling_params.generator
         height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor
         width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor
+        latents = req.sampling_params.latents
+        prompt_embeds = None
+        negative_prompt_embeds = None
+        output_type = req.sampling_params.output_type or "pil"
 
-        if prompt is not None:
-            batch_size = 1 if isinstance(prompt, str) else len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
+        batch_size = 1 if isinstance(prompt, str) else len(prompt)
 
         if not isinstance(first_prompt, str) and "preprocessed_image" in (
             additional_information := first_prompt.get("additional_information", {})
@@ -600,13 +578,7 @@ class LongCatImageEditPipeline(
             calculated_height = additional_information.get("calculated_height", height)
             calculated_width = additional_information.get("calculated_width", width)
         else:
-            image_size = image[0].size if isinstance(image, list) else image.size
-            calculated_width, calculated_height = calculate_dimensions(1024 * 1024, image_size[0] * 1.0 / image_size[1])
-
-            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 = self.image_processor.resize(image, calculated_height // 2, calculated_width // 2)
-                image = self.image_processor.preprocess(image, calculated_height, calculated_width)
+            raise RuntimeError("Missing preprocess image that should have been created by the preprocess function.")
 
         self.check_inputs(
             prompt,
@@ -621,7 +593,8 @@ class LongCatImageEditPipeline(
             prompt=prompt, image=prompt_image, prompt_embeds=prompt_embeds, num_images_per_prompt=num_images_per_prompt
         )
 
-        if guidance_scale > 1:
+        do_true_cfg = guidance_scale > 1 and negative_prompt is not None
+        if do_true_cfg:
             (negative_prompt_embeds, negative_text_ids) = self.encode_prompt(
                 prompt=negative_prompt,
                 image=prompt_image,
@@ -680,7 +653,6 @@ class LongCatImageEditPipeline(
                 latent_model_input = torch.cat([latents, image_latents], dim=1)
 
             timestep = t.expand(latent_model_input.shape[0]).to(latents.dtype)
-            do_true_cfg = guidance_scale > 1
             positive_kwargs = {
                 "hidden_states": latent_model_input,
                 "timestep": timestep / 1000,

vllm_omni/diffusion/models/ltx2/__init__.py

@@ -10,10 +10,8 @@ from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
     get_ltx2_post_process_func,
     load_transformer_config,
 )
-from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import (
-    LTX23ImageToVideoPipeline,
-    LTX23Pipeline,
-)
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
 from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_image2video import (
     LTX2I2VDMD2Pipeline,
     LTX2ImageToVideoPipeline,

vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3_image2video.py

@@ -1,16 +1,452 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""Re-exports for LTX-2.3 I2V pipeline variants.
+"""LTX-2.3 image-to-video pipeline."""
 
-The registry loads pipeline classes by (mod_folder, mod_relname, cls_name).
-This module exposes the I2V class names so the registry can find them.
-"""
+from __future__ import annotations
+
+from typing import Any
+
+import PIL.Image
+import torch
+from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_latents
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.video_processor import VideoProcessor
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 
 from .pipeline_ltx2_3 import (
-    LTX23ImageToVideoPipeline,
-    get_ltx2_post_process_func,  # noqa: F401 - loaded by registry via getattr
+    LTX23Pipeline,
+    _LTX23DenoiseContext,
+    _LTX23ForwardContext,
+    _LTX23PromptContext,
+    _LTX23RequestInputs,
+    get_ltx2_post_process_func,
 )
+from .pipeline_ltx2_image2video import LTX2ImageToVideoPipeline, _I2VVideoAudioScheduler
+
+
+class LTX23ImageToVideoPipeline(LTX23Pipeline):
+    """LTX-2.3 image-to-video pipeline.
+
+    This keeps the LTX-2.3 prompt connector, x0-space CFG, sigma prompt
+    modulation, and audio branch semantics from ``LTX23Pipeline`` while
+    reusing the existing LTX image-conditioning contract: the first video
+    latent frame is encoded from the input image and remains fixed during
+    denoising.
+    """
+
+    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
+        super().__init__(od_config=od_config, prefix=prefix)
+        self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio, resample="bilinear")
+
+    support_image_input = True
+
+    _normalize_latents = staticmethod(LTX2ImageToVideoPipeline._normalize_latents)
+    _create_noised_state = staticmethod(LTX2ImageToVideoPipeline._create_noised_state)
+
+    @staticmethod
+    def _resolve_single_prompt_image(raw_image: Any) -> Any:
+        if isinstance(raw_image, list):
+            if len(raw_image) != 1:
+                raise ValueError(
+                    "LTX-2.3 I2V prompt dictionaries support exactly one image per prompt. "
+                    "Pass one image per prompt for batched I2V requests."
+                )
+            return raw_image[0]
+        return raw_image
+
+    @staticmethod
+    def _resolve_additional_image(additional: dict[str, Any]) -> Any:
+        raw_image = additional.get("preprocessed_image")
+        if raw_image is None:
+            raw_image = additional.get("pixel_values")
+        if raw_image is None:
+            raw_image = additional.get("image")
+        return raw_image
+
+    def prepare_latents(
+        self,
+        image: torch.Tensor | None = None,
+        batch_size: int = 1,
+        num_channels_latents: int = 128,
+        height: int = 512,
+        width: int = 768,
+        num_frames: int = 121,
+        noise_scale: float = 0.0,
+        dtype: torch.dtype | None = None,
+        device: torch.device | None = None,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Prepare I2V latents and the first-frame conditioning mask.
+
+        If caller-provided latents are used without an image, the latents must
+        already represent the full video state including the conditioning first
+        frame. Packed 3D latents are assumed to be in transformer token layout.
+        """
+        height = height // self.vae_spatial_compression_ratio
+        width = width // self.vae_spatial_compression_ratio
+        num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1
+
+        shape = (batch_size, num_channels_latents, num_frames, height, width)
+        mask_shape = (batch_size, 1, num_frames, height, width)
+
+        if latents is not None:
+            if latents.ndim == 5:
+                batch_size, _, num_frames, height, width = latents.shape
+                mask_shape = (batch_size, 1, num_frames, height, width)
+                conditioning_mask = latents.new_zeros(mask_shape)
+                conditioning_mask[:, :, 0] = 1.0
+
+                latents = self._normalize_latents(
+                    latents,
+                    self.vae.latents_mean,
+                    self.vae.latents_std,
+                    self.vae.config.scaling_factor,
+                )
+                latents = self._create_noised_state(latents, noise_scale * (1 - conditioning_mask), generator)
+                latents = self._pack_latents(
+                    latents,
+                    self.transformer_spatial_patch_size,
+                    self.transformer_temporal_patch_size,
+                )
+            else:
+                conditioning_mask = latents.new_zeros(mask_shape)
+                conditioning_mask[:, :, 0] = 1.0
+
+            conditioning_mask = self._pack_latents(
+                conditioning_mask,
+                self.transformer_spatial_patch_size,
+                self.transformer_temporal_patch_size,
+            ).squeeze(-1)
+            if latents.ndim != 3 or latents.shape[:2] != conditioning_mask.shape:
+                raise ValueError(
+                    "Provided `latents` tensor has shape"
+                    f" {latents.shape}, but the expected shape is {conditioning_mask.shape + (num_channels_latents,)}."
+                )
+            return latents.to(device=device, dtype=dtype), conditioning_mask
+
+        if image is None:
+            raise ValueError("`image` must be provided when `latents` is None.")
+
+        image_batch_size = image.shape[0]
+        if image_batch_size == 0:
+            raise ValueError("`image` batch is empty.")
+        if batch_size % image_batch_size != 0:
+            raise ValueError(
+                f"`batch_size` ({batch_size}) must be divisible by image batch size ({image_batch_size}) "
+                "for image-to-video outputs."
+            )
+        num_videos_per_prompt = batch_size // image_batch_size
+
+        if isinstance(generator, list):
+            if len(generator) != batch_size:
+                raise ValueError(
+                    f"You have passed a list of generators of length {len(generator)}, but requested an effective"
+                    f" batch size of {batch_size}. Make sure the batch size matches the length of the generators."
+                )
+            image_generators = [generator[i * num_videos_per_prompt] for i in range(image_batch_size)]
+            init_latents = [
+                retrieve_latents(self.vae.encode(image[i].unsqueeze(0).unsqueeze(2)), image_generators[i], "argmax")
+                for i in range(image_batch_size)
+            ]
+        else:
+            init_latents = [
+                retrieve_latents(self.vae.encode(img.unsqueeze(0).unsqueeze(2)), generator, "argmax") for img in image
+            ]
+
+        init_latents = torch.cat(init_latents, dim=0).to(dtype)
+        if num_videos_per_prompt > 1:
+            init_latents = init_latents.repeat_interleave(num_videos_per_prompt, dim=0)
+        init_latents = self._normalize_latents(
+            init_latents,
+            self.vae.latents_mean,
+            self.vae.latents_std,
+        )
+        init_latents = init_latents.repeat(1, 1, num_frames, 1, 1)
+
+        conditioning_mask = torch.zeros(mask_shape, device=device, dtype=dtype)
+        conditioning_mask[:, :, 0] = 1.0
+
+        noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        latents = init_latents * conditioning_mask + noise * (1 - conditioning_mask)
+
+        conditioning_mask = self._pack_latents(
+            conditioning_mask,
+            self.transformer_spatial_patch_size,
+            self.transformer_temporal_patch_size,
+        ).squeeze(-1)
+        latents = self._pack_latents(latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size)
+
+        return latents, conditioning_mask
+
+    def check_inputs(
+        self,
+        image,
+        height,
+        width,
+        prompt,
+        latents=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        prompt_attention_mask=None,
+        negative_prompt_attention_mask=None,
+    ):
+        if image is None and latents is None:
+            raise ValueError("Provide either `image` or `latents`. Cannot leave both undefined.")
+        super().check_inputs(
+            prompt=prompt,
+            height=height,
+            width=width,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            prompt_attention_mask=prompt_attention_mask,
+            negative_prompt_attention_mask=negative_prompt_attention_mask,
+        )
+
+    _step_video_latents_i2v = LTX2ImageToVideoPipeline._step_video_latents_i2v
+
+    def _resolve_request_image(
+        self,
+        req: DiffusionRequestBatch,
+        image: PIL.Image.Image | torch.Tensor | list[PIL.Image.Image | torch.Tensor] | None,
+        request_inputs: _LTX23RequestInputs,
+    ) -> PIL.Image.Image | torch.Tensor | list[PIL.Image.Image | torch.Tensor] | None:
+        if image is not None or not req.prompts:
+            return image
+
+        raw_images = []
+        for prompt_item in req.prompts:
+            if isinstance(prompt_item, str):
+                raw_image = None
+            else:
+                multi_modal_data = prompt_item.get("multi_modal_data") or {}
+                raw_image = multi_modal_data.get("image")
+                if raw_image is None:
+                    additional = prompt_item.get("additional_information") or {}
+                    raw_image = self._resolve_additional_image(additional)
+            raw_image = self._resolve_single_prompt_image(raw_image)
+            if isinstance(raw_image, str):
+                raw_image = PIL.Image.open(raw_image).convert("RGB")
+            raw_images.append(raw_image)
+
+        if any(img is None for img in raw_images) and request_inputs.latents is None:
+            raise ValueError("Image is required for LTX-2.3 I2V generation.")
+        if len(raw_images) == 1:
+            return raw_images[0]
+        if raw_images:
+            return raw_images
+        return image
+
+    def _check_forward_inputs(
+        self,
+        request_inputs: _LTX23RequestInputs,
+        image: Any | None = None,
+    ) -> None:
+        self.check_inputs(
+            image=image,
+            height=request_inputs.height,
+            width=request_inputs.width,
+            prompt=request_inputs.prompt,
+            latents=request_inputs.latents,
+            prompt_embeds=request_inputs.prompt_embeds,
+            negative_prompt_embeds=request_inputs.negative_prompt_embeds,
+            prompt_attention_mask=request_inputs.prompt_attention_mask,
+            negative_prompt_attention_mask=request_inputs.negative_prompt_attention_mask,
+        )
+
+    def _resolve_video_latent_dimensions(
+        self,
+        request_inputs: _LTX23RequestInputs,
+    ) -> tuple[int, int, int]:
+        latent_num_frames = (request_inputs.num_frames - 1) // self.vae_temporal_compression_ratio + 1
+        latent_height = request_inputs.height // self.vae_spatial_compression_ratio
+        latent_width = request_inputs.width // self.vae_spatial_compression_ratio
+        if request_inputs.latents is not None:
+            if request_inputs.latents.ndim == 5:
+                _, _, latent_num_frames, latent_height, latent_width = request_inputs.latents.shape
+            elif request_inputs.latents.ndim != 3:
+                raise ValueError(
+                    f"Provided `latents` tensor has shape {request_inputs.latents.shape}, but the expected shape is "
+                    "either [batch_size, seq_len, num_features] or "
+                    "[batch_size, latent_dim, latent_frames, latent_height, latent_width]."
+                )
+        return latent_num_frames, latent_height, latent_width
+
+    def _prepare_video_latents_stage(
+        self,
+        request_inputs: _LTX23RequestInputs,
+        prompt_context: _LTX23PromptContext,
+        *,
+        device: torch.device,
+        noise_scale: float,
+        image: Any | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if request_inputs.latents is None:
+            if isinstance(image, torch.Tensor):
+                if image.ndim == 3:
+                    image = image.unsqueeze(0)
+            elif isinstance(image, list) and image and isinstance(image[0], torch.Tensor):
+                image = torch.stack(image, dim=0)
+            else:
+                image = self.video_processor.preprocess(
+                    image,
+                    height=request_inputs.height,
+                    width=request_inputs.width,
+                )
+            image = image.to(device=device, dtype=prompt_context.positive_connector_prompt_embeds.dtype)
+
+        num_channels_latents = self.transformer.config.in_channels
+        return self.prepare_latents(
+            image,
+            prompt_context.batch_size * request_inputs.num_videos_per_prompt,
+            num_channels_latents,
+            request_inputs.height,
+            request_inputs.width,
+            request_inputs.num_frames,
+            noise_scale,
+            torch.float32,
+            device,
+            request_inputs.generator,
+            request_inputs.latents,
+        )
+
+    def _make_video_audio_scheduler(
+        self,
+        audio_scheduler: Any,
+        latent_num_frames: int,
+        latent_height: int,
+        latent_width: int,
+    ) -> Any:
+        return _I2VVideoAudioScheduler(
+            self,
+            audio_scheduler,
+            latent_num_frames,
+            latent_height,
+            latent_width,
+        )
+
+    def _prepare_denoise_context_for_cfg(
+        self,
+        forward_ctx: _LTX23ForwardContext,
+        denoise_ctx: _LTX23DenoiseContext,
+    ) -> _LTX23DenoiseContext:
+        denoise_ctx = super()._prepare_denoise_context_for_cfg(forward_ctx, denoise_ctx)
+        if denoise_ctx.conditioning_mask is None:
+            raise ValueError("I2V denoising requires a conditioning mask.")
+        if self.do_classifier_free_guidance and not forward_ctx.cfg_parallel_ready:
+            denoise_ctx.conditioning_mask_for_model = torch.cat(
+                [denoise_ctx.conditioning_mask, denoise_ctx.conditioning_mask]
+            )
+        else:
+            denoise_ctx.conditioning_mask_for_model = denoise_ctx.conditioning_mask
+        return denoise_ctx
+
+    def _denoise_timestep_kwargs(
+        self,
+        ts: torch.Tensor,
+        forward_ctx: _LTX23ForwardContext,
+        denoise_ctx: _LTX23DenoiseContext,
+    ) -> dict[str, torch.Tensor]:
+        conditioning_mask = (
+            denoise_ctx.conditioning_mask if forward_ctx.cfg_parallel_ready else denoise_ctx.conditioning_mask_for_model
+        )
+        if conditioning_mask is None:
+            raise ValueError("I2V denoising requires a conditioning mask.")
+        return {
+            "timestep": ts.unsqueeze(-1) * (1 - conditioning_mask),
+            "audio_timestep": ts,
+            "sigma": ts,
+        }
+
+    def _step_denoised_latents(
+        self,
+        forward_ctx: _LTX23ForwardContext,
+        denoise_ctx: _LTX23DenoiseContext,
+        noise_pred_video: torch.Tensor,
+        noise_pred_audio: torch.Tensor,
+        t: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        latents, audio_latents = self.scheduler_step_maybe_with_cfg(
+            (noise_pred_video, noise_pred_audio),
+            (t, t),
+            (denoise_ctx.latents, denoise_ctx.audio_latents),
+            do_true_cfg=self.do_classifier_free_guidance,
+            per_request_scheduler=forward_ctx.video_audio_scheduler,
+        )
+        return self._synchronize_cfg_parallel_step_output(
+            (latents, audio_latents),
+            do_true_cfg=self.do_classifier_free_guidance,
+        )
+
+    @torch.no_grad()
+    def forward(
+        self,
+        req: DiffusionRequestBatch,
+        image: PIL.Image.Image | torch.Tensor | list[PIL.Image.Image | torch.Tensor] | None = None,
+        prompt: str | list[str] | None = None,
+        negative_prompt: str | list[str] | None = None,
+        height: int | None = None,
+        width: int | None = None,
+        num_frames: int | None = None,
+        frame_rate: float | None = None,
+        num_inference_steps: int | None = None,
+        sigmas: list[float] | None = None,
+        timesteps: list[int] | None = None,
+        guidance_scale: float = 4.0,
+        noise_scale: float = 0.0,
+        num_videos_per_prompt: int | None = 1,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+        audio_latents: torch.Tensor | None = None,
+        prompt_embeds: torch.Tensor | None = None,
+        negative_prompt_embeds: torch.Tensor | None = None,
+        prompt_attention_mask: torch.Tensor | None = None,
+        negative_prompt_attention_mask: torch.Tensor | None = None,
+        decode_timestep: float | list[float] = 0.0,
+        decode_noise_scale: float | list[float] | None = None,
+        output_type: str = "np",
+        return_dict: bool = True,
+        attention_kwargs: dict[str, Any] | None = None,
+        max_sequence_length: int | None = None,
+    ) -> list[DiffusionOutput]:
+        request_inputs = self._resolve_request_inputs(
+            req,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            frame_rate=frame_rate,
+            num_inference_steps=num_inference_steps,
+            timesteps=timesteps,
+            guidance_scale=guidance_scale,
+            num_videos_per_prompt=num_videos_per_prompt,
+            generator=generator,
+            latents=latents,
+            audio_latents=audio_latents,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            prompt_attention_mask=prompt_attention_mask,
+            negative_prompt_attention_mask=negative_prompt_attention_mask,
+            decode_timestep=decode_timestep,
+            decode_noise_scale=decode_noise_scale,
+            output_type=output_type,
+            max_sequence_length=max_sequence_length,
+        )
+        image = self._resolve_request_image(req, image, request_inputs)
+        return self._forward_impl(
+            req,
+            request_inputs,
+            noise_scale=noise_scale,
+            sigmas=sigmas,
+            timesteps=timesteps,
+            attention_kwargs=attention_kwargs,
+            image=image,
+        )
+
 
 __all__ = [
     "LTX23ImageToVideoPipeline",

vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py

@@ -446,7 +446,7 @@ class MingImagePipeline(ZImagePipeline):
             None if ref_latent is None else tuple(ref_latent.shape),
         )
         try:
-            outputs = super().forward(
+            outputs: DiffusionOutput = super().forward(
                 z_req,
                 prompt=None,
                 height=height,
@@ -462,13 +462,7 @@ class MingImagePipeline(ZImagePipeline):
         finally:
             set_forward_context_ref_latent(None)
 
-        output = outputs[0]
-        if hasattr(output, "output") and output.output is not None:
-            raw = output.output
-        elif hasattr(output, "images"):
-            raw = output.images
-        else:
-            raw = output
+        raw = outputs.output
         if not isinstance(raw, torch.Tensor):
             raise RuntimeError(f"ZImagePipeline returned non-tensor output: {type(raw).__name__}")
         if logger.isEnabledFor(logging.DEBUG):
@@ -478,7 +472,7 @@ class MingImagePipeline(ZImagePipeline):
                 raw.float().min().item(),
                 raw.float().max().item(),
             )
-        return DiffusionOutput(output=raw)
+        return outputs
 
 
 # ----------------------------------------------------------------------

vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py

@@ -18,7 +18,7 @@
 import inspect
 import json
 import os
-from collections.abc import Callable, Iterable
+from collections.abc import Iterable
 from typing import Any, ClassVar
 
 import numpy as np
@@ -551,117 +551,33 @@ class OvisImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMi
     def interrupt(self):
         return self._interrupt
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        guidance_scale: float = 5.0,
-        height: int | None = None,
-        width: int | None = None,
-        num_inference_steps: int = 50,
-        sigmas: list[float] | None = None,
-        num_images_per_prompt: int | None = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.FloatTensor | None = None,
-        prompt_embeds: torch.FloatTensor | None = None,
-        negative_prompt_embeds: torch.FloatTensor | None = None,
-        output_type: str | None = "pil",
-        return_dict: bool = True,
-        joint_attention_kwargs: dict[str, Any] | None = None,
-        callback_on_step_end: Callable[[int, int, dict], None] | None = None,
-        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
-        max_sequence_length: int = 256,
-    ) -> DiffusionOutput:
-        r"""
-        Function invoked when calling the pipeline for generation.
-
-        Args:
-            prompt (`str` or `list[str]`, *optional*):
-                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
-                instead.
-            negative_prompt (`str` or `list[str]`, *optional*):
-                The prompt or prompts not to guide the image generation. If not defined, one has to pass
-                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
-                not greater than `1`).
-            guidance_scale (`float`, *optional*, defaults to 1.0):
-                True classifier-free guidance (guidance scale) is enabled when `guidance_scale` > 1 and
-                `negative_prompt` is provided.
-            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
-                The height in pixels of the generated image. This is set to 1024 by default for the best results.
-            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
-                The width in pixels of the generated image. This is set to 1024 by default for the best results.
-            num_inference_steps (`int`, *optional*, defaults to 50):
-                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
-                expense of slower inference.
-            sigmas (`list[float]`, *optional*):
-                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
-                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
-                will be used.
-            num_images_per_prompt (`int`, *optional*, defaults to 1):
-                The number of images to generate per prompt.
-            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
-                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
-                to make generation deterministic.
-            latents (`torch.FloatTensor`, *optional*):
-                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
-                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
-                tensor will be generated by sampling using the supplied random `generator`.
-            prompt_embeds (`torch.FloatTensor`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
-                provided, text embeddings will be generated from `prompt` input argument.
-            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
-                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
-                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
-                argument.
-            output_type (`str`, *optional*, defaults to `"pil"`):
-                The output format of the generate image. Choose between
-                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`~pipelines.flux.FluxPipelineOutput`] instead of a plain tuple.
-            joint_attention_kwargs (`dict`, *optional*):
-                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
-                `self.processor` in
-                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
-            callback_on_step_end (`Callable`, *optional*):
-                A function that calls at the end of each denoising steps during the inference. The function is called
-                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
-                callback_kwargs: dict)`. `callback_kwargs` will include a list of all tensors as specified by
-                `callback_on_step_end_tensor_inputs`.
-            callback_on_step_end_tensor_inputs (`list`, *optional*):
-                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
-                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
-                `._callback_tensor_inputs` attribute of your pipeline class.
-            max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`.
-
-        Examples:
-
-        Returns:
-            [`~pipelines.ovis_image.OvisImagePipelineOutput`] or `tuple`:
-            [`~pipelines.ovis_image.OvisImagePipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When
-            returning a tuple, the first element is a list with the generated images.
-        """
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
         # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
         # TODO: May be some data formatting operations on the API side. Hack for now.
-        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts]
         if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
             negative_prompt = None
         elif req.prompts:
             negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
 
+        prompt_embeds = None
+        negative_prompt_embeds = None
+
         height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor
         width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor
-        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
-        sigmas = req.sampling_params.sigmas or sigmas
-        guidance_scale = (
-            req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale
-        )
-        generator = req.sampling_params.generator or generator
+        num_inference_steps = req.sampling_params.num_inference_steps or 50
+        sigmas = req.sampling_params.sigmas
+        guidance_scale = req.sampling_params.guidance_scale
+        generator = req.sampling_params.generator
         num_images_per_prompt = (
-            req.sampling_params.num_outputs_per_prompt
-            if req.sampling_params.num_outputs_per_prompt > 0
-            else num_images_per_prompt
+            req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1
         )
+        latents = req.sampling_params.latents
+        max_sequence_length = req.sampling_params.max_sequence_length or 256
+
+        joint_attention_kwargs: dict[str, Any] | None = None
+        output_type = req.sampling_params.output_type or "pil"
+        callback_on_step_end_tensor_inputs = ["latents"]
 
         # Steps:
         # 1. Check Inputs
@@ -691,12 +607,7 @@ class OvisImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMi
         self._interrupt = False
 
         # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
+        batch_size = len(prompt)
 
         do_classifier_free_guidance = guidance_scale > 1.0
         prompt_embeds, text_ids = self.encode_prompt(

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py

@@ -978,58 +978,34 @@ class QwenImagePipeline(
 
         height = state.sampling.height or self.default_sample_size * self.vae_scale_factor
         width = state.sampling.width or self.default_sample_size * self.vae_scale_factor
-        output_type = kwargs.get("output_type", "pil")
+        output_type = kwargs.get("output_type") or state.sampling.output_type or "pil"
 
         return self._decode_latents(state.latents, height, width, output_type)
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        true_cfg_scale: float = 4.0,
-        height: int | None = None,
-        width: int | None = None,
-        num_inference_steps: int = 50,
-        sigmas: list[float] | None = None,
-        guidance_scale: float = 1.0,
-        num_images_per_prompt: int = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.Tensor | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        prompt_embeds_mask: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds_mask: torch.Tensor | None = None,
-        output_type: str | None = "pil",
-        attention_kwargs: dict[str, Any] | None = None,
-        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
-        max_sequence_length: int = 1024,
-    ) -> list[DiffusionOutput]:
+    def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]:
         sampling_params_list = req.sampling_params_list
         common_sampling_params = sampling_params_list[0]
         extracted_prompt, negative_prompt = self._extract_prompts(req.prompts)
-        prompt = extracted_prompt or prompt
+        prompt = extracted_prompt
 
         height = common_sampling_params.height or self.default_sample_size * self.vae_scale_factor
         width = common_sampling_params.width or self.default_sample_size * self.vae_scale_factor
         height, width = normalize_min_aligned_size(height, width, self.vae_scale_factor * 2)
-        num_inference_steps = common_sampling_params.num_inference_steps or num_inference_steps
-        sigmas = common_sampling_params.sigmas or sigmas
-        max_sequence_length = common_sampling_params.max_sequence_length or max_sequence_length
+        num_inference_steps = common_sampling_params.num_inference_steps or 50
+        sigmas = common_sampling_params.sigmas
+        max_sequence_length = common_sampling_params.max_sequence_length or 1024
         num_images_per_prompt = (
-            common_sampling_params.num_outputs_per_prompt
-            if common_sampling_params.num_outputs_per_prompt > 0
-            else num_images_per_prompt
+            common_sampling_params.num_outputs_per_prompt if common_sampling_params.num_outputs_per_prompt > 0 else 1
         )
-        generator = req.collate_request_generators(num_images_per_prompt, generator)
-        latents = req.collate_request_tensors("latents", latents)
+        generator = req.collate_request_generators(num_images_per_prompt, None)
+        latents = req.collate_request_tensors("latents", None)
         prompt_fields = DiffusionRequestBatch.collate_prompt_field_map(
             req.prompts,
             {
-                "prompt_embeds": prompt_embeds,
-                "prompt_embeds_mask": prompt_embeds_mask,
-                "negative_prompt_embeds": negative_prompt_embeds,
-                "negative_prompt_embeds_mask": negative_prompt_embeds_mask,
+                "prompt_embeds": None,
+                "prompt_embeds_mask": None,
+                "negative_prompt_embeds": None,
+                "negative_prompt_embeds_mask": None,
             },
         )
         prompt_embeds = prompt_fields["prompt_embeds"]
@@ -1040,9 +1016,16 @@ class QwenImagePipeline(
             prompt = None
         if negative_prompt_embeds is not None:
             negative_prompt = None
-        true_cfg_scale = common_sampling_params.true_cfg_scale or true_cfg_scale
+        true_cfg_scale = common_sampling_params.true_cfg_scale or 4.0
         if common_sampling_params.guidance_scale_provided:
             guidance_scale = common_sampling_params.guidance_scale
+        else:
+            guidance_scale = 1.0
+
+        latents = req.collate_request_tensors("latents", None)
+        output_type = common_sampling_params.output_type or "pil"
+        attention_kwargs = None
+        callback_on_step_end_tensor_inputs = ["latents"]
 
         ctx = self._prepare_generation_context(
             prompt=prompt,
@@ -1087,6 +1070,7 @@ class QwenImagePipeline(
         )
 
         self._current_timestep = None
+
         result = self._decode_latents(latents, height, width, output_type)
         return split_diffusion_output_by_request(
             result,

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py

@@ -7,7 +7,7 @@ import logging
 import math
 import os
 from collections.abc import Iterable
-from typing import Any, ClassVar, cast
+from typing import ClassVar, cast
 
 import numpy as np
 import PIL.Image
@@ -677,30 +677,7 @@ class QwenImageEditPipeline(
     def interrupt(self):
         return self._interrupt
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        image: PIL.Image.Image | torch.Tensor | None = None,
-        true_cfg_scale: float = 4.0,
-        height: int | None = None,
-        width: int | None = None,
-        num_inference_steps: int = 50,
-        sigmas: list[float] | None = None,
-        guidance_scale: float = 1.0,
-        num_images_per_prompt: int = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.Tensor | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        prompt_embeds_mask: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds_mask: torch.Tensor | None = None,
-        output_type: str | None = "pil",
-        attention_kwargs: dict[str, Any] | None = None,
-        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
-        max_sequence_length: int = 1024,
-    ) -> DiffusionOutput:
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
         """Forward pass for image editing."""
         # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
         # TODO: May be some data formatting operations on the API side. Hack for now.
@@ -731,32 +708,28 @@ class QwenImageEditPipeline(
             height = req.sampling_params.height
             width = req.sampling_params.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
+            raise RuntimeError("Missing preprocess image that should have been created by the preprocess function.")
 
-            height, width = normalize_min_aligned_size(height, width, self.vae_scale_factor * 2)
-
-            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.sampling_params.num_inference_steps or num_inference_steps
-        sigmas = req.sampling_params.sigmas or sigmas
-        max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length
-        generator = req.sampling_params.generator or generator
-        true_cfg_scale = req.sampling_params.true_cfg_scale or true_cfg_scale
+        num_inference_steps = req.sampling_params.num_inference_steps or 50
+        sigmas = req.sampling_params.sigmas
+        max_sequence_length = req.sampling_params.max_sequence_length or 1024
+        generator = req.sampling_params.generator
+        true_cfg_scale = req.sampling_params.true_cfg_scale or 4.0
         if req.sampling_params.guidance_scale_provided:
             guidance_scale = req.sampling_params.guidance_scale
+        else:
+            guidance_scale = 1.0
         num_images_per_prompt = (
-            req.sampling_params.num_outputs_per_prompt
-            if req.sampling_params.num_outputs_per_prompt > 0
-            else num_images_per_prompt
+            req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1
         )
+        latents = req.sampling_params.latents
+        prompt_embeds = None
+        negative_prompt_embeds = None
+        prompt_embeds_mask = None
+        negative_prompt_embeds_mask = None
+        output_type = req.sampling_params.output_type or "pil"
+        attention_kwargs = None
+        callback_on_step_end_tensor_inputs = ["latents"]
 
         # 1. check inputs
         # 2. encode prompts
@@ -783,16 +756,9 @@ class QwenImageEditPipeline(
         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]
+        batch_size = 1
 
-        has_neg_prompt = negative_prompt is not None or (
-            negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None
-        )
+        has_neg_prompt = negative_prompt is not None
 
         do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
         self.check_cfg_parallel_validity(true_cfg_scale, has_neg_prompt)

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py

@@ -5,7 +5,7 @@ import json
 import logging
 import os
 from collections.abc import Iterable
-from typing import Any, ClassVar, cast
+from typing import ClassVar, cast
 
 import numpy as np
 import PIL.Image
@@ -623,30 +623,7 @@ class QwenImageEditPlusPipeline(
     def interrupt(self):
         return self._interrupt
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        image: PIL.Image.Image | list[PIL.Image.Image] | torch.Tensor | None = None,
-        true_cfg_scale: float = 4.0,
-        height: int | None = None,
-        width: int | None = None,
-        num_inference_steps: int = 50,
-        sigmas: list[float] | None = None,
-        guidance_scale: float = 1.0,
-        num_images_per_prompt: int = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.Tensor | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        prompt_embeds_mask: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds_mask: torch.Tensor | None = None,
-        output_type: str | None = "pil",
-        attention_kwargs: dict[str, Any] | None = None,
-        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
-        max_sequence_length: int = 1024,
-    ) -> DiffusionOutput:
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
         """Forward pass for image editing with support for multiple images."""
         # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
         # TODO: May be some data formatting operations on the API side. Hack for now.
@@ -674,55 +651,35 @@ class QwenImageEditPlusPipeline(
         ):
             condition_images = additional_information.get("condition_images")
             vae_images = additional_information.get("vae_images")
-            condition_image_sizes = additional_information.get("condition_image_sizes")
             vae_image_sizes = additional_information.get("vae_image_sizes")
             calculated_height = additional_information.get("calculated_height")
             calculated_width = additional_information.get("calculated_width")
-            height = req.sampling_params.height
-            width = req.sampling_params.width
+            height = req.sampling_params.height or calculated_height
+            width = req.sampling_params.width or calculated_width
+            image = None
         else:
-            # fallback to run pre-processing in pipeline (debug only)
-            if image is None:
-                raise ValueError("Image is required for QwenImageEditPlusPipeline")
+            raise RuntimeError("Missing preprocess images that should have been created by the preprocess function.")
 
-            if not isinstance(image, list):
-                image = [image]
-
-            image_size = image[0].size
-            calculated_width, calculated_height = calculate_dimensions(VAE_IMAGE_SIZE, image_size[0] / image_size[1])
-            height = height or calculated_height
-            width = width or calculated_width
-
-            height, width = normalize_min_aligned_size(height, width, self.vae_scale_factor * 2)
-
-            condition_images = []
-            vae_images = []
-            condition_image_sizes = []
-            vae_image_sizes = []
-
-            for img in image:
-                image_width, image_height = img.size
-                condition_width, condition_height = calculate_dimensions(
-                    CONDITION_IMAGE_SIZE, image_width / image_height
-                )
-                vae_width, vae_height = calculate_dimensions(VAE_IMAGE_SIZE, image_width / image_height)
-                condition_image_sizes.append((condition_width, condition_height))
-                vae_image_sizes.append((vae_width, vae_height))
-                condition_images.append(self.image_processor.resize(img, condition_height, condition_width))
-                vae_images.append(self.image_processor.preprocess(img, vae_height, vae_width).unsqueeze(2))
-
-        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
-        sigmas = req.sampling_params.sigmas or sigmas
-        max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length
-        generator = req.sampling_params.generator or generator
-        true_cfg_scale = req.sampling_params.true_cfg_scale or true_cfg_scale
+        num_inference_steps = req.sampling_params.num_inference_steps or 50
+        sigmas = req.sampling_params.sigmas
+        max_sequence_length = req.sampling_params.max_sequence_length or 1024
+        generator = req.sampling_params.generator
+        true_cfg_scale = req.sampling_params.true_cfg_scale or 4.0
         if req.sampling_params.guidance_scale_provided:
             guidance_scale = req.sampling_params.guidance_scale
+        else:
+            guidance_scale = 1.0
         num_images_per_prompt = (
-            req.sampling_params.num_outputs_per_prompt
-            if req.sampling_params.num_outputs_per_prompt > 0
-            else num_images_per_prompt
+            req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1
         )
+        latents = req.sampling_params.latents
+        prompt_embeds = None
+        negative_prompt_embeds = None
+        prompt_embeds_mask = None
+        negative_prompt_embeds_mask = None
+        output_type = req.sampling_params.output_type or "pil"
+        attention_kwargs = None
+        callback_on_step_end_tensor_inputs = ["latents"]
 
         # 1. check inputs
         # 2. encode prompts
@@ -749,16 +706,9 @@ class QwenImageEditPlusPipeline(
         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]
+        batch_size = 1
 
-        has_neg_prompt = negative_prompt is not None or (
-            negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None
-        )
+        has_neg_prompt = negative_prompt is not None
 
         do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
         self.check_cfg_parallel_validity(true_cfg_scale, has_neg_prompt)

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py

@@ -7,7 +7,7 @@ import logging
 import math
 import os
 from collections.abc import Iterable
-from typing import Any, ClassVar, cast
+from typing import ClassVar, cast
 
 import numpy as np
 import PIL.Image
@@ -648,31 +648,7 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
     def interrupt(self):
         return self._interrupt
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        image: PIL.Image.Image | torch.Tensor | None = None,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        true_cfg_scale: float = 4.0,
-        layers: int | None = 4,
-        num_inference_steps: int = 50,
-        sigmas: list[float] | None = None,
-        guidance_scale: float | None = None,
-        num_images_per_prompt: int = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.Tensor | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        prompt_embeds_mask: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds_mask: torch.Tensor | None = None,
-        output_type: str | None = "pil",
-        attention_kwargs: dict[str, Any] | None = None,
-        max_sequence_length: int = 1024,
-        resolution: int = 640,
-        cfg_normalize: bool = False,
-        use_en_prompt: bool = False,
-    ) -> DiffusionOutput:
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
         """Forward pass for image layered."""
 
         # 1. Get preprocessed image from request (pre-processing is done in DiffusionEngine)
@@ -688,26 +664,27 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
         prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
         negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt")
 
-        layers = req.sampling_params.layers if req.sampling_params.layers is not None else layers
-        resolution = req.sampling_params.resolution if req.sampling_params.resolution is not None else resolution
-        max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length
-        cfg_normalize = (
-            req.sampling_params.cfg_normalize if req.sampling_params.cfg_normalize is not None else cfg_normalize
-        )
-        use_en_prompt = (
-            req.sampling_params.use_en_prompt if req.sampling_params.use_en_prompt is not None else use_en_prompt
-        )
-        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
-        sigmas = req.sampling_params.sigmas or sigmas
-        generator = req.sampling_params.generator or generator
-        true_cfg_scale = req.sampling_params.true_cfg_scale or true_cfg_scale
+        layers = req.sampling_params.layers
+        max_sequence_length = req.sampling_params.max_sequence_length or 1024
+        cfg_normalize = req.sampling_params.cfg_normalize
+        use_en_prompt = req.sampling_params.use_en_prompt
+        num_inference_steps = req.sampling_params.num_inference_steps or 50
+        sigmas = req.sampling_params.sigmas
+        generator = req.sampling_params.generator
+        true_cfg_scale = req.sampling_params.true_cfg_scale or 4.0
+        guidance_scale = None
         if req.sampling_params.guidance_scale_provided:
             guidance_scale = req.sampling_params.guidance_scale
         num_images_per_prompt = (
-            req.sampling_params.num_outputs_per_prompt
-            if req.sampling_params.num_outputs_per_prompt > 0
-            else num_images_per_prompt
+            req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1
         )
+        latents = req.sampling_params.latents
+        prompt_embeds = None
+        negative_prompt_embeds = None
+        prompt_embeds_mask = None
+        negative_prompt_embeds_mask = None
+        output_type = req.sampling_params.output_type or "pil"
+        attention_kwargs = None
 
         if not isinstance(first_prompt, str) and "preprocessed_image" in (
             additional_information := first_prompt.get("additional_information", {})
@@ -720,25 +697,7 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
             height = req.sampling_params.height
             width = req.sampling_params.width
         else:
-            # fallback to run pre-processing in pipeline (debug only)
-            if isinstance(image, PIL.Image.Image) and image.mode != "RGBA":
-                image = image.convert("RGBA")
-            image_size = image[0].size if isinstance(image, list) else image.size
-            assert resolution in [640, 1024], f"resolution must be either 640 or 1024, but got {resolution}"
-            calculated_width, calculated_height = calculate_dimensions(
-                resolution * resolution, image_size[0] / image_size[1]
-            )
-            height = calculated_height
-            width = calculated_width
-
-            height, width = normalize_min_aligned_size(height, width, self.vae_scale_factor * 2)
-
-            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)
-                image = image.to(dtype=self.text_encoder.dtype)
+            raise RuntimeError("Missing preprocess image that should have been created by the preprocess function.")
 
         # 2. check inputs
         self.check_inputs(
@@ -761,16 +720,9 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
         # 3. encode prompot & negative prompt
         if prompt is None or prompt == "" or prompt == " ":
             prompt = self.get_image_caption(prompt_image, use_en_prompt=use_en_prompt, device=self.device)
-        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]
+        batch_size = 1
 
-        has_neg_prompt = negative_prompt is not None or (
-            negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None
-        )
+        has_neg_prompt = negative_prompt is not None
 
         if true_cfg_scale > 1 and not has_neg_prompt:
             logger.warning(

vllm_omni/diffusion/models/sd3/pipeline_sd3.py

@@ -628,57 +628,39 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro
 
         return latents
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | list[str] = "",
-        prompt_2: str | list[str] = "",
-        prompt_3: str | list[str] = "",
-        negative_prompt: str | list[str] = "",
-        negative_prompt_2: str | list[str] = "",
-        negative_prompt_3: str | list[str] = "",
-        height: int | None = None,
-        width: int | None = None,
-        num_inference_steps: int = 28,
-        sigmas: list[float] | None = None,
-        num_images_per_prompt: int = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.Tensor | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        pooled_prompt_embeds: torch.Tensor | None = None,
-        negative_pooled_prompt_embeds: torch.Tensor | None = None,
-        max_sequence_length: int = 256,
-    ) -> list[DiffusionOutput]:
+    def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]:
         # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
         # TODO: May be some data formatting operations on the API side. Hack for now.
         sampling_params_list = req.sampling_params_list
         common_sampling_params = sampling_params_list[0]
-        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
-        negative_prompt = [
-            "" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts
-        ] or negative_prompt
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts]
+        negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
+        prompt_2 = ""
+        prompt_3 = ""
+        negative_prompt_2 = ""
+        negative_prompt_3 = ""
+        prompt_embeds = None
+        negative_prompt_embeds = None
 
         height = common_sampling_params.height or self.default_sample_size * self.vae_scale_factor
         width = common_sampling_params.width or self.default_sample_size * self.vae_scale_factor
-        sigmas = common_sampling_params.sigmas or sigmas
-        max_sequence_length = common_sampling_params.max_sequence_length or max_sequence_length
-        num_inference_steps = common_sampling_params.num_inference_steps or num_inference_steps
+        sigmas = common_sampling_params.sigmas
+        max_sequence_length = common_sampling_params.max_sequence_length or 256
+        num_inference_steps = common_sampling_params.num_inference_steps or 28
         num_images_per_prompt = (
-            common_sampling_params.num_outputs_per_prompt
-            if common_sampling_params.num_outputs_per_prompt > 0
-            else num_images_per_prompt
+            common_sampling_params.num_outputs_per_prompt if common_sampling_params.num_outputs_per_prompt > 0 else 1
         )
-        if generator is None:
-            generator = req.collate_request_generators(num_images_per_prompt, generator)
-        latents = req.collate_request_tensors("latents", latents)
+        generator = req.collate_request_generators(num_images_per_prompt, None)
+        latents = req.collate_request_tensors("latents", None)
+        output_type = common_sampling_params.output_type or "pil"
+
         prompt_fields = DiffusionRequestBatch.collate_prompt_field_map(
             req.prompts,
             {
                 "prompt_embeds": prompt_embeds,
                 "negative_prompt_embeds": negative_prompt_embeds,
-                "pooled_prompt_embeds": pooled_prompt_embeds,
-                "negative_pooled_prompt_embeds": negative_pooled_prompt_embeds,
+                "pooled_prompt_embeds": None,
+                "negative_pooled_prompt_embeds": None,
             },
         )
         prompt_embeds = prompt_fields["prompt_embeds"]
@@ -717,9 +699,7 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro
         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):
+        if prompt is not None:
             batch_size = len(prompt)
         else:
             batch_size = prompt_embeds.shape[0]
@@ -771,7 +751,7 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro
         )
 
         self._current_timestep = None
-        if self.output_type == "latent":
+        if output_type == "latent":
             image = latents
         else:
             # Ensure the latents are the same dtype as the VAE for decode

vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py

@@ -385,37 +385,21 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, SupportsComponentDiscov
         latents = latents * self.scheduler.init_noise_sigma
         return latents
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        audio_end_in_s: float | None = None,
-        audio_start_in_s: float = 0.0,
-        guidance_scale: float = 7.0,
-        num_waveforms_per_prompt: int = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.Tensor | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        output_type: str = "np",
-    ) -> DiffusionOutput:
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
         """
         Generate audio from text prompt.
 
         Args:
-            req: OmniDiffusionRequest containing generation parameters
-            prompt: Text prompt for audio generation
-            negative_prompt: Negative prompt for CFG
-            audio_end_in_s: Audio end time in seconds (max ~47s for stable-audio-open-1.0)
-            audio_start_in_s: Audio start time in seconds
-            guidance_scale: CFG scale
-            num_waveforms_per_prompt: Number of audio outputs per prompt
-            generator: Random generator for reproducibility
-            latents: Pre-generated latents
-            prompt_embeds: Pre-computed prompt embeddings
-            negative_prompt_embeds: Pre-computed negative prompt embeddings
-            output_type: Output format ("np", "pt", or "latent")
+            req: OmniDiffusionRequest containing generation parameters.
+                The `req.sampling_params.extra_args` can include the following keys:
+                - audio_start_in_s (`float`, *optional*, defaults to 0.0):
+                    Start time of the audio in seconds.
+                - audio_end_in_s (`float`, *optional*):
+                    End time of the audio in seconds.
+                - num_waveforms_per_prompt (`int`, *optional*, defaults to 1):
+                    Number of audio outputs per prompt.
+                - output_type (`str`, *optional*, defaults to "np"):
+                    Output format ("np", "pt", or "latent").
 
         Returns:
             DiffusionOutput containing generated audio
@@ -423,28 +407,33 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, SupportsComponentDiscov
         # Extract from request
         # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
         # TODO: May be some data formatting operations on the API side. Hack for now.
-        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts]
         if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
             negative_prompt = None
         elif req.prompts:
             negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
-        num_inference_steps = req.sampling_params.num_inference_steps
-        if num_inference_steps is None:
-            num_inference_steps = 100  # Default steps
-        num_inference_steps = req.sampling_params.num_inference_steps
-        if num_inference_steps is None:
-            num_inference_steps = 50  # Default steps
+
+        prompt_embeds = None
+        negative_prompt_embeds = None
+
+        num_inference_steps = req.sampling_params.num_inference_steps or 50
         if req.sampling_params.guidance_scale_provided:
             guidance_scale = req.sampling_params.guidance_scale
+        else:
+            guidance_scale = 7.0
 
-        if generator is None:
-            generator = req.sampling_params.generator
+        generator = req.sampling_params.generator
         if generator is None and req.sampling_params.seed is not None:
             generator = torch.Generator(device=self.device).manual_seed(req.sampling_params.seed)
+        latents = req.sampling_params.latents
 
         # Get audio duration from request extra params or defaults
-        audio_start_in_s = req.sampling_params.extra_args.get("audio_start_in_s", audio_start_in_s)
-        audio_end_in_s = req.sampling_params.extra_args.get("audio_end_in_s", audio_end_in_s)
+        audio_start_in_s: float = req.sampling_params.extra_args.get("audio_start_in_s", 0.0)
+        audio_end_in_s: float | None = req.sampling_params.extra_args.get("audio_end_in_s", None)
+        num_waveforms_per_prompt: int = req.sampling_params.extra_args.get(
+            "num_waveforms_per_prompt", req.sampling_params.num_outputs_per_prompt or 1
+        )
+        output_type = req.sampling_params.output_type or "np"
 
         # Calculate audio length
         downsample_ratio = self.vae.hop_length
@@ -474,12 +463,7 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, SupportsComponentDiscov
         )
 
         # Determine batch size
-        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]
+        batch_size = len(prompt)
 
         device = self.device
         do_classifier_free_guidance = guidance_scale > 1.0
@@ -500,7 +484,7 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, SupportsComponentDiscov
             audio_start_in_s,
             audio_end_in_s,
             device,
-            do_classifier_free_guidance and (negative_prompt is not None or negative_prompt_embeds is not None),
+            do_classifier_free_guidance and negative_prompt is not None,
             batch_size,
         )
 
@@ -515,7 +499,7 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, SupportsComponentDiscov
         )
 
         # Handle CFG without negative prompt
-        if do_classifier_free_guidance and negative_prompt_embeds is None and negative_prompt is None:
+        if do_classifier_free_guidance and negative_prompt is None:
             negative_text_audio_duration_embeds = torch.zeros_like(text_audio_duration_embeds)
             text_audio_duration_embeds = torch.cat(
                 [negative_text_audio_duration_embeds, text_audio_duration_embeds],

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py

@@ -540,38 +540,27 @@ class Wan22Pipeline(
 
         return latents
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | None = None,
-        negative_prompt: str | None = None,
-        height: int = 480,
-        width: int = 832,
-        num_inference_steps: int = 40,
-        guidance_scale: float | tuple[float, float] = 4.0,
-        frame_num: int = 81,
-        output_type: str | None = "np",
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        attention_kwargs: dict | None = None,
-        **kwargs,
-    ) -> DiffusionOutput:
-        # Get parameters from request or arguments
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
+        prompt: str | None = None
+        negative_prompt: str | None = None
+        prompt_embeds: torch.Tensor | None = None
+        negative_prompt_embeds: torch.Tensor | None = None
         if len(req.prompts) > 1:
             raise ValueError(
                 """This model only supports a single prompt, not a batched request.""",
                 """Please pass in a single prompt object or string, or a single-item list.""",
             )
-        if len(req.prompts) == 1:  # If req.prompt is empty, default to prompt & neg_prompt in param list
-            prompt = req.prompts[0] if isinstance(req.prompts[0], str) else req.prompts[0].get("prompt")
-            negative_prompt = None if isinstance(req.prompts[0], str) else req.prompts[0].get("negative_prompt")
-        if prompt is None and prompt_embeds is None:
-            raise ValueError("Prompt or prompt_embeds is required for Wan2.2 generation.")
+        if len(req.prompts) == 1:
+            first_prompt = req.prompts[0]
+            prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
+            negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt")
 
-        height = req.sampling_params.height or height
-        width = req.sampling_params.width or width
-        num_frames = req.sampling_params.num_frames if req.sampling_params.num_frames else frame_num
+        if not prompt:
+            raise ValueError("Prompt is required for Wan2.2 generation.")
+
+        height = req.sampling_params.height or 480
+        width = req.sampling_params.width or 832
+        num_frames = req.sampling_params.num_frames or 81
 
         # Ensure dimensions are compatible with VAE and patch size
         # For expand_timesteps mode, we need latent dims to be even (divisible by patch_size)
@@ -579,11 +568,16 @@ class Wan22Pipeline(
         mod_value = self.vae_scale_factor_spatial * patch_size[1]  # 16*2=32 for TI2V, 8*2=16 for I2V
         height = (height // mod_value) * mod_value
         width = (width // mod_value) * mod_value
-        num_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        num_steps = 40 if req.sampling_params.num_inference_steps is None else req.sampling_params.num_inference_steps
 
         # Respect per-request guidance_scale when explicitly provided.
         if req.sampling_params.guidance_scale_provided:
             guidance_scale = req.sampling_params.guidance_scale
+        else:
+            guidance_scale = 4.0
+
+        output_type = req.sampling_params.output_type or "np"
+        attention_kwargs: dict | None = None
 
         guidance_low = guidance_scale if isinstance(guidance_scale, (int, float)) else guidance_scale[0]
         guidance_high = (
@@ -634,8 +628,7 @@ class Wan22Pipeline(
             dtype = self.text_encoder.dtype
 
         # Seed / generator
-        if generator is None:
-            generator = req.sampling_params.generator
+        generator = req.sampling_params.generator
         if generator is None and req.sampling_params.seed is not None:
             generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed)
 
@@ -644,24 +637,16 @@ class Wan22Pipeline(
             current_omni_platform.synchronize()
             _t_pipeline_start = time.perf_counter()
             _t_text_enc_start = _t_pipeline_start
-        if prompt_embeds is None:
-            prompt_embeds, negative_prompt_embeds = self.encode_prompt(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0,
-                num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1,
-                max_sequence_length=req.sampling_params.max_sequence_length or 512,
-                device=device,
-                dtype=dtype,
-            )
-        else:
-            prompt_embeds = prompt_embeds.to(device=device, dtype=dtype)
-            if negative_prompt_embeds is not None:
-                negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype)
-            elif guidance_low > 1.0 or guidance_high > 1.0:
-                raise ValueError(
-                    "negative_prompt_embeds must be provided when prompt_embeds are given and guidance > 1."
-                )
+        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0,
+            num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1,
+            max_sequence_length=req.sampling_params.max_sequence_length or 512,
+            device=device,
+            dtype=dtype,
+        )
+
         if DEBUG_PERF:
             current_omni_platform.synchronize()
             _t_text_enc_ms = (time.perf_counter() - _t_text_enc_start) * 1000

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py

@@ -429,66 +429,59 @@ class Wan22I2VPipeline(
         quant_config = getattr(self.od_config, "quantization_config", None)
         return create_transformer_from_config(config, quant_config=quant_config)
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | None = None,
-        negative_prompt: str | None = None,
-        image: PIL.Image.Image | torch.Tensor | None = None,
-        height: int = 480,
-        width: int = 832,
-        num_inference_steps: int = 40,
-        guidance_scale: float | tuple[float, float] = 5.0,
-        frame_num: int = 81,
-        output_type: str | None = "np",
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        image_embeds: torch.Tensor | None = None,
-        last_image: PIL.Image.Image | torch.Tensor | None = None,
-        attention_kwargs: dict | None = None,
-        **kwargs,
-    ) -> DiffusionOutput:
-        # Get parameters from request or arguments
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
+        prompt: str | None = None
+        negative_prompt: str | None = None
+        prompt_embeds: torch.Tensor | None = None
+        negative_prompt_embeds: torch.Tensor | None = None
+        image_embeds: torch.Tensor | None = None
         if len(req.prompts) > 1:
             raise ValueError(
                 """This model only supports a single prompt, not a batched request.""",
                 """Please pass in a single prompt object or string, or a single-item list.""",
             )
-        if len(req.prompts) == 1:  # If req.prompt is empty, default to prompt & neg_prompt in param list
-            prompt = req.prompts[0] if isinstance(req.prompts[0], str) else req.prompts[0].get("prompt")
-            negative_prompt = None if isinstance(req.prompts[0], str) else req.prompts[0].get("negative_prompt")
-        if prompt is None and prompt_embeds is None:
-            raise ValueError("Prompt or prompt_embeds is required for Wan2.2 generation.")
+        if len(req.prompts) == 1:
+            first_prompt = req.prompts[0]
+            prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
+            negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt")
+
+        if not prompt:
+            raise ValueError("Prompt is required for Wan2.2 generation.")
 
         # Get image from request
-        if image is None:
-            multi_modal_data = (
-                req.prompts[0].get("multi_modal_data", {}) if not isinstance(req.prompts[0], str) else None
-            )
-            raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None
-            if raw_image is None:
-                raise ValueError("Image is required for I2V generation.")
-            if isinstance(raw_image, list):
-                if len(raw_image) > 1:
-                    logger.warning(
-                        """Received a list of image. Only a single image is supported by this model."""
-                        """Taking only the first image for now."""
-                    )
-                raw_image = raw_image[0]
-            if isinstance(raw_image, str):
-                image = PIL.Image.open(raw_image)
-            else:
-                image = cast(PIL.Image.Image | torch.Tensor, raw_image)
+        multi_modal_data = req.prompts[0].get("multi_modal_data", {}) if not isinstance(req.prompts[0], str) else None
+        raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None
+        if raw_image is None:
+            raise ValueError("Image is required for I2V generation.")
+        if isinstance(raw_image, list):
+            if len(raw_image) > 1:
+                logger.warning(
+                    """Received a list of image. Only a single image is supported by this model."""
+                    """Taking only the first image for now."""
+                )
+            raw_image = raw_image[0]
+        if isinstance(raw_image, str):
+            image = PIL.Image.open(raw_image)
+        else:
+            image = cast(PIL.Image.Image | torch.Tensor, raw_image)
 
-        height = req.sampling_params.height or height
-        width = req.sampling_params.width or width
-        num_frames = req.sampling_params.num_frames or frame_num
-        num_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        last_image: PIL.Image.Image | torch.Tensor | None = None
+        if multi_modal_data is not None:
+            last_image = multi_modal_data.get("last_image", None)
+
+        height = req.sampling_params.height or 480
+        width = req.sampling_params.width or 832
+        num_frames = req.sampling_params.num_frames or 81
+        num_steps = 40 if req.sampling_params.num_inference_steps is None else req.sampling_params.num_inference_steps
 
         # Respect per-request guidance_scale when explicitly provided.
         if req.sampling_params.guidance_scale_provided:
             guidance_scale = req.sampling_params.guidance_scale
+        else:
+            guidance_scale = 5.0
+
+        output_type = req.sampling_params.output_type or "np"
+        attention_kwargs: dict | None = None
 
         # Handle guidance scales
         guidance_low = guidance_scale if isinstance(guidance_scale, (int, float)) else guidance_scale[0]
@@ -533,8 +526,7 @@ class Wan22I2VPipeline(
         dtype = self.transformer.dtype
 
         # Generator setup
-        if generator is None:
-            generator = req.sampling_params.generator
+        generator = req.sampling_params.generator
         if generator is None and req.sampling_params.seed is not None:
             generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed)
 
@@ -544,20 +536,15 @@ class Wan22I2VPipeline(
             _t_pipeline_start = time.perf_counter()
             _t_text_enc_start = _t_pipeline_start
 
-        if prompt_embeds is None:
-            prompt_embeds, negative_prompt_embeds = self.encode_prompt(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0,
-                num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1,
-                max_sequence_length=req.sampling_params.max_sequence_length or 512,
-                device=device,
-                dtype=dtype,
-            )
-        else:
-            prompt_embeds = prompt_embeds.to(device=device, dtype=dtype)
-            if negative_prompt_embeds is not None:
-                negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype)
+        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0,
+            num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1,
+            max_sequence_length=req.sampling_params.max_sequence_length or 512,
+            device=device,
+            dtype=dtype,
+        )
 
         if DEBUG_PERF:
             current_omni_platform.synchronize()

vllm_omni/diffusion/models/z_image/pipeline_z_image.py

@@ -24,7 +24,7 @@ from typing import Any, ClassVar
 import PIL.Image
 import torch
 import torch.nn as nn
-from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
+from diffusers.image_processor import VaeImageProcessor
 from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
 from diffusers.utils import logging
 from diffusers.utils.torch_utils import randn_tensor
@@ -410,124 +410,21 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsComponen
     def interrupt(self):
         return self._interrupt
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | list[str] | None = None,
-        image: PipelineImageInput = None,
-        strength: float = 0.6,
-        height: int = 1024,
-        width: int = 1024,
-        num_inference_steps: int = 50,
-        sigmas: list[float] | None = None,
-        guidance_scale: float = 5.0,
-        cfg_normalization: bool = False,
-        cfg_truncation: float = 1.0,
-        negative_prompt: str | list[str] | None = None,
-        num_images_per_prompt: int = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.FloatTensor | None = None,
-        prompt_embeds: list[torch.FloatTensor] | None = None,
-        negative_prompt_embeds: list[torch.FloatTensor] | None = None,
-        output_type: str | None = "pil",
-        return_dict: bool = True,
-        joint_attention_kwargs: dict[str, Any] | None = None,
-        callback_on_step_end: Callable[[int, int, dict], None] | None = None,
-        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
-        max_sequence_length: int = 512,
-    ) -> DiffusionOutput:
-        r"""
-        Function invoked when calling the pipeline for generation.
-
-        Args:
-            prompt (`str` or `list[str]`, *optional*):
-                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
-                instead.
-            image (`PipelineImageInput`, *optional*):
-                The image to use for img2img generation. If provided, the pipeline
-                will perform img2img instead of text-to-image.
-            strength (`float`, *optional*, defaults to 0.6):
-                Indicates extent to transform the reference `image`. Must be between 0 and 1.
-            height (`int`, *optional*, defaults to 1024):
-                The height in pixels of the generated image.
-            width (`int`, *optional*, defaults to 1024):
-                The width in pixels of the generated image.
-            num_inference_steps (`int`, *optional*, defaults to 50):
-                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
-                expense of slower inference.
-            sigmas (`list[float]`, *optional*):
-                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
-                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
-                will be used.
-            guidance_scale (`float`, *optional*, defaults to 5.0):
-                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
-                `guidance_scale` is defined as `w` of equation 2. of [Imagen
-                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
-                0`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
-                usually at the expense of lower image quality.
-            cfg_normalization (`bool`, *optional*, defaults to False):
-                Whether to apply configuration normalization.
-            cfg_truncation (`float`, *optional*, defaults to 1.0):
-                The truncation value for configuration.
-            negative_prompt (`str` or `list[str]`, *optional*):
-                The prompt or prompts not to guide the image generation. If not defined, one has to pass
-                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
-                less than or equal to `0`).
-            num_images_per_prompt (`int`, *optional*, defaults to 1):
-                The number of images to generate per prompt.
-            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
-                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
-                to make generation deterministic.
-            latents (`torch.FloatTensor`, *optional*):
-                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
-                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
-                tensor will be generated by sampling using the supplied random `generator`.
-            prompt_embeds (`list[torch.FloatTensor]`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
-                provided, text embeddings will be generated from `prompt` input argument.
-            negative_prompt_embeds (`list[torch.FloatTensor]`, *optional*):
-                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
-                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
-                argument.
-            output_type (`str`, *optional*, defaults to `"pil"`):
-                The output format of the generate image. Choose between
-                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`~pipelines.stable_diffusion.ZImagePipelineOutput`] instead of a plain
-                tuple.
-            joint_attention_kwargs (`dict`, *optional*):
-                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
-                `self.processor` in
-                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
-            callback_on_step_end (`Callable`, *optional*):
-                A function that calls at the end of each denoising steps during the inference. The function is called
-                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
-                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
-                `callback_on_step_end_tensor_inputs`.
-            callback_on_step_end_tensor_inputs (`list`, *optional*):
-                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
-                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
-                `._callback_tensor_inputs` attribute of your pipeline class.
-            max_sequence_length (`int`, *optional*, defaults to 512):
-                Maximum sequence length to use with the `prompt`.
-
-        Examples:
-
-        Returns:
-            [`~pipelines.z_image.ZImagePipelineOutput`] or `tuple`: [`~pipelines.z_image.ZImagePipelineOutput`] if
-            `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the
-            generated images.
-        """
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
         # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
         # TODO: May be some data formatting operations on the API side. Hack for now.
-        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts]
+
         if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
             negative_prompt = None
         elif req.prompts:
             negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
 
-        # Handle img2img: extract image from request
-        if image is None and req.prompts:
+        prompt_embeds = None
+        negative_prompt_embeds = None
+
+        image = None
+        if req.prompts:
             if len(req.prompts) > 1:
                 logger.warning(
                     "This model only supports a single prompt for img2img, not a batched request. "
@@ -542,10 +439,8 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsComponen
                     else:
                         image = PIL.Image.open(raw_image) if isinstance(raw_image, str) else raw_image
 
-        # strength is currently only applicable for Z-Image I2I; other pipelines ignore this parameter
         explicit_strength = req.sampling_params.strength is not None
-        if explicit_strength:
-            strength = req.sampling_params.strength
+        strength = req.sampling_params.strength if explicit_strength else 0.6
         if explicit_strength and image is None:
             logger.warning(
                 "strength parameter (%.2f) is only applicable for image-to-image (I2I) generation. "
@@ -556,20 +451,24 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsComponen
         if image is not None and strength is not None and (strength < 0 or strength > 1):
             raise ValueError(f"The value of strength should be in [0.0, 1.0] but is {strength}")
 
-        height = req.sampling_params.height or height
-        width = req.sampling_params.width or width
-        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        height = req.sampling_params.height or 1024
+        width = req.sampling_params.width or 1024
+        num_inference_steps = req.sampling_params.num_inference_steps or 50
         generator = req.sampling_params.generator
-        sigmas = req.sampling_params.sigmas or sigmas
-        max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length
-        guidance_scale = (
-            req.sampling_params.guidance_scale if req.sampling_params.guidance_rescale is not None else guidance_scale
-        )
+        sigmas = req.sampling_params.sigmas
+        max_sequence_length = req.sampling_params.max_sequence_length or 512
+        guidance_scale = req.sampling_params.guidance_scale
         num_images_per_prompt = (
-            req.sampling_params.num_outputs_per_prompt
-            if req.sampling_params.num_outputs_per_prompt > 0
-            else num_images_per_prompt
+            req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1
         )
+        latents = req.sampling_params.latents
+
+        cfg_normalization = req.sampling_params.cfg_normalize
+        cfg_truncation = req.sampling_params.extra_args.get("cfg_truncation", 1.0)
+        joint_attention_kwargs: dict[str, Any] | None = None
+        callback_on_step_end: Callable[[int, int, dict], None] | None = None
+        callback_on_step_end_tensor_inputs = ["latents"]
+        output_type = req.sampling_params.output_type or "pil"
 
         vae_scale = self.vae_scale_factor * 2
         if height % vae_scale != 0:
@@ -591,33 +490,20 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsComponen
         self._cfg_normalization = cfg_normalization
         self._cfg_truncation = cfg_truncation
         # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = len(prompt_embeds)
+        batch_size = len(prompt)
 
-        # If prompt_embeds is provided and prompt is None, skip encoding
-        if prompt_embeds is not None and prompt is None:
-            if self.do_classifier_free_guidance and negative_prompt_embeds is None:
-                raise ValueError(
-                    "When `prompt_embeds` is provided without `prompt`, "
-                    "`negative_prompt_embeds` must also be provided for classifier-free guidance."
-                )
-        else:
-            (
-                prompt_embeds,
-                negative_prompt_embeds,
-            ) = self.encode_prompt(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                do_classifier_free_guidance=self.do_classifier_free_guidance,
-                prompt_embeds=prompt_embeds,
-                negative_prompt_embeds=negative_prompt_embeds,
-                device=device,
-                max_sequence_length=max_sequence_length,
-            )
+        (
+            prompt_embeds,
+            negative_prompt_embeds,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            do_classifier_free_guidance=self.do_classifier_free_guidance,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            device=device,
+            max_sequence_length=max_sequence_length,
+        )
 
         # 4. Prepare latent variables
         num_channels_latents = self.transformer.in_channels

vllm_omni/diffusion/profiler/diffusion_pipeline_profiler.py

@@ -14,6 +14,7 @@ logger = init_logger(__name__)
 
 def profiler(name: str, func: Callable, instance: Any) -> Callable:
     """Timing a function execution."""
+    metric_name = f"{name.rsplit('.', 1)[0]}.diffuse" if name.endswith(".denoise_step") else name
 
     @functools.wraps(func)
     def wrapper(*args, **kwargs) -> Any:
@@ -28,10 +29,10 @@ def profiler(name: str, func: Callable, instance: Any) -> Callable:
             if current_omni_platform.is_available():
                 current_omni_platform.synchronize()
             duration = time.perf_counter() - start_time
-            logger.info(f"[DiffusionPipelineProfiler] {name} took {duration:.6f}s")
+            logger.info(f"[DiffusionPipelineProfiler] {metric_name} took {duration:.6f}s")
             # record the profiling data: duration of stages
             with instance._profiler_lock:
-                instance._stage_durations[name] = instance._stage_durations.get(name, 0.0) + duration
+                instance._stage_durations[metric_name] = instance._stage_durations.get(metric_name, 0.0) + duration
 
     return wrapper
 
@@ -93,7 +94,11 @@ class DiffusionPipelineProfilerMixin:
         self._stage_durations: dict[str, float] = {}
         targets = profiler_targets if profiler_targets is not None else self._PROFILER_TARGETS
         if not targets:
-            return
+            targets = []
+        else:
+            targets = list(targets)
+        if not profiler_targets and hasattr(self, "denoise_step"):
+            targets.append("denoise_step")
 
         targets = ["forward"] + [
             t for t in targets if t != "forward"

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -39,7 +39,15 @@ from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
 from vllm_omni.diffusion.worker.input_batch import InputBatch, scatter_latents
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
-from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, DiffusionRequestState, RunnerOutput
+from vllm_omni.diffusion.worker.utils import (
+    BatchRunnerOutput,
+    DiffusionRequestState,
+    RunnerOutput,
+    attach_stage_durations,
+    clear_pipeline_stage_durations,
+    consume_pipeline_stage_durations,
+    merge_stage_durations,
+)
 from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager
 from vllm_omni.platforms import current_omni_platform
 from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
@@ -487,6 +495,19 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
 
         return self._runner_output_from_outputs(reqs, outputs)
 
+    def _attach_stepwise_metrics(
+        self,
+        state: DiffusionRequestState,
+        output: DiffusionOutput,
+        *,
+        is_primary: bool,
+    ) -> None:
+        merge_stage_durations(
+            state,
+            consume_pipeline_stage_durations(self.pipeline),
+        )
+        attach_stage_durations(state, output)
+
     def execute_model(self, req: OmniDiffusionRequest, kv_prefetch_jobs: dict | None = None) -> DiffusionOutput:
         """
         Execute a forward pass for the given requests.
@@ -601,8 +622,13 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                     else:
                         gen_device = self.device
                     state.sampling.generator = torch.Generator(device=gen_device).manual_seed(state.sampling.seed)
+                clear_pipeline_stage_durations(self.pipeline)
                 # encode
                 self.pipeline.prepare_encode(state)
+                merge_stage_durations(
+                    state,
+                    consume_pipeline_stage_durations(self.pipeline),
+                )
 
         input_batch = InputBatch.make_batch(
             states,
@@ -658,19 +684,27 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         use_hsdp = self.od_config.parallel_config.use_hsdp
         grad_context = torch.no_grad() if use_hsdp else torch.inference_mode()
         with grad_context:
+            had_active_states = bool(self.state_cache)
             states, new_request_ids = self._update_states(scheduler_output)
+            is_primary = not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
+            if new_request_ids and not had_active_states and is_primary and current_omni_platform.is_available():
+                current_omni_platform.reset_peak_memory_stats()
             input_batch = self._prepare_batch_inputs(states, new_request_ids)
             attn_metadata = self._prepare_attn_metadata(input_batch)
-            is_primary = not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
-            if is_primary:
-                current_omni_platform.reset_peak_memory_stats()
 
             with set_forward_context(
                 vllm_config=self.vllm_config,
                 omni_diffusion_config=self.od_config,
                 attn_metadata=attn_metadata,
             ):
+                clear_pipeline_stage_durations(self.pipeline)
                 noise_pred = self.pipeline.denoise_step(input_batch, states=states)
+                denoise_stage_durations = consume_pipeline_stage_durations(self.pipeline)
+                for state in states:
+                    merge_stage_durations(
+                        state,
+                        denoise_stage_durations,
+                    )
 
                 runner_output_list = []
                 pipeline_interrupted = getattr(self.pipeline, "interrupt", False)
@@ -693,14 +727,20 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                             self.pipeline.step_scheduler(
                                 req, noise_pred[offset : offset + row_num] if noise_pred is not None else None
                             )
-                            offset = offset + row_num
                             if self.od_config.streaming_output:
                                 should_decode = req.chunk_denoise_completed
                             else:
                                 should_decode = req.denoise_completed
 
                             if should_decode:
+                                clear_pipeline_stage_durations(self.pipeline)
                                 result = self.pipeline.post_decode(req)
+                                if result is not None:
+                                    self._attach_stepwise_metrics(
+                                        req,
+                                        result,
+                                        is_primary=is_primary,
+                                    )
                             else:
                                 result = None
                             # finished should be computed after post_decode() advanced chunk_index
@@ -717,6 +757,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                                     result=result,
                                 )
                             )
+                            offset = offset + row_num
                         except Exception as per_req_exc:
                             offset = offset + row_num
                             logger.error(

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -246,9 +246,18 @@ class DiffusionWorker:
         # Create vllm_config for parallel configuration. Pass explicit device_config
         # so DeviceConfig does not rely on current_platform in worker subprocesses.
         vllm_config = _create_diffusion_worker_vllm_config(self.device, self.od_config)
-        vllm_config.parallel_config.tensor_parallel_size = self.od_config.parallel_config.tensor_parallel_size
-        vllm_config.parallel_config.data_parallel_size = self.od_config.parallel_config.data_parallel_size
-        vllm_config.parallel_config.enable_expert_parallel = self.od_config.parallel_config.enable_expert_parallel
+        parallel_config = self.od_config.parallel_config
+        vllm_config.parallel_config.tensor_parallel_size = parallel_config.tensor_parallel_size
+        vllm_config.parallel_config.data_parallel_size = parallel_config.data_parallel_size
+        if parallel_config.enable_expert_parallel and self.od_config.is_moe:
+            # Diffusion uses its own DP/CFG/SP groups normally. vLLM groups are
+            # only remapped for expert-parallel runtimes that consume vLLM's
+            # FusedMoE/EP semantics.
+            vllm_config.parallel_config.data_parallel_size = (
+                parallel_config.data_parallel_size * parallel_config.cfg_parallel_size
+            )
+            vllm_config.parallel_config.prefill_context_parallel_size = parallel_config.sequence_parallel_size
+        vllm_config.parallel_config.enable_expert_parallel = parallel_config.enable_expert_parallel
         vllm_config.profiler_config = self.od_config.profiler_config
         vllm_config.model_config = _make_diffusion_vllm_model_config(self.od_config)  # type: ignore[assignment]
         vllm_config.quant_config = self.od_config.quantization_config

vllm_omni/diffusion/worker/utils.py

@@ -15,6 +15,41 @@ if TYPE_CHECKING:
     from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
 
 
+def clear_pipeline_stage_durations(pipeline: Any) -> None:
+    clear_records = getattr(pipeline, "clear_profiler_records", None)
+    if getattr(pipeline, "enable_diffusion_pipeline_profiler", False) and callable(clear_records):
+        clear_records()
+
+
+def consume_pipeline_stage_durations(pipeline: Any) -> dict[str, float]:
+    if not getattr(pipeline, "enable_diffusion_pipeline_profiler", False):
+        return {}
+    stage_durations = getattr(pipeline, "stage_durations", None)
+    if not isinstance(stage_durations, dict):
+        return {}
+    result = {stage: float(duration) for stage, duration in stage_durations.items()}
+    clear_pipeline_stage_durations(pipeline)
+    return result
+
+
+def merge_stage_durations(
+    state: DiffusionRequestState,
+    stage_durations: dict[str, float],
+) -> None:
+    if not stage_durations:
+        return
+    for stage, duration in stage_durations.items():
+        state.stage_durations[stage] = float(state.stage_durations.get(stage, 0.0)) + float(duration)
+
+
+def attach_stage_durations(
+    state: DiffusionRequestState,
+    output: DiffusionOutput,
+) -> None:
+    if state.stage_durations:
+        output.stage_durations = dict(state.stage_durations)
+
+
 @dataclass
 class DiffusionRequestState:
     """Per-request mutable state across all pipeline stages.
@@ -79,6 +114,9 @@ class DiffusionRequestState:
     # For example: Wan condition tensors / masks, or Bagel KV contexts.
     extra: dict[str, Any] = field(default_factory=dict)
 
+    # ── Runner-owned profiling metadata ──
+    stage_durations: dict[str, float] = field(default_factory=dict)
+
     # Peak device memory observed while this request is active in step mode.
     peak_memory_mb: float = 0.0
 

vllm_omni/engine/orchestrator.py

@@ -1287,7 +1287,7 @@ class Orchestrator:
                             [req_id, *self._cfg_tracker.cleanup_parent(req_id)],
                         )
                         return
-                    if already_submitted and len(diffusion_prompt) == 1:
+                    if len(diffusion_prompt) == 1:
                         diffusion_prompt = diffusion_prompt[0]
             else:
                 diffusion_prompt = req_state.prompt

vllm_omni/engine/stage_pool.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import asyncio
 import time as _time
+from collections.abc import Mapping
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, cast
 
@@ -597,7 +598,6 @@ class StagePool:
         batch_id = metrics.batch_seq
         metrics.agg_total_tokens += num_tokens_out
         metrics.agg_total_gen_time_ms += stage_gen_time_ms
-        audio_rtf = defs.compute_audio_rtf(stage_gen_time_ms / 1000.0, audio_duration_s)
 
         return StageRequestMetrics(
             num_tokens_in=num_tokens_in,
@@ -616,7 +616,6 @@ class StagePool:
             audio_generated_frames=audio_generated_frames,
             audio_sample_rate=audio_sample_rate,
             audio_duration_s=audio_duration_s,
-            audio_rtf=audio_rtf,
             image_pixels=image_pixels,
             denoise_step_latency_ms=denoise_step_latency_ms,
             output_unit_type=output_unit_type,
@@ -686,7 +685,7 @@ class StagePool:
     def _has_audio_output(self, request_outputs: list[Any]) -> bool:
         for ro in request_outputs:
             for mm_output in self._iter_multimodal_outputs(ro):
-                if isinstance(mm_output, dict) and mm_output.get("audio") is not None:
+                if isinstance(mm_output, Mapping) and mm_output.get("audio") is not None:
                     return True
         return False
 
@@ -695,7 +694,7 @@ class StagePool:
         sample_rate = 0
         for ro in request_outputs:
             for mm_output in self._iter_multimodal_outputs(ro):
-                if not isinstance(mm_output, dict):
+                if not isinstance(mm_output, Mapping):
                     continue
                 if sample_rate <= 0:
                     sample_rate = self._infer_audio_sample_rate(mm_output)
@@ -779,7 +778,7 @@ class StagePool:
     def _count_images(self, request_output: Any) -> int:
         total_images = self._count_value_units(getattr(request_output, "images", None))
         for mm_output in self._iter_multimodal_outputs(request_output):
-            if isinstance(mm_output, dict):
+            if isinstance(mm_output, Mapping):
                 total_images += self._count_value_units(mm_output.get("image"))
                 total_images += self._count_value_units(mm_output.get("images"))
         return total_images
@@ -789,7 +788,7 @@ class StagePool:
         for ro in request_outputs:
             total_pixels += self._count_image_value_pixels(getattr(ro, "images", None))
             for mm_output in self._iter_multimodal_outputs(ro):
-                if isinstance(mm_output, dict):
+                if isinstance(mm_output, Mapping):
                     total_pixels += self._count_image_value_pixels(mm_output.get("image"))
                     total_pixels += self._count_image_value_pixels(mm_output.get("images"))
         return total_pixels
@@ -824,7 +823,7 @@ class StagePool:
         total_videos = self._count_video_units(getattr(request_output, "video", None))
         total_videos += self._count_video_units(getattr(request_output, "videos", None))
         for mm_output in self._iter_multimodal_outputs(request_output):
-            if isinstance(mm_output, dict):
+            if isinstance(mm_output, Mapping):
                 total_videos += self._count_video_units(mm_output.get("video"))
                 total_videos += self._count_video_units(mm_output.get("videos"))
         return total_videos
@@ -849,7 +848,7 @@ class StagePool:
         if isinstance(custom_output, dict) and custom_output:
             return True
         for mm_output in self._iter_multimodal_outputs(request_output):
-            if isinstance(mm_output, dict) and any(self._is_non_empty_value(value) for value in mm_output.values()):
+            if isinstance(mm_output, Mapping) and any(self._is_non_empty_value(value) for value in mm_output.values()):
                 return True
         for output in getattr(request_output, "outputs", None) or []:
             if getattr(output, "text", None):
@@ -917,14 +916,14 @@ class StagePool:
             if self._audio_sample_rate_by_request.get(rid, 0) <= 0 and audio_sample_rate > 0:
                 self._audio_sample_rate_by_request[rid] = audio_sample_rate
 
-    def _iter_multimodal_outputs(self, request_output: Any) -> list[dict[str, Any]]:
-        multimodal_outputs: list[dict[str, Any]] = []
+    def _iter_multimodal_outputs(self, request_output: object) -> list[Mapping[str, Any]]:
+        multimodal_outputs: list[Mapping[str, Any]] = []
         outer_mm = getattr(request_output, "multimodal_output", None)
-        if isinstance(outer_mm, dict) and outer_mm:
+        if isinstance(outer_mm, Mapping) and outer_mm:
             multimodal_outputs.append(outer_mm)
         for output in getattr(request_output, "outputs", None) or []:
             inner_mm = getattr(output, "multimodal_output", None)
-            if isinstance(inner_mm, dict) and inner_mm:
+            if isinstance(inner_mm, Mapping) and inner_mm:
                 multimodal_outputs.append(inner_mm)
         return multimodal_outputs
 

vllm_omni/entrypoints/omni_base.py

@@ -221,6 +221,13 @@ class OmniBase(PDDisaggregationMixin):
         """Expose engine stage configs for PD disaggregation detection and validation."""
         return self.engine.stage_configs
 
+    def _consumed_metric_message_ids(self, request_id: str) -> set[int]:
+        consumed_by_request = getattr(self, "_consumed_metric_messages", None)
+        if consumed_by_request is None:
+            consumed_by_request = {}
+            self._consumed_metric_messages = consumed_by_request
+        return consumed_by_request.setdefault(request_id, set())
+
     def _has_dead_stage(self) -> bool:
         for stage_client in self.engine.stage_clients:
             if getattr(stage_client, "_engine_dead", False):
@@ -310,7 +317,9 @@ class OmniBase(PDDisaggregationMixin):
             )
         finally:
             self.request_states.pop(request_id, None)
-            self._consumed_metric_messages.pop(request_id, None)
+            consumed_by_request = getattr(self, "_consumed_metric_messages", None)
+            if consumed_by_request is not None:
+                consumed_by_request.pop(request_id, None)
             # Republish gauges so any stale value left by the per-stage
             # publish in _process_single_result (which runs while the request
             # is still in self.request_states) is corrected after the pop.
@@ -397,7 +406,7 @@ class OmniBase(PDDisaggregationMixin):
             stage_meta = self.engine.get_stage_metadata(stage_id)
             output_type = getattr(msg.engine_outputs, "final_output_type", stage_meta.final_output_type)
             msg_id = id(msg)
-            consumed = self._consumed_metric_messages.setdefault(req_id, set())
+            consumed = self._consumed_metric_message_ids(req_id)
             if msg_id not in consumed:
                 req_state.metrics.on_stage_metrics(stage_id, req_id, msg.metrics, output_type)
                 submit_ts = msg.stage_submit_ts
@@ -505,7 +514,11 @@ class OmniBase(PDDisaggregationMixin):
         stage_meta = self.engine.get_stage_metadata(stage_id)
         output_type = getattr(engine_outputs, "final_output_type", stage_meta.final_output_type)
         if finished and _m is not None:
-            metrics.on_stage_metrics(stage_id, req_id, _m, output_type)
+            msg_id = id(result)
+            consumed = self._consumed_metric_message_ids(req_id)
+            if msg_id not in consumed:
+                metrics.on_stage_metrics(stage_id, req_id, _m, output_type)
+                consumed.add(msg_id)
 
         if not stage_meta.final_output:
             return None

vllm_omni/entrypoints/openai/api_server.py

@@ -751,7 +751,7 @@ async def omni_init_app_state(
         state.diffusion_engine = engine_client
         state.openai_serving_models = _DiffusionServingModels(base_model_paths)
         # OMNI: tokenization endpoints are not supported in pure diffusion mode.
-        state.openai_serving_tokenization = None
+        state.serving_tokenization = None
 
         # Use for_diffusion method to create chat handler
         state.openai_serving_chat = OmniOpenAIServingChat.for_diffusion(
@@ -1014,7 +1014,7 @@ async def omni_init_app_state(
         if any(t in supported_tasks for t in ("embed", "score", "token_embed"))
         else None
     )
-    state.openai_serving_tokenization = ServingTokenization(
+    state.serving_tokenization = ServingTokenization(
         state.openai_serving_models,
         state.openai_serving_render,
         request_logger=request_logger,
@@ -1154,7 +1154,7 @@ async def create_chat_completion(request: ChatCompletionRequest, raw_request: Re
     metrics_header_format = raw_request.headers.get(ENDPOINT_LOAD_METRICS_FORMAT_HEADER_LABEL, "")
     handler = Omnichat(raw_request)
     if handler is None:
-        base_server = getattr(raw_request.app.state, "openai_serving_tokenization", None)
+        base_server = getattr(raw_request.app.state, "serving_tokenization", None)
         if base_server is None:
             raise HTTPException(
                 status_code=HTTPStatus.NOT_FOUND.value,
@@ -1245,7 +1245,7 @@ async def create_speech(request: OpenAICreateSpeechRequest, raw_request: Request
     """
     handler = Omnispeech(raw_request)
     if handler is None:
-        base_server = getattr(raw_request.app.state, "openai_serving_tokenization", None)
+        base_server = getattr(raw_request.app.state, "serving_tokenization", None)
         if base_server is None:
             raise HTTPException(
                 status_code=HTTPStatus.NOT_FOUND.value,
@@ -1283,7 +1283,7 @@ async def create_speech(request: OpenAICreateSpeechRequest, raw_request: Request
 async def create_speech_batch(request: BatchSpeechRequest, raw_request: Request):
     handler = Omnispeech(raw_request)
     if handler is None:
-        base_server = getattr(raw_request.app.state, "openai_serving_tokenization", None)
+        base_server = getattr(raw_request.app.state, "serving_tokenization", None)
         if base_server is None:
             raise HTTPException(
                 status_code=HTTPStatus.NOT_FOUND.value,
@@ -1326,7 +1326,7 @@ async def create_speech_batch(request: BatchSpeechRequest, raw_request: Request)
 async def create_audio_generate(request: OpenAICreateAudioGenerateRequest, raw_request: Request):
     handler = OmniAudioGenerate(raw_request)
     if handler is None:
-        base_server = getattr(raw_request.app.state, "openai_serving_tokenization", None)
+        base_server = getattr(raw_request.app.state, "serving_tokenization", None)
         if base_server is None:
             raise HTTPException(
                 status_code=HTTPStatus.NOT_FOUND.value,

vllm_omni/entrypoints/openai/serving_speech.py

@@ -2620,16 +2620,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         self,
         request,
         *,
-        include_language=False,
         include_voice=False,
         plain_text_passthrough=False,
     ):
         instruction_text = request.instructions.strip() if isinstance(request.instructions, str) else None
         instruction_dict: dict[str, Any] = {}
 
-        if include_language and request.language not in (None, "", "Auto"):
-            instruction_dict["方言"] = request.language
-
         voice_lower = request.voice.lower() if isinstance(request.voice, str) else None
         if include_voice and request.voice and not (voice_lower and voice_lower in self.uploaded_speakers):
             instruction_dict["IP"] = request.voice
@@ -2652,7 +2648,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         """Build a Ming instruction payload from OpenAI speech fields."""
         return self._parse_ming_instruction_fields(
             request,
-            include_language=True,
             include_voice=True,
             plain_text_passthrough=True,
         )
@@ -3705,15 +3700,14 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 sampling_params_list[0].max_tokens = int(request.max_new_tokens) + 1
 
         if request.seed is not None and sampling_params_list:
-            if sampling_params_list is self.engine_client.default_sampling_params_list:
-                import copy
+            import copy
 
-                sampling_params_list = copy.deepcopy(sampling_params_list)
-            sampling_params_list[0].seed = request.seed
-            if self._tts_model_type == "qwen3_tts":
-                if sampling_params_list[0].extra_args is None:
-                    sampling_params_list[0].extra_args = {}
-                sampling_params_list[0].extra_args["qwen3_tts_request_seed"] = request.seed
+            sampling_params_list = copy.deepcopy(sampling_params_list)
+            stage0_params = sampling_params_list[0]
+            stage0_params.seed = request.seed
+            if stage0_params.extra_args is None:
+                stage0_params.extra_args = {}
+            stage0_params.extra_args["tts_local_seed"] = request.seed
 
         generator = self.engine_client.generate(
             prompt=prompt,

vllm_omni/metrics/definitions.py

@@ -11,11 +11,14 @@ time-bearing metrics use the ``_s`` suffix (values in seconds), counters use
 ``_total`` (auto-suffixed by the prometheus client), sizes use ``_bytes``.
 """
 
+import logging
+import os
 from typing import Any
 
 # vllm_omni: namespace for omni-specific Prometheus families, distinct from
 # the upstream vllm:* families.
 METRIC_PREFIX = "vllm_omni:"
+logger = logging.getLogger(__name__)
 
 
 # ============================================================================
@@ -76,10 +79,8 @@ VLLM_TTFTS = "vllm_ttfts"
 VLLM_TPOTS = "vllm_tpots"
 VLLM_ITLS = "vllm_itls"
 AUDIO_TTFPS = "audio_ttfps"
-AUDIO_RTFS = "audio_rtfs"
 AUDIO_DURATIONS = "audio_durations"
 MISSING_AUDIO_DURATION_COUNT = "missing_audio_duration_count"
-MISSING_AUDIO_RTF_COUNT = "missing_audio_rtf_count"
 STAGE_GEN_TIME = "stage_gen_time"
 STAGE_GEN_TIME_MS = f"{STAGE_GEN_TIME}_ms"
 STAGE_GEN_TIMES_MS = f"{STAGE_GEN_TIME}s_ms"
@@ -282,6 +283,9 @@ def compute_denoise_step_latency(stage_gen_time: float, num_inference_steps: int
 # ming_flash 16000 — these models populate multimodal_output["audio_sample_rate"]
 # at runtime so this default only kicks in when the field is missing.
 DEFAULT_AUDIO_SAMPLE_RATE = 24000
+DEFAULT_AUDIO_CHANNELS = 1
+AUDIO_SAMPLE_RATE_ENV = "VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE"
+AUDIO_CHANNELS_ENV = "VLLM_OMNI_BENCH_AUDIO_CHANNELS"
 
 _SAMPLE_RATE_KEYS = ("output_sample_rate", "audio_sample_rate", "sample_rate", "sampling_rate", "sr")
 
@@ -308,3 +312,26 @@ def resolve_audio_sample_rate(source: dict[str, Any] | Any | None) -> int:
         if value > 0:
             return value
     return DEFAULT_AUDIO_SAMPLE_RATE
+
+
+def stream_pcm_format_from_env(
+    *,
+    default_sample_rate: int = DEFAULT_AUDIO_SAMPLE_RATE,
+    default_channels: int = DEFAULT_AUDIO_CHANNELS,
+) -> tuple[int, int]:
+    """Return the sample rate and channel count for streamed raw PCM."""
+    sample_rate = default_sample_rate
+    channels = default_channels
+    raw_sr = os.environ.get(AUDIO_SAMPLE_RATE_ENV)
+    if raw_sr:
+        try:
+            sample_rate = int(raw_sr)
+        except ValueError:
+            logger.warning("Invalid %s=%r; using default %d", AUDIO_SAMPLE_RATE_ENV, raw_sr, sample_rate)
+    raw_ch = os.environ.get(AUDIO_CHANNELS_ENV)
+    if raw_ch:
+        try:
+            channels = int(raw_ch)
+        except ValueError:
+            logger.warning("Invalid %s=%r; using default %d", AUDIO_CHANNELS_ENV, raw_ch, channels)
+    return max(sample_rate, 1), max(channels, 1)

vllm_omni/metrics/stats.py

@@ -48,7 +48,6 @@ class StageRequestStats:
     audio_generated_frames: int = 0
     audio_sample_rate: int = 0
     audio_duration_s: float = 0.0
-    audio_rtf: float = 0.0
     image_pixels: int = 0
     denoise_step_latency_ms: float = 0.0
     pipeline_timings: dict[str, float] | None = None
@@ -471,7 +470,6 @@ class OrchestratorAggregator:
                 defs.AUDIO_FRAMES: int(evt.audio_generated_frames),
                 defs.AUDIO_SAMPLE_RATE: int(evt.audio_sample_rate),
                 f"{defs.AUDIO_DURATION}_s": float(evt.audio_duration_s),
-                defs.AUDIO_RTF: float(evt.audio_rtf),
                 defs.IMAGE_PIXELS: int(evt.image_pixels),
                 defs.DENOISE_STEP_LATENCY_MS: float(evt.denoise_step_latency_ms),
                 "output_unit_type": evt.output_unit_type,
@@ -541,10 +539,6 @@ class OrchestratorAggregator:
             0.0,
         )
         current[defs.TIME_PER_OUTPUT_UNIT_MS] = remaining_ms / float(output_count - 1) if output_count > 1 else 0.0
-        duration_s = float(current.get(f"{defs.AUDIO_DURATION}_s", 0.0))
-        current[defs.AUDIO_RTF] = defs.compute_audio_rtf(
-            float(current.get(defs.STAGE_GEN_TIME_MS, 0.0)) / 1000.0, duration_s
-        )
         return current
 
     def _build_stage_metrics_snapshot(self, req_id: str) -> dict[str, dict[str, Any]]:

vllm_omni/model_executor/models/common/ming/audio_vae.py

@@ -138,8 +138,8 @@ class Encoder(nn.Module):
         self.norm = nn.LayerNorm(config.hidden_size)
         self.patch_size = patch_size
         if patch_size != -1:
-            aggregator_config = _qwen2_config({**encoder_args, "num_hidden_layers": 4})
-            self.aggregator = Qwen2Model(aggregator_config)
+            config.num_hidden_layers = 4
+            self.aggregator = Qwen2Model(config)
             self.cls_embed = nn.Parameter(torch.rand(1, 1, config.hidden_size))
             self.cls_embed.data.normal_(0, 0.02)
 
@@ -346,8 +346,9 @@ class AudioVAE(PreTrainedModel):
         h, _ = self.encoder(waveform)
         h = h.transpose(1, 2)
 
-        mean = torch.chunk(h, 2, dim=1)[0]
-        latent = mean.transpose(1, 2)
+        mean, scale = torch.chunk(h, 2, dim=1)
+        std = F.softplus(scale) + 1e-4
+        latent = (mean + std * torch.randn_like(mean)).transpose(1, 2)
         return latent, frame_num
 
     @torch.inference_mode()

vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py

@@ -222,14 +222,8 @@ class HunyuanModel(HunYuanModel):
             if self.config.tie_word_embeddings and "lm_head.weight" in name:
                 continue
 
-            if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)):
-                # Loading kv cache scales for compressed-tensors quantization
-                param = params_dict[scale_name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                loaded_weight = loaded_weight[0]
-                weight_loader(param, loaded_weight)
-                continue
-
+            # KV-cache scales are renamed to .attn.{k,v,q}_scale by the
+            # quant_config mapper in the outer AutoWeightsLoader.
             is_found = False
             for param_name, weight_name, shard_id in stacked_params_mapping:
                 if weight_name not in name:

vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py

@@ -483,19 +483,15 @@ class MammothModa2Qwen2ForCausalLM(nn.Module, SupportsPP):
         ]
         params_dict = dict(self.named_parameters(remove_duplicate=False))
         loaded_params: set[str] = set()
+        if self.quant_config is not None:
+            cache_scale_mapper = self.quant_config.get_cache_scale_mapper()
+            if cache_scale_mapper is not None:
+                weights = cache_scale_mapper.apply(weights)
 
         for name, loaded_weight in weights:
             if "rotary_emb.inv_freq" in name:
                 continue
 
-            if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)):
-                param = params_dict[scale_name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                loaded_weight = loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
-                weight_loader(param, loaded_weight)
-                loaded_params.add(scale_name)
-                continue
-
             for param_name, weight_name, shard_id in stacked_params_mapping:
                 if weight_name not in name:
                     continue