cherry-pick pr3067 (#3087)

Signed-off-by: lyj-jjj <liuyingjun5@huawei.com>

docs/.nav.yml

@@ -62,6 +62,7 @@ nav:
         - FP8: user_guide/diffusion/quantization/fp8.md
         - Int8: user_guide/diffusion/quantization/int8.md
         - GGUF: user_guide/diffusion/quantization/gguf.md
+      - Frame Interpolation: user_guide/diffusion/frame_interpolation.md
       - Step Execution: user_guide/diffusion/step_execution.md
       - Parallelism Acceleration: user_guide/diffusion/parallelism_acceleration.md
       - CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md

docs/design/feature/vae_parallel.md

@@ -1,14 +1,15 @@
 # VAE Patch Parallelism
 
 This document describes how to add **VAE Patch Parallelism** support to a diffusion model.
-We use **Qwen-Image** as the reference implementation.
+We use **Qwen-Image** as the reference implementation for decode parallel, and **Wan2.2** for encode parallel.
 
 ---
 
 ## Table of Contents
 
 - [Overview](#overview)
-- [Step-by-Step Implementation](#step-by-step-implementation)
+- [Step-by-Step Implementation (Decode)](#step-by-step-implementation-decode)
+- [Encode Parallel Implementation](#encode-parallel-implementation)
 - [Testing](#testing)
 - [Reference Implementations](#reference-implementations)
 - [Summary](#summary)
@@ -19,13 +20,13 @@ We use **Qwen-Image** as the reference implementation.
 
 ### What is Vae Patch parallel?
 
-**VAE Patch Parallelism** is a decoding acceleration technique. Instead of decoding the entire latent tensor at once, the latent tensor is:
+**VAE Patch Parallelism** is an acceleration technique for both **encoding** and **decoding**. Instead of processing the entire tensor at once, the tensor is:
 
 + Split into multiple spatial tiles
 
 + Distributed across multiple ranks
 
-+ Decoded in parallel
++ Encoded/Decoded in parallel
 
 + Merged to reconstruct the final output
 
@@ -35,10 +36,17 @@ This approach:
 
 + Reduces peak memory usage per device
 
-+ Accelerates decoding latency
++ Accelerates encoding/decoding latency
+
+### When to Use Encode vs Decode Parallel
+
+| Operation | Use Case | Example |
+|-----------|----------|---------|
+| **Decode Parallel** | Text-to-Image, Text-to-Video | Latent → Image/Video |
+| **Encode Parallel** | Image-to-Video (I2V) | Image → Latent (for conditioning) |
 
 ### Architecture
-We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE decoding.
+We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE encoding/decoding.
 
 The executor is model-agnostic and accepts three function parameters:
 
@@ -84,7 +92,7 @@ Therefore:
 
 + Merge must perform blending to avoid seams
 
-## Step-by-Step Implementation
+## Step-by-Step Implementation (Decode)
 
 ### Step 1: Implement DistributedAutoencoderKLQwenImage
 `QwenImagePipeline` use `AutoencoderKLQwenImage` for vae, so implement a distributed version:
@@ -205,14 +213,14 @@ def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid
 We need to override tiled_decode, the main logic is:
 + check distributed is enabled
 + select split/exec/merge
-+ Invoke self.distributed_decoder.execute to decode
++ Invoke self.distributed_executor.execute to decode
 ```
 def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
     if not self.is_distributed_enabled():
         return super().tiled_decode(z, return_dict=return_dict)
 
     logger.info("Decode run with distributed executor")
-    result = self.distributed_decoder.execute(
+    result = self.distributed_executor.execute(
         z,
         DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
         broadcast_result=True,
@@ -243,6 +251,166 @@ class YourModelPipeline(nn.Module):
 +       ).to(self.device)
 ```
 
+## Encode Parallel Implementation
+
+For models that require VAE encoding (e.g., Image-to-Video), you can also parallelize the encode operation. We use **Wan2.2** as the reference implementation.
+
+### Step 1: Implement encode_tile_split
+
+Similar to decode, split the input tensor into tiles. Key considerations:
+
++ **Patchify handling**: If the model uses `patch_size`, scale tile parameters accordingly
++ **Temporal chunking**: Video VAEs may have temporal compression (e.g., 4x)
+
+```python
+def encode_tile_split(self, x: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
+    _, _, num_frames, height, width = x.shape
+    encode_spatial_compression_ratio = self.spatial_compression_ratio
+
+    # Scale tile parameters for patchified coordinate system
+    tile_sample_min_height = self.tile_sample_min_height
+    tile_sample_min_width = self.tile_sample_min_width
+    tile_sample_stride_height = self.tile_sample_stride_height
+    tile_sample_stride_width = self.tile_sample_stride_width
+
+    if self.config.patch_size is not None:
+        # When input is patchified, scale tile parameters accordingly
+        encode_spatial_compression_ratio = self.spatial_compression_ratio // self.config.patch_size
+        tile_sample_min_height = tile_sample_min_height // self.config.patch_size
+        tile_sample_min_width = tile_sample_min_width // self.config.patch_size
+        tile_sample_stride_height = tile_sample_stride_height // self.config.patch_size
+        tile_sample_stride_width = tile_sample_stride_width // self.config.patch_size
+
+    latent_height = height // encode_spatial_compression_ratio
+    latent_width = width // encode_spatial_compression_ratio
+
+    tile_latent_min_height = tile_sample_min_height // encode_spatial_compression_ratio
+    tile_latent_min_width = tile_sample_min_width // encode_spatial_compression_ratio
+    tile_latent_stride_height = tile_sample_stride_height // encode_spatial_compression_ratio
+    tile_latent_stride_width = tile_sample_stride_width // encode_spatial_compression_ratio
+
+    blend_height = tile_latent_min_height - tile_latent_stride_height
+    blend_width = tile_latent_min_width - tile_latent_stride_width
+
+    tiletask_list = []
+    # Use temporal compression ratio from config instead of hardcoding
+    temporal_compression = self.config.scale_factor_temporal
+
+    for i in range(0, height, tile_sample_stride_height):
+        for j in range(0, width, tile_sample_stride_width):
+            time_list = []
+            frame_range = 1 + (num_frames - 1) // temporal_compression
+            for k in range(frame_range):
+                if k == 0:
+                    tile = x[:, :, :1, i : i + tile_sample_min_height, j : j + tile_sample_min_width]
+                else:
+                    tile = x[
+                        :, :,
+                        1 + temporal_compression * (k - 1) : 1 + temporal_compression * k,
+                        i : i + tile_sample_min_height,
+                        j : j + tile_sample_min_width,
+                    ]
+                time_list.append(tile)
+            tiletask_list.append(
+                TileTask(len(tiletask_list), (i // tile_sample_stride_height, j // tile_sample_stride_width),
+                         time_list, workload=time_list[0].shape[3] * time_list[0].shape[4])
+            )
+
+    grid_spec = GridSpec(
+        split_dims=(3, 4),
+        grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
+        tile_spec={
+            "latent_height": latent_height, "latent_width": latent_width,
+            "blend_height": blend_height, "blend_width": blend_width,
+            "tile_latent_stride_height": tile_latent_stride_height,
+            "tile_latent_stride_width": tile_latent_stride_width,
+        },
+        output_dtype=self.dtype,
+    )
+    return tiletask_list, grid_spec
+```
+
+### Step 2: Implement encode_tile_exec
+
+```python
+def encode_tile_exec(self, task: TileTask) -> torch.Tensor:
+    """Encode a single sample tile into latent space."""
+    self.clear_cache()
+    time = []
+    for k, tile in enumerate(task.tensor):
+        self._enc_conv_idx = [0]
+        encoded = self.encoder(tile, feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx)
+        encoded = self.quant_conv(encoded)
+        time.append(encoded)
+    result = torch.cat(time, dim=2)
+    self.clear_cache()
+    return result
+```
+
+### Step 3: Implement encode_tile_merge
+
+```python
+def encode_tile_merge(
+    self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec
+) -> torch.Tensor:
+    """Merge encoded tiles into a full latent tensor."""
+    grid_h, grid_w = grid_spec.grid_shape
+    result_rows = []
+    for i in range(grid_h):
+        result_row = []
+        for j in range(grid_w):
+            tile = coord_tensor_map[(i, j)]
+            if i > 0:
+                tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
+            if j > 0:
+                tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
+            result_row.append(tile[:, :, :,
+                : grid_spec.tile_spec["tile_latent_stride_height"],
+                : grid_spec.tile_spec["tile_latent_stride_width"]])
+        result_rows.append(torch.cat(result_row, dim=-1))
+
+    enc = torch.cat(result_rows, dim=3)[
+        :, :, :, : grid_spec.tile_spec["latent_height"], : grid_spec.tile_spec["latent_width"]
+    ]
+    return enc
+```
+
+### Step 4: Override tiled_encode method
+
+Override `tiled_encode` instead of `encode`. The parent's `_encode()` handles patchify before calling `tiled_encode()`, so input `x` is already patchified.
+
+```python
+def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
+    """
+    Encode using distributed VAE executor.
+
+    Note: x is already patchified by parent's _encode() before calling this method.
+    """
+    if not self.is_distributed_enabled():
+        return super().tiled_encode(x)
+
+    self.clear_cache()
+    result = self.distributed_executor.execute(
+        x,
+        DistributedOperator(
+            split=self.encode_tile_split,
+            exec=self.encode_tile_exec,
+            merge=self.encode_tile_merge,
+        ),
+        broadcast_result=True,  # Latents needed by all ranks for diffusion
+    )
+    self.clear_cache()
+    return result
+```
+
+**Key differences from decode parallel:**
+
+| Aspect | Decode Parallel | Encode Parallel |
+|--------|-----------------|-----------------|
+| `broadcast_result` | Often `False` (only rank 0 needs output) | `True` (all ranks need latents for diffusion) |
+| Patchify | Applied in merge (unpatchify) | Handled by parent `_encode()` before `tiled_encode()` |
+| Temporal chunking | Frame-by-frame | Chunk-based (e.g., 1 + 4n frames) |
+
 ## Testing
 Verify numerical consistency between:
 + vae_patch_parallel_size = 1
@@ -272,18 +440,20 @@ When vae_patch_parallel_size is larger than the DiT world size, it will automati
 
 Complete examples in the codebase:
 
-| Model | Path | Notes |
-|-------|------|-------|
-| **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | Distributed AutoencoderKL |
-| **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | Distributed AutoencoderKLWan |
-| **Qwen-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py` | Distributed AutoencoderKLQwenImage |
+| Model | Path | Decode Parallel | Encode Parallel |
+|-------|------|-----------------|-----------------|
+| **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | ✅ | ❌ |
+| **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | ✅ | ✅ |
+| **Qwen-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py` | ✅ | ❌ |
 
 ---
 
 ## Summary
 
-Adding Vae Patch Parallel support to diffusion model:
+Adding VAE Patch Parallel support to diffusion model:
 
-1. **Implement Distributed Vae** - mainly copy from `diffusers` tiled_decode, and refactor into split/exec/merge
-2. **Change vae model in pipeline to Distributed Vae**
-3. **Test** - Verify with `tensor_parallel_size=N` quality
+1. **Implement Distributed VAE** - Inherit from base VAE class and `DistributedVaeMixin`
+2. **Decode Parallel** - Refactor `tiled_decode` into `tile_split`/`tile_exec`/`tile_merge`
+3. **Encode Parallel** (optional) - Implement `encode_tile_split`/`encode_tile_exec`/`encode_tile_merge` for I2V models
+4. **Change VAE model in pipeline** - Use the distributed version
+5. **Test** - Verify numerical consistency with `vae_patch_parallel_size=1` vs `N`

docs/user_guide/diffusion/frame_interpolation.md

@@ -0,0 +1,92 @@
+# Frame Interpolation
+
+## Overview
+
+vLLM-Omni supports post-generation frame interpolation for supported video
+diffusion pipelines. This feature inserts synthesized intermediate frames
+between adjacent generated frames to improve temporal smoothness without
+rerunning the diffusion denoising loop.
+
+Frame interpolation runs in the diffusion worker post-processing path instead
+of the API server encoding path. This allows the interpolation step to reuse
+the worker's current accelerator device and keeps the FastAPI event loop free
+from heavy synchronous PyTorch work.
+
+For an input video with `N` generated frames and interpolation exponent `exp`,
+the output frame count is:
+
+```text
+(N - 1) * 2**exp + 1
+```
+
+The output FPS is multiplied by `2**exp` so the clip duration remains close to
+the original generated video.
+
+## Supported Pipelines
+
+Frame interpolation is currently supported for:
+
+- `WanPipeline` (Wan2.2 text-to-video)
+- `WanImageToVideoPipeline`
+- `Wan22TI2VPipeline`
+
+## Request Parameters
+
+The video APIs `/v1/videos` and `/v1/videos/sync` accept:
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `enable_frame_interpolation` | bool | `false` | Enable post-generation frame interpolation |
+| `frame_interpolation_exp` | int | `1` | Interpolation exponent. `1=2x`, `2=4x`, etc. |
+| `frame_interpolation_scale` | float | `1.0` | RIFE inference scale |
+| `frame_interpolation_model_path` | str | `None` | Local directory or Hugging Face repo ID containing `flownet.pkl` |
+
+## Execution Flow
+
+For supported Wan2.2 pipelines, the execution order is:
+
+1. Diffusion worker finishes denoising and decodes the raw video tensor.
+2. Worker-side model-specific post-processing runs.
+3. If frame interpolation is enabled, RIFE interpolates the decoded video
+   tensor on the worker side and records a FPS multiplier in `custom_output`.
+4. The API server receives the already-interpolated video and only performs
+   MP4 export.
+
+This design keeps interpolation close to the generated tensor and avoids
+introducing another heavyweight GPU context in the API server process.
+
+## Example
+
+Start the server:
+
+```bash
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091
+```
+
+Run a sync request with interpolation enabled:
+
+```bash
+curl -X POST http://localhost:8091/v1/videos/sync \
+  -F "prompt=A dog running through a park" \
+  -F "num_frames=81" \
+  -F "width=832" \
+  -F "height=480" \
+  -F "fps=16" \
+  -F "num_inference_steps=40" \
+  -F "guidance_scale=1.0" \
+  -F "guidance_scale_2=1.0" \
+  -F "enable_frame_interpolation=true" \
+  -F "frame_interpolation_exp=1" \
+  -F "frame_interpolation_scale=1.0" \
+  -F "seed=42" \
+  -o sync_t2v_interpolated.mp4
+```
+
+## Notes
+
+- This is a post-processing feature. It does not modify the diffusion denoising
+  schedule.
+- Higher interpolation exponents increase post-processing time and memory usage.
+- If the interpolation model weights are not available locally,
+  `frame_interpolation_model_path` may point to a Hugging Face repo containing
+  `flownet.pkl`.

docs/user_guide/examples/online_serving/text_to_video.md

@@ -123,6 +123,9 @@ curl -X POST http://localhost:8091/v1/videos \
   -F "guidance_scale_2=4.0" \
   -F "boundary_ratio=0.875" \
   -F "flow_shift=5.0" \
+  -F "enable_frame_interpolation=true" \
+  -F "frame_interpolation_exp=1" \
+  -F "frame_interpolation_scale=1.0" \
   -F "seed=42"
 ```
 
@@ -145,6 +148,35 @@ curl -X POST http://localhost:8091/v1/videos \
 | `flow_shift`          | float  | None    | Scheduler flow shift (Wan2.2)                    |
 | `seed`                | int    | None    | Random seed (reproducible)                       |
 | `lora`                | object | None    | LoRA configuration                               |
+| `enable_frame_interpolation` | bool | false | Enable RIFE frame interpolation before MP4 encoding |
+| `frame_interpolation_exp` | int | 1 | Interpolation exponent; 1=2x temporal resolution, 2=4x |
+| `frame_interpolation_scale` | float | 1.0 | RIFE inference scale; use 0.5 for high-resolution inputs |
+| `frame_interpolation_model_path` | str | None | Local directory or Hugging Face repo ID with `flownet.pkl`; defaults to `elfgum/RIFE-4.22.lite` |
+
+## Frame Interpolation
+
+Frame interpolation is an optional post-processing step for `/v1/videos` and
+`/v1/videos/sync`. It synthesizes intermediate frames between generated frames
+without rerunning the diffusion model. If the generated video has `N` frames,
+the interpolated output frame count is `(N - 1) * 2**exp + 1`. The encoder FPS
+is multiplied by `2**exp` so the output duration remains close to the original.
+
+Frame interpolation runs in the diffusion worker post-processing path instead of
+the API server encoding path, so it can reuse the worker's current accelerator
+device without blocking the FastAPI event loop.
+
+Example: generate 5 frames and interpolate to 9 frames:
+
+```bash
+curl -X POST http://localhost:8091/v1/videos/sync \
+  -F "prompt=A dog running through a park" \
+  -F "num_frames=5" \
+  -F "fps=8" \
+  -F "enable_frame_interpolation=true" \
+  -F "frame_interpolation_exp=1" \
+  -F "frame_interpolation_scale=1.0" \
+  -o sync_t2v_interpolated.mp4
+```
 
 ## Create Response Format
 

tests/diffusion/distributed/test_autoencoder_kl_wan.py

@@ -0,0 +1,43 @@
+import pytest
+import torch
+
+from vllm_omni.diffusion.distributed.autoencoders import autoencoder_kl_wan as wan_vae_module
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _DummyOmniAutoencoderKLWan(OmniAutoencoderKLWan):
+    def __init__(self, *, dtype: torch.dtype):
+        torch.nn.Module.__init__(self)
+        self.register_parameter("dummy_weight", torch.nn.Parameter(torch.ones(1, dtype=dtype)))
+
+
+def test_wan_vae_execution_context_handles_fp32():
+    model = _DummyOmniAutoencoderKLWan(dtype=torch.float32)
+    with model._execution_context():
+        output = model.dummy_weight + 1
+    assert output.dtype == torch.float32
+
+
+def test_wan_vae_execution_context_handles_bf16():
+    model = _DummyOmniAutoencoderKLWan(dtype=torch.bfloat16)
+    with model._execution_context():
+        output = model.dummy_weight + 1
+    assert output.dtype == torch.bfloat16
+
+
+def test_wan_vae_execution_context_uses_platform_autocast(mocker):
+    sentinel = object()
+    platform = mocker.Mock()
+    platform.create_autocast_context.return_value = sentinel
+    mocker.patch.object(wan_vae_module, "current_omni_platform", platform)
+
+    model = _DummyOmniAutoencoderKLWan(dtype=torch.bfloat16)
+
+    assert model._execution_context() is sentinel
+    platform.create_autocast_context.assert_called_once_with(
+        device_type=model.dummy_weight.device.type,
+        dtype=torch.bfloat16,
+        enabled=True,
+    )

tests/diffusion/distributed/test_autoencoder_kl_wan_encode.py

@@ -0,0 +1,273 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Unit tests for DistributedAutoencoderKLWan encode parallel (CPU-only)."""
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.cpu, pytest.mark.core_model]
+
+
+class _DummyConfig:
+    def __init__(self, patch_size=None, scale_factor_temporal=4):
+        self.patch_size = patch_size
+        self.scale_factor_temporal = scale_factor_temporal
+
+
+class _DummyWanVae:
+    """Minimal mock of DistributedAutoencoderKLWan for testing encode_tile_split."""
+
+    def __init__(
+        self,
+        config=None,
+        spatial_compression_ratio=8,
+        tile_sample_min_height=256,
+        tile_sample_min_width=256,
+        tile_sample_stride_height=192,
+        tile_sample_stride_width=192,
+    ):
+        self.config = config or _DummyConfig()
+        self.spatial_compression_ratio = spatial_compression_ratio
+        self.tile_sample_min_height = tile_sample_min_height
+        self.tile_sample_min_width = tile_sample_min_width
+        self.tile_sample_stride_height = tile_sample_stride_height
+        self.tile_sample_stride_width = tile_sample_stride_width
+        self.dtype = torch.float32
+
+        # Mock caches
+        self._enc_feat_map = None
+        self._enc_conv_idx = [0]
+
+    def clear_cache(self):
+        self._enc_feat_map = None
+        self._enc_conv_idx = [0]
+
+    def encoder(self, x, feat_cache=None, feat_idx=None):  # noqa: ARG002
+        # Simple mock: just return the input
+        return x
+
+    def quant_conv(self, x):
+        return x
+
+    def blend_v(self, _a, b, _blend_extent):
+        return b
+
+    def blend_h(self, _a, b, _blend_extent):
+        return b
+
+
+def _import_encode_tile_split():
+    """Import the encode_tile_split method from the module."""
+    from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
+        DistributedAutoencoderKLWan,
+    )
+
+    return DistributedAutoencoderKLWan.encode_tile_split
+
+
+def _import_encode_tile_exec():
+    from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
+        DistributedAutoencoderKLWan,
+    )
+
+    return DistributedAutoencoderKLWan.encode_tile_exec
+
+
+def _import_encode_tile_merge():
+    from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
+        DistributedAutoencoderKLWan,
+    )
+
+    return DistributedAutoencoderKLWan.encode_tile_merge
+
+
+class TestEncodeTileSplit:
+    """Tests for encode_tile_split method."""
+
+    def test_basic_split_without_patch_size(self):
+        """Test basic tile splitting without patch_size."""
+        encode_tile_split = _import_encode_tile_split()
+
+        vae = _DummyWanVae(
+            config=_DummyConfig(patch_size=None, scale_factor_temporal=4),
+            spatial_compression_ratio=8,
+            tile_sample_min_height=256,
+            tile_sample_min_width=256,
+            tile_sample_stride_height=192,
+            tile_sample_stride_width=192,
+        )
+
+        # Input: (B, C, T, H, W) = (1, 3, 5, 256, 256)
+        x = torch.randn(1, 3, 5, 256, 256)
+
+        tiletask_list, grid_spec = encode_tile_split(vae, x)
+
+        # With stride 192 and input size 256, we should get:
+        # Height: ceil(256/192) = 2 positions (0, 192) but 192+256 > 256, so only 1
+        # Actually for i in range(0, 256, 192): i = 0, 192 but 192 is out of bounds
+        # So we get 1x1 grid
+        assert len(tiletask_list) >= 1
+        assert grid_spec.grid_shape[0] >= 1
+        assert grid_spec.grid_shape[1] >= 1
+
+        # Check temporal chunking: 5 frames -> 1 + (5-1)//4 = 2 chunks
+        first_task = tiletask_list[0]
+        assert len(first_task.tensor) == 2  # 2 temporal chunks
+
+    def test_split_with_patch_size_scales_coordinates(self):
+        """Test that patch_size properly scales tile coordinates."""
+        encode_tile_split = _import_encode_tile_split()
+
+        # Without patch_size
+        vae_no_patch = _DummyWanVae(
+            config=_DummyConfig(patch_size=None, scale_factor_temporal=4),
+            spatial_compression_ratio=8,
+            tile_sample_min_height=256,
+            tile_sample_min_width=256,
+            tile_sample_stride_height=128,
+            tile_sample_stride_width=128,
+        )
+
+        # With patch_size=2 (simulating patchified input)
+        vae_with_patch = _DummyWanVae(
+            config=_DummyConfig(patch_size=2, scale_factor_temporal=4),
+            spatial_compression_ratio=8,
+            tile_sample_min_height=256,
+            tile_sample_min_width=256,
+            tile_sample_stride_height=128,
+            tile_sample_stride_width=128,
+        )
+
+        # Same patchified input size
+        x = torch.randn(1, 3, 5, 256, 256)
+
+        tasks_no_patch, _ = encode_tile_split(vae_no_patch, x)
+        tasks_with_patch, _ = encode_tile_split(vae_with_patch, x)
+
+        # With patch_size=2, stride becomes 128//2=64, so more tiles
+        assert len(tasks_with_patch) >= len(tasks_no_patch)
+
+    def test_temporal_compression_from_config(self):
+        """Test that temporal compression ratio is read from config."""
+        encode_tile_split = _import_encode_tile_split()
+
+        # temporal_compression=4 (default)
+        vae_4x = _DummyWanVae(
+            config=_DummyConfig(scale_factor_temporal=4),
+            tile_sample_min_height=512,
+            tile_sample_min_width=512,
+            tile_sample_stride_height=512,
+            tile_sample_stride_width=512,
+        )
+
+        # temporal_compression=2
+        vae_2x = _DummyWanVae(
+            config=_DummyConfig(scale_factor_temporal=2),
+            tile_sample_min_height=512,
+            tile_sample_min_width=512,
+            tile_sample_stride_height=512,
+            tile_sample_stride_width=512,
+        )
+
+        # 9 frames input
+        x = torch.randn(1, 3, 9, 512, 512)
+
+        tasks_4x, _ = encode_tile_split(vae_4x, x)
+        tasks_2x, _ = encode_tile_split(vae_2x, x)
+
+        # With 4x compression: 1 + (9-1)//4 = 3 chunks
+        assert len(tasks_4x[0].tensor) == 3
+
+        # With 2x compression: 1 + (9-1)//2 = 5 chunks
+        assert len(tasks_2x[0].tensor) == 5
+
+    def test_grid_spec_latent_dimensions(self):
+        """Test that grid_spec contains correct latent dimensions."""
+        encode_tile_split = _import_encode_tile_split()
+
+        vae = _DummyWanVae(
+            config=_DummyConfig(patch_size=None),
+            spatial_compression_ratio=8,
+            tile_sample_min_height=512,
+            tile_sample_min_width=512,
+            tile_sample_stride_height=512,
+            tile_sample_stride_width=512,
+        )
+
+        # Input: 512x512 with compression 8 -> 64x64 latent
+        x = torch.randn(1, 3, 5, 512, 512)
+
+        _, grid_spec = encode_tile_split(vae, x)
+
+        assert grid_spec.tile_spec["latent_height"] == 64
+        assert grid_spec.tile_spec["latent_width"] == 64
+
+
+class TestEncodeTileExec:
+    """Tests for encode_tile_exec method."""
+
+    def test_basic_exec(self):
+        """Test basic tile execution."""
+        encode_tile_exec = _import_encode_tile_exec()
+
+        vae = _DummyWanVae()
+
+        from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+            TileTask,
+        )
+
+        # Create a simple task with 2 temporal chunks
+        tile1 = torch.randn(1, 3, 1, 32, 32)
+        tile2 = torch.randn(1, 3, 4, 32, 32)
+        task = TileTask(tile_id=0, grid_coord=(0, 0), tensor=[tile1, tile2])
+
+        result = encode_tile_exec(vae, task)
+
+        # Result should concatenate temporal dimension
+        assert result.shape[2] == 5  # 1 + 4 frames
+
+
+class TestEncodeTileMerge:
+    """Tests for encode_tile_merge method."""
+
+    def test_basic_merge(self):
+        """Test basic tile merging."""
+        encode_tile_merge = _import_encode_tile_merge()
+
+        vae = _DummyWanVae()
+
+        from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+            GridSpec,
+        )
+
+        # Create 2x2 grid of tiles
+        tile_00 = torch.ones(1, 16, 2, 32, 32) * 0
+        tile_01 = torch.ones(1, 16, 2, 32, 32) * 1
+        tile_10 = torch.ones(1, 16, 2, 32, 32) * 2
+        tile_11 = torch.ones(1, 16, 2, 32, 32) * 3
+
+        coord_tensor_map = {
+            (0, 0): tile_00,
+            (0, 1): tile_01,
+            (1, 0): tile_10,
+            (1, 1): tile_11,
+        }
+
+        grid_spec = GridSpec(
+            split_dims=(3, 4),
+            grid_shape=(2, 2),
+            tile_spec={
+                "latent_height": 48,
+                "latent_width": 48,
+                "blend_height": 8,
+                "blend_width": 8,
+                "tile_latent_stride_height": 24,
+                "tile_latent_stride_width": 24,
+            },
+        )
+
+        result = encode_tile_merge(vae, coord_tensor_map, grid_spec)
+
+        # Output should be (1, 16, 2, 48, 48)
+        assert result.shape == (1, 16, 2, 48, 48)

tests/diffusion/distributed/test_distributed_vae_executor.py

@@ -59,9 +59,9 @@ class E2EOperator:
 class DummyMixin(DistributedVaeMixin):
     def __init__(self):
         self.use_tiling = True
-        self.distributed_decoder = MagicMock()
-        self.distributed_decoder.parallel_size = 2
-        self.distributed_decoder.group = None
+        self.distributed_executor = MagicMock()
+        self.distributed_executor.parallel_size = 2
+        self.distributed_executor.group = None
 
 
 @pytest.fixture(autouse=True)

tests/diffusion/models/qwen_image/test_qwen_image_edit_plus.py

@@ -0,0 +1,38 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+from pathlib import Path
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+from PIL import Image
+
+from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit_plus import (
+    get_qwen_image_edit_plus_pre_process_func,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def test_qwen_image_edit_plus_rejects_too_many_input_images(tmp_path: Path):
+    vae_dir = tmp_path / "vae"
+    vae_dir.mkdir()
+    # Keep the mock config intentionally minimal: this test only needs the
+    # fields touched during pre-process initialization.
+    (vae_dir / "config.json").write_text(json.dumps({"z_dim": 16}))
+
+    pre_process = get_qwen_image_edit_plus_pre_process_func(SimpleNamespace(model=str(tmp_path)))
+    image = Image.fromarray(np.zeros((32, 32, 3), dtype=np.uint8))
+    request = SimpleNamespace(
+        prompts=[
+            {
+                "prompt": "combine",
+                "multi_modal_data": {"image": [image, image, image, image, image]},
+            }
+        ],
+        sampling_params=SimpleNamespace(height=None, width=None),
+    )
+
+    with pytest.raises(ValueError, match=r"At most 4 images are supported by this model"):
+        pre_process(request)

tests/diffusion/models/qwen_image/test_qwen_image_max_sequence_length.py

@@ -0,0 +1,260 @@
+import inspect
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import (
+    QwenImagePipeline,
+)
+from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit import (
+    QwenImageEditPipeline,
+)
+from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit_plus import (
+    QwenImageEditPlusPipeline,
+)
+from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_layered import (
+    QwenImageLayeredPipeline,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _RejectingTextEncoder:
+    dtype = torch.float32
+
+    def __call__(self, *args, **kwargs):
+        raise AssertionError("text encoder should not run for prompts that exceed max_sequence_length")
+
+
+class _FakeModelInputs:
+    def __init__(self, total_sequence_length: int):
+        attention_mask = torch.ones((1, total_sequence_length), dtype=torch.long)
+        self.input_ids = attention_mask.clone()
+        self.attention_mask = attention_mask
+        self.pixel_values = None
+        self.image_grid_thw = None
+
+    def to(self, device):
+        return self
+
+
+class _FakeTokenizer:
+    def __init__(self, total_sequence_length: int | list[int]):
+        if isinstance(total_sequence_length, list):
+            self.total_sequence_lengths = list(total_sequence_length)
+        else:
+            self.total_sequence_lengths = [total_sequence_length]
+
+    def __call__(self, *args, **kwargs):
+        if len(self.total_sequence_lengths) > 1:
+            total_sequence_length = self.total_sequence_lengths.pop(0)
+        else:
+            total_sequence_length = self.total_sequence_lengths[0]
+        return _FakeModelInputs(total_sequence_length)
+
+
+class _FakeProcessor(_FakeTokenizer):
+    pass
+
+
+class _FakeScheduler:
+    def __init__(self):
+        self.begin_index = None
+
+    def set_begin_index(self, begin_index: int):
+        self.begin_index = begin_index
+
+
+PIPELINE_CASES = [
+    pytest.param(QwenImagePipeline, 34, "tokenizer", id="qwen-image"),
+    pytest.param(QwenImageLayeredPipeline, 34, "tokenizer", id="qwen-image-layered"),
+    pytest.param(QwenImageEditPipeline, 64, "processor", id="qwen-image-edit"),
+    pytest.param(QwenImageEditPlusPipeline, 64, "processor", id="qwen-image-edit-plus"),
+]
+
+
+def _make_pipeline(
+    pipeline_class: type,
+    *,
+    total_sequence_length: int,
+    drop_idx: int,
+    input_kind: str,
+):
+    pipeline = object.__new__(pipeline_class)
+    nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.text_encoder = _RejectingTextEncoder()
+    pipeline.tokenizer_max_length = 1024
+    pipeline.prompt_template_encode = "{}"
+    pipeline.prompt_template_encode_start_idx = drop_idx
+    pipeline.tokenizer = _FakeTokenizer([total_sequence_length, 0])
+    if input_kind == "processor":
+        pipeline.processor = _FakeProcessor(total_sequence_length)
+    return pipeline
+
+
+@pytest.mark.parametrize(("pipeline_class", "drop_idx", "input_kind"), PIPELINE_CASES)
+def test_encode_prompt_rejects_prompt_longer_than_default_max_sequence_length(
+    pipeline_class: type,
+    drop_idx: int,
+    input_kind: str,
+):
+    pipeline = _make_pipeline(
+        pipeline_class,
+        total_sequence_length=1025,
+        drop_idx=drop_idx,
+        input_kind=input_kind,
+    )
+
+    with pytest.raises(ValueError, match=r"got 1025 tokens, but `max_sequence_length` is 1024"):
+        pipeline.encode_prompt(prompt="prompt")
+
+
+@pytest.mark.parametrize(("pipeline_class", "drop_idx", "input_kind"), PIPELINE_CASES)
+def test_encode_prompt_rejects_prompt_longer_than_explicit_max_sequence_length(
+    pipeline_class: type,
+    drop_idx: int,
+    input_kind: str,
+):
+    pipeline = _make_pipeline(
+        pipeline_class,
+        total_sequence_length=17,
+        drop_idx=drop_idx,
+        input_kind=input_kind,
+    )
+
+    with pytest.raises(ValueError, match=r"got 17 tokens, but `max_sequence_length` is 16"):
+        pipeline.encode_prompt(prompt="prompt", max_sequence_length=16)
+
+
+def test_prepare_encode_defaults_to_tokenizer_max_length():
+    pipeline = object.__new__(QwenImagePipeline)
+    nn.Module.__init__(pipeline)
+    pipeline.tokenizer_max_length = 1024
+    pipeline.vae_scale_factor = 8
+    pipeline.default_sample_size = 128
+    pipeline.scheduler = _FakeScheduler()
+    pipeline._extract_prompts = lambda prompts: (["prompt"], None)
+
+    captured = {}
+
+    def _fake_prepare_generation_context(**kwargs):
+        captured["max_sequence_length"] = kwargs["max_sequence_length"]
+        embeds = torch.ones((1, 1, 1))
+        mask = torch.ones((1, 1), dtype=torch.long)
+        return {
+            "prompt_embeds": embeds,
+            "prompt_embeds_mask": mask,
+            "negative_prompt_embeds": None,
+            "negative_prompt_embeds_mask": None,
+            "latents": embeds,
+            "timesteps": torch.tensor([1]),
+            "do_true_cfg": False,
+            "guidance": None,
+            "img_shapes": [[(1, 1, 1)]],
+            "txt_seq_lens": [1],
+            "negative_txt_seq_lens": None,
+        }
+
+    pipeline._prepare_generation_context = _fake_prepare_generation_context
+    state = SimpleNamespace(
+        prompts=["prompt"],
+        sampling=SimpleNamespace(
+            height=None,
+            width=None,
+            num_inference_steps=None,
+            sigmas=None,
+            guidance_scale_provided=False,
+            num_outputs_per_prompt=0,
+            generator=None,
+            true_cfg_scale=None,
+            max_sequence_length=None,
+        ),
+    )
+
+    pipeline.prepare_encode(state)
+
+    assert captured["max_sequence_length"] == 1024
+
+
+@pytest.mark.parametrize(
+    ("pipeline_class", "drop_idx"),
+    [
+        pytest.param(QwenImageEditPipeline, 64, id="qwen-image-edit"),
+        pytest.param(QwenImageEditPlusPipeline, 64, id="qwen-image-edit-plus"),
+    ],
+)
+def test_edit_pipelines_validate_text_prompt_length_before_image_token_expansion(
+    pipeline_class: type,
+    drop_idx: int,
+):
+    pipeline = object.__new__(pipeline_class)
+    nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.text_encoder = _RejectingTextEncoder()
+    pipeline.tokenizer_max_length = 1024
+    pipeline.prompt_template_encode = "{}"
+    pipeline.prompt_template_encode_start_idx = drop_idx
+    pipeline.tokenizer = _FakeTokenizer([8, 0])
+    pipeline.processor = _FakeProcessor(drop_idx + 1500)
+
+    with pytest.raises(AssertionError, match="text encoder should not run"):
+        pipeline.encode_prompt(prompt="short prompt")
+
+
+@pytest.mark.parametrize(
+    "pipeline_class",
+    [
+        pytest.param(QwenImagePipeline, id="qwen-image"),
+        pytest.param(QwenImageLayeredPipeline, id="qwen-image-layered"),
+    ],
+)
+def test_qwen_generation_validator_excludes_template_suffix_from_budget(pipeline_class: type):
+    pipeline = object.__new__(pipeline_class)
+    nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.text_encoder = _RejectingTextEncoder()
+    pipeline.tokenizer_max_length = 1024
+    pipeline.prompt_template_encode = "{}"
+    pipeline.prompt_template_encode_start_idx = 34
+    pipeline.tokenizer = _FakeTokenizer([1029, 5])
+
+    with pytest.raises(AssertionError, match="text encoder should not run"):
+        pipeline.encode_prompt(prompt="boundary prompt")
+
+
+@pytest.mark.parametrize(
+    "pipeline_class",
+    [
+        pytest.param(QwenImageEditPipeline, id="qwen-image-edit"),
+        pytest.param(QwenImageEditPlusPipeline, id="qwen-image-edit-plus"),
+    ],
+)
+def test_qwen_edit_validator_excludes_image_placeholders_from_budget(pipeline_class: type):
+    pipeline = object.__new__(pipeline_class)
+    nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.text_encoder = _RejectingTextEncoder()
+    pipeline.tokenizer_max_length = 1024
+    pipeline.prompt_template_encode = "{}"
+    pipeline.prompt_template_encode_start_idx = 64
+    pipeline.tokenizer = _FakeTokenizer([30, 20])
+    pipeline.processor = _FakeProcessor(1500)
+
+    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/qwen_image/test_qwen_image_size_utils.py

@@ -0,0 +1,26 @@
+import pytest
+
+from vllm_omni.diffusion.utils.size_utils import (
+    normalize_min_aligned_size,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize(
+    ("height", "width", "expected"),
+    [
+        (1, 1, (16, 16)),
+        (15, 15, (16, 16)),
+        (17, 17, (16, 16)),
+        (31, 33, (16, 32)),
+        (64, 80, (64, 80)),
+    ],
+)
+def test_normalize_min_aligned_size_clamps_to_minimum_aligned_shape(height, width, expected):
+    assert normalize_min_aligned_size(height, width, alignment=16) == expected
+
+
+def test_normalize_min_aligned_size_rejects_invalid_alignment():
+    with pytest.raises(ValueError, match="positive alignment"):
+        normalize_min_aligned_size(16, 16, alignment=0)

tests/diffusion/test_diffusion_scheduler.py

@@ -6,7 +6,7 @@ from unittest.mock import Mock, patch
 
 import pytest
 
-from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.data import DiffusionOutput, OmniRequestError, normalize_omni_error
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched import (
@@ -245,6 +245,39 @@ class TestDiffusionEngine:
         assert output is expected
         engine.executor.add_req.assert_called_once_with(request)
 
+    def test_add_req_and_wait_for_response_wraps_executor_errors(self) -> None:
+        engine = DiffusionEngine.__new__(DiffusionEngine)
+        engine.scheduler = RequestScheduler()
+        engine.scheduler.initialize(Mock())
+        engine.executor = Mock()
+        engine.executor.add_req.side_effect = RuntimeError("worker oom")
+        engine._rpc_lock = threading.Lock()
+
+        request = _make_request("engine-error")
+
+        with pytest.raises(OmniRequestError, match="worker oom") as exc_info:
+            engine.add_req_and_wait_for_response(request)
+
+        assert exc_info.value.error_type == "RuntimeError"
+
+    def test_add_req_and_wait_for_response_raises_scheduler_error_when_empty_without_requests(self) -> None:
+        engine = DiffusionEngine.__new__(DiffusionEngine)
+        engine.scheduler = Mock()
+        engine.scheduler.add_request.return_value = "sched-1"
+        engine.scheduler.schedule.return_value = Mock(is_empty=True)
+        engine.scheduler.has_requests.return_value = False
+        engine.scheduler.pop_request_state.return_value = None
+        engine.executor = Mock()
+        engine._rpc_lock = threading.Lock()
+
+        request = _make_request("engine-empty")
+
+        with pytest.raises(OmniRequestError, match="no runnable requests") as exc_info:
+            engine.add_req_and_wait_for_response(request)
+
+        assert exc_info.value.error_type == "SchedulerError"
+        engine.scheduler.pop_request_state.assert_called_once_with("sched-1")
+
     def test_supports_scheduler_interface_injection(self) -> None:
         request = _make_request("engine_iface")
         expected = DiffusionOutput(output=None)
@@ -297,3 +330,15 @@ class TestDiffusionEngine:
 
         with pytest.raises(RuntimeError, match="Dummy run failed: boom"):
             engine._dummy_run()
+
+
+def test_normalize_omni_error_preserves_existing_metadata() -> None:
+    err = OmniRequestError("boom", status_code=409, error_type="Conflict")
+
+    normalized = normalize_omni_error(err, request_id="req-1", stage_id=3)
+
+    assert normalized is err
+    assert normalized.request_id == "req-1"
+    assert normalized.stage_id == 3
+    assert normalized.status_code == 409
+    assert normalized.error_type == "Conflict"

tests/diffusion/test_diffusion_step_pipeline.py

@@ -443,14 +443,42 @@ class TestWorker:
 class TestIPC:
     def test_pack_unpack_runner_output_shm(self):
         tensor = torch.zeros(300_000, dtype=torch.float32)
-        output = RunnerOutput(req_id="req-1", finished=True, result=DiffusionOutput(output=tensor))
+        log_probs = torch.ones(300_000, dtype=torch.float32)
+        output = RunnerOutput(
+            req_id="req-1",
+            finished=True,
+            result=DiffusionOutput(output=tensor, trajectory_log_probs=log_probs),
+        )
 
         packed = pack_diffusion_output_shm(output)
         assert isinstance(packed.result.output, dict)
         assert packed.result.output["__tensor_shm__"] is True
+        assert isinstance(packed.result.trajectory_log_probs, dict)
+        assert packed.result.trajectory_log_probs["__tensor_shm__"] is True
 
         unpacked = unpack_diffusion_output_shm(packed)
         assert isinstance(unpacked.result.output, torch.Tensor)
+        assert isinstance(unpacked.result.trajectory_log_probs, torch.Tensor)
+        torch.testing.assert_close(unpacked.result.output, tensor)
+        torch.testing.assert_close(unpacked.result.trajectory_log_probs, log_probs)
+
+    def test_pack_unpack_runner_output_shm_bfloat16(self):
+        tensor = torch.arange(600_000, dtype=torch.bfloat16)
+        output = RunnerOutput(
+            req_id="req-1",
+            finished=True,
+            result=DiffusionOutput(output=tensor),
+        )
+
+        packed = pack_diffusion_output_shm(output)
+        assert isinstance(packed.result.output, dict)
+        assert packed.result.output["__tensor_shm__"] is True
+        assert packed.result.output["torch_dtype"] == str(torch.bfloat16)
+        assert packed.result.output["numpy_dtype"] == "uint16"
+
+        unpacked = unpack_diffusion_output_shm(packed)
+        assert isinstance(unpacked.result.output, torch.Tensor)
+        assert unpacked.result.output.dtype == torch.bfloat16
         torch.testing.assert_close(unpacked.result.output, tensor)
 
 

tests/diffusion/test_inline_stage_diffusion_client.py

@@ -0,0 +1,122 @@
+from __future__ import annotations
+
+import asyncio
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient
+from vllm_omni.engine.stage_init_utils import StageMetadata
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def mock_engine():
+    with patch("vllm_omni.diffusion.inline_stage_diffusion_client.DiffusionEngine") as mock:
+        engine_instance = MagicMock()
+        mock.make_engine.return_value = engine_instance
+        yield engine_instance
+
+
+@pytest.fixture
+def client(mock_engine):
+    metadata = StageMetadata(
+        stage_id=0,
+        stage_type="diffusion",
+        engine_output_type="image",
+        is_comprehension=False,
+        requires_multimodal_data=False,
+        engine_input_source="prompt",
+        final_output=True,
+        final_output_type="image",
+        default_sampling_params={},
+        custom_process_input_func=None,
+        model_stage=None,
+        runtime_cfg=None,
+    )
+    with patch.object(InlineStageDiffusionClient, "_enrich_config"):
+        od_config = MagicMock(spec=OmniDiffusionConfig)
+        c = InlineStageDiffusionClient(model="test_model", od_config=od_config, metadata=metadata, batch_size=1)
+        yield c
+        c.shutdown()
+
+
+@pytest.mark.asyncio
+async def test_inline_dispatch_request_success(client, mock_engine):
+    mock_result = OmniRequestOutput.from_diffusion(request_id="req-1", images=[MagicMock()])
+    mock_engine.step.return_value = [mock_result]
+
+    sampling_params = OmniDiffusionSamplingParams()
+    await client.add_request_async("req-1", "A test prompt", sampling_params)
+
+    # Wait for the task to be processed
+    for _ in range(10):
+        output = client.get_diffusion_output_nowait()
+        if output is not None:
+            break
+        await asyncio.sleep(0.01)
+
+    assert output is not None
+    assert output.request_id == "req-1"
+    mock_engine.step.assert_called_once()
+    request = mock_engine.step.call_args.args[0]
+    assert request.prompts == ["A test prompt"]
+    assert request.request_ids == ["req-1"]
+
+
+@pytest.mark.asyncio
+async def test_inline_dispatch_batch_success(client, mock_engine):
+    mock_engine.step.return_value = [
+        OmniRequestOutput.from_diffusion(request_id="req-batch-0", images=[MagicMock()]),
+        OmniRequestOutput.from_diffusion(request_id="req-batch-1", images=[MagicMock()]),
+    ]
+
+    sampling_params = OmniDiffusionSamplingParams()
+    prompts = ["prompt-0", "prompt-1"]
+    await client.add_batch_request_async("req-batch", prompts, sampling_params)
+
+    for _ in range(10):
+        output = client.get_diffusion_output_nowait()
+        if output is not None:
+            break
+        await asyncio.sleep(0.01)
+
+    assert output is not None
+    assert output.request_id == "req-batch"
+    assert len(output.images) == 2
+    request = mock_engine.step.call_args.args[0]
+    assert request.prompts == prompts
+    assert request.request_ids == ["req-batch-0", "req-batch-1"]
+
+
+@pytest.mark.asyncio
+async def test_inline_dispatch_request_error(client, mock_engine):
+    mock_engine.step.side_effect = RuntimeError("Engine failure")
+
+    sampling_params = OmniDiffusionSamplingParams()
+    await client.add_request_async("req-err", "A test prompt", sampling_params)
+
+    for _ in range(10):
+        output = client.get_diffusion_output_nowait()
+        if output is not None:
+            break
+        await asyncio.sleep(0.01)
+
+    assert output is not None
+    assert output["request_id"] == "req-err"
+    assert output["error"] == "Engine failure"
+    assert output["error_type"] == "RuntimeError"
+
+
+def test_inline_shutdown(client, mock_engine):
+    assert not client._shutting_down
+
+    # Shutting down should cleanly cancel anything queued and close engine
+    client.shutdown()
+
+    assert client._shutting_down
+    mock_engine.close.assert_called_once()

tests/engine/test_async_omni_engine_stage_init.py

@@ -71,6 +71,67 @@ def test_initialize_stages_restores_device_visibility_after_diffusion_init(monke
             os.environ[env_var] = old_env
 
 
+def test_initialize_stages_uses_inline_diffusion_client_for_single_stage(monkeypatch):
+    """Single-stage diffusion init should request the inline client path."""
+    import vllm_omni.engine.async_omni_engine as engine_mod
+    from vllm_omni.platforms import current_omni_platform
+
+    engine = object.__new__(AsyncOmniEngine)
+    engine.model = "dummy-model"
+    engine.config_path = "dummy-config"
+    engine.num_stages = 1
+    engine.async_chunk = False
+    engine.diffusion_batch_size = 1
+    engine.stage_configs = [types.SimpleNamespace(stage_id=0, stage_type="diffusion")]
+
+    env_var = current_omni_platform.device_control_env_var
+    old_env = os.environ.get(env_var)
+    os.environ[env_var] = "0"
+
+    metadata = types.SimpleNamespace(
+        stage_id=0,
+        stage_type="diffusion",
+        runtime_cfg={"devices": "0"},
+        prompt_expand_func=None,
+    )
+    captured_use_inline = None
+    diffusion_client = types.SimpleNamespace(is_comprehension=False)
+
+    monkeypatch.setattr(engine_mod, "prepare_engine_environment", lambda: None)
+    monkeypatch.setattr(engine_mod, "load_omni_transfer_config_for_model", lambda *_: None)
+    monkeypatch.setattr(engine_mod, "extract_stage_metadata", lambda _cfg: metadata)
+    monkeypatch.setattr(engine_mod, "get_stage_connector_spec", lambda **_: {})
+    monkeypatch.setattr(engine_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None))
+    monkeypatch.setattr(engine_mod, "setup_stage_devices", lambda *_: None)
+    monkeypatch.setattr(engine_mod, "_inject_kv_stage_info", lambda *_: None)
+
+    def _fake_initialize_diffusion_stage(*args, **kwargs):
+        nonlocal captured_use_inline
+        captured_use_inline = kwargs.get("use_inline")
+        return diffusion_client
+
+    monkeypatch.setattr(engine_mod, "initialize_diffusion_stage", _fake_initialize_diffusion_stage)
+    monkeypatch.setattr(
+        engine_mod,
+        "finalize_initialized_stages",
+        lambda stage_clients, _input_processor: (
+            stage_clients,
+            [types.SimpleNamespace()],
+            [{"final_output_type": "image"}],
+        ),
+    )
+
+    try:
+        engine._initialize_stages(stage_init_timeout=1)
+    finally:
+        if old_env is None:
+            os.environ.pop(env_var, None)
+        else:
+            os.environ[env_var] = old_env
+
+    assert captured_use_inline is True
+
+
 def test_attach_llm_stage_uses_omni_input_preprocessor(monkeypatch):
     """Regression test for GLM-Image t2i preprocessing path.
 

tests/entrypoints/openai_api/test_image_server.py

@@ -115,7 +115,7 @@ class MockGenerationResult:
 class FakeAsyncOmni:
     """Fake AsyncOmni that yields a single diffusion output."""
 
-    def __init__(self):
+    def __init__(self, images=None):
         self.stage_configs = [
             SimpleNamespace(stage_type="llm"),
             SimpleNamespace(stage_type="diffusion"),
@@ -123,11 +123,12 @@ class FakeAsyncOmni:
         self.default_sampling_params_list = [SamplingParams(temperature=0.1), OmniDiffusionSamplingParams()]
         self.captured_sampling_params_list = None
         self.captured_prompt = None
+        self._images = images or [Image.new("RGB", (64, 64), color="green")]
 
     async def generate(self, prompt, request_id, sampling_params_list):
         self.captured_sampling_params_list = sampling_params_list
         self.captured_prompt = prompt
-        images = [Image.new("RGB", (64, 64), color="green")]
+        images = [img.copy() for img in self._images]
         yield MockGenerationResult(images)
 
 
@@ -170,8 +171,8 @@ def test_client(mock_async_diffusion):
     app.state.stage_configs = [SimpleNamespace(stage_type="diffusion")]
     app.state.diffusion_model_name = "Qwen/Qwen-Image"  # For models endpoint
     app.state.args = Namespace(
-        default_sampling_params='{"0": {"num_inference_steps":4, "guidance_scale":7.5}}',
-        max_generated_image_size=4096,  # 64*64
+        default_sampling_params='{"0": {"num_inference_steps":4, "guidance_scale":7.5, "generator_device":"cpu"}}',
+        max_generated_image_size=1024 * 1792,
     )
 
     return TestClient(app)
@@ -193,12 +194,34 @@ def async_omni_test_client():
         SimpleNamespace(stage_type="diffusion"),
     ]
     app.state.args = Namespace(
-        default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5}}',
+        default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5, "generator_device":"cpu"}}',
         max_generated_image_size=1048576,  # 1024*1024 to support resolution tests
     )
     return TestClient(app)
 
 
+@pytest.fixture
+def async_omni_rgba_test_client():
+    """Create test client with mocked AsyncOmni engine returning RGBA output."""
+    from fastapi import FastAPI
+
+    from vllm_omni.entrypoints.openai.api_server import router
+
+    app = FastAPI()
+    app.include_router(router)
+
+    app.state.engine_client = FakeAsyncOmni(images=[Image.new("RGBA", (64, 64), color=(0, 255, 0, 128))])
+    app.state.stage_configs = [
+        SimpleNamespace(stage_type="llm"),
+        SimpleNamespace(stage_type="diffusion"),
+    ]
+    app.state.args = Namespace(
+        default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5, "generator_device":"cpu"}}',
+        max_generated_image_size=1048576,
+    )
+    return TestClient(app)
+
+
 @pytest.fixture
 def async_omni_stage_configs_only_client():
     """Create test client with refactored AsyncOmni compatibility surface only."""
@@ -215,8 +238,8 @@ def async_omni_stage_configs_only_client():
     # Intentionally do not populate app.state.stage_configs. Refactored
     # AsyncOmni exposes stage_configs on the engine instance.
     app.state.args = Namespace(
-        default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5}}',
-        max_generated_image_size=4096,  # 64*64
+        default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5, "generator_device":"cpu"}}',
+        max_generated_image_size=1024 * 1792,
     )
     return TestClient(app)
 
@@ -363,6 +386,18 @@ def test_image_edits_async_omni_stage_configs_only(async_omni_stage_configs_only
     assert len(captured) == 2
 
 
+def test_generate_images_max_size_rejected(async_omni_test_client):
+    """Test that a size exceeding max_generated_image_size returns 400."""
+    response = async_omni_test_client.post(
+        "/v1/images/generations",
+        json={
+            "prompt": "a cat",
+            "size": "2048x2048",  # 4,194,304 pixels > max_generated_image_size (1,048,576)
+        },
+    )
+    assert response.status_code == 400
+
+
 def test_generate_multiple_images(test_client):
     """Test generating multiple images"""
     response = test_client.post(
@@ -592,6 +627,13 @@ def test_parameter_validation():
     with pytest.raises(ValueError):
         ImageGenerationRequest(prompt="test", guidance_scale=21.0)
 
+    # Invalid layers for layered models (must stay within the backend-supported range)
+    with pytest.raises(ValueError):
+        ImageGenerationRequest(prompt="test", layers=2)
+
+    with pytest.raises(ValueError):
+        ImageGenerationRequest(prompt="test", layers=11)
+
 
 # Pass-Through Tests
 
@@ -711,6 +753,86 @@ def test_image_edit_images_processing(async_omni_test_client):
     assert processed_images[1].size == (24, 24)
 
 
+def test_image_edit_rejects_multiple_images_when_model_does_not_support_them(async_omni_test_client):
+    img_bytes_1 = make_test_image_bytes((16, 16))
+    img_bytes_2 = make_test_image_bytes((32, 32))
+
+    engine = async_omni_test_client.app.state.engine_client
+    engine.get_diffusion_od_config = lambda: SimpleNamespace(supports_multimodal_inputs=False)
+
+    response = async_omni_test_client.post(
+        "/v1/images/edits",
+        files=[
+            ("image", img_bytes_1),
+            ("image", img_bytes_2),
+        ],
+        data={"prompt": "hello world."},
+    )
+
+    assert response.status_code == 400
+    assert (
+        response.json()["detail"] == "Received multiple input images. Only a single image is supported by this model."
+    )
+    assert engine.captured_prompt is None
+
+
+def test_image_edit_rejects_too_many_images_for_qwen_image_edit_2511(async_omni_test_client):
+    engine = async_omni_test_client.app.state.engine_client
+    engine.get_diffusion_od_config = lambda: SimpleNamespace(
+        supports_multimodal_inputs=True,
+        max_multimodal_image_inputs=4,
+    )
+
+    response = async_omni_test_client.post(
+        "/v1/images/edits",
+        files=[
+            ("image", make_test_image_bytes((16, 16))),
+            ("image", make_test_image_bytes((16, 16))),
+            ("image", make_test_image_bytes((16, 16))),
+            ("image", make_test_image_bytes((16, 16))),
+            ("image", make_test_image_bytes((16, 16))),
+        ],
+        data={"prompt": "hello world."},
+    )
+
+    assert response.status_code == 400
+    assert response.json()["detail"] == "Received 5 input images. At most 4 images are supported by this model."
+    assert engine.captured_prompt is None
+
+
+def test_image_edit_rejects_too_many_images_for_qwen_image_edit_2511_before_loading(
+    async_omni_test_client, monkeypatch: pytest.MonkeyPatch
+):
+    import vllm_omni.entrypoints.openai.api_server as api_server_module
+
+    engine = async_omni_test_client.app.state.engine_client
+    engine.get_diffusion_od_config = lambda: SimpleNamespace(
+        supports_multimodal_inputs=True,
+        max_multimodal_image_inputs=4,
+    )
+
+    def _fail_load(*args, **kwargs):
+        raise AssertionError("_load_input_images should not run for over-limit requests")
+
+    monkeypatch.setattr(api_server_module, "_load_input_images", _fail_load)
+
+    response = async_omni_test_client.post(
+        "/v1/images/edits",
+        files=[
+            ("image", make_test_image_bytes((16, 16))),
+            ("image", make_test_image_bytes((16, 16))),
+            ("image", make_test_image_bytes((16, 16))),
+            ("image", make_test_image_bytes((16, 16))),
+            ("image", make_test_image_bytes((16, 16))),
+        ],
+        data={"prompt": "hello world."},
+    )
+
+    assert response.status_code == 400
+    assert response.json()["detail"] == "Received 5 input images. At most 4 images are supported by this model."
+    assert engine.captured_prompt is None
+
+
 def test_image_edit_parameter_pass(async_omni_test_client):
     img_bytes_1 = make_test_image_bytes((16, 16))
 
@@ -816,35 +938,36 @@ def test_image_edit_invalid_resolution(async_omni_test_client):
 
 
 def test_image_edit_invalid_layers(async_omni_test_client):
-    """Test that invalid layers values (< 1) are rejected with 400."""
+    """Test that layered image edits reject out-of-range layers with 400."""
     img_bytes = make_test_image_bytes((16, 16))
 
-    # Test layers = 0
+    # Test layers below the supported range
     response = async_omni_test_client.post(
         "/v1/images/edits",
         files=[("image", img_bytes)],
         data={
             "prompt": "test",
-            "layers": 0,
+            "layers": 2,
         },
     )
     assert response.status_code == 400
     detail = response.json()["detail"]
     assert "Invalid layers" in detail
-    assert "layers must be >= 1" in detail
+    assert "layers must be between 3 and 10 inclusive" in detail
 
-    # Test layers = -1
+    # Test layers above the supported range
     response = async_omni_test_client.post(
         "/v1/images/edits",
         files=[("image", img_bytes)],
         data={
             "prompt": "test",
-            "layers": -1,
+            "layers": 11,
         },
     )
     assert response.status_code == 400
     detail = response.json()["detail"]
     assert "Invalid layers" in detail
+    assert "layers must be between 3 and 10 inclusive" in detail
 
 
 def test_image_edit_resolution_and_size_conflict(async_omni_test_client):
@@ -888,6 +1011,7 @@ def test_image_edit_parameter_default(async_omni_test_client):
     assert captured_sampling_params.num_outputs_per_prompt == 1
     assert captured_sampling_params.num_inference_steps == 4
     assert captured_sampling_params.guidance_scale == 7.5
+    assert captured_sampling_params.generator_device == "cpu"
 
     # Test that a size exceeding max_generated_image_size returns 400
     response = async_omni_test_client.post(
@@ -921,13 +1045,15 @@ def test_image_edit_parameter_default_single_stage(test_client):
     assert captured_sampling_params.num_outputs_per_prompt == 1
     assert captured_sampling_params.num_inference_steps == 4
     assert captured_sampling_params.guidance_scale == 7.5
+    assert captured_sampling_params.generator_device == "cpu"
 
+    # Size exceeding max_generated_image_size (1024*1792) returns 400
     response = test_client.post(
         "/v1/images/edits",
         files=[("image", img_bytes_1)],
         data={
             "prompt": "hello world.",
-            "size": "96x96",
+            "size": "2048x2048",
         },
     )
     assert response.status_code == 400
@@ -977,6 +1103,27 @@ def test_image_edit_compression_jpeg(test_client):
     assert len(img_bytes_50) < len(img_bytes_100)
 
 
+def test_image_edit_rgba_output_converts_to_jpeg(async_omni_rgba_test_client):
+    img_bytes_1 = make_test_image_bytes((16, 16))
+
+    response = async_omni_rgba_test_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes_1)],
+        data={
+            "prompt": "hello world.",
+            "output_format": "jpeg",
+        },
+    )
+    assert response.status_code == 200
+
+    data = response.json()
+    img_bytes = base64.b64decode(data["data"][0]["b64_json"])
+    img = Image.open(io.BytesIO(img_bytes))
+    assert img.format.lower() == "jpeg"
+    assert img.mode == "RGB"
+    assert data["output_format"] == "jpeg"
+
+
 def test_image_edit_compression_png(async_omni_test_client):
     img_bytes_1 = make_test_image_bytes((16, 16))
     # uploadfile with image key

tests/entrypoints/openai_api/test_video_api_utils.py

@@ -0,0 +1,146 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for OpenAI-compatible video API encoding helpers."""
+
+import numpy as np
+import pytest
+import torch
+
+from vllm_omni.diffusion.postprocess import rife_interpolator
+from vllm_omni.entrypoints.openai import video_api_utils
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _install_fake_video_mux(monkeypatch, mux_calls):
+    def _fake_mux_video_audio_bytes(frames, audio, fps, audio_sample_rate, video_codec_options=None):
+        mux_calls.append(
+            {
+                "frames": frames,
+                "audio": audio,
+                "fps": fps,
+                "audio_sample_rate": audio_sample_rate,
+                "video_codec_options": video_codec_options,
+            }
+        )
+        return b"fake-video"
+
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.utils.media_utils.mux_video_audio_bytes",
+        _fake_mux_video_audio_bytes,
+    )
+
+
+def _install_fake_export_to_video(monkeypatch, export_calls):
+    def _fake_export_to_video(frames, output_path, fps):
+        export_calls.append(
+            {
+                "frames": frames,
+                "output_path": output_path,
+                "fps": fps,
+            }
+        )
+        with open(output_path, "wb") as f:
+            f.write(b"fake-video")
+
+    monkeypatch.setattr("diffusers.utils.export_to_video", _fake_export_to_video)
+
+
+def test_encode_video_bytes_exports_frames_without_interpolation(monkeypatch):
+    export_calls = []
+    _install_fake_export_to_video(monkeypatch, export_calls)
+
+    frames = [np.full((2, 2, 3), fill_value=i / 5, dtype=np.float32) for i in range(5)]
+    video_bytes = video_api_utils._encode_video_bytes(
+        frames,
+        fps=8,
+    )
+
+    assert video_bytes == b"fake-video"
+    assert len(export_calls) == 1
+    assert len(export_calls[0]["frames"]) == 5
+    assert export_calls[0]["frames"][0].shape == (2, 2, 3)
+    assert export_calls[0]["fps"] == 8
+
+
+def test_encode_video_bytes_uses_array_fast_path(monkeypatch):
+    export_calls = []
+    _install_fake_export_to_video(monkeypatch, export_calls)
+    monkeypatch.setattr(
+        video_api_utils,
+        "_normalize_frames",
+        lambda frames: (_ for _ in ()).throw(AssertionError("ndarray input should not go through frame list path")),
+    )
+
+    video = np.linspace(0.0, 1.0, num=5 * 2 * 2 * 3, dtype=np.float32).reshape(5, 2, 2, 3)
+    video_bytes = video_api_utils._encode_video_bytes(video, fps=12)
+
+    assert video_bytes == b"fake-video"
+    assert len(export_calls) == 1
+    assert len(export_calls[0]["frames"]) == 5
+    assert export_calls[0]["frames"][0].flags["C_CONTIGUOUS"]
+    assert export_calls[0]["fps"] == 12.0
+
+
+def test_encode_video_bytes_without_audio_uses_diffusers_export(monkeypatch):
+    export_calls = []
+    _install_fake_export_to_video(monkeypatch, export_calls)
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.utils.media_utils.mux_video_audio_bytes",
+        lambda *args, **kwargs: (_ for _ in ()).throw(
+            AssertionError("no-audio path should not use mux_video_audio_bytes")
+        ),
+    )
+
+    video = np.linspace(0.0, 1.0, num=4 * 2 * 2 * 3, dtype=np.float32).reshape(4, 2, 2, 3)
+    video_bytes = video_api_utils._encode_video_bytes(video, fps=10)
+
+    assert video_bytes == b"fake-video"
+    assert len(export_calls) == 1
+    assert export_calls[0]["fps"] == 10
+    assert len(export_calls[0]["frames"]) == 4
+
+
+def test_rife_model_inference_runs_on_dummy_tensors():
+    model = rife_interpolator.Model().eval()
+    img0 = torch.rand(1, 3, 32, 32)
+    img1 = torch.rand(1, 3, 32, 32)
+
+    output = model.inference(img0, img1, scale=1.0)
+
+    assert output.shape == (1, 3, 32, 32)
+    assert torch.isfinite(output).all()
+
+
+def test_frame_interpolator_runs_actual_torch_tensor_path(monkeypatch):
+    model = rife_interpolator.Model().eval()
+    interpolator = rife_interpolator.FrameInterpolator()
+    monkeypatch.setattr(interpolator, "_ensure_model_loaded", lambda preferred_device=None: model)
+
+    video = torch.zeros(1, 3, 2, 32, 32)
+    output_video, multiplier = interpolator.interpolate_tensor(video, exp=1, scale=1.0)
+
+    assert multiplier == 2
+    assert output_video.shape == (1, 3, 3, 32, 32)
+    assert torch.isfinite(output_video).all()
+
+
+def test_frame_interpolator_uses_platform_device_when_tensor_is_cpu(monkeypatch):
+    chosen_devices = []
+    model = rife_interpolator.Model().eval()
+
+    def _fake_ensure_model_loaded(*, preferred_device=None):
+        chosen_devices.append(preferred_device)
+        return model
+
+    interpolator = rife_interpolator.FrameInterpolator()
+    monkeypatch.setattr(interpolator, "_ensure_model_loaded", _fake_ensure_model_loaded)
+    monkeypatch.setattr(model.flownet, "to", lambda device: model.flownet)
+    monkeypatch.setattr(rife_interpolator, "_select_torch_device", lambda: torch.device("cuda"))
+
+    video = torch.zeros(1, 3, 2, 32, 32)
+    output_video, multiplier = interpolator.interpolate_tensor(video, exp=1, scale=1.0)
+
+    assert chosen_devices == [torch.device("cuda")]
+    assert multiplier == 2
+    assert output_video.shape == (1, 3, 3, 32, 32)

tests/entrypoints/openai_api/test_video_server.py

@@ -14,11 +14,12 @@ import time
 from types import SimpleNamespace
 
 import pytest
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
 from fastapi.testclient import TestClient
 from PIL import Image
 from pytest_mock import MockerFixture
 
+from vllm_omni.diffusion.data import OmniRequestError
 from vllm_omni.entrypoints.openai import api_server
 from vllm_omni.entrypoints.openai.api_server import router
 from vllm_omni.entrypoints.openai.protocol.videos import (
@@ -29,22 +30,39 @@ from vllm_omni.entrypoints.openai.protocol.videos import (
 from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo
 from vllm_omni.entrypoints.openai.storage import LocalStorageManager
 from vllm_omni.entrypoints.openai.stores import AsyncDictStore, TaskRegistry
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 class MockVideoResult:
-    def __init__(self, videos, audios=None, sample_rate=None):
+    def __init__(
+        self,
+        videos,
+        audios=None,
+        sample_rate=None,
+        custom_output=None,
+        stage_durations=None,
+        peak_memory_mb=0.0,
+    ):
         self.multimodal_output = {"video": videos}
         if audios is not None:
             self.multimodal_output["audio"] = audios
         if sample_rate is not None:
             self.multimodal_output["audio_sample_rate"] = sample_rate
+        self._custom_output = custom_output or {}
+        self.stage_durations = stage_durations or {}
+        self.peak_memory_mb = peak_memory_mb
+
+    @property
+    def custom_output(self):
+        return self._custom_output
 
 
 class FakeAsyncOmni:
     def __init__(self):
         self.stage_configs = [SimpleNamespace(stage_type="diffusion")]
+        self.default_sampling_params_list = [OmniDiffusionSamplingParams()]
         self.captured_prompt = None
         self.captured_sampling_params_list = None
 
@@ -56,6 +74,18 @@ class FakeAsyncOmni:
         yield MockVideoResult(videos)
 
 
+class ErrorAsyncOmni(FakeAsyncOmni):
+    def __init__(self, exc: Exception):
+        super().__init__()
+        self._exc = exc
+
+    async def generate(self, prompt, request_id, sampling_params_list):
+        self.captured_prompt = prompt
+        self.captured_sampling_params_list = sampling_params_list
+        raise self._exc
+        yield  # pragma: no cover
+
+
 class BlockingVideoHandler:
     def __init__(self):
         self.model_name = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
@@ -67,7 +97,7 @@ class BlockingVideoHandler:
         if self.stage_configs is None:
             self.stage_configs = stage_configs
 
-    async def generate_videos(self, request, reference_id, *, reference_image=None):
+    async def generate_video_bytes(self, request, reference_id, *, reference_image=None):
         self.started.set()
         try:
             await asyncio.Future()
@@ -135,15 +165,81 @@ def _wait_until(predicate, timeout_s: float = 2.0, interval_s: float = 0.02):
     raise AssertionError("Timed out waiting for condition")
 
 
+def test_async_video_generation_bypasses_base64(test_client, mocker: MockerFixture):
+    """Regression test: Ensure async video generation saves raw bytes directly
+    without bouncing through base64 encoding."""
+    # We mock _encode_video_bytes (the correct path)
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"raw-mp4-bytes",
+    )
+
+    # We assert that encode_video_base64 is never called
+    mock_base64 = mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
+        side_effect=RuntimeError("Regression: async video path should not base64 encode"),
+    )
+
+    response = test_client.post(
+        "/v1/videos",
+        data={"prompt": "A base64 test."},
+    )
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+
+    # Wait for completion. If it used base64, the RuntimeError would fail the task
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+    mock_base64.assert_not_called()
+
+
+def test_async_video_generation_with_audio_bypasses_base64(test_client, mocker: MockerFixture):
+    """Regression test: Ensure async video generation passes audio through
+    generate_video_bytes without bouncing through base64 encoding."""
+    mock_encode = mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"raw-mp4-bytes",
+    )
+
+    mock_base64 = mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
+        side_effect=RuntimeError("Regression: async video path should not base64 encode"),
+    )
+
+    engine = test_client.app.state.openai_serving_video._engine_client
+
+    async def _generate(prompt, request_id, sampling_params_list):
+        engine.captured_prompt = prompt
+        engine.captured_sampling_params_list = sampling_params_list
+        yield MockVideoResult([object()], audios=[object()], sample_rate=48000)
+
+    engine.generate = _generate
+
+    response = test_client.post(
+        "/v1/videos",
+        data={"prompt": "A base64 test with audio."},
+    )
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+    mock_base64.assert_not_called()
+
+    mock_encode.assert_called_once()
+    kwargs = mock_encode.call_args.kwargs
+    assert "audio" in kwargs
+    assert kwargs["audio"] is not None
+    assert kwargs["audio_sample_rate"] == 48000
+
+
 def test_t2v_video_generation_form(test_client, mocker: MockerFixture):
     fps_values = []
 
-    def _fake_encode(video, fps):
+    def _fake_encode(video, fps, audio=None, audio_sample_rate=None, **kwargs):
         fps_values.append(fps)
-        return "Zg=="
+        return b"fake-video"
 
     mocker.patch(
-        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
         side_effect=_fake_encode,
     )
     response = test_client.post(
@@ -175,8 +271,8 @@ def test_i2v_video_generation_form(test_client, mocker: MockerFixture):
     image_bytes = _make_test_image_bytes((48, 32))
 
     mocker.patch(
-        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
-        return_value="Zg==",
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
     )
     response = test_client.post(
         "/v1/videos",
@@ -201,8 +297,8 @@ def test_i2v_video_generation_resizes_input_to_requested_dimensions(test_client,
     image_bytes = _make_test_image_bytes((48, 32))
 
     mocker.patch(
-        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
-        return_value="Zg==",
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
     )
     response = test_client.post(
         "/v1/videos",
@@ -227,8 +323,8 @@ def test_i2v_video_generation_resizes_input_to_requested_dimensions(test_client,
 
 def test_i2v_video_generation_with_image_reference_form(test_client, mocker: MockerFixture):
     mocker.patch(
-        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
-        return_value="Zg==",
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
     )
     response = test_client.post(
         "/v1/videos",
@@ -252,12 +348,12 @@ def test_i2v_video_generation_with_image_reference_form(test_client, mocker: Moc
 def test_seconds_defaults_fps_and_frames(test_client, mocker: MockerFixture):
     fps_values = []
 
-    def _fake_encode(video, fps):
+    def _fake_encode(video, fps, audio=None, audio_sample_rate=None, **kwargs):
         fps_values.append(fps)
-        return "Zg=="
+        return b"fake-video"
 
     mocker.patch(
-        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
         side_effect=_fake_encode,
     )
     response = test_client.post(
@@ -281,8 +377,8 @@ def test_seconds_defaults_fps_and_frames(test_client, mocker: MockerFixture):
 
 def test_size_param_sets_width_height(test_client, mocker: MockerFixture):
     mocker.patch(
-        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
-        return_value="Zg==",
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
     )
     response = test_client.post(
         "/v1/videos",
@@ -303,8 +399,8 @@ def test_size_param_sets_width_height(test_client, mocker: MockerFixture):
 
 def test_sampling_params_pass_through(test_client, mocker: MockerFixture):
     mocker.patch(
-        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
-        return_value="Zg==",
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
     )
     response = test_client.post(
         "/v1/videos",
@@ -332,13 +428,155 @@ def test_sampling_params_pass_through(test_client, mocker: MockerFixture):
     assert captured.extra_args["flow_shift"] == 0.25
 
 
+def test_frame_interpolation_params_pass_to_diffusion_sampling_params(test_client, mocker: MockerFixture):
+    """Frame interpolation parameters should be forwarded to diffusion worker sampling params."""
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
+    )
+    response = test_client.post(
+        "/v1/videos",
+        data={
+            "prompt": "smooth motion",
+            "fps": "8",
+            "enable_frame_interpolation": "true",
+            "frame_interpolation_exp": "2",
+            "frame_interpolation_scale": "0.5",
+            "frame_interpolation_model_path": "local-rife",
+        },
+    )
+
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+
+    engine = test_client.app.state.openai_serving_video._engine_client
+    captured = engine.captured_sampling_params_list[0]
+    assert captured.enable_frame_interpolation is True
+    assert captured.frame_interpolation_exp == 2
+    assert captured.frame_interpolation_scale == 0.5
+    assert captured.frame_interpolation_model_path == "local-rife"
+
+
+def test_default_sampling_params_apply_to_video_requests(test_client, mocker: MockerFixture):
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
+    )
+    engine = test_client.app.state.openai_serving_video._engine_client
+    engine.default_sampling_params_list = [
+        OmniDiffusionSamplingParams(
+            num_inference_steps=4,
+            guidance_scale=7.5,
+            generator_device="cpu",
+            enable_frame_interpolation=True,
+            frame_interpolation_exp=2,
+            frame_interpolation_scale=0.5,
+            frame_interpolation_model_path="default-rife",
+        )
+    ]
+
+    response = test_client.post(
+        "/v1/videos",
+        data={
+            "prompt": "default param pass-through",
+        },
+    )
+
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+
+    captured = engine.captured_sampling_params_list[0]
+    assert captured.num_inference_steps == 4
+    assert captured.guidance_scale == 7.5
+    assert captured.generator_device == "cpu"
+    assert captured.enable_frame_interpolation is True
+    assert captured.frame_interpolation_exp == 2
+    assert captured.frame_interpolation_scale == 0.5
+    assert captured.frame_interpolation_model_path == "default-rife"
+
+
+def test_request_params_override_default_video_sampling_params(test_client, mocker: MockerFixture):
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
+    )
+    engine = test_client.app.state.openai_serving_video._engine_client
+    engine.default_sampling_params_list = [
+        OmniDiffusionSamplingParams(
+            num_inference_steps=4,
+            guidance_scale=7.5,
+            enable_frame_interpolation=True,
+            frame_interpolation_exp=2,
+            frame_interpolation_scale=0.5,
+            frame_interpolation_model_path="default-rife",
+        )
+    ]
+
+    response = test_client.post(
+        "/v1/videos",
+        data={
+            "prompt": "explicit override",
+            "num_inference_steps": "8",
+            "enable_frame_interpolation": "false",
+            "frame_interpolation_exp": "1",
+            "frame_interpolation_scale": "1.0",
+            "frame_interpolation_model_path": "custom-rife",
+        },
+    )
+
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+
+    captured = engine.captured_sampling_params_list[0]
+    assert captured.num_inference_steps == 8
+    assert captured.guidance_scale == 7.5
+    assert captured.enable_frame_interpolation is False
+    assert captured.frame_interpolation_exp == 1
+    assert captured.frame_interpolation_scale == 1.0
+    assert captured.frame_interpolation_model_path == "custom-rife"
+
+
+def test_worker_fps_multiplier_is_applied_to_async_encoding(test_client, mocker: MockerFixture):
+    fps_values = []
+    engine = test_client.app.state.openai_serving_video._engine_client
+
+    async def _generate(prompt, request_id, sampling_params_list):
+        engine.captured_prompt = prompt
+        engine.captured_sampling_params_list = sampling_params_list
+        import numpy as np
+
+        yield MockVideoResult([np.zeros((1, 64, 64, 3), dtype=np.uint8)], custom_output={"video_fps_multiplier": 2})
+
+    engine.generate = _generate
+
+    def _fake_encode(video, fps, **kwargs):
+        del video, kwargs
+        fps_values.append(fps)
+        return b"fake-video"
+
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        side_effect=_fake_encode,
+    )
+
+    response = test_client.post("/v1/videos", data={"prompt": "fps multiplier", "fps": "8"})
+
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+    assert fps_values == [16]
+
+
 def test_audio_sample_rate_comes_from_model_config(test_client, mocker: MockerFixture):
     audio_sample_rates = []
 
-    def _fake_encode(video, fps, audio=None, audio_sample_rate=None):
-        del video, fps, audio
+    def _fake_encode(video, fps, audio=None, audio_sample_rate=None, video_codec_options=None):
+        del video, fps, audio, video_codec_options
         audio_sample_rates.append(audio_sample_rate)
-        return "Zg=="
+        return b"fake-video"
 
     engine = test_client.app.state.openai_serving_video._engine_client
     engine.model_config = SimpleNamespace(
@@ -352,12 +590,14 @@ def test_audio_sample_rate_comes_from_model_config(test_client, mocker: MockerFi
     async def _generate(prompt, request_id, sampling_params_list):
         engine.captured_prompt = prompt
         engine.captured_sampling_params_list = sampling_params_list
-        yield MockVideoResult([object()], audios=[object()])
+        import numpy as np
+
+        yield MockVideoResult([np.zeros((1, 64, 64, 3), dtype=np.uint8)], audios=[object()])
 
     engine.generate = _generate
 
     mocker.patch(
-        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
         side_effect=_fake_encode,
     )
     response = test_client.post(
@@ -371,6 +611,33 @@ def test_audio_sample_rate_comes_from_model_config(test_client, mocker: MockerFi
     assert audio_sample_rates == [16000]
 
 
+def test_video_job_persists_profiler_metadata(test_client, mocker: MockerFixture):
+    engine = test_client.app.state.openai_serving_video._engine_client
+
+    async def _generate(prompt, request_id, sampling_params_list):
+        engine.captured_prompt = prompt
+        engine.captured_sampling_params_list = sampling_params_list
+        yield MockVideoResult(
+            [object()],
+            stage_durations={"diffuse": 2.5, "vae.decode": 0.3},
+            peak_memory_mb=4096.5,
+        )
+
+    engine.generate = _generate
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
+    )
+
+    response = test_client.post("/v1/videos", data={"prompt": "profile me"})
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+    completed = _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+
+    assert completed["stage_durations"] == {"diffuse": 2.5, "vae.decode": 0.3}
+    assert completed["peak_memory_mb"] == 4096.5
+
+
 def test_missing_handler_returns_503():
     app = FastAPI()
     app.include_router(router)
@@ -428,8 +695,8 @@ def test_invalid_seconds_returns_422(test_client):
 
 def test_negative_prompt_and_seed_pass_through(test_client, mocker: MockerFixture):
     mocker.patch(
-        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
-        return_value="Zg==",
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
     )
     response = test_client.post(
         "/v1/videos",
@@ -498,12 +765,16 @@ def test_video_request_validation():
 
     with pytest.raises(ValueError):
         VideoGenerationRequest(prompt="test", image_reference={"file_id": "file-1", "image_url": "https://example.com"})
+    with pytest.raises(ValueError):
+        VideoGenerationRequest(prompt="test", frame_interpolation_exp=0)
+    with pytest.raises(ValueError):
+        VideoGenerationRequest(prompt="test", frame_interpolation_scale=0)
 
 
 def test_list_videos_supports_order_after_and_limit(test_client, mocker: MockerFixture):
     mocker.patch(
-        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
-        return_value="Zg==",
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
     )
     ids = []
     for i in range(3):
@@ -571,8 +842,8 @@ def test_list_videos_supports_order_after_and_limit(test_client, mocker: MockerF
 
 def test_delete_completed_job_removes_file_and_metadata(test_client, mocker: MockerFixture):
     mocker.patch(
-        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
-        return_value="Zg==",
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
     )
     create_resp = test_client.post("/v1/videos", data={"prompt": "Delete this video"})
     assert create_resp.status_code == 200
@@ -643,8 +914,8 @@ def test_video_response_file_extension_is_robust():
 def test_extra_params_merged_into_extra_args(test_client, mocker: MockerFixture):
     """extra_params JSON object is merged into sampling_params.extra_args."""
     mocker.patch(
-        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
-        return_value="Zg==",
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
     )
     extra_params = {
         "is_enable_stage2": True,
@@ -674,8 +945,8 @@ def test_extra_params_merged_into_extra_args(test_client, mocker: MockerFixture)
 def test_extra_params_none_by_default(test_client, mocker: MockerFixture):
     """When extra_params is omitted, extra_args stays empty."""
     mocker.patch(
-        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
-        return_value="Zg==",
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
     )
     response = test_client.post(
         "/v1/videos",
@@ -715,8 +986,8 @@ def test_extra_params_invalid_json(test_client):
 def test_extra_params_merged_with_existing_extra_args(test_client, mocker: MockerFixture):
     """extra_params is merged on top of existing extra_args (e.g. flow_shift)."""
     mocker.patch(
-        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
-        return_value="Zg==",
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
     )
     response = test_client.post(
         "/v1/videos",
@@ -735,3 +1006,59 @@ def test_extra_params_merged_with_existing_extra_args(test_client, mocker: Mocke
     assert captured.extra_args["flow_shift"] == 0.5
     assert captured.extra_args["use_zero_init"] is True
     assert captured.extra_args["zero_steps"] == 2
+
+
+def test_sample_solver_forwarded_via_extra_params(test_client, mocker: MockerFixture):
+    """sample_solver can be passed through existing extra_params for Wan2.2 online serving."""
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
+    )
+    response = test_client.post(
+        "/v1/videos",
+        data={
+            "prompt": "A fox running through snow.",
+            "extra_params": json.dumps({"sample_solver": "euler"}),
+        },
+    )
+
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+    engine = test_client.app.state.openai_serving_video._engine_client
+    captured = engine.captured_sampling_params_list[0]
+    assert captured.extra_args["sample_solver"] == "euler"
+
+
+@pytest.mark.asyncio
+async def test_run_generation_maps_omni_request_error_to_http_exception():
+    serving = OmniOpenAIServingVideo.for_diffusion(
+        diffusion_engine=ErrorAsyncOmni(
+            OmniRequestError(
+                "oom",
+                status_code=507,
+                request_id="req-oom",
+                stage_id=2,
+                error_type="OutOfMemoryError",
+                detail={"retryable": False},
+            )
+        ),
+        model_name="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+    )
+
+    with pytest.raises(HTTPException) as exc_info:
+        await serving._run_generation(
+            {"prompt": "test"},
+            OmniDiffusionSamplingParams(),
+            "req-oom",
+        )
+
+    exc = exc_info.value
+    assert exc.status_code == 507
+    assert exc.detail == {
+        "message": "oom",
+        "request_id": "req-oom",
+        "stage_id": 2,
+        "error_type": "OutOfMemoryError",
+        "detail": {"retryable": False},
+    }

tests/entrypoints/test_async_omni_diffusion.py

@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm_omni.entrypoints.async_omni_diffusion import AsyncOmniDiffusion
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_get_diffusion_od_config_returns_direct_config():
+    diffusion = object.__new__(AsyncOmniDiffusion)
+    diffusion.od_config = object()
+
+    assert diffusion.get_diffusion_od_config() is diffusion.od_config

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -69,6 +69,20 @@ def test_default_stage_config_propagates_ulysses_mode():
     assert parallel_config.ulysses_mode == "advanced_uaa"
 
 
+def test_default_stage_config_includes_default_sampling_params():
+    """Ensure default sampling params survive the default diffusion-stage builder."""
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+        {
+            "default_sampling_params": '{"0": {"generator_device":"cpu", "guidance_scale":7.5}}',
+        }
+    )[0]
+
+    assert stage_cfg["default_sampling_params"] == {
+        "generator_device": "cpu",
+        "guidance_scale": 7.5,
+    }
+
+
 def test_serve_cli_accepts_ulysses_mode():
     """Ensure diffusion serve CLI exposes ulysses_mode and wires it to parallel_config."""
     parser = FlexibleArgumentParser()

tests/entrypoints/test_omni_base_profiler.py

@@ -4,6 +4,8 @@ from unittest.mock import MagicMock, patch
 
 import pytest
 
+from vllm_omni.diffusion.data import OmniRequestError
+
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
@@ -204,6 +206,45 @@ class TestOmniBaseProfilerSignatureConsistency:
         # stages should default to None
         assert sig.parameters["stages"].default is None
 
+
+class TestOmniBaseOutputHandling:
+    def test_handle_output_message_raises_structured_omni_error(self):
+        mock_engine = MagicMock()
+        mock_engine.num_stages = 1
+        mock_engine.default_sampling_params_list = [MagicMock()]
+        mock_engine.get_stage_metadata.return_value = {
+            "final_output_type": "image",
+            "final_output": True,
+        }
+
+        with (
+            patch("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", return_value=mock_engine),
+            patch("vllm_omni.entrypoints.omni_base.omni_snapshot_download", side_effect=lambda x: x),
+            patch("vllm_omni.entrypoints.omni_base.weakref.finalize"),
+        ):
+            from vllm_omni.entrypoints.omni_base import OmniBase
+
+            omni_base = OmniBase(model="test-model")
+
+        with pytest.raises(OmniRequestError, match="oom") as exc_info:
+            omni_base._handle_output_message(
+                {
+                    "type": "error",
+                    "error": "oom",
+                    "status_code": 507,
+                    "request_id": "req-1",
+                    "stage_id": 2,
+                    "error_type": "OutOfMemoryError",
+                    "detail": {"retryable": False},
+                }
+            )
+
+        assert exc_info.value.status_code == 507
+        assert exc_info.value.request_id == "req-1"
+        assert exc_info.value.stage_id == 2
+        assert exc_info.value.error_type == "OutOfMemoryError"
+        assert exc_info.value.detail == {"retryable": False}
+
     def test_stop_profile_default_values(self):
         """Verify stop_profile has correct default parameter values."""
         import inspect

tests/entrypoints/test_omni_entrypoints.py

@@ -10,6 +10,7 @@ from vllm.entrypoints.openai.models.protocol import BaseModelPath
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
 from vllm.sampling_params import RequestOutputKind, SamplingParams
 
+from vllm_omni.diffusion.data import OmniRequestError
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.omni import Omni
 
@@ -379,6 +380,19 @@ def test_openai_serving_models_can_consume_async_omni_compat_attrs():
     assert serving_models.input_processor is input_processor
 
 
+def test_get_diffusion_od_config_returns_diffusion_stage_config():
+    diffusion_od_config = object()
+    omni = object.__new__(AsyncOmni)
+    omni.engine = SimpleNamespace(
+        stage_clients=[
+            SimpleNamespace(stage_type="llm"),
+            SimpleNamespace(stage_type="diffusion", _engine=SimpleNamespace(od_config=diffusion_od_config)),
+        ]
+    )
+
+    assert omni.get_diffusion_od_config() is diffusion_od_config
+
+
 @pytest.mark.asyncio
 async def test_async_omni_yields_only_final_stage_outputs(monkeypatch: pytest.MonkeyPatch):
     engine = FakeAsyncOmniEngine(
@@ -527,12 +541,16 @@ async def test_async_omni_propagates_engine_error(monkeypatch: pytest.MonkeyPatc
 
     app = AsyncOmni("dummy-model")
     try:
-        with pytest.raises(RuntimeError, match="engine boom"):
+        with pytest.raises(OmniRequestError, match="engine boom") as exc_info:
             async for _ in app.generate(prompt="hello", request_id="req-1"):
                 pass
     finally:
         app.shutdown()
 
+    assert exc_info.value.request_id == "req-1"
+    assert exc_info.value.stage_id == 0
+    assert exc_info.value.status_code == 500
+
 
 def test_omni_generate_py_generator_yields_final_outputs_for_each_request(monkeypatch: pytest.MonkeyPatch):
     sampling_params = [SamplingParams(max_tokens=8) for _ in range(3)]

tests/test_diffusion_config_propagation.py

@@ -14,6 +14,7 @@ from vllm_omni.diffusion.data import (
     DiffusionParallelConfig,
     OmniDiffusionConfig,
 )
+from vllm_omni.diffusion.model_metadata import QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES
 
 
 def _roundtrip_diffusion_config(**kwargs) -> OmniDiffusionConfig:
@@ -106,3 +107,12 @@ class TestCreateDefaultDiffusion:
         ea = stages[0]["engine_args"]
         assert ea["enforce_eager"] is True
         assert ea["lora_path"] == "/tmp/lora"
+
+
+def test_qwen_image_edit_plus_sets_generic_multimodal_limit():
+    od_config = OmniDiffusionConfig(model="Qwen/Qwen-Image-Edit-2511", model_class_name="QwenImageEditPlusPipeline")
+
+    od_config.update_multimodal_support()
+
+    assert od_config.supports_multimodal_inputs is True
+    assert od_config.max_multimodal_image_inputs == QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES

vllm_omni/diffusion/attention/layer.py

@@ -36,6 +36,7 @@ class Attention(nn.Module):
         scatter_idx: int = 2,
         gather_idx: int = 1,
         use_sync: bool = False,
+        skip_sequence_parallel: bool = False,
     ):
         super().__init__()
         self.attn_backend = get_attn_backend(-1)
@@ -62,6 +63,7 @@ class Attention(nn.Module):
         self.gather_idx = gather_idx
         self.use_sync = use_sync
         self.causal = causal
+        self.skip_sequence_parallel = skip_sequence_parallel
 
         self.use_ring = False
         self.ring_pg = None
@@ -98,6 +100,8 @@ class Attention(nn.Module):
         (e.g., in noise_refiner/context_refiner before unified_prepare in Z-Image).
         This avoids unnecessary SP communication for layers not covered by _sp_plan.
         """
+        if self.skip_sequence_parallel:
+            return self._no_parallel_strategy
         if is_forward_context_available():
             ctx = get_forward_context()
             if not ctx.sp_active:

vllm_omni/diffusion/cache/teacache/extractors.py

@@ -221,7 +221,8 @@ def extract_qwen_context(
     block = module.transformer_blocks[0]
     img_mod_params = block.img_mod(temb)
     img_mod1, _ = img_mod_params.chunk(2, dim=-1)
-    img_modulated, _ = block.img_norm1(hidden_states, img_mod1)
+    img_scale1, img_shift1, _ = block._modulate(img_mod1)
+    img_modulated = block.img_norm1(hidden_states, img_scale1, img_shift1)
 
     # ============================================================================
     # DEFINE TRANSFORMER EXECUTION (Qwen-specific)

vllm_omni/diffusion/data.py

@@ -17,6 +17,7 @@ from vllm.model_executor.layers.quantization.base_config import (
     QuantizationConfig,
 )
 
+from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
 from vllm_omni.diffusion.utils.network_utils import is_port_available
 from vllm_omni.quantization import build_quant_config
 
@@ -468,8 +469,10 @@ class OmniDiffusionConfig:
     # Scheduler flow_shift for Wan2.2 (12.0 for 480p, 5.0 for 720p)
     flow_shift: float | None = None
 
-    # support multi images input
+    # Support multi-image inputs and expose any model-specific request limit
+    # through a generic config field so serving code stays model-agnostic.
     supports_multimodal_inputs: bool = False
+    max_multimodal_image_inputs: int | None = None
 
     log_level: str = "info"
 
@@ -615,7 +618,54 @@ class OmniDiffusionConfig:
             raise ValueError("max_cpu_loras must be >= 1 for diffusion LoRA")
 
     def update_multimodal_support(self) -> None:
-        self.supports_multimodal_inputs = self.model_class_name in {"QwenImageEditPlusPipeline"}
+        # Resolve serving-visible multimodal behavior from shared metadata
+        # instead of importing concrete pipeline modules into the config layer.
+        metadata = get_diffusion_model_metadata(self.model_class_name)
+        self.supports_multimodal_inputs = metadata.supports_multimodal_inputs
+        self.max_multimodal_image_inputs = metadata.max_multimodal_image_inputs
+
+    def enrich_config(self) -> None:
+        """Load model metadata from HuggingFace and populate config fields.
+
+        Diffusers-style models expose ``model_index.json`` with ``_class_name``.
+        Non-diffusers models (e.g. Bagel, NextStep) only have ``config.json``,
+        so we fall back to reading that and mapping model_type manually.
+        """
+        from vllm.transformers_utils.config import get_hf_file_to_dict
+
+        try:
+            config_dict = get_hf_file_to_dict("model_index.json", self.model)
+            if config_dict is not None:
+                if self.model_class_name is None:
+                    self.model_class_name = config_dict.get("_class_name", None)
+                self.update_multimodal_support()
+
+                tf_config_dict = get_hf_file_to_dict("transformer/config.json", self.model)
+                self.tf_model_config = TransformerConfig.from_dict(tf_config_dict)
+            else:
+                raise FileNotFoundError("model_index.json not found")
+        except (AttributeError, OSError, ValueError, FileNotFoundError):
+            cfg = get_hf_file_to_dict("config.json", self.model)
+            if cfg is None:
+                raise ValueError(f"Could not find config.json or model_index.json for model {self.model}")
+
+            self.tf_model_config = TransformerConfig.from_dict(cfg)
+            model_type = cfg.get("model_type")
+            architectures = cfg.get("architectures") or []
+
+            if model_type == "bagel" or "BagelForConditionalGeneration" in architectures:
+                self.model_class_name = "BagelPipeline"
+                self.tf_model_config = TransformerConfig()
+                self.update_multimodal_support()
+            elif model_type == "nextstep":
+                if self.model_class_name is None:
+                    self.model_class_name = "NextStep11Pipeline"
+                self.tf_model_config = TransformerConfig()
+                self.update_multimodal_support()
+            elif architectures and len(architectures) == 1:
+                self.model_class_name = architectures[0]
+            else:
+                raise
 
     @classmethod
     def from_kwargs(cls, **kwargs: Any) -> "OmniDiffusionConfig":
@@ -656,6 +706,7 @@ class DiffusionOutput:
     output: torch.Tensor | None = None
     trajectory_timesteps: list[torch.Tensor] | None = None
     trajectory_latents: torch.Tensor | None = None
+    trajectory_log_probs: torch.Tensor | None = None
     trajectory_decoded: list[torch.Tensor] | None = None
     error: str | None = None
 
@@ -675,6 +726,49 @@ class DiffusionOutput:
     peak_memory_mb: float = 0.0
 
 
+@dataclass
+class OmniRequestError(RuntimeError):
+    def __init__(
+        self,
+        message: str,
+        *,
+        status_code: int = 500,
+        request_id: str | None = None,
+        stage_id: int | None = None,
+        error_type: str | None = None,
+        detail: dict[str, Any] | None = None,
+    ):
+        super().__init__(message)
+        self.status_code = status_code
+        self.request_id = request_id
+        self.stage_id = stage_id
+        self.error_type = error_type or self.__class__.__name__
+        self.detail = detail or {}
+
+
+def normalize_omni_error(
+    exc: Exception,
+    *,
+    status_code: int = 500,
+    request_id: str | None = None,
+    stage_id: int | None = None,
+) -> OmniRequestError:
+    if isinstance(exc, OmniRequestError):
+        if exc.request_id is None:
+            exc.request_id = request_id
+        if exc.stage_id is None:
+            exc.stage_id = stage_id
+        return exc
+
+    return OmniRequestError(
+        str(exc),
+        status_code=status_code,
+        request_id=request_id,
+        stage_id=stage_id,
+        error_type=type(exc).__name__,
+    )
+
+
 class AttentionBackendEnum(enum.Enum):
     FA = enum.auto()
     SLIDING_TILE_ATTN = enum.auto()

vllm_omni/diffusion/diffusion_engine.py

@@ -1,6 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from __future__ import annotations
+
+import inspect
 import threading
 import time
 from collections.abc import Iterable
@@ -11,7 +14,7 @@ import PIL.Image
 import torch
 from vllm.logger import init_logger
 
-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig, OmniRequestError, normalize_omni_error
 from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
 from vllm_omni.diffusion.registry import (
     DiffusionModelRegistry,
@@ -69,6 +72,12 @@ class DiffusionEngine:
 
         self.post_process_func = get_diffusion_post_process_func(od_config)
         self.pre_process_func = get_diffusion_pre_process_func(od_config)
+        # Cache whether the model-specific postprocess accepts request-level
+        # sampling params so step() can support both legacy and extended hooks.
+        self._post_process_accepts_sampling_params = bool(
+            self.post_process_func is not None
+            and "sampling_params" in inspect.signature(self.post_process_func).parameters
+        )
 
         executor_class = DiffusionExecutor.get_class(od_config)
         self.executor = executor_class(od_config)
@@ -99,7 +108,12 @@ class DiffusionEngine:
         exec_total_time = time.perf_counter() - exec_start_time
 
         if output.error:
-            raise Exception(f"{output.error}")
+            # raise Exception(f"{output.error}")
+            raise OmniRequestError(
+                output.error,
+                status_code=500,
+                error_type="DiffusionExecutionError",
+            )
         logger.info("Generation completed successfully.")
 
         if output.output is None:
@@ -127,10 +141,24 @@ class DiffusionEngine:
             output_data = output_data.cpu()
 
         postprocess_start_time = time.perf_counter()
-        outputs = self.post_process_func(output_data) if self.post_process_func is not None else output_data
+        if self.post_process_func is not None:
+            # Some video pipelines need request-level controls during
+            # postprocess (for example worker-side frame interpolation).
+            if self._post_process_accepts_sampling_params:
+                outputs = self.post_process_func(output_data, sampling_params=request.sampling_params)
+            else:
+                outputs = self.post_process_func(output_data)
+        else:
+            outputs = output_data
         audio_payload = None
+        custom_output = output.custom_output or {}
+        model_audio_sample_rate = None
+        model_fps = None
         if isinstance(outputs, dict):
             audio_payload = outputs.get("audio")
+            custom_output.update(outputs.get("custom_output") or {})
+            model_audio_sample_rate = outputs.get("audio_sample_rate")
+            model_fps = outputs.get("fps")
             outputs = outputs.get("video", outputs)
         postprocess_time = time.perf_counter() - postprocess_start_time
         logger.info(f"Post-processing completed in {postprocess_time:.4f} seconds")
@@ -186,6 +214,10 @@ class DiffusionEngine:
                 mm_output = {}
                 if audio_payload is not None:
                     mm_output["audio"] = audio_payload
+                if model_audio_sample_rate is not None:
+                    mm_output["audio_sample_rate"] = model_audio_sample_rate
+                if model_fps is not None:
+                    mm_output["fps"] = model_fps
                 return [
                     OmniRequestOutput.from_diffusion(
                         request_id=request_id,
@@ -193,7 +225,11 @@ class DiffusionEngine:
                         prompt=prompt,
                         metrics=metrics,
                         latents=output.trajectory_latents,
-                        custom_output=output.custom_output or {},
+                        trajectory_latents=output.trajectory_latents,
+                        trajectory_timesteps=output.trajectory_timesteps,
+                        trajectory_log_probs=output.trajectory_log_probs,
+                        trajectory_decoded=output.trajectory_decoded,
+                        custom_output=custom_output,
                         multimodal_output=mm_output,
                         stage_durations=output.stage_durations,
                         peak_memory_mb=output.peak_memory_mb,
@@ -244,6 +280,10 @@ class DiffusionEngine:
                                 if num_outputs == 1:
                                     sliced_audio = sliced_audio[0]
                         mm_output["audio"] = sliced_audio
+                    if model_audio_sample_rate is not None:
+                        mm_output["audio_sample_rate"] = model_audio_sample_rate
+                    if model_fps is not None:
+                        mm_output["fps"] = model_fps
                     results.append(
                         OmniRequestOutput.from_diffusion(
                             request_id=request_id,
@@ -251,7 +291,11 @@ class DiffusionEngine:
                             prompt=prompt,
                             metrics=metrics,
                             latents=output.trajectory_latents,
-                            custom_output=output.custom_output or {},
+                            trajectory_latents=output.trajectory_latents,
+                            trajectory_timesteps=output.trajectory_timesteps,
+                            trajectory_log_probs=output.trajectory_log_probs,
+                            trajectory_decoded=output.trajectory_decoded,
+                            custom_output=custom_output,
                             multimodal_output=mm_output,
                             stage_durations=output.stage_durations,
                             peak_memory_mb=output.peak_memory_mb,
@@ -264,7 +308,7 @@ class DiffusionEngine:
     def make_engine(
         config: OmniDiffusionConfig,
         scheduler: SchedulerInterface | None = None,
-    ) -> "DiffusionEngine":
+    ) -> DiffusionEngine:
         """Factory method to create a DiffusionEngine instance.
 
         Args:
@@ -280,33 +324,53 @@ class DiffusionEngine:
             target_sched_req_id = self.scheduler.add_request(request)
 
             # keep scheduling and executing until the target request is finished
-            while True:
-                sched_output = self.scheduler.schedule()
-                if sched_output.is_empty:
-                    if not self.scheduler.has_requests():
-                        raise RuntimeError("Diffusion scheduler has no runnable requests.")
-                    continue
+            try:
+                while True:
+                    sched_output = self.scheduler.schedule()
+                    if sched_output.is_empty:
+                        if not self.scheduler.has_requests():
+                            # raise RuntimeError("Diffusion scheduler has no runnable requests.")
+                            raise OmniRequestError(
+                                "Diffusion scheduler has no runnable requests.",
+                                status_code=500,
+                                error_type="SchedulerError",
+                            )
+                        continue
 
-                # NOTE: add_req_and_wait_for_response() is synchronous, and
-                # the scheduler currently enforces _max_batch_size = 1 (see
-                # vllm_omni/diffusion/sched/base_scheduler.py), so we directly
-                # take the single scheduled request here.
-                sched_req_id = sched_output.scheduled_req_ids[0]
-                req = sched_output.scheduled_new_reqs[0].req
+                    # NOTE: add_req_and_wait_for_response() is synchronous, and
+                    # the scheduler currently enforces _max_batch_size = 1 (see
+                    # vllm_omni/diffusion/sched/base_scheduler.py), so we directly
+                    # take the single scheduled request here.
+                    sched_req_id = sched_output.scheduled_req_ids[0]
+                    req = sched_output.scheduled_new_reqs[0].req
+                    try:
+                        output = self.executor.add_req(req)
+                    except Exception as exc:
+                        logger.error(
+                            "Execution failed for diffusion request %s",
+                            sched_req_id,
+                            exc_info=True,
+                        )
+                        # output = DiffusionOutput(error=str(exc))
+                        raise normalize_omni_error(exc) from exc
+
+                    finished_req_ids = self.scheduler.update_from_output(sched_output, output)
+
+                    if output.error:
+                        status_code = 400 if "`max_sequence_length`" in output.error else 500
+                        raise OmniRequestError(
+                            output.error,
+                            status_code=status_code,
+                            error_type="DiffusionExecutionError",
+                        )
+                    if target_sched_req_id in finished_req_ids:
+                        # self.scheduler.pop_request_state(target_sched_req_id)
+                        return output
+            finally:
                 try:
-                    output = self.executor.add_req(req)
-                except Exception as exc:
-                    logger.error(
-                        "Execution failed for diffusion request %s",
-                        sched_req_id,
-                        exc_info=True,
-                    )
-                    output = DiffusionOutput(error=str(exc))
-
-                finished_req_ids = self.scheduler.update_from_output(sched_output, output)
-                if target_sched_req_id in finished_req_ids:
                     self.scheduler.pop_request_state(target_sched_req_id)
-                    return output
+                except Exception:
+                    logger.debug("Request state already removed: %s", target_sched_req_id, exc_info=True)
 
     def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None:
         """Start or stop torch profiling on all diffusion workers.

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py

@@ -93,7 +93,7 @@ class DistributedAutoencoderKL_base(DistributedVaeMixin):
 
         _, _, latent_h, latent_w = z.shape
         scale = int(2 ** (len(self.config.block_out_channels) - 1))
-        max_parallel_size = self.distributed_decoder.parallel_size
+        max_parallel_size = self.distributed_executor.parallel_size
 
         root = int(math.sqrt(max_parallel_size))
         for rows in range(root, 0, -1):
@@ -187,7 +187,7 @@ class DistributedAutoencoderKL_base(DistributedVaeMixin):
         if split is not None:
             strategy = "tile" if split == self.tile_split else "patch"
             logger.info(f"Decode run with distributed executor, split strategy is {strategy}")
-            result = self.distributed_decoder.execute(
+            result = self.distributed_executor.execute(
                 z, DistributedOperator(split=split, exec=exec, merge=merge), broadcast_result=False
             )
             if not return_dict:

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py

@@ -108,8 +108,8 @@ class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMi
         if not self.is_distributed_enabled():
             return super().tiled_decode(z, return_dict=return_dict)
 
-        logger.info("Decode run with distributed executor")
-        result = self.distributed_decoder.execute(
+        logger.debug("Decode running with distributed executor")
+        result = self.distributed_executor.execute(
             z,
             DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
             broadcast_result=True,

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from contextlib import nullcontext
 from typing import Any
 
 import torch
@@ -15,11 +16,38 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor
     GridSpec,
     TileTask,
 )
+from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
 
-class DistributedAutoencoderKLWan(AutoencoderKLWan, DistributedVaeMixin):
+class OmniAutoencoderKLWan(AutoencoderKLWan):
+    def _execution_context(self):
+        try:
+            first_param = next(self.parameters())
+        except StopIteration:
+            return nullcontext()
+
+        dtype = first_param.dtype
+        if dtype not in (torch.float16, torch.bfloat16):
+            return nullcontext()
+
+        return current_omni_platform.create_autocast_context(
+            device_type=first_param.device.type,
+            dtype=dtype,
+            enabled=True,
+        )
+
+    def encode(self, x: torch.Tensor, return_dict: bool = True):
+        with self._execution_context():
+            return super().encode(x, return_dict=return_dict)
+
+    def decode(self, z: torch.Tensor, return_dict: bool = True):
+        with self._execution_context():
+            return super().decode(z, return_dict=return_dict)
+
+
+class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin):
     @classmethod
     def from_pretrained(cls, *args: Any, **kwargs: Any):
         model = super().from_pretrained(*args, **kwargs)
@@ -84,14 +112,128 @@ class DistributedAutoencoderKLWan(AutoencoderKLWan, DistributedVaeMixin):
         """Decode a single latent tile into RGB space."""
         self.clear_cache()
         time = []
-        for k in range(len(task.tensor)):
-            self._conv_idx = [0]
-            tile = self.post_quant_conv(task.tensor[k])
-            decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=(k == 0))
-            time.append(decoded)
+        with self._execution_context():
+            for k in range(len(task.tensor)):
+                self._conv_idx = [0]
+                tile = self.post_quant_conv(task.tensor[k])
+                decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=(k == 0))
+                time.append(decoded)
         result = torch.cat(time, dim=2)
         return result
 
+    def encode_tile_split(self, x: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
+        _, _, num_frames, height, width = x.shape
+        encode_spatial_compression_ratio = self.spatial_compression_ratio
+        # Scale tile parameters for patchified coordinate system
+        tile_sample_min_height = self.tile_sample_min_height
+        tile_sample_min_width = self.tile_sample_min_width
+        tile_sample_stride_height = self.tile_sample_stride_height
+        tile_sample_stride_width = self.tile_sample_stride_width
+        if self.config.patch_size is not None:
+            assert encode_spatial_compression_ratio % self.config.patch_size == 0
+            encode_spatial_compression_ratio = self.spatial_compression_ratio // self.config.patch_size
+            # When input is patchified, scale tile parameters accordingly
+            tile_sample_min_height = tile_sample_min_height // self.config.patch_size
+            tile_sample_min_width = tile_sample_min_width // self.config.patch_size
+            tile_sample_stride_height = tile_sample_stride_height // self.config.patch_size
+            tile_sample_stride_width = tile_sample_stride_width // self.config.patch_size
+
+        latent_height = height // encode_spatial_compression_ratio
+        latent_width = width // encode_spatial_compression_ratio
+
+        tile_latent_min_height = tile_sample_min_height // encode_spatial_compression_ratio
+        tile_latent_min_width = tile_sample_min_width // encode_spatial_compression_ratio
+        tile_latent_stride_height = tile_sample_stride_height // encode_spatial_compression_ratio
+        tile_latent_stride_width = tile_sample_stride_width // encode_spatial_compression_ratio
+
+        blend_height = tile_latent_min_height - tile_latent_stride_height
+        blend_width = tile_latent_min_width - tile_latent_stride_width
+
+        tiletask_list = []
+        temporal_compression = self.config.scale_factor_temporal
+        for i in range(0, height, tile_sample_stride_height):
+            for j in range(0, width, tile_sample_stride_width):
+                time_list = []
+                frame_range = 1 + (num_frames - 1) // temporal_compression
+                for k in range(frame_range):
+                    if k == 0:
+                        tile = x[:, :, :1, i : i + tile_sample_min_height, j : j + tile_sample_min_width]
+                    else:
+                        tile = x[
+                            :,
+                            :,
+                            1 + temporal_compression * (k - 1) : 1 + temporal_compression * k,
+                            i : i + tile_sample_min_height,
+                            j : j + tile_sample_min_width,
+                        ]
+                    time_list.append(tile)
+                tiletask_list.append(
+                    TileTask(
+                        len(tiletask_list),
+                        (i // tile_sample_stride_height, j // tile_sample_stride_width),
+                        time_list,
+                        workload=time_list[0].shape[3] * time_list[0].shape[4],
+                    )
+                )
+
+        grid_spec = GridSpec(
+            split_dims=(3, 4),
+            grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
+            tile_spec={
+                "latent_height": latent_height,
+                "latent_width": latent_width,
+                "blend_height": blend_height,
+                "blend_width": blend_width,
+                "tile_latent_stride_height": tile_latent_stride_height,
+                "tile_latent_stride_width": tile_latent_stride_width,
+            },
+            output_dtype=self.dtype,
+        )
+        return tiletask_list, grid_spec
+
+    def encode_tile_exec(self, task: TileTask) -> torch.Tensor:
+        """Encode a single sample tile into latent space."""
+        self.clear_cache()
+        time = []
+        for k, tile in enumerate(task.tensor):
+            self._enc_conv_idx = [0]
+            encoded = self.encoder(tile, feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx)
+            encoded = self.quant_conv(encoded)
+            time.append(encoded)
+        result = torch.cat(time, dim=2)
+        self.clear_cache()
+        return result
+
+    def encode_tile_merge(
+        self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec
+    ) -> torch.Tensor:
+        """Merge encoded tiles into a full latent tensor."""
+        grid_h, grid_w = grid_spec.grid_shape
+        result_rows = []
+        for i in range(grid_h):
+            result_row = []
+            for j in range(grid_w):
+                tile = coord_tensor_map[(i, j)]
+                if i > 0:
+                    tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
+                if j > 0:
+                    tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
+                result_row.append(
+                    tile[
+                        :,
+                        :,
+                        :,
+                        : grid_spec.tile_spec["tile_latent_stride_height"],
+                        : grid_spec.tile_spec["tile_latent_stride_width"],
+                    ]
+                )
+            result_rows.append(torch.cat(result_row, dim=-1))
+
+        enc = torch.cat(result_rows, dim=3)[
+            :, :, :, : grid_spec.tile_spec["latent_height"], : grid_spec.tile_spec["latent_width"]
+        ]
+        return enc
+
     def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor:
         """Merge decoded tiles into a full image."""
         grid_h, grid_w = grid_spec.grid_shape
@@ -130,8 +272,8 @@ class DistributedAutoencoderKLWan(AutoencoderKLWan, DistributedVaeMixin):
         if not self.is_distributed_enabled():
             return super().tiled_decode(z, return_dict=return_dict)
 
-        logger.info("Decode run with distributed executor")
-        result = self.distributed_decoder.execute(
+        logger.debug("Decode running with distributed executor")
+        result = self.distributed_executor.execute(
             z,
             DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
             broadcast_result=False,
@@ -140,3 +282,26 @@ class DistributedAutoencoderKLWan(AutoencoderKLWan, DistributedVaeMixin):
             return (result,)
 
         return DecoderOutput(sample=result)
+
+    def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Encode using distributed VAE executor.
+
+        Note: x is already patchified by parent's _encode() before calling this method.
+        """
+        if not self.is_distributed_enabled():
+            return super().tiled_encode(x)
+
+        logger.debug("Encode running with distributed executor")
+        self.clear_cache()
+        result = self.distributed_executor.execute(
+            x,
+            DistributedOperator(
+                split=self.encode_tile_split,
+                exec=self.encode_tile_exec,
+                merge=self.encode_tile_merge,
+            ),
+            broadcast_result=True,
+        )
+        self.clear_cache()
+        return result

vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py

@@ -54,9 +54,9 @@ class DistributedVaeExecutor:
         self.parallel_size = parallel_size
 
     def gather_tensors(self, tensor: torch.Tensor):
-        gather_list = [torch.empty_like(tensor) for _ in range(self.world_size)] if self.rank == 0 else None
-        dist.gather(tensor, gather_list=gather_list, dst=0, group=self.group)
-        return gather_list
+        gather_list = [torch.empty_like(tensor) for _ in range(self.world_size)]
+        dist.all_gather(gather_list, tensor, group=self.group)
+        return gather_list if self.rank == 0 else None
 
     def broadcast_tensor(self, tensor: torch.Tensor):
         dist.broadcast(tensor, src=0, group=self.group)
@@ -168,25 +168,25 @@ class DistributedVaeExecutor:
 
 class DistributedVaeMixin:
     def init_distributed(self):
-        self.distributed_decoder = DistributedVaeExecutor()
+        self.distributed_executor = DistributedVaeExecutor()
 
-    def set_parallel_size(self, parallel_size: int) -> bool:
-        return self.distributed_decoder.set_parallel_size(parallel_size)
+    def set_parallel_size(self, parallel_size: int) -> None:
+        self.distributed_executor.set_parallel_size(parallel_size)
 
     def is_distributed_enabled(self) -> bool:
         if (
-            self.distributed_decoder.parallel_size <= 1
+            self.distributed_executor.parallel_size <= 1
             or not dist.is_initialized()
             or not getattr(self, "use_tiling", False)
         ):
             return False
-        world_size = dist.get_world_size(group=self.distributed_decoder.group)
-        pp_size = min(int(self.distributed_decoder.parallel_size), int(world_size))
+        world_size = dist.get_world_size(group=self.distributed_executor.group)
+        pp_size = min(int(self.distributed_executor.parallel_size), int(world_size))
         if pp_size <= 1:
             return False
-        if self.distributed_decoder.parallel_size > pp_size:
+        if self.distributed_executor.parallel_size > pp_size:
             logger.warning(
-                f"vae_patch_parallel_size={self.distributed_decoder.parallel_size} "
+                f"vae_patch_parallel_size={self.distributed_executor.parallel_size} "
                 f"is greater than dit_group={world_size};"
                 f" using dit_group size={world_size}"
             )

vllm_omni/diffusion/inline_stage_diffusion_client.py

@@ -0,0 +1,362 @@
+"""Inline Stage Diffusion Client for vLLM-Omni multi-stage runtime.
+
+Runs DiffusionEngine in a ThreadPoolExecutor inside the Orchestrator process
+instead of spawning a separate StageDiffusionProc subprocess, eliminating ZMQ
+IPC overhead. Used when there is only a single diffusion stage.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from concurrent.futures import ThreadPoolExecutor
+from typing import TYPE_CHECKING, Any
+
+import torch
+from PIL import Image
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.data import normalize_omni_error
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.engine.stage_init_utils import StageMetadata
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+    from vllm_omni.inputs.data import OmniPromptType
+
+logger = init_logger(__name__)
+
+
+class InlineStageDiffusionClient:
+    """Runs DiffusionEngine in a thread executor inside the Orchestrator."""
+
+    stage_type: str = "diffusion"
+
+    def __init__(
+        self,
+        model: str,
+        od_config: OmniDiffusionConfig,
+        metadata: StageMetadata,
+        batch_size: int = 1,
+    ) -> None:
+        self.model = model
+        self.od_config = od_config
+        self.stage_id = metadata.stage_id
+        self.final_output = metadata.final_output
+        self.final_output_type = metadata.final_output_type
+        self.default_sampling_params = metadata.default_sampling_params
+        self.custom_process_input_func = metadata.custom_process_input_func
+        self.engine_input_source = metadata.engine_input_source
+        self.batch_size = batch_size
+
+        self._enrich_config()
+        self._engine = DiffusionEngine.make_engine(self.od_config)
+        self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="inline-diffusion")
+
+        self._output_queue: asyncio.Queue[OmniRequestOutput | dict[str, Any]] = asyncio.Queue()
+        self._tasks: dict[str, asyncio.Task[Any]] = {}
+        self._shutting_down = False
+
+        logger.info(
+            "[InlineStageDiffusionClient] Stage-%s initialized inline (batch_size=%d)",
+            self.stage_id,
+            self.batch_size,
+        )
+
+    def _enrich_config(self) -> None:
+        """Load model metadata from HuggingFace and populate od_config fields."""
+        self.od_config.enrich_config()
+
+    def _queue_error(self, request_id: str, exc: Exception, log_prefix: str) -> None:
+        err = normalize_omni_error(exc, request_id=request_id, stage_id=self.stage_id)
+        logger.exception(
+            "[InlineStageDiffusionClient] Stage-%s %s req=%s failed: %s",
+            self.stage_id,
+            log_prefix,
+            request_id,
+            err,
+        )
+        self._output_queue.put_nowait(
+            {
+                "type": "error",
+                "request_id": request_id,
+                "stage_id": self.stage_id,
+                "status_code": err.status_code,
+                "error": str(err),
+                "error_type": err.error_type,
+                "detail": err.detail,
+                "finished": True,
+            }
+        )
+
+    # ------------------------------------------------------------------
+    # Request processing
+    # ------------------------------------------------------------------
+
+    async def add_request_async(
+        self,
+        request_id: str,
+        prompt: OmniPromptType,
+        sampling_params: OmniDiffusionSamplingParams,
+        kv_sender_info: dict[int, dict[str, Any]] | None = None,
+    ) -> None:
+        task = asyncio.create_task(
+            self._dispatch_request(
+                request_id,
+                prompt,
+                sampling_params,
+                kv_sender_info,
+            )
+        )
+        self._tasks[request_id] = task
+
+    async def _dispatch_request(
+        self,
+        request_id: str,
+        prompt: Any,
+        sampling_params: OmniDiffusionSamplingParams,
+        kv_sender_info: dict[str, Any] | None = None,
+    ) -> None:
+        try:
+            # Inline diffusion uses the same request envelope as AsyncOmniDiffusion.
+            # kv_sender_info is accepted at the interface layer for compatibility
+            # with other stage clients but is not consumed by DiffusionEngine here.
+            _ = kv_sender_info
+            request = OmniDiffusionRequest(
+                prompts=[prompt],
+                sampling_params=sampling_params,
+                request_ids=[request_id],
+            )
+
+            loop = asyncio.get_running_loop()
+            results = await loop.run_in_executor(self._executor, self._engine.step, request)
+            result = results[0]
+            if not result.request_id:
+                result.request_id = request_id
+
+            self._output_queue.put_nowait(result)
+        except asyncio.CancelledError:
+            logger.info("request_id: %s cancelled", request_id)
+        except Exception as e:
+            self._queue_error(request_id, e, "request")
+        finally:
+            self._tasks.pop(request_id, None)
+
+    async def add_batch_request_async(
+        self,
+        request_id: str,
+        prompts: list[OmniPromptType],
+        sampling_params: OmniDiffusionSamplingParams,
+        kv_sender_info: dict[int, dict[str, Any]] | None = None,
+    ) -> None:
+        task = asyncio.create_task(
+            self._dispatch_batch(
+                request_id,
+                prompts,
+                sampling_params,
+                kv_sender_info,
+            )
+        )
+        self._tasks[request_id] = task
+
+    async def _dispatch_batch(
+        self,
+        request_id: str,
+        prompts: list[Any],
+        sampling_params: OmniDiffusionSamplingParams,
+        kv_sender_info: dict[str, Any] | None = None,
+    ) -> None:
+        try:
+            _ = kv_sender_info
+            request = OmniDiffusionRequest(
+                prompts=prompts,
+                sampling_params=sampling_params,
+                request_ids=[f"{request_id}-{i}" for i in range(len(prompts))],
+            )
+
+            loop = asyncio.get_running_loop()
+            results = await loop.run_in_executor(self._executor, self._engine.step, request)
+
+            all_images: list = []
+            merged_mm: dict[str, Any] = {}
+            merged_metrics: dict[str, Any] = {}
+            merged_durations: dict[str, float] = {}
+            merged_custom: dict[str, Any] = {}
+            peak_mem = 0.0
+            latents = None
+            trajectory_latents: torch.Tensor | None = None
+            trajectory_timesteps: list[torch.Tensor] | None = None
+            trajectory_log_probs: torch.Tensor | None = None
+            trajectory_decoded: list[torch.Tensor] | list[Image.Image] | None = None
+            final_output_type = "image"
+
+            for r in results:
+                all_images.extend(r.images)
+                merged_mm.update(r._multimodal_output)
+                merged_metrics.update(r.metrics)
+                merged_durations.update(r.stage_durations)
+                merged_custom.update(r._custom_output)
+                peak_mem = max(peak_mem, r.peak_memory_mb)
+                if latents is None and r.latents is not None:
+                    latents = r.latents
+                if trajectory_latents is None:
+                    trajectory_latents = r.trajectory_latents
+                if trajectory_timesteps is None:
+                    trajectory_timesteps = r.trajectory_timesteps
+                if trajectory_log_probs is None:
+                    trajectory_log_probs = r.trajectory_log_probs
+                if trajectory_decoded is None:
+                    trajectory_decoded = r.trajectory_decoded
+                if r.final_output_type != "image":
+                    final_output_type = r.final_output_type
+
+            result = OmniRequestOutput.from_diffusion(
+                request_id=request_id,
+                images=all_images,
+                prompt=prompts[0] if len(prompts) == 1 else None,
+                metrics=merged_metrics,
+                latents=latents,
+                trajectory_latents=trajectory_latents,
+                trajectory_timesteps=trajectory_timesteps,
+                trajectory_log_probs=trajectory_log_probs,
+                trajectory_decoded=trajectory_decoded,
+                custom_output=merged_custom or None,
+                multimodal_output=merged_mm or None,
+                final_output_type=final_output_type,
+                stage_durations=merged_durations,
+                peak_memory_mb=peak_mem,
+            )
+
+            self._output_queue.put_nowait(result)
+        except asyncio.CancelledError:
+            logger.info("request_id: %s cancelled", request_id)
+        except Exception as e:
+            self._queue_error(request_id, e, "batch request")
+        finally:
+            self._tasks.pop(request_id, None)
+
+    def get_diffusion_output_async(self) -> OmniRequestOutput | dict[str, Any] | None:
+        try:
+            return self._output_queue.get_nowait()
+        except asyncio.QueueEmpty:
+            return None
+
+    def get_diffusion_output_nowait(self) -> OmniRequestOutput | dict[str, Any] | None:
+        return self.get_diffusion_output_async()
+
+    async def abort_requests_async(self, request_ids: list[str]) -> None:
+        for rid in request_ids:
+            task = self._tasks.pop(rid, None)
+            if task:
+                task.cancel()
+            self._engine.abort(rid)
+
+    async def collective_rpc_async(
+        self,
+        method: str,
+        timeout: float | None = None,
+        args: tuple[Any, ...] = (),
+        kwargs: dict[str, Any] | None = None,
+    ) -> Any:
+        loop = asyncio.get_running_loop()
+
+        if method == "profile":
+            is_start = args[0] if args else True
+            profile_prefix = args[1] if len(args) > 1 else None
+            if is_start and profile_prefix is None:
+                profile_prefix = f"stage_{self.stage_id}_diffusion_{int(time.time())}"
+            return await loop.run_in_executor(
+                self._executor,
+                self._engine.profile,
+                is_start,
+                profile_prefix,
+            )
+
+        kwargs = kwargs or {}
+
+        # LoRA methods
+        if method == "add_lora":
+            lora_request = args[0] if args else kwargs.get("lora_request")
+            results = await loop.run_in_executor(
+                self._executor,
+                self._engine.collective_rpc,
+                "add_lora",
+                timeout,
+                (),
+                {"lora_request": lora_request},
+                None,
+            )
+            return all(results) if isinstance(results, list) else results
+
+        if method == "remove_lora":
+            results = await loop.run_in_executor(
+                self._executor,
+                self._engine.collective_rpc,
+                "remove_lora",
+                timeout,
+                args,
+                kwargs,
+                None,
+            )
+            return all(results) if isinstance(results, list) else results
+
+        if method == "list_loras":
+            results = await loop.run_in_executor(
+                self._executor,
+                self._engine.collective_rpc,
+                "list_loras",
+                timeout,
+                (),
+                {},
+                None,
+            )
+            if not isinstance(results, list):
+                return results or []
+            merged: set[int] = set()
+            for part in results:
+                merged.update(part or [])
+            return sorted(merged)
+
+        if method == "pin_lora":
+            lora_id = args[0] if args else kwargs.get("adapter_id")
+            results = await loop.run_in_executor(
+                self._executor,
+                self._engine.collective_rpc,
+                "pin_lora",
+                timeout,
+                (),
+                {"adapter_id": lora_id},
+                None,
+            )
+            return all(results) if isinstance(results, list) else results
+
+        return await loop.run_in_executor(
+            self._executor,
+            self._engine.collective_rpc,
+            method,
+            timeout,
+            args,
+            kwargs,
+            None,
+        )
+
+    def shutdown(self) -> None:
+        self._shutting_down = True
+
+        # Cancel all pending tasks
+        for task in self._tasks.values():
+            task.cancel()
+
+        try:
+            # Cancel queued futures and wait for the running one to complete deterministically
+            self._executor.shutdown(wait=True, cancel_futures=True)
+        except Exception:
+            pass
+
+        try:
+            self._engine.close()
+        except Exception:
+            pass

vllm_omni/diffusion/ipc.py

@@ -31,7 +31,13 @@ def _tensor_to_shm(tensor: torch.Tensor) -> dict[str, Any]:
     import numpy as np
 
     tensor = tensor.detach().cpu().contiguous()
-    arr = tensor.numpy()
+    if tensor.dtype == torch.bfloat16:
+        # NumPy cannot materialize BF16 tensors directly, so store the raw
+        # 16-bit payload and reconstruct the original dtype on unpack.
+        tensor_for_storage = tensor.view(torch.uint16)
+        arr = tensor_for_storage.numpy()
+    else:
+        arr = tensor.numpy()
     nbytes = arr.nbytes
     shm = shared_memory.SharedMemory(create=True, size=nbytes)
     shm_arr = np.ndarray(arr.shape, dtype=arr.dtype, buffer=shm.buf[:nbytes])
@@ -59,6 +65,8 @@ def _tensor_from_shm(handle: dict[str, Any]) -> torch.Tensor:
         np_dtype = np.dtype(handle["numpy_dtype"])
         arr = np.ndarray(handle["shape"], dtype=np_dtype, buffer=shm.buf[: handle["nbytes"]])
         tensor = torch.from_numpy(arr.copy())
+        if handle["torch_dtype"] == str(torch.bfloat16):
+            tensor = tensor.view(torch.bfloat16)
     finally:
         shm.close()
         shm.unlink()
@@ -72,6 +80,9 @@ def _pack_diffusion_fields(output: DiffusionOutput) -> DiffusionOutput:
     if output.trajectory_latents is not None and isinstance(output.trajectory_latents, torch.Tensor):
         if output.trajectory_latents.nelement() * output.trajectory_latents.element_size() > _SHM_TENSOR_THRESHOLD:
             output.trajectory_latents = _tensor_to_shm(output.trajectory_latents)
+    if output.trajectory_log_probs is not None and isinstance(output.trajectory_log_probs, torch.Tensor):
+        if output.trajectory_log_probs.nelement() * output.trajectory_log_probs.element_size() > _SHM_TENSOR_THRESHOLD:
+            output.trajectory_log_probs = _tensor_to_shm(output.trajectory_log_probs)
     return output
 
 
@@ -95,6 +106,8 @@ def _unpack_diffusion_fields(output: DiffusionOutput) -> DiffusionOutput:
         output.output = _tensor_from_shm(output.output)
     if isinstance(output.trajectory_latents, dict) and output.trajectory_latents.get("__tensor_shm__"):
         output.trajectory_latents = _tensor_from_shm(output.trajectory_latents)
+    if isinstance(output.trajectory_log_probs, dict) and output.trajectory_log_probs.get("__tensor_shm__"):
+        output.trajectory_log_probs = _tensor_from_shm(output.trajectory_log_probs)
     return output
 
 

vllm_omni/diffusion/layers/adalayernorm.py

@@ -1,10 +1,10 @@
 from importlib.util import find_spec
 
 import torch
-import torch.nn as nn
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.layers.custom_op import CustomOp
+from vllm_omni.diffusion.layers.norm import LayerNorm
 
 logger = init_logger(__name__)
 
@@ -22,104 +22,60 @@ class AdaLayerNorm(CustomOp):
         self.eps = eps
         self.elementwise_affine = elementwise_affine
         self.hidden_size = hidden_size
-        self.layernorm = nn.LayerNorm(self.hidden_size, elementwise_affine=self.elementwise_affine, eps=self.eps)
-
-    def preprocess(
-        self,
-        mod_params: torch.Tensor,
-        index: torch.Tensor = None,
-    ) -> torch.Tensor:
-        # shift: b d, scale: b d, gate: b d
-        shift, scale, gate = mod_params.chunk(3, dim=-1)
-
-        if index is not None:
-            # Assuming mod_params batch dim is 2*actual_batch (chunked into 2 parts)
-            # So shift, scale, gate have shape [2*actual_batch, d]
-            actual_batch = shift.size(0) // 2
-            shift_0, shift_1 = shift[:actual_batch], shift[actual_batch:]  # each: [actual_batch, d]
-            scale_0, scale_1 = scale[:actual_batch], scale[actual_batch:]
-            gate_0, gate_1 = gate[:actual_batch], gate[actual_batch:]
-
-            # index: [b, l] where b is actual batch size
-            # Expand to [b, l, 1] to match feature dimension
-            index_expanded = index.unsqueeze(-1)  # [b, l, 1]
-
-            # Expand chunks to [b, 1, d] then broadcast to [b, l, d]
-            shift_0_exp = shift_0.unsqueeze(1)  # [b, 1, d]
-            shift_1_exp = shift_1.unsqueeze(1)  # [b, 1, d]
-            scale_0_exp = scale_0.unsqueeze(1)
-            scale_1_exp = scale_1.unsqueeze(1)
-            gate_0_exp = gate_0.unsqueeze(1)
-            gate_1_exp = gate_1.unsqueeze(1)
-
-            # Use torch.where to select based on index
-            shift_result = torch.where(index_expanded == 0, shift_0_exp, shift_1_exp)
-            scale_result = torch.where(index_expanded == 0, scale_0_exp, scale_1_exp)
-            gate_result = torch.where(index_expanded == 0, gate_0_exp, gate_1_exp)
-        else:
-            shift_result = shift.unsqueeze(1)
-            scale_result = scale.unsqueeze(1)
-            gate_result = gate.unsqueeze(1)
-
-        return shift_result, scale_result, gate_result
+        self.layernorm = LayerNorm(self.hidden_size, elementwise_affine=self.elementwise_affine, eps=self.eps)
 
     def forward_cuda(
         self,
         x: torch.Tensor,
-        mod_params: torch.Tensor,
-        index: torch.Tensor = None,
+        scale: torch.Tensor,
+        shift: torch.Tensor,
     ) -> torch.Tensor:
-        return self.forward_native(x, mod_params, index)
+        return self.forward_native(x, scale, shift)
 
     def forward_hip(
         self,
         x: torch.Tensor,
-        mod_params: torch.Tensor,
-        index: torch.Tensor = None,
+        scale: torch.Tensor,
+        shift: torch.Tensor,
     ) -> torch.Tensor:
-        return self.forward_native(x, mod_params, index)
+        return self.forward_native(x, scale, shift)
 
     def forward_npu(
         self,
         x: torch.Tensor,
-        mod_params: torch.Tensor,
-        index: torch.Tensor = None,
+        scale: torch.Tensor,
+        shift: torch.Tensor,
     ) -> torch.Tensor:
-        shift_result, scale_result, gate_result = self.preprocess(mod_params, index)
-
         if _HAS_MINDIESD:
             try:
                 from mindiesd import layernorm_scale_shift
 
-                output = layernorm_scale_shift(self.layernorm, x, scale_result, shift_result, fused=True)
+                output = layernorm_scale_shift(self.layernorm, x, scale, shift, fused=True)
 
-                return output, gate_result
+                return output
             except ImportError as e:
                 logger.warning_once(f"mindiesd import failed, falling back to torch_npu: {e}")
 
         import torch_npu
 
         output = (
-            torch_npu.npu_layer_norm_eval(x, normalized_shape=[self.hidden_size], eps=self.eps) * (1 + scale_result)
-            + shift_result
+            torch_npu.npu_layer_norm_eval(x, normalized_shape=[self.hidden_size], eps=self.eps) * (1 + scale) + shift
         )
 
-        return output, gate_result
+        return output
 
     def forward_xpu(
         self,
         x: torch.Tensor,
-        mod_params: torch.Tensor,
-        index: torch.Tensor = None,
+        scale: torch.Tensor,
+        shift: torch.Tensor,
     ) -> torch.Tensor:
-        return self.forward_native(x, mod_params, index)
+        return self.forward_native(x, scale, shift)
 
     def forward_native(
         self,
         x: torch.Tensor,
-        mod_params: torch.Tensor,
-        index: torch.Tensor = None,
+        scale: torch.Tensor,
+        shift: torch.Tensor,
     ) -> torch.Tensor:
-        shift_result, scale_result, gate_result = self.preprocess(mod_params, index)
-
-        return self.layernorm(x) * (1 + scale_result) + shift_result, gate_result
+        return self.layernorm(x) * (1 + scale) + shift

vllm_omni/diffusion/layers/norm.py

@@ -0,0 +1,185 @@
+from importlib.util import find_spec
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.layers.custom_op import CustomOp
+
+logger = init_logger(__name__)
+
+_HAS_MINDIESD = find_spec("mindiesd") is not None
+
+
+class LayerNorm(nn.LayerNorm, CustomOp):
+    """
+    LayerNorm implementation that inherits from both ``nn.LayerNorm`` and ``CustomOp``.
+    NPU:
+        Uses ``mindiesd.fast_layernorm(self, x)`` when MindIE-SD is installed.
+    CUDA / HIP / XPU / native:
+        Falls back to FP32 nn.LayerNorm implementation.
+    """
+
+    def __init__(self, dim: int, eps: float = 1e-6, elementwise_affine: bool = True):
+        super().__init__(normalized_shape=dim, eps=eps, elementwise_affine=elementwise_affine)
+        self._forward_method = CustomOp.dispatch_forward(self)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self._forward_method(x)
+
+    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
+        return self.forward_native(x)
+
+    def forward_hip(self, x: torch.Tensor) -> torch.Tensor:
+        return self.forward_native(x)
+
+    def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
+        return self.forward_native(x)
+
+    def forward_npu(self, x: torch.Tensor) -> torch.Tensor:
+        if _HAS_MINDIESD:
+            try:
+                from mindiesd import fast_layernorm
+
+                return fast_layernorm(self, x)
+            except ImportError as e:
+                logger.warning_once(
+                    "mindiesd.fast_layernorm import failed, falling back to FP32 layer_norm: %s",
+                    e,
+                )
+
+        return self.forward_native(x)
+
+    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
+        origin_dtype = x.dtype
+        return F.layer_norm(
+            x.float(),
+            self.normalized_shape,
+            self.weight.float() if self.weight is not None else None,
+            self.bias.float() if self.bias is not None else None,
+            self.eps,
+        ).to(origin_dtype)
+
+
+class RMSNorm(CustomOp):
+    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward_cuda(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.forward_native(x)
+
+    def forward_hip(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.forward_native(x)
+
+    def forward_npu(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        import torch_npu
+
+        output = torch_npu.npu_rms_norm(x, gamma=self.weight, epsilon=self.variance_epsilon)[0]
+
+        return output
+
+    def forward_xpu(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.forward_native(x)
+
+    def forward_native(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        input_dtype = x.dtype
+        x = x.to(torch.float32)
+        variance = x.pow(2).mean(-1, keepdim=True)
+        out = x * torch.rsqrt(variance + self.variance_epsilon)
+        out = self.weight.to(torch.float32) * out
+        return out.to(input_dtype)
+
+class RMSNormVAE(CustomOp):
+    """Root Mean Square Layer Normalization for Channel-First or Last"""
+
+    def __init__(
+        self,
+        dim: int,
+        channel_first: bool = True,
+        images: bool = True,
+        bias: bool = False,
+        epsilon: float = 1e-6,
+    ) -> None:
+        super().__init__()
+        broadcastable_dims = (1, 1, 1) if not images else (1, 1)
+        shape = (dim, *broadcastable_dims) if channel_first else (dim,)
+
+        self.channel_first = channel_first
+        self.scale = dim**0.5
+        self.gamma = nn.Parameter(torch.ones(shape))
+        self.bias = nn.Parameter(torch.zeros(shape)) if bias else None
+        self.epsilon = epsilon
+
+        self.gamma_rmsnorm = None
+
+    def forward_cuda(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.forward_native(x)
+
+    def forward_hip(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.forward_native(x)
+
+    def forward_npu(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        import torch_npu
+
+        if self.gamma_rmsnorm is None:
+            self.gamma_rmsnorm = self.gamma.reshape(-1)
+
+        if self.channel_first:
+            x = x.transpose(1, -1)
+            out = torch_npu.npu_rms_norm(x, self.gamma_rmsnorm, epsilon=self.epsilon)[0].transpose(1, -1)
+        else:
+            out = torch_npu.npu_rms_norm(x, self.gamma_rmsnorm, epsilon=self.epsilon)[0]
+
+        if self.bias is not None:
+            out = out + self.bias
+        return out
+
+    def forward_xpu(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.forward_native(x)
+
+    def forward_native(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        out = (
+            F.normalize(
+                x,
+                dim=(1 if self.channel_first else -1),
+                eps=self.epsilon,
+            )
+            * self.scale
+            * self.gamma
+        )
+        if self.bias is not None:
+            out = out + self.bias
+        return out

vllm_omni/diffusion/model_metadata.py

@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class DiffusionModelMetadata:
+    # Keep serving-facing capability metadata in a lightweight shared module so
+    # config/model plumbing can read it without importing concrete pipelines.
+    supports_multimodal_inputs: bool = False
+    max_multimodal_image_inputs: int | None = None
+
+
+QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES = 4
+
+
+_DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = {
+    "QwenImageEditPlusPipeline": DiffusionModelMetadata(
+        supports_multimodal_inputs=True,
+        max_multimodal_image_inputs=QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES,
+    ),
+}
+
+
+def get_diffusion_model_metadata(model_class_name: str | None) -> DiffusionModelMetadata:
+    # Unknown models fall back to "no special multimodal capabilities" so new
+    # pipelines do not accidentally inherit limits meant for other models.
+    if model_class_name is None:
+        return DiffusionModelMetadata()
+    return _DIFFUSION_MODEL_METADATA.get(model_class_name, DiffusionModelMetadata())

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py

@@ -34,6 +34,12 @@ from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
 )
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.utils.prompt_utils import (
+    validate_prompt_sequence_lengths,
+)
+from vllm_omni.diffusion.utils.size_utils import (
+    normalize_min_aligned_size,
+)
 from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
 
 if TYPE_CHECKING:
@@ -360,8 +366,10 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
                 "that was used to generate `negative_prompt_embeds`."
             )
 
-        if max_sequence_length is not None and max_sequence_length > 1024:
-            raise ValueError(f"`max_sequence_length` cannot be greater than 1024 but is {max_sequence_length}")
+        if max_sequence_length is not None and max_sequence_length > self.tokenizer_max_length:
+            raise ValueError(
+                f"`max_sequence_length` cannot be greater than {self.tokenizer_max_length} but is {max_sequence_length}"
+            )
 
     def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor):
         bool_mask = mask.bool()
@@ -375,6 +383,8 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
         self,
         prompt: str | list[str] = None,
         dtype: torch.dtype | None = None,
+        max_sequence_length: int | None = None,
+        prompt_name: str = "prompt",
     ):
         dtype = dtype or self.text_encoder.dtype
 
@@ -385,12 +395,27 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
         txt = [template.format(e) for e in prompt]
         txt_tokens = self.tokenizer(
             txt,
-            max_length=self.tokenizer_max_length + drop_idx,
             padding=True,
-            truncation=True,
+            truncation=False,
             return_tensors="pt",
         ).to(self.device)
-        # print(f"attention mask: {txt_tokens.attention_mask}")
+        # Validate only the user prompt contribution. The Qwen template also
+        # adds a fixed suffix after the user text, so subtracting only
+        # prompt_template_encode_start_idx would overcount near-limit prompts.
+        template_tokens = self.tokenizer(
+            [template.format("")],
+            padding=True,
+            truncation=False,
+            return_tensors="pt",
+        ).to(self.device)
+        validate_prompt_sequence_lengths(
+            txt_tokens.attention_mask,
+            max_sequence_length=max_sequence_length or self.tokenizer_max_length,
+            supported_max_sequence_length=self.tokenizer_max_length,
+            prompt_name=prompt_name,
+            baseline_attention_mask=template_tokens.attention_mask,
+            error_context="after applying the Qwen prompt template",
+        )
         encoder_hidden_states = self.text_encoder(
             input_ids=txt_tokens.input_ids,
             attention_mask=txt_tokens.attention_mask,
@@ -419,6 +444,7 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
         prompt_embeds: torch.Tensor | None = None,
         prompt_embeds_mask: torch.Tensor | None = None,
         max_sequence_length: int = 1024,
+        prompt_name: str = "prompt",
     ):
         r"""
 
@@ -436,7 +462,11 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
         batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0]
 
         if prompt_embeds is None:
-            prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt)
+            prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(
+                prompt,
+                max_sequence_length=max_sequence_length,
+                prompt_name=prompt_name,
+            )
 
         prompt_embeds = prompt_embeds[:, :max_sequence_length]
         prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length]
@@ -629,6 +659,7 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
                 prompt_embeds_mask=negative_prompt_embeds_mask,
                 num_images_per_prompt=num_images_per_prompt,
                 max_sequence_length=max_sequence_length,
+                prompt_name="negative_prompt",
             )
         else:
             negative_prompt_embeds = None
@@ -700,7 +731,7 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
             num_images_per_prompt=sampling.num_outputs_per_prompt if sampling.num_outputs_per_prompt > 0 else 1,
             generator=sampling.generator,
             true_cfg_scale=sampling.true_cfg_scale or 4.0,
-            max_sequence_length=sampling.max_sequence_length or 512,
+            max_sequence_length=sampling.max_sequence_length or self.tokenizer_max_length,
             attention_kwargs=kwargs.get("attention_kwargs"),
         )
 
@@ -931,13 +962,14 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
         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,
+        max_sequence_length: int = 1024,
     ) -> DiffusionOutput:
         extracted_prompt, negative_prompt = self._extract_prompts(req.prompts)
         prompt = extracted_prompt or prompt
 
         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
+        height, width = normalize_min_aligned_size(height, width, self.vae_scale_factor * 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

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py

@@ -37,6 +37,12 @@ from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
 )
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.utils.prompt_utils import (
+    validate_prompt_sequence_lengths,
+)
+from vllm_omni.diffusion.utils.size_utils import (
+    normalize_min_aligned_size,
+)
 from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
 from vllm_omni.inputs.data import OmniTextPrompt
 from vllm_omni.model_executor.model_loader.weight_utils import (
@@ -97,9 +103,7 @@ def get_qwen_image_edit_pre_process_func(
             width = request.sampling_params.width or calculated_width
 
             # Ensure dimensions are multiples of vae_scale_factor * 2
-            multiple_of = vae_scale_factor * 2
-            height = height // multiple_of * multiple_of
-            width = width // multiple_of * multiple_of
+            height, width = normalize_min_aligned_size(height, width, vae_scale_factor * 2)
 
             # Store calculated dimensions in request
             prompt["additional_information"]["calculated_height"] = calculated_height
@@ -322,8 +326,10 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi
                 "that was used to generate `negative_prompt_embeds`."
             )
 
-        if max_sequence_length is not None and max_sequence_length > 1024:
-            raise ValueError(f"`max_sequence_length` cannot be greater than 1024 but is {max_sequence_length}")
+        if max_sequence_length is not None and max_sequence_length > self.tokenizer_max_length:
+            raise ValueError(
+                f"`max_sequence_length` cannot be greater than {self.tokenizer_max_length} but is {max_sequence_length}"
+            )
 
     def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor):
         bool_mask = mask.bool()
@@ -383,6 +389,8 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi
         prompt: str | list[str] = None,
         image: PIL.Image.Image | torch.Tensor | None = None,
         dtype: torch.dtype | None = None,
+        max_sequence_length: int | None = None,
+        prompt_name: str = "prompt",
     ):
         """Get prompt embeddings with image support for editing."""
         dtype = dtype or self.text_encoder.dtype
@@ -392,6 +400,33 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi
         template = self.prompt_template_encode
         drop_idx = self.prompt_template_encode_start_idx
         txt = [template.format(e) for e in prompt]
+        txt_tokens = self.tokenizer(
+            txt,
+            padding=True,
+            truncation=False,
+            return_tensors="pt",
+        ).to(self.device)
+        # The edit template contains fixed multimodal scaffolding around the
+        # instruction. Validate against the empty-template baseline so image
+        # placeholder text does not consume the user's text budget.
+        template_tokens = self.tokenizer(
+            [template.format("")],
+            padding=True,
+            truncation=False,
+            return_tensors="pt",
+        ).to(self.device)
+        # Qwen-Image-Edit expands image placeholders into many vision tokens
+        # inside the processor. `max_sequence_length` is meant to constrain the
+        # prompt text length, so validate on the text template before image
+        # token expansion.
+        validate_prompt_sequence_lengths(
+            txt_tokens.attention_mask,
+            max_sequence_length=max_sequence_length or self.tokenizer_max_length,
+            supported_max_sequence_length=self.tokenizer_max_length,
+            prompt_name=prompt_name,
+            baseline_attention_mask=template_tokens.attention_mask,
+            error_context="after applying the Qwen prompt template",
+        )
 
         # Use processor to handle both text and image inputs
         model_inputs = self.processor(
@@ -433,6 +468,7 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi
         prompt_embeds: torch.Tensor | None = None,
         prompt_embeds_mask: torch.Tensor | None = None,
         max_sequence_length: int = 1024,
+        prompt_name: str = "prompt",
     ):
         r"""
 
@@ -452,7 +488,12 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi
         batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0]
 
         if prompt_embeds is None:
-            prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt, image)
+            prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(
+                prompt,
+                image,
+                max_sequence_length=max_sequence_length,
+                prompt_name=prompt_name,
+            )
 
         _, seq_len, _ = prompt_embeds.shape
         prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
@@ -623,7 +664,7 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi
         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,
+        max_sequence_length: int = 1024,
     ) -> DiffusionOutput:
         """Forward pass for image editing."""
         # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
@@ -661,9 +702,7 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi
             height = height or calculated_height
             width = width or calculated_width
 
-            multiple_of = self.vae_scale_factor * 2
-            width = width // multiple_of * multiple_of
-            height = height // multiple_of * multiple_of
+            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)
@@ -740,6 +779,7 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi
                 prompt_embeds_mask=negative_prompt_embeds_mask,
                 num_images_per_prompt=num_images_per_prompt,
                 max_sequence_length=max_sequence_length,
+                prompt_name="negative_prompt",
             )
 
         num_channels_latents = self.transformer.in_channels // 4

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py

@@ -25,6 +25,7 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_metadata import QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES
 from vllm_omni.diffusion.models.interface import SupportImageInput
 from vllm_omni.diffusion.models.qwen_image.cfg_parallel import (
     QwenImageCFGParallelMixin,
@@ -40,6 +41,12 @@ from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
 )
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.utils.prompt_utils import (
+    validate_prompt_sequence_lengths,
+)
+from vllm_omni.diffusion.utils.size_utils import (
+    normalize_min_aligned_size,
+)
 from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
 from vllm_omni.inputs.data import OmniTextPrompt
 from vllm_omni.model_executor.model_loader.weight_utils import (
@@ -50,6 +57,12 @@ logger = logging.getLogger(__name__)
 
 CONDITION_IMAGE_SIZE = 384 * 384
 VAE_IMAGE_SIZE = 1024 * 1024
+# Keep this in sync with the practical conditioning-token budget for
+# Qwen-Image-Edit-2511. Empirically, 4 images stays within the supported range
+# while 5 images overflows the prompt/conditioning path and fails downstream.
+# Re-export the shared metadata value locally so this pipeline keeps a nearby,
+# descriptive constant for validation and tests without becoming the source of truth.
+MAX_QWEN_IMAGE_EDIT_PLUS_INPUT_IMAGES = QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES
 
 
 def get_qwen_image_edit_plus_pre_process_func(
@@ -87,6 +100,11 @@ def get_qwen_image_edit_plus_pre_process_func(
 
             if not isinstance(raw_image, list):
                 raw_image = [raw_image]
+            if len(raw_image) > MAX_QWEN_IMAGE_EDIT_PLUS_INPUT_IMAGES:
+                raise ValueError(
+                    f"Received {len(raw_image)} input images. "
+                    f"At most {MAX_QWEN_IMAGE_EDIT_PLUS_INPUT_IMAGES} images are supported by this model."
+                )
             image = [
                 PIL.Image.open(im) if isinstance(im, str) else cast(PIL.Image.Image | np.ndarray | torch.Tensor, im)
                 for im in raw_image
@@ -99,9 +117,7 @@ def get_qwen_image_edit_plus_pre_process_func(
             width = request.sampling_params.width or calculated_width
 
             # Ensure dimensions are multiples of vae_scale_factor * 2
-            multiple_of = vae_scale_factor * 2
-            height = height // multiple_of * multiple_of
-            width = width // multiple_of * multiple_of
+            height, width = normalize_min_aligned_size(height, width, vae_scale_factor * 2)
 
             # Store calculated dimensions in request
             prompt["additional_information"]["calculated_height"] = calculated_height
@@ -282,8 +298,10 @@ class QwenImageEditPlusPipeline(
                 "that was used to generate `negative_prompt_embeds`."
             )
 
-        if max_sequence_length is not None and max_sequence_length > 1024:
-            raise ValueError(f"`max_sequence_length` cannot be greater than 1024 but is {max_sequence_length}")
+        if max_sequence_length is not None and max_sequence_length > self.tokenizer_max_length:
+            raise ValueError(
+                f"`max_sequence_length` cannot be greater than {self.tokenizer_max_length} but is {max_sequence_length}"
+            )
 
     def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor):
         bool_mask = mask.bool()
@@ -298,6 +316,8 @@ class QwenImageEditPlusPipeline(
         prompt: str | list[str],
         image: list[torch.Tensor] | torch.Tensor | None = None,
         dtype: torch.dtype | None = None,
+        max_sequence_length: int | None = None,
+        prompt_name: str = "prompt",
     ):
         """Get prompt embeddings with support for multiple images."""
         dtype = dtype or self.text_encoder.dtype
@@ -318,6 +338,32 @@ class QwenImageEditPlusPipeline(
         template = self.prompt_template_encode
         drop_idx = self.prompt_template_encode_start_idx
         txt = [template.format(base_img_prompt + e) for e in prompt]
+        txt_tokens = self.tokenizer(
+            txt,
+            padding=True,
+            truncation=False,
+            return_tensors="pt",
+        ).to(self.device)
+        # Multi-image edit prepends "Picture N" placeholders before the user
+        # instruction. Subtract the placeholder-aware baseline so attached
+        # images do not reduce the remaining prompt budget.
+        template_tokens = self.tokenizer(
+            [template.format(base_img_prompt)],
+            padding=True,
+            truncation=False,
+            return_tensors="pt",
+        ).to(self.device)
+        # The processor expands image placeholders into many vision tokens.
+        # `max_sequence_length` should guard the prompt text length before that
+        # multimodal expansion happens.
+        validate_prompt_sequence_lengths(
+            txt_tokens.attention_mask,
+            max_sequence_length=max_sequence_length or self.tokenizer_max_length,
+            supported_max_sequence_length=self.tokenizer_max_length,
+            prompt_name=prompt_name,
+            baseline_attention_mask=template_tokens.attention_mask,
+            error_context="after applying the Qwen prompt template",
+        )
 
         # Use processor to handle both text and image inputs
         model_inputs = self.processor(
@@ -359,6 +405,7 @@ class QwenImageEditPlusPipeline(
         prompt_embeds: torch.Tensor | None = None,
         prompt_embeds_mask: torch.Tensor | None = None,
         max_sequence_length: int = 1024,
+        prompt_name: str = "prompt",
     ):
         r"""
 
@@ -378,7 +425,12 @@ class QwenImageEditPlusPipeline(
         batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0]
 
         if prompt_embeds is None:
-            prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt, image)
+            prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(
+                prompt,
+                image,
+                max_sequence_length=max_sequence_length,
+                prompt_name=prompt_name,
+            )
 
         _, seq_len, _ = prompt_embeds.shape
         prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
@@ -556,7 +608,7 @@ class QwenImageEditPlusPipeline(
         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,
+        max_sequence_length: int = 1024,
     ) -> 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("...", "")
@@ -604,9 +656,7 @@ class QwenImageEditPlusPipeline(
             height = height or calculated_height
             width = width or calculated_width
 
-            multiple_of = self.vae_scale_factor * 2
-            width = width // multiple_of * multiple_of
-            height = height // multiple_of * multiple_of
+            height, width = normalize_min_aligned_size(height, width, self.vae_scale_factor * 2)
 
             condition_images = []
             vae_images = []
@@ -693,6 +743,7 @@ class QwenImageEditPlusPipeline(
                 prompt_embeds_mask=negative_prompt_embeds_mask,
                 num_images_per_prompt=num_images_per_prompt,
                 max_sequence_length=max_sequence_length,
+                prompt_name="negative_prompt",
             )
 
         num_channels_latents = self.transformer.in_channels // 4

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py

@@ -36,6 +36,12 @@ from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
 )
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.utils.prompt_utils import (
+    validate_prompt_sequence_lengths,
+)
+from vllm_omni.diffusion.utils.size_utils import (
+    normalize_min_aligned_size,
+)
 from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
 from vllm_omni.inputs.data import OmniTextPrompt
 from vllm_omni.model_executor.model_loader.weight_utils import (
@@ -109,9 +115,7 @@ def get_qwen_image_layered_pre_process_func(
             height = calculated_height
             width = calculated_width
 
-            multiple_of = vae_scale_factor * 2
-            width = width // multiple_of * multiple_of
-            height = height // multiple_of * multiple_of
+            height, width = normalize_min_aligned_size(height, width, vae_scale_factor * 2)
 
             # Store calculated dimensions in request
             prompt["additional_information"]["calculated_height"] = calculated_height
@@ -339,8 +343,10 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
                 "generate `negative_prompt_embeds`."
             )
 
-        if max_sequence_length is not None and max_sequence_length > 1024:
-            raise ValueError(f"`max_sequence_length` cannot be greater than 1024 but is {max_sequence_length}")
+        if max_sequence_length is not None and max_sequence_length > self.tokenizer_max_length:
+            raise ValueError(
+                f"`max_sequence_length` cannot be greater than {self.tokenizer_max_length} but is {max_sequence_length}"
+            )
 
     def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor):
         bool_mask = mask.bool()
@@ -355,6 +361,8 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
         prompt: str | list[str] | None = None,
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
+        max_sequence_length: int | None = None,
+        prompt_name: str = "prompt",
     ):
         device = device or self.device
         dtype = dtype or self.text_encoder.dtype
@@ -367,8 +375,26 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
         txt_tokens = self.tokenizer(
             txt,
             padding=True,
+            truncation=False,
             return_tensors="pt",
         ).to(device)
+        # The layered template also appends fixed non-user tokens after the
+        # editable text, so use the empty-template tokenized baseline instead of
+        # counting everything after prompt_template_encode_start_idx.
+        template_tokens = self.tokenizer(
+            [template.format("")],
+            padding=True,
+            truncation=False,
+            return_tensors="pt",
+        ).to(device)
+        validate_prompt_sequence_lengths(
+            txt_tokens.attention_mask,
+            max_sequence_length=max_sequence_length or self.tokenizer_max_length,
+            supported_max_sequence_length=self.tokenizer_max_length,
+            prompt_name=prompt_name,
+            baseline_attention_mask=template_tokens.attention_mask,
+            error_context="after applying the Qwen prompt template",
+        )
         encoder_hidden_states = self.text_encoder(
             input_ids=txt_tokens.input_ids,
             attention_mask=txt_tokens.attention_mask,
@@ -398,6 +424,7 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
         prompt_embeds: torch.Tensor | None = None,
         prompt_embeds_mask: torch.Tensor | None = None,
         max_sequence_length: int = 1024,
+        prompt_name: str = "prompt",
     ):
         r"""
 
@@ -418,7 +445,11 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
         batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0]
 
         if prompt_embeds is None:
-            prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt)
+            prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(
+                prompt,
+                max_sequence_length=max_sequence_length,
+                prompt_name=prompt_name,
+            )
 
         prompt_embeds = prompt_embeds[:, :max_sequence_length]
         prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length]
@@ -602,7 +633,7 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
         negative_prompt_embeds_mask: torch.Tensor | None = None,
         output_type: str | None = "pil",
         attention_kwargs: dict[str, Any] | None = None,
-        max_sequence_length: int = 512,
+        max_sequence_length: int = 1024,
         resolution: int = 640,
         cfg_normalize: bool = False,
         use_en_prompt: bool = False,
@@ -665,9 +696,7 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
             height = calculated_height
             width = calculated_width
 
-            multiple_of = self.vae_scale_factor * 2
-            width = width // multiple_of * multiple_of
-            height = height // multiple_of * multiple_of
+            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)
@@ -737,6 +766,7 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
                 device=self.device,
                 num_images_per_prompt=num_images_per_prompt,
                 max_sequence_length=max_sequence_length,
+                prompt_name="negative_prompt",
             )
 
         # 4. Prepare latent variables

vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py

@@ -743,9 +743,9 @@ class QwenImageTransformerBlock(nn.Module):
 
         self.zero_cond_t = zero_cond_t
 
-    def _modulate(self, x, mod_params, index=None):
+    def _modulate(self, mod_params, index=None):
         """Apply modulation to input tensor"""
-        # x: b l d, shift: b d, scale: b d, gate: b d
+        # shift: b d, scale: b d, gate: b d
         shift, scale, gate = mod_params.chunk(3, dim=-1)
 
         if index is not None:
@@ -777,7 +777,7 @@ class QwenImageTransformerBlock(nn.Module):
             scale_result = scale.unsqueeze(1)
             gate_result = gate.unsqueeze(1)
 
-        return x * (1 + scale_result) + shift_result, gate_result
+        return scale_result, shift_result, gate_result
 
     def forward(
         self,
@@ -803,10 +803,12 @@ class QwenImageTransformerBlock(nn.Module):
         txt_mod1, txt_mod2 = txt_mod_params.chunk(2, dim=-1)  # Each [B, 3*dim]
 
         # Process image stream - norm1 + modulation
-        img_modulated, img_gate1 = self.img_norm1(hidden_states, img_mod1, modulate_index)
+        img_scale1, img_shift1, img_gate1 = self._modulate(img_mod1, modulate_index)
+        img_modulated = self.img_norm1(hidden_states, img_scale1, img_shift1)
 
         # Process text stream - norm1 + modulation
-        txt_modulated, txt_gate1 = self.txt_norm1(encoder_hidden_states, txt_mod1)
+        txt_scale1, txt_shift1, txt_gate1 = self._modulate(txt_mod1)
+        txt_modulated = self.txt_norm1(encoder_hidden_states, txt_scale1, txt_shift1)
 
         # Use QwenAttnProcessor2_0 for joint attention computation
         # This directly implements the DoubleStreamLayerMegatron logic:
@@ -831,13 +833,16 @@ class QwenImageTransformerBlock(nn.Module):
         encoder_hidden_states = encoder_hidden_states + txt_gate1 * txt_attn_output
 
         # Process image stream - norm2 + MLP
-        img_modulated2, img_gate2 = self.img_norm2(hidden_states, img_mod2, modulate_index)
+        img_scale2, img_shift2, img_gate2 = self._modulate(img_mod2, modulate_index)
+        img_modulated2 = self.img_norm2(hidden_states, img_scale2, img_shift2)
 
         img_mlp_output = self.img_mlp(img_modulated2)
         hidden_states = hidden_states + img_gate2 * img_mlp_output
 
         # Process text stream - norm2 + MLP
-        txt_modulated2, txt_gate2 = self.txt_norm2(encoder_hidden_states, txt_mod2)
+        txt_scale2, txt_shift2, txt_gate2 = self._modulate(txt_mod2)
+        txt_modulated2 = self.txt_norm2(encoder_hidden_states, txt_scale2, txt_shift2)
+
         txt_mlp_output = self.txt_mlp(txt_modulated2)
         encoder_hidden_states = encoder_hidden_states + txt_gate2 * txt_mlp_output
 

vllm_omni/diffusion/models/wan2_2/__init__.py

@@ -1,3 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from .patch_diffusers import patch_wan_rms_norm
 from .pipeline_wan2_2 import (
     Wan22Pipeline,
     create_transformer_from_config,
@@ -33,3 +37,5 @@ __all__ = [
     "get_wan22_ti2v_pre_process_func",
     "WanTransformer3DModel",
 ]
+
+patch_wan_rms_norm()

vllm_omni/diffusion/models/wan2_2/patch_diffusers.py

@@ -0,0 +1,14 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import sys
+
+from vllm_omni.diffusion.layers.norm import RMSNormVAE
+
+
+def patch_wan_rms_norm():
+    """Patch diffusers Wan RMSNorm implementation with RMSNormVAE."""
+
+    for module_name, module in sys.modules.items():
+        if hasattr(module, "WanRMS_norm"):
+            setattr(module, "WanRMS_norm", RMSNormVAE)

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py

@@ -25,6 +25,7 @@ from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineL
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero
 from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler
 from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel
+from vllm_omni.diffusion.postprocess import interpolate_video_tensor
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.inputs.data import OmniTextPrompt
@@ -121,10 +122,23 @@ def get_wan22_post_process_func(
     def post_process_func(
         video: torch.Tensor,
         output_type: str = "np",
+        sampling_params=None,
     ):
         if output_type == "latent":
             return video
-        return video_processor.postprocess_video(video, output_type=output_type)
+        custom_output = {}
+        if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False):
+            video, multiplier = interpolate_video_tensor(
+                video,
+                exp=sampling_params.frame_interpolation_exp,
+                scale=sampling_params.frame_interpolation_scale,
+                model_path=sampling_params.frame_interpolation_model_path,
+            )
+            custom_output["video_fps_multiplier"] = multiplier
+        return {
+            "video": video_processor.postprocess_video(video, output_type=output_type),
+            "custom_output": custom_output,
+        }
 
     return post_process_func
 
@@ -134,9 +148,6 @@ def get_wan22_pre_process_func(
 ):
     """Pre-process function for Wan2.2: optionally load and resize input image for I2V mode."""
     import numpy as np
-    from diffusers.video_processor import VideoProcessor
-
-    video_processor = VideoProcessor(vae_scale_factor=8)
 
     def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest:
         for i, prompt in enumerate(request.prompts):
@@ -180,10 +191,6 @@ def get_wan22_pre_process_func(
             )
             prompt["multi_modal_data"]["image"] = image  # type: ignore # key existence already checked above
 
-            # Preprocess for VAE
-            prompt["additional_information"]["preprocessed_image"] = video_processor.preprocess(
-                image, height=request.sampling_params.height, width=request.sampling_params.width
-            )
             request.prompts[i] = prompt
         return request
 
@@ -272,7 +279,7 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe
             model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only
         ).to(self.device)
         self.vae = DistributedAutoencoderKLWan.from_pretrained(
-            model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only
+            model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only
         ).to(self.device)
 
         # Initialize transformers with correct config (weights loaded via load_weights)

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py

@@ -12,6 +12,7 @@ from typing import Any, cast
 import numpy as np
 import PIL.Image
 import torch
+import torchvision.transforms.functional as TF
 from diffusers.utils.torch_utils import randn_tensor
 from torch import nn
 from transformers import AutoTokenizer, CLIPImageProcessor, CLIPVisionModel, UMT5EncoderModel
@@ -30,6 +31,7 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
     load_transformer_config,
     retrieve_latents,
 )
+from vllm_omni.diffusion.postprocess import interpolate_video_tensor
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.inputs.data import OmniTextPrompt
@@ -72,10 +74,23 @@ def get_wan22_i2v_post_process_func(
     def post_process_func(
         video: torch.Tensor,
         output_type: str = "np",
+        sampling_params=None,
     ):
         if output_type == "latent":
             return video
-        return video_processor.postprocess_video(video, output_type=output_type)
+        custom_output = {}
+        if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False):
+            video, multiplier = interpolate_video_tensor(
+                video,
+                exp=sampling_params.frame_interpolation_exp,
+                scale=sampling_params.frame_interpolation_scale,
+                model_path=sampling_params.frame_interpolation_model_path,
+            )
+            custom_output["video_fps_multiplier"] = multiplier
+        return {
+            "video": video_processor.postprocess_video(video, output_type=output_type),
+            "custom_output": custom_output,
+        }
 
     return post_process_func
 
@@ -84,9 +99,6 @@ def get_wan22_i2v_pre_process_func(
     od_config: OmniDiffusionConfig,
 ):
     """Pre-process function for I2V: load and resize input image."""
-    from diffusers.video_processor import VideoProcessor
-
-    video_processor = VideoProcessor(vae_scale_factor=8)
 
     def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest:
         for i, prompt in enumerate(request.prompts):
@@ -132,10 +144,6 @@ def get_wan22_i2v_pre_process_func(
             )
             prompt["multi_modal_data"]["image"] = image  # type: ignore # key existence already checked above
 
-            # Preprocess for VAE
-            prompt["additional_information"]["preprocessed_image"] = video_processor.preprocess(
-                image, height=request.sampling_params.height, width=request.sampling_params.width
-            )
             request.prompts[i] = prompt
         return request
 
@@ -217,7 +225,7 @@ class Wan22I2VPipeline(
 
         # VAE
         self.vae = DistributedAutoencoderKLWan.from_pretrained(
-            model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only
+            model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only
         ).to(self.device)
 
         # Transformers (weights loaded via load_weights)
@@ -459,6 +467,7 @@ class Wan22I2VPipeline(
         video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)
 
         if isinstance(image, PIL.Image.Image):
+            image = TF.to_tensor(image).to(device)
             image_tensor = video_processor.preprocess(image, height=height, width=width)
         else:
             image_tensor = image
@@ -467,6 +476,7 @@ class Wan22I2VPipeline(
         # Handle last_image if provided
         if last_image is not None:
             if isinstance(last_image, PIL.Image.Image):
+                last_image = TF.to_tensor(last_image).to(device)
                 last_image_tensor = video_processor.preprocess(last_image, height=height, width=width)
             else:
                 last_image_tensor = last_image
@@ -789,12 +799,14 @@ class Wan22I2VPipeline(
             return latents, latent_condition, first_frame_mask
 
         # Wan2.1 style: create mask and concatenate with condition
-        mask_lat_size = torch.ones(batch_size, 1, num_frames, latent_height, latent_width)
+        mask_lat_size = torch.ones(
+            batch_size, 1, num_frames, latent_height, latent_width, device=latent_condition.device
+        )
 
         if last_image is None:
-            mask_lat_size[:, :, list(range(1, num_frames))] = 0
+            mask_lat_size[:, :, 1:] = 0
         else:
-            mask_lat_size[:, :, list(range(1, num_frames - 1))] = 0
+            mask_lat_size[:, :, 1 : num_frames - 1] = 0
 
         first_frame_mask = mask_lat_size[:, :, 0:1]
         first_frame_mask = torch.repeat_interleave(first_frame_mask, dim=2, repeats=self.vae_scale_factor_temporal)

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py

@@ -24,13 +24,13 @@ from typing import Any, cast
 import numpy as np
 import PIL.Image
 import torch
-from diffusers import AutoencoderKLWan
 from diffusers.utils.torch_utils import randn_tensor
 from torch import nn
 from transformers import AutoTokenizer, UMT5EncoderModel
 from vllm.model_executor.models.utils import AutoWeightsLoader
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
@@ -42,6 +42,7 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
     load_transformer_config,
     retrieve_latents,
 )
+from vllm_omni.diffusion.postprocess import interpolate_video_tensor
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.inputs.data import OmniTextPrompt
 from vllm_omni.platforms import current_omni_platform
@@ -59,10 +60,23 @@ def get_wan22_ti2v_post_process_func(
     def post_process_func(
         video: torch.Tensor,
         output_type: str = "np",
+        sampling_params=None,
     ):
         if output_type == "latent":
             return video
-        return video_processor.postprocess_video(video, output_type=output_type)
+        custom_output = {}
+        if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False):
+            video, multiplier = interpolate_video_tensor(
+                video,
+                exp=sampling_params.frame_interpolation_exp,
+                scale=sampling_params.frame_interpolation_scale,
+                model_path=sampling_params.frame_interpolation_model_path,
+            )
+            custom_output["video_fps_multiplier"] = multiplier
+        return {
+            "video": video_processor.postprocess_video(video, output_type=output_type),
+            "custom_output": custom_output,
+        }
 
     return post_process_func
 
@@ -71,9 +85,6 @@ def get_wan22_ti2v_pre_process_func(
     od_config: OmniDiffusionConfig,
 ):
     """Pre-process function for TI2V: optionally load and resize input image."""
-    from diffusers.video_processor import VideoProcessor
-
-    video_processor = VideoProcessor(vae_scale_factor=8)
 
     def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest:
         for i, prompt in enumerate(request.prompts):
@@ -119,10 +130,6 @@ def get_wan22_ti2v_pre_process_func(
             )
             prompt["multi_modal_data"]["image"] = image  # type: ignore # key existence already checked above
 
-            # Preprocess for VAE
-            prompt["additional_information"]["preprocessed_image"] = video_processor.preprocess(
-                image, height=request.sampling_params.height, width=request.sampling_params.width
-            )
             request.prompts[i] = prompt
         return request
 
@@ -174,8 +181,8 @@ class Wan22TI2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, Progress
         ).to(self.device)
 
         # VAE
-        self.vae = AutoencoderKLWan.from_pretrained(
-            model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only
+        self.vae = OmniAutoencoderKLWan.from_pretrained(
+            model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only
         ).to(self.device)
 
         # Single transformer (TI2V uses dense 5B model, not MoE)

vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py

@@ -11,7 +11,6 @@ import torch.nn.functional as F
 from diffusers.models.attention import FeedForward
 from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbedding, Timesteps
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
-from diffusers.models.normalization import FP32LayerNorm
 from vllm.distributed import (
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
@@ -29,6 +28,8 @@ from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelOutput,
 )
 from vllm_omni.diffusion.forward_context import get_forward_context
+from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm
+from vllm_omni.diffusion.layers.norm import LayerNorm, RMSNorm
 
 logger = init_logger(__name__)
 
@@ -52,10 +53,14 @@ def apply_rotary_emb_wan(
     x1, x2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1)
     cos = freqs_cos[..., 0::2]
     sin = freqs_sin[..., 1::2]
-    out = torch.empty_like(hidden_states)
-    out[..., 0::2] = x1 * cos - x2 * sin
-    out[..., 1::2] = x1 * sin + x2 * cos
-    return out.type_as(hidden_states)
+    rotated = torch.stack(
+        (
+            x1 * cos - x2 * sin,
+            x1 * sin + x2 * cos,
+        ),
+        dim=-1,
+    )
+    return rotated.flatten(-2, -1).to(hidden_states.dtype)
 
 
 class DistributedRMSNorm(nn.Module):
@@ -230,9 +235,9 @@ class WanImageEmbedding(nn.Module):
     def __init__(self, in_features: int, out_features: int, pos_embed_seq_len: int | None = None):
         super().__init__()
 
-        self.norm1 = FP32LayerNorm(in_features)
+        self.norm1 = LayerNorm(in_features)
         self.ff = FeedForward(in_features, out_features, mult=1, activation_fn="gelu")
-        self.norm2 = FP32LayerNorm(out_features)
+        self.norm2 = LayerNorm(out_features)
         if pos_embed_seq_len is not None:
             self.pos_embed = nn.Parameter(torch.zeros(1, pos_embed_seq_len, in_features))
         else:
@@ -372,8 +377,12 @@ class WanSelfAttention(nn.Module):
         self.tp_inner_dim = self.num_heads * head_dim
 
         # QK normalization using vLLM's RMSNorm
-        self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
-        self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
+        if get_tensor_model_parallel_world_size() > 1:
+            self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
+            self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
+        else:
+            self.norm_q = RMSNorm(self.tp_inner_dim, eps=eps)
+            self.norm_k = RMSNorm(self.tp_inner_dim, eps=eps)
 
         self.to_out = RowParallelLinear(
             self.inner_dim,
@@ -492,8 +501,12 @@ class WanCrossAttention(nn.Module):
         self.tp_inner_dim = self.num_heads * head_dim
 
         # QK normalization
-        self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
-        self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
+        if get_tensor_model_parallel_world_size() > 1:
+            self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
+            self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
+        else:
+            self.norm_q = RMSNorm(self.tp_inner_dim, eps=eps)
+            self.norm_k = RMSNorm(self.tp_inner_dim, eps=eps)
 
         # Optional added KV projections for I2V (image embeddings)
         self.added_kv_proj_dim = added_kv_proj_dim
@@ -512,7 +525,10 @@ class WanCrossAttention(nn.Module):
                 gather_output=False,
                 return_bias=False,
             )
-            self.norm_added_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
+            if get_tensor_model_parallel_world_size() > 1:
+                self.norm_added_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
+            else:
+                self.norm_added_k = RMSNorm(self.tp_inner_dim, eps=eps)
         else:
             self.add_k_proj = None
             self.add_v_proj = None
@@ -535,6 +551,7 @@ class WanCrossAttention(nn.Module):
             num_kv_heads=self.num_heads,
             softmax_scale=1.0 / (head_dim**0.5),
             causal=False,
+            skip_sequence_parallel=True,
         )
 
     def forward(
@@ -614,7 +631,8 @@ class WanTransformerBlock(nn.Module):
         head_dim = dim // num_heads
 
         # 1. Self-attention
-        self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
+        self.norm1 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps)
+
         self.attn1 = WanSelfAttention(
             dim=dim,
             num_heads=num_heads,
@@ -630,11 +648,11 @@ class WanTransformerBlock(nn.Module):
             eps=eps,
             added_kv_proj_dim=added_kv_proj_dim,
         )
-        self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()
+        self.norm2 = LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()
 
         # 3. Feed-forward
         self.ffn = WanFeedForward(dim=dim, inner_dim=ffn_dim, dim_out=dim)
-        self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False)
+        self.norm3 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps)
 
         # Scale-shift table for modulation
         self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
@@ -650,7 +668,7 @@ class WanTransformerBlock(nn.Module):
         if temb.ndim == 4:
             # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v)
             shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
-                self.scale_shift_table.unsqueeze(0) + temb.float()
+                self.scale_shift_table.unsqueeze(0) + temb
             ).chunk(6, dim=2)
             shift_msa = shift_msa.squeeze(2)
             scale_msa = scale_msa.squeeze(2)
@@ -661,25 +679,23 @@ class WanTransformerBlock(nn.Module):
         else:
             # temb: batch_size, 6, inner_dim (wan2.1/wan2.2 14B)
             shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
-                self.scale_shift_table + temb.float()
+                self.scale_shift_table + temb
             ).chunk(6, dim=1)
 
         # 1. Self-attention
-        norm_hidden_states = (self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa).type_as(hidden_states)
+        norm_hidden_states = self.norm1(hidden_states, scale_msa, shift_msa).type_as(hidden_states)
         attn_output = self.attn1(norm_hidden_states, rotary_emb, hidden_states_mask)
-        hidden_states = (hidden_states.float() + attn_output * gate_msa).type_as(hidden_states)
+        hidden_states = (hidden_states + attn_output * gate_msa).type_as(hidden_states)
 
         # 2. Cross-attention
-        norm_hidden_states = self.norm2(hidden_states.float()).type_as(hidden_states)
+        norm_hidden_states = self.norm2(hidden_states).type_as(hidden_states)
         attn_output = self.attn2(norm_hidden_states, encoder_hidden_states)
         hidden_states = hidden_states + attn_output
 
         # 3. Feed-forward
-        norm_hidden_states = (self.norm3(hidden_states.float()) * (1 + c_scale_msa) + c_shift_msa).type_as(
-            hidden_states
-        )
+        norm_hidden_states = self.norm3(hidden_states, c_scale_msa, c_shift_msa).type_as(hidden_states)
         ff_output = self.ffn(norm_hidden_states)
-        hidden_states = (hidden_states.float() + ff_output.float() * c_gate_msa).type_as(hidden_states)
+        hidden_states = (hidden_states + ff_output * c_gate_msa).type_as(hidden_states)
 
         return hidden_states
 
@@ -848,7 +864,7 @@ class WanTransformer3DModel(nn.Module):
         )
 
         # 4. Output norm & projection
-        self.norm_out = FP32LayerNorm(inner_dim, eps, elementwise_affine=False)
+        self.norm_out = AdaLayerNorm(inner_dim, elementwise_affine=False, eps=eps)
         self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size))
 
         # SP helper modules
@@ -936,7 +952,7 @@ class WanTransformer3DModel(nn.Module):
             shift = shift.unsqueeze(1)
             scale = scale.unsqueeze(1)
 
-        hidden_states = (self.norm_out(hidden_states.float()) * (1 + scale) + shift).type_as(hidden_states)
+        hidden_states = self.norm_out(hidden_states, scale, shift).type_as(hidden_states)
         hidden_states = self.proj_out(hidden_states)
 
         hidden_states = hidden_states.reshape(

vllm_omni/diffusion/postprocess/__init__.py

@@ -0,0 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Diffusion post-processing helpers."""
+
+from vllm_omni.diffusion.postprocess.rife_interpolator import (
+    FrameInterpolator,
+    interpolate_video_tensor,
+)
+
+__all__ = ["FrameInterpolator", "interpolate_video_tensor"]

vllm_omni/diffusion/postprocess/rife_interpolator.py

@@ -0,0 +1,443 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+RIFE 4.22.lite frame interpolation for vLLM-Omni video generation.
+
+RIFE model code is vendored and adapted from:
+  - https://github.com/hzwer/ECCV2022-RIFE  (MIT License)
+  - https://github.com/hzwer/Practical-RIFE  (MIT License)
+  Copyright (c) 2021 Zhewei Huang
+
+The FrameInterpolator wrapper and vLLM-Omni integration code are original work.
+"""
+
+from __future__ import annotations
+
+import os
+import threading
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+_DEFAULT_RIFE_HF_REPO = "elfgum/RIFE-4.22.lite"
+_MODEL_CACHE: dict[tuple[str, str], Model] = {}
+_MODEL_CACHE_LOCK = threading.Lock()
+
+
+def warp(ten_input: torch.Tensor, ten_flow: torch.Tensor) -> torch.Tensor:
+    """Warp input tensor by optical flow using grid_sample."""
+    ten_horizontal = (
+        torch.linspace(-1.0, 1.0, ten_flow.shape[3], device=ten_flow.device)
+        .view(1, 1, 1, ten_flow.shape[3])
+        .expand(ten_flow.shape[0], -1, ten_flow.shape[2], -1)
+    )
+    ten_vertical = (
+        torch.linspace(-1.0, 1.0, ten_flow.shape[2], device=ten_flow.device)
+        .view(1, 1, ten_flow.shape[2], 1)
+        .expand(ten_flow.shape[0], -1, -1, ten_flow.shape[3])
+    )
+    ten_grid = torch.cat([ten_horizontal, ten_vertical], dim=1)
+
+    ten_flow = torch.cat(
+        [
+            ten_flow[:, 0:1, :, :] / ((ten_input.shape[3] - 1.0) / 2.0),
+            ten_flow[:, 1:2, :, :] / ((ten_input.shape[2] - 1.0) / 2.0),
+        ],
+        dim=1,
+    )
+    grid = (ten_grid + ten_flow).permute(0, 2, 3, 1)
+    return F.grid_sample(
+        input=ten_input,
+        grid=grid,
+        mode="bilinear",
+        padding_mode="border",
+        align_corners=True,
+    )
+
+
+def _conv(
+    in_planes: int,
+    out_planes: int,
+    kernel_size: int = 3,
+    stride: int = 1,
+    padding: int = 1,
+    dilation: int = 1,
+) -> nn.Sequential:
+    return nn.Sequential(
+        nn.Conv2d(
+            in_planes,
+            out_planes,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            bias=True,
+        ),
+        nn.LeakyReLU(0.2, True),
+    )
+
+
+class ResConv(nn.Module):
+    """Residual convolution block with learnable beta scaling."""
+
+    def __init__(self, c: int, dilation: int = 1):
+        super().__init__()
+        self.conv = nn.Conv2d(c, c, 3, 1, dilation, dilation=dilation, groups=1)
+        self.beta = nn.Parameter(torch.ones((1, c, 1, 1)), requires_grad=True)
+        self.relu = nn.LeakyReLU(0.2, True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.relu(self.conv(x) * self.beta + x)
+
+
+class IFBlock(nn.Module):
+    """Single-scale optical flow, mask, and feature block."""
+
+    def __init__(self, in_planes: int, c: int = 64):
+        super().__init__()
+        self.conv0 = nn.Sequential(
+            _conv(in_planes, c // 2, 3, 2, 1),
+            _conv(c // 2, c, 3, 2, 1),
+        )
+        self.convblock = nn.Sequential(
+            ResConv(c),
+            ResConv(c),
+            ResConv(c),
+            ResConv(c),
+            ResConv(c),
+            ResConv(c),
+            ResConv(c),
+            ResConv(c),
+        )
+        self.lastconv = nn.Sequential(
+            nn.ConvTranspose2d(c, 4 * 13, 4, 2, 1),
+            nn.PixelShuffle(2),
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        flow: torch.Tensor | None = None,
+        scale: float = 1.0,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        x = F.interpolate(x, scale_factor=1.0 / scale, mode="bilinear", align_corners=False)
+        if flow is not None:
+            flow = (
+                F.interpolate(
+                    flow,
+                    scale_factor=1.0 / scale,
+                    mode="bilinear",
+                    align_corners=False,
+                )
+                * 1.0
+                / scale
+            )
+            x = torch.cat((x, flow), 1)
+        feat = self.conv0(x)
+        feat = self.convblock(feat)
+        tmp = self.lastconv(feat)
+        tmp = F.interpolate(tmp, scale_factor=scale, mode="bilinear", align_corners=False)
+        flow = tmp[:, :4] * scale
+        mask = tmp[:, 4:5]
+        feat = tmp[:, 5:]
+        return flow, mask, feat
+
+
+class Head(nn.Module):
+    """Feature encoder producing four-channel features at full resolution."""
+
+    def __init__(self):
+        super().__init__()
+        self.cnn0 = nn.Conv2d(3, 16, 3, 2, 1)
+        self.cnn1 = nn.Conv2d(16, 16, 3, 1, 1)
+        self.cnn2 = nn.Conv2d(16, 16, 3, 1, 1)
+        self.cnn3 = nn.ConvTranspose2d(16, 4, 4, 2, 1)
+        self.relu = nn.LeakyReLU(0.2, True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x0 = self.cnn0(x)
+        x = self.relu(x0)
+        x1 = self.cnn1(x)
+        x = self.relu(x1)
+        x2 = self.cnn2(x)
+        x = self.relu(x2)
+        x3 = self.cnn3(x)
+        return x3
+
+
+class IFNet(nn.Module):
+    """Four-scale IFNet optical flow network."""
+
+    def __init__(self):
+        super().__init__()
+        self.block0 = IFBlock(7 + 8, c=192)
+        self.block1 = IFBlock(8 + 4 + 8 + 8, c=128)
+        self.block2 = IFBlock(8 + 4 + 8 + 8, c=64)
+        self.block3 = IFBlock(8 + 4 + 8 + 8, c=32)
+        self.encode = Head()
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        timestep: float = 0.5,
+        scale_list: list[float] | None = None,
+    ) -> tuple[list[torch.Tensor], torch.Tensor, list[tuple[torch.Tensor, torch.Tensor] | torch.Tensor]]:
+        if scale_list is None:
+            scale_list = [8, 4, 2, 1]
+
+        channel = x.shape[1] // 2
+        img0 = x[:, :channel]
+        img1 = x[:, channel:]
+
+        if not torch.is_tensor(timestep):
+            timestep = (x[:, :1].clone() * 0 + 1) * timestep
+        else:
+            timestep = timestep.repeat(1, 1, img0.shape[2], img0.shape[3])
+
+        f0 = self.encode(img0[:, :3])
+        f1 = self.encode(img1[:, :3])
+
+        flow_list: list[torch.Tensor] = []
+        merged: list[tuple[torch.Tensor, torch.Tensor] | torch.Tensor] = []
+        mask_list: list[torch.Tensor] = []
+        warped_img0 = img0
+        warped_img1 = img1
+        flow = None
+        mask = None
+
+        for i, block in enumerate([self.block0, self.block1, self.block2, self.block3]):
+            if flow is None:
+                flow, mask, feat = block(
+                    torch.cat((img0[:, :3], img1[:, :3], f0, f1, timestep), 1),
+                    None,
+                    scale=scale_list[i],
+                )
+            else:
+                wf0 = warp(f0, flow[:, :2])
+                wf1 = warp(f1, flow[:, 2:4])
+                fd, m0, feat = block(
+                    torch.cat(
+                        (
+                            warped_img0[:, :3],
+                            warped_img1[:, :3],
+                            wf0,
+                            wf1,
+                            timestep,
+                            mask,
+                            feat,
+                        ),
+                        1,
+                    ),
+                    flow,
+                    scale=scale_list[i],
+                )
+                mask = m0
+                flow = flow + fd
+
+            mask_list.append(mask)
+            flow_list.append(flow)
+            warped_img0 = warp(img0, flow[:, :2])
+            warped_img1 = warp(img1, flow[:, 2:4])
+            merged.append((warped_img0, warped_img1))
+
+        mask = torch.sigmoid(mask)
+        merged[3] = warped_img0 * mask + warped_img1 * (1 - mask)
+        return flow_list, mask_list[3], merged
+
+
+class Model:
+    """Wraps IFNet and exposes RIFE-compatible load/inference helpers."""
+
+    def __init__(self):
+        self.flownet = IFNet()
+
+    def eval(self) -> Model:
+        self.flownet.eval()
+        return self
+
+    def device(self) -> torch.device:
+        return next(self.flownet.parameters()).device
+
+    def load_model(self, path: str) -> None:
+        flownet_path = os.path.join(path, "flownet.pkl")
+        if not os.path.isfile(flownet_path):
+            raise FileNotFoundError(
+                f"RIFE weight file not found: {flownet_path}. Expected layout: <model_path>/flownet.pkl"
+            )
+
+        state = torch.load(flownet_path, map_location="cpu", weights_only=False)
+        state = {k.removeprefix("module."): v for k, v in state.items()}
+        self.flownet.load_state_dict(state, strict=False)
+        logger.info("Loaded RIFE weights from %s", flownet_path)
+
+    def inference(
+        self,
+        img0: torch.Tensor,
+        img1: torch.Tensor,
+        scale: float = 1.0,
+        timestep: float = 0.5,
+    ) -> torch.Tensor:
+        _n, _c, h, w = img0.shape
+        ph = ((h - 1) // 32 + 1) * 32
+        pw = ((w - 1) // 32 + 1) * 32
+        pad = (0, pw - w, 0, ph - h)
+        img0 = F.pad(img0, pad)
+        img1 = F.pad(img1, pad)
+
+        imgs = torch.cat((img0, img1), 1)
+        scale_list = [8 / scale, 4 / scale, 2 / scale, 1 / scale]
+        with torch.no_grad():
+            _flow_list, _mask, merged = self.flownet(
+                imgs,
+                timestep=timestep,
+                scale_list=scale_list,
+            )
+        return merged[3][:, :, :h, :w]
+
+
+def _resolve_rife_model_path(model_path: str | None) -> str:
+    model_path = model_path or _DEFAULT_RIFE_HF_REPO
+    if os.path.isdir(model_path):
+        return model_path
+    from vllm_omni.model_executor.model_loader.weight_utils import (
+        download_weights_from_hf_specific,
+    )
+
+    return download_weights_from_hf_specific(
+        model_path,
+        cache_dir=None,
+        allow_patterns=["flownet.pkl"],
+        require_all=True,
+    )
+
+
+def _select_torch_device() -> torch.device:
+    try:
+        from vllm_omni.platforms import current_omni_platform
+
+        return current_omni_platform.get_torch_device()
+    except Exception as exc:
+        logger.warning("Failed to resolve current vLLM-Omni torch device: %s", exc)
+
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    return torch.device("cpu")
+
+
+def _normalize_video_tensor_layout(video: torch.Tensor) -> tuple[torch.Tensor, Any]:
+    if video.ndim == 5:
+        if video.shape[1] in (3, 4):
+            return video, lambda out: out
+        if video.shape[2] in (3, 4):
+            return video.permute(0, 2, 1, 3, 4), lambda out: out.permute(0, 2, 1, 3, 4)
+    elif video.ndim == 4:
+        if video.shape[0] in (3, 4):
+            return video.unsqueeze(0), lambda out: out.squeeze(0)
+        if video.shape[1] in (3, 4):
+            return video.permute(1, 0, 2, 3).unsqueeze(0), lambda out: out.squeeze(0).permute(1, 0, 2, 3)
+    raise ValueError(f"Unsupported video tensor shape for interpolation: {tuple(video.shape)}")
+
+
+def _normalize_video_tensor_range(video: torch.Tensor) -> tuple[torch.Tensor, Any]:
+    original_dtype = video.dtype
+    video = video.detach()
+    if video.is_floating_point():
+        video = video.to(torch.float32)
+        if torch.amin(video) < 0.0 or torch.amax(video) > 1.0:
+            return video.clamp(-1.0, 1.0) * 0.5 + 0.5, lambda out: (out * 2.0 - 1.0).to(original_dtype)
+        return video.clamp(0.0, 1.0), lambda out: out.to(original_dtype)
+    return video.to(torch.float32) / 255.0, lambda out: (out * 255.0).round().clamp(0, 255).to(original_dtype)
+
+
+class FrameInterpolator:
+    """Lazy-loaded RIFE 4.22.lite frame interpolator."""
+
+    def __init__(self, model_path: str | None = None):
+        self._model_path = model_path
+        self._resolved_path: str | None = None
+
+    def _ensure_model_loaded(self, preferred_device: torch.device | None = None) -> Model:
+        resolved_path = _resolve_rife_model_path(self._model_path)
+        self._resolved_path = resolved_path
+        device = preferred_device or _select_torch_device()
+        cache_key = (resolved_path, str(device))
+
+        with _MODEL_CACHE_LOCK:
+            if cache_key in _MODEL_CACHE:
+                return _MODEL_CACHE[cache_key]
+
+            model = Model()
+            model.load_model(resolved_path)
+            model.eval()
+            model.flownet = model.flownet.to(device)
+            _MODEL_CACHE[cache_key] = model
+            logger.info("RIFE model loaded on device: %s", device)
+            return model
+
+    def _make_inference(
+        self,
+        model: Model,
+        img0: torch.Tensor,
+        img1: torch.Tensor,
+        n: int,
+        scale: float,
+    ) -> list[torch.Tensor]:
+        if n == 1:
+            return [model.inference(img0, img1, scale=scale)]
+        mid = model.inference(img0, img1, scale=scale)
+        return (
+            self._make_inference(model, img0, mid, n // 2, scale)
+            + [mid]
+            + self._make_inference(model, mid, img1, n // 2, scale)
+        )
+
+    def interpolate_tensor(
+        self,
+        video: torch.Tensor,
+        exp: int = 1,
+        scale: float = 1.0,
+    ) -> tuple[torch.Tensor, int]:
+        if exp < 1:
+            raise ValueError(f"frame interpolation exp must be >= 1, got {exp}")
+        if scale <= 0:
+            raise ValueError(f"frame interpolation scale must be > 0, got {scale}")
+
+        video, restore_layout = _normalize_video_tensor_layout(video)
+        if video.shape[2] < 2:
+            return restore_layout(video), 1
+
+        video, restore_range = _normalize_video_tensor_range(video)
+        # A CPU tensor may be transport/offload state rather than an execution
+        # choice, so only trust it when it is already on an accelerator.
+        preferred_device = video.device
+        if preferred_device.type == "cpu":
+            preferred_device = _select_torch_device()
+        model = self._ensure_model_loaded(preferred_device=preferred_device)
+        video = video.to(model.device())
+        intermediates_per_pair = 2**exp // 2
+
+        result_frames: list[torch.Tensor] = []
+        for idx in range(video.shape[2] - 1):
+            img0 = video[:, :, idx, :, :]
+            img1 = video[:, :, idx + 1, :, :]
+            result_frames.append(img0)
+            result_frames.extend(self._make_inference(model, img0, img1, intermediates_per_pair, scale))
+        result_frames.append(video[:, :, -1, :, :])
+        result = torch.stack(result_frames, dim=2)
+        return restore_layout(restore_range(result)), 2**exp
+
+
+def interpolate_video_tensor(
+    video: torch.Tensor,
+    exp: int = 1,
+    scale: float = 1.0,
+    model_path: str | None = None,
+) -> tuple[torch.Tensor, int]:
+    """Interpolate a video tensor and return the FPS multiplier."""
+    interpolator = FrameInterpolator(model_path=model_path)
+    return interpolator.interpolate_tensor(video, exp=exp, scale=scale)

vllm_omni/diffusion/stage_diffusion_client.py

@@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any
 
 from vllm.logger import init_logger
 
+from vllm_omni.diffusion.data import normalize_omni_error
 from vllm_omni.engine.stage_init_utils import StageMetadata
 from vllm_omni.entrypoints.async_omni_diffusion import AsyncOmniDiffusion
 from vllm_omni.outputs import OmniRequestOutput
@@ -23,6 +24,21 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 
+def create_diffusion_client(
+    model: str,
+    od_config: OmniDiffusionConfig,
+    metadata: StageMetadata,
+    batch_size: int = 1,
+    use_inline: bool = False,
+) -> Any:
+    """Factory to create either an inline or the default diffusion client."""
+    if use_inline:
+        from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient
+
+        return InlineStageDiffusionClient(model, od_config, metadata, batch_size=batch_size)
+    return StageDiffusionClient(model, od_config, metadata, batch_size=batch_size)
+
+
 class StageDiffusionClient:
     """Wraps AsyncOmniDiffusion for use inside the Orchestrator.
 
@@ -75,11 +91,24 @@ class StageDiffusionClient:
             result = await self._engine.generate(prompt, sampling_params, request_id)
             await self._output_queue.put(result)
         except Exception as e:
+            err = normalize_omni_error(e, request_id=request_id, stage_id=self.stage_id)
             logger.exception(
                 "[StageDiffusionClient] Stage-%s req=%s failed: %s",
                 self.stage_id,
                 request_id,
-                e,
+                err,
+            )
+            await self._output_queue.put(
+                {
+                    "type": "error",
+                    "request_id": request_id,
+                    "stage_id": self.stage_id,
+                    "status_code": err.status_code,
+                    "error": str(err),
+                    "error_type": err.error_type,
+                    "detail": err.detail,
+                    "finished": True,
+                }
             )
         finally:
             self._tasks.pop(request_id, None)
@@ -162,10 +191,8 @@ class StageDiffusionClient:
                     "todo": True,
                     "reason": f"AsyncOmniDiffusion.{method} is not implemented",
                 }
-            # Extract is_start and profile_prefix from args
             is_start = args[0] if args else True
             profile_prefix = args[1] if len(args) > 1 else None
-            # Generate profile_prefix with stage_id if starting and no prefix provided
             if is_start and profile_prefix is None:
                 profile_prefix = f"stage_{self.stage_id}_diffusion_{int(time.time())}"
             result = target(is_start, profile_prefix)
@@ -186,7 +213,6 @@ class StageDiffusionClient:
                 return await asyncio.wait_for(result, timeout=timeout)
             return await result
 
-        # Fall back to collective RPC for other methods
         loop = asyncio.get_running_loop()
         return await loop.run_in_executor(
             self._engine._executor,

vllm_omni/diffusion/utils/media_utils.py

@@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Video/audio muxing utilities using PyAV (no ffmpeg binary dependency)."""
+
+from __future__ import annotations
+
+import io
+from fractions import Fraction
+
+import av
+import numpy as np
+
+
+def mux_video_audio_bytes(
+    video_frames: np.ndarray,
+    audio_waveform: np.ndarray | None = None,
+    *,
+    fps: float = 25.0,
+    audio_sample_rate: int = 44100,
+    video_codec: str = "h264",
+    audio_codec: str = "aac",
+    crf: str = "18",
+    video_codec_options: dict[str, str] | None = None,
+) -> bytes:
+    """Mux video frames and optional audio waveform into MP4 bytes.
+
+    Args:
+        video_frames: uint8 array of shape ``(T, H, W, 3)`` (RGB).
+        audio_waveform: float32 array – mono ``(N,)`` or ``(N, C)`` / ``(C, N)``.
+        fps: Video frame rate.
+        audio_sample_rate: Audio sample rate in Hz.
+        video_codec: Video codec name.
+        audio_codec: Audio codec name.
+        crf: Constant rate factor for the video encoder.
+
+    Returns:
+        Raw MP4 bytes ready to be written to disk or streamed.
+    """
+    buf = io.BytesIO()
+    container = av.open(buf, mode="w", format="mp4")
+
+    v_stream = container.add_stream(video_codec, rate=Fraction(fps).limit_denominator(10000))
+    v_stream.width = video_frames.shape[2]
+    v_stream.height = video_frames.shape[1]
+    v_stream.pix_fmt = "yuv420p"
+
+    options = {"crf": str(crf)}
+    if video_codec_options:
+        options.update(video_codec_options)
+    v_stream.options = options
+
+    a_stream = None
+    if audio_waveform is not None:
+        samples = audio_waveform.astype(np.float32)
+        if samples.ndim == 1:
+            samples = samples.reshape(1, -1)
+        elif samples.ndim == 2 and samples.shape[0] > samples.shape[1]:
+            samples = np.ascontiguousarray(samples.T)
+        num_channels = samples.shape[0]
+        layout = "stereo" if num_channels >= 2 else "mono"
+        a_stream = container.add_stream(audio_codec, rate=audio_sample_rate)
+        a_stream.layout = layout
+
+    for frame_data in video_frames:
+        frame = av.VideoFrame.from_ndarray(frame_data, format="rgb24")
+        for packet in v_stream.encode(frame):
+            container.mux(packet)
+    for packet in v_stream.encode():
+        container.mux(packet)
+
+    if a_stream is not None and audio_waveform is not None:
+        audio_frame = av.AudioFrame.from_ndarray(samples, format="fltp", layout=layout)
+        audio_frame.sample_rate = audio_sample_rate
+        for packet in a_stream.encode(audio_frame):
+            container.mux(packet)
+        for packet in a_stream.encode():
+            container.mux(packet)
+
+    container.close()
+    return buf.getvalue()

vllm_omni/diffusion/utils/prompt_utils.py

@@ -0,0 +1,38 @@
+import torch
+
+
+def validate_prompt_sequence_lengths(
+    attention_mask: torch.Tensor,
+    *,
+    max_sequence_length: int,
+    supported_max_sequence_length: int,
+    prompt_name: str = "prompt",
+    length_offset: int = 0,
+    baseline_attention_mask: torch.Tensor | None = None,
+    error_context: str,
+) -> None:
+    sequence_lengths = attention_mask.sum(dim=1)
+    if baseline_attention_mask is not None:
+        # Some callers need to validate only the user-controlled portion of a
+        # templated prompt. In those cases we subtract the fully-tokenized
+        # template baseline instead of only removing a fixed prefix length,
+        # because the template may also contribute a suffix or image markers.
+        baseline_lengths = baseline_attention_mask.sum(dim=1)
+        if baseline_lengths.shape[0] == 1 and sequence_lengths.shape[0] > 1:
+            baseline_lengths = baseline_lengths.expand(sequence_lengths.shape[0])
+        sequence_lengths = sequence_lengths - baseline_lengths
+    if length_offset:
+        sequence_lengths = sequence_lengths - length_offset
+    sequence_lengths = torch.clamp(sequence_lengths, min=0)
+    too_long = torch.nonzero(sequence_lengths > max_sequence_length, as_tuple=False)
+    if too_long.numel() == 0:
+        return
+
+    batch_idx = int(too_long[0].item())
+    actual_length = int(sequence_lengths[batch_idx].item())
+    prompt_ref = f"`{prompt_name}` at batch index {batch_idx}" if attention_mask.shape[0] > 1 else f"`{prompt_name}`"
+    raise ValueError(
+        f"{prompt_ref} is too long {error_context}: got {actual_length} tokens, but "
+        f"`max_sequence_length` is {max_sequence_length}. Shorten the prompt or increase "
+        f"`max_sequence_length` up to {supported_max_sequence_length}."
+    )

vllm_omni/diffusion/utils/size_utils.py

@@ -0,0 +1,20 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Shared size normalization helpers for diffusion pipelines."""
+
+
+def normalize_min_aligned_size(height: int, width: int, alignment: int) -> tuple[int, int]:
+    """Clamp dimensions to the minimum valid aligned size.
+
+    This preserves floor-to-alignment behavior for normal requests while
+    preventing very small dimensions from collapsing to zero after alignment.
+    """
+
+    alignment = int(alignment)
+    if alignment <= 0:
+        raise ValueError(f"Expected positive alignment, got {alignment}")
+
+    normalized_height = max(alignment, (int(height) // alignment) * alignment)
+    normalized_width = max(alignment, (int(width) // alignment) * alignment)
+    return normalized_height, normalized_width

vllm_omni/engine/async_omni_engine.py

@@ -487,11 +487,13 @@ class AsyncOmniEngine:
 
                                     inject_omni_kv_config(stage_cfg, omni_conn_cfg, omni_from, omni_to)
                                 _inject_kv_stage_info(stage_cfg, stage_id)
+                                use_inline = self.num_stages == 1
                                 stage_clients[stage_id] = initialize_diffusion_stage(
                                     self.model,
                                     stage_cfg,
                                     metadata,
                                     batch_size=self.diffusion_batch_size,
+                                    use_inline=use_inline,
                                 )
                                 logger.info(
                                     "[AsyncOmniEngine] Stage %s initialized (diffusion, batch_size=%d)",
@@ -798,6 +800,16 @@ class AsyncOmniEngine:
         # We temporally create a default config for diffusion stage.
         # In the future, we should merge the default config with the user-provided config.
         normalized_kwargs = dict(kwargs)
+        default_sampling_params = normalized_kwargs.get("default_sampling_params")
+        if isinstance(default_sampling_params, str):
+            try:
+                default_sampling_params = json.loads(default_sampling_params)
+            except json.JSONDecodeError:
+                logger.warning("Invalid default_sampling_params JSON, ignoring stage defaults.")
+                default_sampling_params = None
+        if not isinstance(default_sampling_params, dict):
+            default_sampling_params = None
+        stage_default_sampling_params = default_sampling_params.get("0", {}) if default_sampling_params else {}
 
         # TODO: hack, convert dtype to string to avoid non-premitive omegaconf create error.
         if "dtype" in normalized_kwargs and not isinstance(normalized_kwargs["dtype"], str):
@@ -886,6 +898,7 @@ class AsyncOmniEngine:
                         else {}
                     ),
                 },
+                "default_sampling_params": stage_default_sampling_params,
                 "final_output": True,
                 "final_output_type": "image",
             }

vllm_omni/engine/orchestrator.py

@@ -242,8 +242,48 @@ class Orchestrator:
                     output = stage_client.get_diffusion_output_async()
                     if output is not None:
                         idle = False
+
+                        if isinstance(output, dict) and output.get("type") == "error":
+                            req_id = output.get("request_id")
+                            req_state = self.request_states.get(req_id)
+
+                            await self.output_async_queue.put(
+                                {
+                                    "type": "error",
+                                    "request_id": req_id,
+                                    "stage_id": stage_id,
+                                    "status_code": output.get("status_code", 500),
+                                    "error": output.get("error", "Unknown diffusion stage error"),
+                                    "error_type": output.get("error_type"),
+                                    "detail": output.get("detail") or {},
+                                    "finished": True,
+                                    "stage_submit_ts": req_state.stage_submit_ts.get(stage_id) if req_state else None,
+                                }
+                            )
+
+                            self._cleanup_companion_state(req_id)
+                            self.request_states.pop(req_id, None)
+                            continue
+
                         req_state = self.request_states.get(output.request_id)
                         if req_state is not None:
+                            if getattr(output, "error", None) is not None:
+                                parent_id = self._companion_to_parent.get(output.request_id, output.request_id)
+                                await self.output_async_queue.put(
+                                    {
+                                        "type": "error",
+                                        "request_id": parent_id,
+                                        "stage_id": stage_id,
+                                        "error": output.error,
+                                    }
+                                )
+                                role_map = self._companion_map.get(parent_id, {})
+                                for cid in role_map.values():
+                                    self.request_states.pop(cid, None)
+                                self._cleanup_companion_state(parent_id)
+                                self.request_states.pop(parent_id, None)
+                                continue
+
                             stage_metrics = self._build_stage_metrics(stage_id, output.request_id, [output], req_state)
                             await self._route_output(stage_id, output, req_state, stage_metrics)
                     continue

vllm_omni/engine/stage_init_utils.py

@@ -437,6 +437,7 @@ def initialize_diffusion_stage(
     stage_cfg: Any,
     metadata: StageMetadata,
     batch_size: int = 1,
+    use_inline: bool = False,
 ) -> Any:
     """Build a diffusion stage client.
 
@@ -447,9 +448,10 @@ def initialize_diffusion_stage(
         batch_size: Maximum number of requests to batch together in the
             diffusion engine.  Passed through to ``StageDiffusionClient``
             and ultimately to ``AsyncOmniDiffusion``.
+        use_inline: If True, uses the inline diffusion client instead of subprocess.
     """
     from vllm_omni.diffusion.data import OmniDiffusionConfig
-    from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
+    from vllm_omni.diffusion.stage_diffusion_client import create_diffusion_client
 
     od_config = OmniDiffusionConfig.from_kwargs(
         model=model,
@@ -457,7 +459,7 @@ def initialize_diffusion_stage(
     )
     if metadata.cfg_kv_collect_func is not None:
         od_config.cfg_kv_collect_func = metadata.cfg_kv_collect_func
-    return StageDiffusionClient(model, od_config, metadata, batch_size=batch_size)
+    return create_diffusion_client(model, od_config, metadata, batch_size=batch_size, use_inline=use_inline)
 
 
 def close_started_llm_stage(started: StartedLlmStage) -> None:

vllm_omni/entrypoints/async_omni.py

@@ -21,6 +21,7 @@ from vllm.pooling_params import PoolingParams
 from vllm.tasks import SupportedTask
 from vllm.v1.engine.exceptions import EngineDeadError
 
+from vllm_omni.diffusion.data import OmniRequestError
 from vllm_omni.entrypoints.client_request_state import ClientRequestState
 from vllm_omni.entrypoints.omni_base import OmniBase
 from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics
@@ -118,6 +119,23 @@ class AsyncOmni(EngineClient, OmniBase):
         """Compatibility helper for call sites expecting async vllm config access."""
         return self.vllm_config
 
+    def get_diffusion_od_config(self) -> Any | None:
+        """Return the diffusion-stage config when the pipeline has one."""
+        for stage_client in self.engine.stage_clients:
+            if getattr(stage_client, "stage_type", None) != "diffusion":
+                continue
+
+            od_config = getattr(stage_client, "od_config", None)
+            if od_config is not None:
+                return od_config
+
+            inner_engine = getattr(stage_client, "_engine", None)
+            od_config = getattr(inner_engine, "od_config", None)
+            if od_config is not None:
+                return od_config
+
+        return None
+
     @property
     def model_config(self):
         """Return the model config for the comprehension stage when present."""
@@ -273,14 +291,15 @@ class AsyncOmni(EngineClient, OmniBase):
             stage_id = result.get("stage_id", 0)
 
             # Check for errors
-            if "error" in result:
-                logger.error(
-                    "[AsyncOmni] Orchestrator error for req=%s stage-%s: %s",
-                    request_id,
-                    stage_id,
-                    result["error"],
+            if isinstance(result, dict) and result.get("type") == "error":
+                raise OmniRequestError(
+                    result.get("error", "Unknown orchestrator error"),
+                    status_code=result.get("status_code", 500),
+                    request_id=result.get("request_id", request_id),
+                    stage_id=result.get("stage_id", stage_id),
+                    error_type=result.get("error_type"),
+                    detail=result.get("detail") or {},
                 )
-                raise RuntimeError(result)
 
             # Process the result (constructs OmniRequestOutput)
             output_to_yield = self._process_single_result(
@@ -327,6 +346,13 @@ class AsyncOmni(EngineClient, OmniBase):
                         await asyncio.sleep(_FINAL_OUTPUT_IDLE_SLEEP_S)
                         continue
 
+                    if isinstance(msg, dict) and msg.get("type") == "error":
+                        req_id = msg.get("request_id")
+                        req_state = self.request_states.get(req_id)
+                        if req_state is not None:
+                            await req_state.queue.put(msg)
+                        continue
+
                     should_continue, _, stage_id, req_state = self._handle_output_message(msg)
                     if should_continue:
                         continue
@@ -341,7 +367,16 @@ class AsyncOmni(EngineClient, OmniBase):
             except Exception as e:
                 logger.exception("[AsyncOmni] final_output_loop failed.")
                 for req_state in list(self.request_states.values()):
-                    error_msg = {"request_id": req_state.request_id, "error": str(e)}
+                    # error_msg = {"request_id": req_state.request_id, "error": str(e)}
+                    error_msg = {
+                        "type": "error",
+                        "request_id": req_state.request_id,
+                        "status_code": 500,
+                        "error": str(e),
+                        "error_type": type(e).__name__,
+                        "detail": {},
+                        "finished": True,
+                    }
                     await req_state.queue.put(error_msg)
                 self.final_output_task = None
 

vllm_omni/entrypoints/async_omni_diffusion.py

@@ -18,7 +18,7 @@ from typing import Any
 from vllm.logger import init_logger
 from vllm.transformers_utils.config import get_hf_file_to_dict
 
-from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig
+from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig, normalize_omni_error
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
@@ -237,7 +237,8 @@ class AsyncOmniDiffusion:
             )
         except Exception as e:
             logger.error("Batch generation failed for request %s: %s", request_id, e)
-            raise RuntimeError(f"Diffusion batch generation failed: {e}") from e
+            # raise RuntimeError(f"Diffusion batch generation failed: {e}") from e
+            raise normalize_omni_error(e, request_id=request_id) from e
 
         # Combine all per-prompt results into a single OmniRequestOutput
         all_images = []
@@ -251,6 +252,10 @@ class AsyncOmniDiffusion:
             finished=True,
         )
 
+    def get_diffusion_od_config(self) -> OmniDiffusionConfig:
+        """Return the diffusion config used by this engine."""
+        return self.od_config
+
     # ------------------------------------------------------------------
     # Public generate API
     # ------------------------------------------------------------------
@@ -306,7 +311,8 @@ class AsyncOmniDiffusion:
             result = result[0]
         except Exception as e:
             logger.error("Generation failed for request %s: %s", request_id, e)
-            raise RuntimeError(f"Diffusion generation failed: {e}") from e
+            # raise RuntimeError(f"Diffusion generation failed: {e}") from e
+            raise normalize_omni_error(e, request_id=request_id) from e
 
         if not result.request_id:
             result.request_id = request_id

vllm_omni/entrypoints/omni_base.py

@@ -12,6 +12,7 @@ import huggingface_hub
 from vllm.logger import init_logger
 from vllm.v1.engine.exceptions import EngineDeadError
 
+from vllm_omni.diffusion.data import OmniRequestError
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.entrypoints.client_request_state import ClientRequestState
 from vllm_omni.entrypoints.utils import get_final_stage_id_for_e2e
@@ -206,7 +207,14 @@ class OmniBase:
             return True, None, None, None
 
         if msg_type == "error":
-            raise RuntimeError(msg.get("error", "Orchestrator returned an error message"))
+            raise OmniRequestError(
+                msg.get("error", "Orchestrator returned an error message"),
+                status_code=msg.get("status_code", 500),
+                request_id=msg.get("request_id"),
+                stage_id=msg.get("stage_id"),
+                error_type=msg.get("error_type"),
+                detail=msg.get("detail") or {},
+            )
 
         if msg_type != "output":
             logger.warning("[%s] got unexpected msg type: %s", self.__class__.__name__, msg_type)

vllm_omni/entrypoints/openai/api_server.py

@@ -81,11 +81,14 @@ from vllm.tool_parsers import ToolParserManager
 from vllm.utils import random_uuid
 from vllm.utils.system_utils import decorate_logs
 
+from vllm_omni.diffusion.data import OmniRequestError
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
 from vllm_omni.entrypoints.openai.image_api_utils import (
+    SUPPORTED_LAYERED_RESOLUTIONS,
     encode_image_base64,
     parse_size,
+    validate_layered_layers,
 )
 from vllm_omni.entrypoints.openai.protocol.audio import BatchSpeechRequest, OpenAICreateSpeechRequest
 from vllm_omni.entrypoints.openai.protocol.images import (
@@ -116,8 +119,6 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParam
 logger = init_logger(__name__)
 router = APIRouter()
 
-# Supported resolution buckets for layered models (e.g., Qwen-Image-Layered)
-SUPPORTED_LAYERED_RESOLUTIONS = (640, 1024)
 profiler_router = APIRouter()
 
 
@@ -1299,6 +1300,10 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
             size_str = f"{width}x{height}"
         else:
             size_str = "model default"
+
+        app_state_args = getattr(raw_request.app.state, "args", None)
+        _check_max_generated_image_size(app_state_args, width, height)
+
         _update_if_not_none(gen_params, "width", width)
         _update_if_not_none(gen_params, "height", height)
 
@@ -1314,6 +1319,7 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
             gen_params, "seed", request.seed if request.seed is not None else random.randint(0, 2**32 - 1)
         )
         _update_if_not_none(gen_params, "generator_device", request.generator_device)
+        _update_if_not_none(gen_params, "layers", request.layers)
 
         request_id = f"img_gen-{random_uuid()}"
 
@@ -1427,6 +1433,23 @@ async def edit_images(
             input_images_list.extend(urls)
         if not input_images_list:
             raise HTTPException(status_code=422, detail="Field 'image' or 'url' is required")
+        # Reject oversized multi-image edit requests before fetching or decoding
+        # any inputs. This keeps over-limit URL requests from burning network,
+        # CPU, and memory on work that will be rejected anyway.
+        max_input_images = _get_max_edit_input_images(raw_request, engine_client)
+        if max_input_images is not None and len(input_images_list) > max_input_images:
+            detail = (
+                "Received multiple input images. Only a single image is supported by this model."
+                if max_input_images == 1
+                else (
+                    f"Received {len(input_images_list)} input images. "
+                    f"At most {max_input_images} images are supported by this model."
+                )
+            )
+            raise HTTPException(
+                status_code=HTTPStatus.BAD_REQUEST.value,
+                detail=detail,
+            )
         pil_images = await _load_input_images(input_images_list)
         prompt["multi_modal_data"] = {}
         prompt["multi_modal_data"]["image"] = pil_images
@@ -1462,11 +1485,13 @@ async def edit_images(
                 detail=f"Invalid resolution {resolution}. Supported resolutions: {SUPPORTED_LAYERED_RESOLUTIONS}.",
             )
         # 3.2.1 Validate layers if provided
-        if layers is not None and layers < 1:
+        try:
+            layers = validate_layered_layers(layers)
+        except ValueError as e:
             raise HTTPException(
                 status_code=HTTPStatus.BAD_REQUEST.value,
-                detail=f"Invalid layers value {layers}. layers must be >= 1.",
-            )
+                detail=str(e),
+            ) from e
         # 3.2.2 Check for conflicting size and resolution parameters
         if resolution is not None and size.lower() != "auto":
             raise HTTPException(
@@ -1476,7 +1501,6 @@ async def edit_images(
             )
 
         # 3.3 Parse and add size if provided
-        max_generated_image_size = getattr(app_state_args, "max_generated_image_size", None)
         width, height = None, None
         if size.lower() == "auto":
             if resolution is None:
@@ -1486,23 +1510,7 @@ async def edit_images(
         else:
             width, height = parse_size(size)
 
-        # Check max_generated_image_size
-        if max_generated_image_size is not None:
-            if width is not None and height is not None:
-                if width * height > max_generated_image_size:
-                    raise HTTPException(
-                        status_code=HTTPStatus.BAD_REQUEST.value,
-                        detail=f"Requested image size {width}x{height} exceeds the maximum allowed "
-                        f"size of {max_generated_image_size} pixels.",
-                    )
-            elif resolution is not None:
-                # When resolution is set, the output size is resolution * resolution
-                if resolution * resolution > max_generated_image_size:
-                    raise HTTPException(
-                        status_code=HTTPStatus.BAD_REQUEST.value,
-                        detail=f"Requested resolution {resolution} (max {resolution}x{resolution} pixels) "
-                        f"exceeds the maximum allowed size of {max_generated_image_size} pixels.",
-                    )
+        _check_max_generated_image_size(app_state_args, width, height, resolution)
 
         size_str = f"{width}x{height}" if width is not None and height is not None else "auto"
         _update_if_not_none(gen_params, "width", width)
@@ -1606,6 +1614,27 @@ def _get_engine_and_model(raw_request: Request):
     return engine_client, model_name, normalized_stage_configs
 
 
+def _get_diffusion_od_config(raw_request: Request, engine_client: Any) -> Any:
+    diffusion_engine = getattr(raw_request.app.state, "diffusion_engine", None) or engine_client
+    get_diffusion_od_config = getattr(diffusion_engine, "get_diffusion_od_config", None)
+    return (
+        get_diffusion_od_config() if callable(get_diffusion_od_config) else getattr(diffusion_engine, "od_config", None)
+    )
+
+
+def _get_max_edit_input_images(raw_request: Request, engine_client: Any) -> int | None:
+    od_config = _get_diffusion_od_config(raw_request, engine_client)
+    if od_config is None:
+        # Preserve the existing compatibility behavior when the diffusion
+        # config is not exposed on the serving surface.
+        return None
+
+    if not bool(getattr(od_config, "supports_multimodal_inputs", False)):
+        return 1
+
+    return getattr(od_config, "max_multimodal_image_inputs", None)
+
+
 def _get_lora_from_json_str(lora_body):
     if lora_body is None:
         return None
@@ -1673,11 +1702,14 @@ async def _generate_with_async_omni(
                 pass
         sampling_params_list.append(default_stage_params)
 
-    async for output in engine_client.generate(
-        sampling_params_list=sampling_params_list,
-        **kwargs,
-    ):
-        result = output
+    try:
+        async for output in engine_client.generate(
+            sampling_params_list=sampling_params_list,
+            **kwargs,
+        ):
+            result = output
+    except OmniRequestError as e:
+        raise HTTPException(status_code=e.status_code, detail=str(e)) from e
 
     if result is None:
         raise HTTPException(
@@ -1687,6 +1719,34 @@ async def _generate_with_async_omni(
     return result
 
 
+def _check_max_generated_image_size(
+    app_state_args: Any,
+    width: int | None,
+    height: int | None,
+    resolution: int | None = None,
+) -> None:
+    """Raise 400 if the requested image size exceeds --max-generated-image-size."""
+    max_generated_image_size = getattr(app_state_args, "max_generated_image_size", None)
+    # Check max_generated_image_size
+    if max_generated_image_size is None:
+        return
+    if width is not None and height is not None:
+        if width * height > max_generated_image_size:
+            raise HTTPException(
+                status_code=HTTPStatus.BAD_REQUEST.value,
+                detail=f"Requested image size {width}x{height} exceeds the maximum allowed "
+                f"size of {max_generated_image_size} pixels.",
+            )
+    elif resolution is not None:
+        # When resolution is set, the output size is resolution * resolution
+        if resolution * resolution > max_generated_image_size:
+            raise HTTPException(
+                status_code=HTTPStatus.BAD_REQUEST.value,
+                detail=f"Requested resolution {resolution} (max {resolution}x{resolution} pixels) "
+                f"exceeds the maximum allowed size of {max_generated_image_size} pixels.",
+            )
+
+
 def _update_if_not_none(object: Any, key: str, val: Any) -> None:
     if val is not None:
         setattr(object, key, val)
@@ -1775,10 +1835,27 @@ def _choose_output_format(output_format: str | None, background: str | None) ->
     return "jpeg"
 
 
+def _prepare_image_for_output_format(image: Image.Image, format: str) -> Image.Image:
+    fmt = format.lower()
+    if fmt not in {"jpg", "jpeg"}:
+        return image
+
+    if image.mode == "RGB":
+        return image
+
+    if image.mode in {"RGBA", "LA"} or (image.mode == "P" and "transparency" in image.info):
+        alpha_image = image.convert("RGBA")
+        flattened = Image.new("RGB", alpha_image.size, (255, 255, 255))
+        flattened.paste(alpha_image, mask=alpha_image.getchannel("A"))
+        return flattened
+
+    return image.convert("RGB")
+
+
 def _encode_image_base64_with_compression(
     image: Image.Image, format: str = "png", output_compression: int = 100
 ) -> str:
-    """Encode PIL Image to base64 PNG string.
+    """Encode PIL Image to a base64 image string.
 
     Args:
         image: PIL Image object
@@ -1788,6 +1865,7 @@ def _encode_image_base64_with_compression(
         Base64-encoded image as string
     """
     buffer = io.BytesIO()
+    image = _prepare_image_for_output_format(image, format)
     save_kwargs = {}
     if format in ("jpg", "jpeg", "webp"):
         save_kwargs["quality"] = output_compression
@@ -1862,18 +1940,6 @@ def video_response_from_request(model_name: str, req: VideoGenerationRequest) ->
     return resp
 
 
-async def decode_and_save_video_output(output: Any, file_name: str) -> str:
-    if not output.b64_json:
-        raise RuntimeError(f"Video output for {file_name} did not include b64_json content.")
-
-    try:
-        video_bytes = base64.b64decode(output.b64_json)
-    except Exception as decode_exc:
-        raise RuntimeError(f"Failed to decode generated video payload for {file_name}") from decode_exc
-
-    return await STORAGE_MANAGER.save(video_bytes, file_name)
-
-
 def _cleanup_video(video_id: str, output_path: str | None):
     try:
         if output_path is not None:
@@ -1897,15 +1963,12 @@ async def _run_video_generation_job(
     started_at = time.perf_counter()
     output_path = None
     try:
-        response = await handler.generate_videos(request, video_id, reference_image=reference_image)
-        if not response.data:
-            raise RuntimeError("Video generation completed but returned no outputs.")
-
-        if (video_count := len(response.data)) > 1:
-            logger.warning("Video request %s generated %s outputs but we only expected one.", video_id, video_count)
+        video_bytes, stage_durations, peak_memory_mb = await handler.generate_video_bytes(
+            request, video_id, reference_image=reference_image
+        )
 
         file_name = f"{video_id}.{job.file_extension}"
-        output_path = await decode_and_save_video_output(response.data[0], file_name)
+        output_path = await STORAGE_MANAGER.save(video_bytes, file_name)
         logger.info("Video request %s persisted %s output file.", video_id, output_path)
 
         await VIDEO_STORE.update_fields(
@@ -1916,6 +1979,8 @@ async def _run_video_generation_job(
                 "file_name": file_name,
                 "completed_at": int(time.time()),
                 "inference_time_s": time.perf_counter() - started_at,
+                "stage_durations": stage_durations,
+                "peak_memory_mb": peak_memory_mb,
             },
         )
     except Exception as exc:
@@ -1968,6 +2033,10 @@ async def create_video(
     true_cfg_scale: float | None = Form(default=None),
     seed: int | None = Form(default=None),
     negative_prompt: str | None = Form(default=None),
+    enable_frame_interpolation: bool | None = Form(default=None),
+    frame_interpolation_exp: int | None = Form(default=None, ge=1),
+    frame_interpolation_scale: float | None = Form(default=None, gt=0.0),
+    frame_interpolation_model_path: str | None = Form(default=None),
     lora: str | None = Form(default=None),
     extra_params: str | None = Form(default=None),
 ) -> VideoResponse:
@@ -2037,6 +2106,10 @@ async def create_video(
         "true_cfg_scale": true_cfg_scale,
         "seed": seed,
         "negative_prompt": negative_prompt,
+        "enable_frame_interpolation": enable_frame_interpolation,
+        "frame_interpolation_exp": frame_interpolation_exp,
+        "frame_interpolation_scale": frame_interpolation_scale,
+        "frame_interpolation_model_path": frame_interpolation_model_path,
         "lora": _parse_form_json(lora, expected_type=dict),
         "extra_params": _parse_form_json(extra_params, expected_type=dict),
     }

vllm_omni/entrypoints/openai/image_api_utils.py

@@ -13,6 +13,9 @@ import io
 
 import PIL.Image
 
+SUPPORTED_LAYERED_RESOLUTIONS = (640, 1024)
+SUPPORTED_LAYERED_LAYERS_RANGE = range(3, 11)
+
 
 def parse_size(size_str: str) -> tuple[int, int]:
     """Parse size string to width and height tuple.
@@ -63,3 +66,16 @@ def encode_image_base64(image: PIL.Image.Image) -> str:
     image.save(buffer, format="PNG")
     buffer.seek(0)
     return base64.b64encode(buffer.read()).decode("utf-8")
+
+
+def validate_layered_layers(layers: int | None) -> int | None:
+    """Validate the Qwen-Image-Layered ``layers`` parameter."""
+    if layers is None:
+        return None
+    if layers not in SUPPORTED_LAYERED_LAYERS_RANGE:
+        raise ValueError(
+            f"Invalid layers value {layers}. layers must be between "
+            f"{SUPPORTED_LAYERED_LAYERS_RANGE.start} and "
+            f"{SUPPORTED_LAYERED_LAYERS_RANGE.stop - 1} inclusive."
+        )
+    return layers

vllm_omni/entrypoints/openai/protocol/images.py

@@ -12,6 +12,8 @@ from typing import Any
 
 from pydantic import BaseModel, Field, field_validator
 
+from vllm_omni.entrypoints.openai.image_api_utils import validate_layered_layers
+
 
 class ResponseFormat(str, Enum):
     """Image response format"""
@@ -43,6 +45,10 @@ class ImageGenerationRequest(BaseModel):
     )
     response_format: ResponseFormat = Field(default=ResponseFormat.B64_JSON, description="Format of the returned image")
     user: str | None = Field(default=None, description="User identifier for tracking")
+    layers: int | None = Field(
+        default=None,
+        description="Number of output layers for layered image models. Supported range: 3-10.",
+    )
 
     @field_validator("size")
     @classmethod
@@ -67,6 +73,12 @@ class ImageGenerationRequest(BaseModel):
             raise ValueError(f"Only 'b64_json' response format is supported, got: {v}")
         return v
 
+    @field_validator("layers")
+    @classmethod
+    def validate_layers(cls, v):
+        """Validate the layers parameter for layered image models."""
+        return validate_layered_layers(v)
+
     # vllm-omni extensions for diffusion control
     negative_prompt: str | None = Field(default=None, description="Text describing what to avoid in the image")
     num_inference_steps: int | None = Field(

vllm_omni/entrypoints/openai/protocol/videos.py

@@ -150,6 +150,29 @@ class VideoGenerationRequest(BaseModel):
     )
     seed: int | None = Field(default=None, description="Random seed for reproducibility")
 
+    # vllm-omni extensions for post-generation frame interpolation.
+    enable_frame_interpolation: bool = Field(
+        default=False,
+        description="Enable post-generation RIFE frame interpolation before MP4 encoding.",
+    )
+    frame_interpolation_exp: int = Field(
+        default=1,
+        ge=1,
+        description="Interpolation exponent: 1=2x temporal resolution, 2=4x, etc.",
+    )
+    frame_interpolation_scale: float = Field(
+        default=1.0,
+        gt=0.0,
+        description="RIFE inference scale. Use 0.5 for high-resolution inputs to save memory.",
+    )
+    frame_interpolation_model_path: str | None = Field(
+        default=None,
+        description=(
+            "Local directory or Hugging Face repo ID containing RIFE flownet.pkl weights. "
+            "Defaults to elfgum/RIFE-4.22.lite."
+        ),
+    )
+
     # vllm-omni extension for per-request LoRA.
     lora: dict[str, Any] | None = Field(
         default=None,

vllm_omni/entrypoints/openai/serving_chat.py

@@ -82,6 +82,7 @@ from vllm.tool_parsers.mistral_tool_parser import MistralToolCall
 from vllm.utils.collection_utils import as_list
 
 from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
+from vllm_omni.entrypoints.openai.image_api_utils import validate_layered_layers
 from vllm_omni.entrypoints.openai.protocol import OmniChatCompletionStreamResponse
 from vllm_omni.entrypoints.openai.protocol.audio import AudioResponse, CreateAudio
 from vllm_omni.entrypoints.openai.utils import parse_lora_request
@@ -2094,6 +2095,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             layers = extra_body.get("layers")
             resolution = extra_body.get("resolution")
 
+            try:
+                layers = validate_layered_layers(layers)
+            except ValueError as e:
+                return self._create_error_response(str(e), status_code=400)
+
             logger.info(
                 "Diffusion chat request %s: prompt=%r, ref_images=%d, params=%s",
                 request_id,

vllm_omni/entrypoints/openai/serving_video.py

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import copy
 import time
 from dataclasses import dataclass
 from http import HTTPStatus
@@ -13,6 +14,7 @@ from PIL import Image
 from vllm.engine.protocol import EngineClient
 from vllm.logger import init_logger
 
+from vllm_omni.diffusion.data import OmniRequestError
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.openai.protocol.videos import (
     VideoData,
@@ -20,7 +22,7 @@ from vllm_omni.entrypoints.openai.protocol.videos import (
     VideoGenerationResponse,
 )
 from vllm_omni.entrypoints.openai.utils import get_stage_type, parse_lora_request
-from vllm_omni.entrypoints.openai.video_api_utils import encode_video_base64
+from vllm_omni.entrypoints.openai.video_api_utils import _encode_video_bytes, encode_video_base64
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams, OmniTextPrompt
 
 logger = init_logger(__name__)
@@ -28,11 +30,23 @@ logger = init_logger(__name__)
 
 @dataclass
 class ReferenceImage:
-    """Reference class for tracking additional metadata if needed"""
+    """Reference class for tracking additional metadata if needed."""
 
     data: Image.Image
 
 
+@dataclass
+class VideoGenerationArtifacts:
+    """Normalized outputs and profiler metadata extracted from one request."""
+
+    videos: list[Any]
+    audios: list[Any | None]
+    audio_sample_rate: int
+    output_fps: int
+    stage_durations: dict[str, float]
+    peak_memory_mb: float
+
+
 class OmniOpenAIServingVideo:
     """OpenAI-style video generation handler for omni diffusion models."""
 
@@ -71,18 +85,19 @@ class OmniOpenAIServingVideo:
             stage_configs=stage_configs,
         )
 
-    async def generate_videos(
+    async def _run_and_extract(
         self,
         request: VideoGenerationRequest,
         reference_id: str,
         *,
         reference_image: ReferenceImage | None = None,
-    ) -> VideoGenerationResponse:
+    ) -> VideoGenerationArtifacts:
+        """Run the generation pipeline and extract video/audio/profiler outputs."""
         prompt: OmniTextPrompt = OmniTextPrompt(prompt=request.prompt)
         if request.negative_prompt is not None:
             prompt["negative_prompt"] = request.negative_prompt
 
-        gen_params = OmniDiffusionSamplingParams()
+        gen_params = self._resolve_default_sampling_params()
 
         input_image = None if reference_image is None else reference_image.data
         vp = request.resolve_video_params()
@@ -100,18 +115,27 @@ class OmniOpenAIServingVideo:
         if vp.fps is not None:
             gen_params.fps = vp.fps
             gen_params.frame_rate = float(vp.fps)
+        provided_fields = request.model_fields_set
+        if "enable_frame_interpolation" in provided_fields:
+            gen_params.enable_frame_interpolation = request.enable_frame_interpolation
+        if "frame_interpolation_exp" in provided_fields:
+            gen_params.frame_interpolation_exp = request.frame_interpolation_exp
+        if "frame_interpolation_scale" in provided_fields:
+            gen_params.frame_interpolation_scale = request.frame_interpolation_scale
+        if "frame_interpolation_model_path" in provided_fields:
+            gen_params.frame_interpolation_model_path = request.frame_interpolation_model_path
 
-        if request.num_inference_steps is not None:
+        if "num_inference_steps" in provided_fields and request.num_inference_steps is not None:
             gen_params.num_inference_steps = request.num_inference_steps
-        if request.guidance_scale is not None:
+        if "guidance_scale" in provided_fields and request.guidance_scale is not None:
             gen_params.guidance_scale = request.guidance_scale
-        if request.guidance_scale_2 is not None:
+        if "guidance_scale_2" in provided_fields and request.guidance_scale_2 is not None:
             gen_params.guidance_scale_2 = request.guidance_scale_2
-        if request.true_cfg_scale is not None:
+        if "true_cfg_scale" in provided_fields and request.true_cfg_scale is not None:
             gen_params.true_cfg_scale = request.true_cfg_scale
-        if request.seed is not None:
+        if "seed" in provided_fields and request.seed is not None:
             gen_params.seed = request.seed
-        if request.boundary_ratio is not None:
+        if "boundary_ratio" in provided_fields and request.boundary_ratio is not None:
             gen_params.boundary_ratio = request.boundary_ratio
 
         logger.info(
@@ -119,17 +143,15 @@ class OmniOpenAIServingVideo:
             request.boundary_ratio,
             gen_params.boundary_ratio,
         )
-        if request.flow_shift is not None:
+        if "flow_shift" in provided_fields and request.flow_shift is not None:
             gen_params.extra_args["flow_shift"] = request.flow_shift
 
-        # Apply model-specific extra parameters
         if request.extra_params is not None:
             if not isinstance(request.extra_params, dict):
                 raise HTTPException(
                     status_code=HTTPStatus.BAD_REQUEST.value,
                     detail="extra_params must be a JSON object/dict.",
                 )
-            # Merge extra_params into extra_args
             gen_params.extra_args.update(request.extra_params)
             logger.info("Applied extra_params: %s", request.extra_params)
 
@@ -144,31 +166,114 @@ class OmniOpenAIServingVideo:
         )
 
         result = await self._run_generation(prompt, gen_params, reference_id)
-        _t_encode_start = time.perf_counter()
         videos = self._extract_video_outputs(result)
         audios = self._extract_audio_outputs(result, expected_count=len(videos))
         audio_sample_rate = self._resolve_audio_sample_rate(result)
-        output_fps = vp.fps or 24
+        output_fps = (vp.fps or self._resolve_fps(result) or 24) * self._resolve_video_fps_multiplier(result)
+        return VideoGenerationArtifacts(
+            videos=videos,
+            audios=audios,
+            audio_sample_rate=audio_sample_rate,
+            output_fps=output_fps,
+            stage_durations=self._extract_stage_durations(result),
+            peak_memory_mb=self._extract_peak_memory_mb(result),
+        )
 
+    async def generate_videos(
+        self,
+        request: VideoGenerationRequest,
+        reference_id: str,
+        *,
+        reference_image: ReferenceImage | None = None,
+    ) -> VideoGenerationResponse:
+        artifacts = await self._run_and_extract(request, reference_id, reference_image=reference_image)
+
+        video_codec_options = self._resolve_video_codec_options(request)
+        _t_encode_start = time.perf_counter()
         video_data = [
             VideoData(
                 b64_json=(
-                    encode_video_base64(video, fps=output_fps)
-                    if audios[idx] is None
+                    encode_video_base64(video, fps=artifacts.output_fps, video_codec_options=video_codec_options)
+                    if artifacts.audios[idx] is None
                     else encode_video_base64(
                         video,
-                        fps=output_fps,
-                        audio=audios[idx],
-                        audio_sample_rate=audio_sample_rate,
+                        fps=artifacts.output_fps,
+                        audio=artifacts.audios[idx],
+                        audio_sample_rate=artifacts.audio_sample_rate,
+                        video_codec_options=video_codec_options,
                     )
                 )
             )
-            for idx, video in enumerate(videos)
+            for idx, video in enumerate(artifacts.videos)
         ]
         _t_encode_ms = (time.perf_counter() - _t_encode_start) * 1000
         logger.info("Video response encoding (MP4+base64): %.2f ms", _t_encode_ms)
         return VideoGenerationResponse(created=int(time.time()), data=video_data)
 
+    async def generate_video_bytes(
+        self,
+        request: VideoGenerationRequest,
+        reference_id: str,
+        *,
+        reference_image: ReferenceImage | None = None,
+    ) -> tuple[bytes, dict[str, float], float]:
+        """Generate a video and return raw MP4 bytes, bypassing base64 encoding."""
+        artifacts = await self._run_and_extract(request, reference_id, reference_image=reference_image)
+        if len(artifacts.videos) > 1:
+            logger.warning(
+                "Video request %s generated %d outputs; returning only the first.",
+                reference_id,
+                len(artifacts.videos),
+            )
+        audio = artifacts.audios[0]
+
+        _t_encode_start = time.perf_counter()
+        video_bytes = _encode_video_bytes(
+            artifacts.videos[0],
+            fps=artifacts.output_fps,
+            **({"audio": audio, "audio_sample_rate": artifacts.audio_sample_rate} if audio is not None else {}),
+            video_codec_options=self._resolve_video_codec_options(request),
+        )
+        _t_encode_ms = (time.perf_counter() - _t_encode_start) * 1000
+        logger.info("Video response encoding (MP4 bytes): %.2f ms", _t_encode_ms)
+        return video_bytes, artifacts.stage_durations, artifacts.peak_memory_mb
+
+    @staticmethod
+    def _resolve_video_codec_options(request: VideoGenerationRequest) -> dict[str, str]:
+        video_codec_options = {"preset": "ultrafast", "threads": "0"}
+        if request.extra_params is not None and isinstance(request.extra_params, dict):
+            extra_codec_options = request.extra_params.get("video_codec_options")
+            if isinstance(extra_codec_options, dict):
+                video_codec_options = {str(k): str(v) for k, v in extra_codec_options.items()}
+        return video_codec_options
+
+    def _resolve_default_sampling_params(self) -> OmniDiffusionSamplingParams:
+        default_sampling_params_list = getattr(self._engine_client, "default_sampling_params_list", None)
+        if default_sampling_params_list:
+            for params in default_sampling_params_list:
+                if isinstance(params, OmniDiffusionSamplingParams):
+                    # Requests mutate sampling params in-place, including
+                    # nested dict fields like extra_args. Deep-copy the stage
+                    # defaults so one request cannot leak state into another.
+                    return copy.deepcopy(params)
+        return OmniDiffusionSamplingParams()
+
+    @staticmethod
+    def _resolve_video_fps_multiplier(result: Any) -> int:
+        custom_output = getattr(result, "custom_output", None)
+        if isinstance(custom_output, dict):
+            multiplier = custom_output.get("video_fps_multiplier")
+            if multiplier is not None:
+                return int(multiplier)
+        request_output = getattr(result, "request_output", None)
+        if request_output is not None:
+            custom_output = getattr(request_output, "custom_output", None)
+            if isinstance(custom_output, dict):
+                multiplier = custom_output.get("video_fps_multiplier")
+                if multiplier is not None:
+                    return int(multiplier)
+        return 1
+
     @staticmethod
     def _apply_lora(lora_body: Any, gen_params: OmniDiffusionSamplingParams) -> None:
         try:
@@ -200,7 +305,6 @@ class OmniOpenAIServingVideo:
                 detail="Stage configs not found. Start server with an omni diffusion model.",
             )
 
-        # Video generation endpoint only supports diffusion stages.
         for stage in stage_configs:
             stage_type = get_stage_type(stage)
             if stage_type != "diffusion":
@@ -209,17 +313,28 @@ class OmniOpenAIServingVideo:
                     detail=f"Video generation only supports diffusion stages, found '{stage_type}' stage.",
                 )
 
-        # Common generation logic for both paths
         engine_client = cast(AsyncOmni, self._engine_client)
         sampling_params_list: list[OmniSamplingParams] = [gen_params for _ in stage_configs]
 
         result = None
-        async for output in engine_client.generate(
-            prompt=prompt,
-            request_id=request_id,
-            sampling_params_list=sampling_params_list,
-        ):
-            result = output
+        try:
+            async for output in engine_client.generate(
+                prompt=prompt,
+                request_id=request_id,
+                sampling_params_list=sampling_params_list,
+            ):
+                result = output
+        except OmniRequestError as e:
+            raise HTTPException(
+                status_code=e.status_code,
+                detail={
+                    "message": str(e),
+                    "request_id": e.request_id,
+                    "stage_id": e.stage_id,
+                    "error_type": e.error_type,
+                    "detail": e.detail,
+                },
+            ) from e
 
         if result is None:
             raise HTTPException(
@@ -323,6 +438,45 @@ class OmniOpenAIServingVideo:
 
         return 24000
 
+    @staticmethod
+    def _resolve_fps(result: Any) -> int | None:
+        multimodal_output = getattr(result, "multimodal_output", None)
+        if isinstance(multimodal_output, dict):
+            fps = multimodal_output.get("fps")
+            if fps is not None:
+                try:
+                    fps_val = fps.item() if hasattr(fps, "item") else int(fps)
+                    if fps_val > 0:
+                        return fps_val
+                except (TypeError, ValueError):
+                    pass
+
+        request_output = getattr(result, "request_output", None)
+        if isinstance(request_output, dict):
+            mm = request_output.get("multimodal_output") or {}
+            if isinstance(mm, dict):
+                fps = mm.get("fps")
+                if fps is not None:
+                    try:
+                        fps_val = fps.item() if hasattr(fps, "item") else int(fps)
+                        if fps_val > 0:
+                            return fps_val
+                    except (TypeError, ValueError):
+                        pass
+        elif hasattr(request_output, "multimodal_output"):
+            mm = getattr(request_output, "multimodal_output", None)
+            if isinstance(mm, dict):
+                fps = mm.get("fps")
+                if fps is not None:
+                    try:
+                        fps_val = fps.item() if hasattr(fps, "item") else int(fps)
+                        if fps_val > 0:
+                            return fps_val
+                    except (TypeError, ValueError):
+                        pass
+
+        return None
+
     @classmethod
     def _extract_audio_sample_rate_from_result(cls, result: Any) -> int | None:
         multimodal_output = getattr(result, "multimodal_output", None)
@@ -401,3 +555,16 @@ class OmniOpenAIServingVideo:
             return None
 
         return sample_rate if sample_rate > 0 else None
+
+    @staticmethod
+    def _extract_stage_durations(result: Any) -> dict[str, float]:
+        stage_durations = getattr(result, "stage_durations", None)
+        return stage_durations if isinstance(stage_durations, dict) else {}
+
+    @staticmethod
+    def _extract_peak_memory_mb(result: Any) -> float:
+        peak_memory_mb = getattr(result, "peak_memory_mb", 0.0)
+        try:
+            return float(peak_memory_mb or 0.0)
+        except (TypeError, ValueError):
+            return 0.0

vllm_omni/entrypoints/openai/video_api_utils.py

@@ -159,111 +159,116 @@ def _normalize_frames(frames: list[Any]) -> list[np.ndarray]:
     return normalized
 
 
-def _coerce_video_to_frames(video: Any) -> list[np.ndarray]:
-    """Convert a video payload into a list of frames for export_to_video."""
+def _coerce_video_to_array(video: Any) -> np.ndarray:
+    """Convert a video payload into a normalized array of frames (F, H, W, C)."""
     if isinstance(video, torch.Tensor):
         video_array = _normalize_video_tensor(video)
-        return list(video_array)
+        if video_array.ndim == 4:
+            return video_array
+        if video_array.ndim == 3:
+            return video_array[None, ...]
+        raise ValueError(f"Unsupported video array shape: {video_array.shape}")
     if isinstance(video, np.ndarray):
         video_array = _normalize_video_array(video)
         if isinstance(video_array, list):
             raise ValueError("Batched video arrays must be split before encoding.")
         if video_array.ndim == 4:
-            return list(video_array)
+            return video_array
         if video_array.ndim == 3:
-            return [video_array]
+            return video_array[None, ...]
         raise ValueError(f"Unsupported video array shape: {video_array.shape}")
     if isinstance(video, list):
         if not video:
-            return []
+            return np.empty((0,), dtype=np.float32)
         # If this looks like a list of frames, normalize directly.
         if all(isinstance(item, (np.ndarray, torch.Tensor, Image.Image)) for item in video):
             # If each item is itself a video (ndim==4), handle elsewhere.
             if all(hasattr(item, "ndim") and item.ndim >= 4 for item in video):
                 raise ValueError("Expected a single video, got a list of video tensors/arrays.")
-            return _normalize_frames(video)
+            return np.stack(_normalize_frames(video), axis=0)
         raise ValueError("Unsupported list contents for video payload.")
     raise ValueError(f"Unsupported video payload type: {type(video)}")
 
 
-def _coerce_audio_to_waveform(audio: Any) -> torch.Tensor:
-    """Convert an audio payload into a 2-channel CPU float tensor for LTX2 export."""
+def _coerce_audio_to_numpy(audio: Any) -> np.ndarray:
+    """Convert an audio payload into a float32 numpy array for muxing."""
     if isinstance(audio, torch.Tensor):
-        waveform = audio.detach().cpu()
+        arr = audio.detach().cpu().float().numpy()
     elif isinstance(audio, np.ndarray):
-        waveform = torch.from_numpy(audio)
+        arr = audio
     elif isinstance(audio, list):
-        waveform = torch.tensor(audio)
+        arr = np.array(audio)
     else:
         raise ValueError(f"Unsupported audio payload type: {type(audio)}")
 
-    waveform = waveform.squeeze()
-
-    if waveform.ndim == 0:
+    arr = np.squeeze(arr)
+    if arr.ndim == 0:
         raise ValueError("Audio payload must contain at least one sample.")
 
-    if waveform.ndim == 1:
-        waveform = waveform.unsqueeze(0)
-    elif waveform.ndim == 2:
-        if waveform.shape[0] in (1, 2):
-            pass
-        elif waveform.shape[1] in (1, 2):
-            waveform = waveform.transpose(0, 1)
-        else:
-            raise ValueError(f"Unsupported audio payload shape: {tuple(waveform.shape)}")
-    else:
-        raise ValueError(f"Unsupported audio payload rank: {waveform.ndim}")
-
-    if waveform.shape[0] == 1:
-        waveform = waveform.repeat(2, 1)
-    elif waveform.shape[0] != 2:
-        raise ValueError(f"Expected mono or stereo audio, got shape {tuple(waveform.shape)}")
-
-    return waveform.float().contiguous()
+    return arr.astype(np.float32)
 
 
-def _encode_video_bytes(video: Any, fps: int, audio: Any | None = None, audio_sample_rate: int | None = None) -> bytes:
+def _encode_video_bytes(
+    video: Any,
+    fps: int,
+    audio: Any | None = None,
+    audio_sample_rate: int | None = None,
+    video_codec_options: dict[str, str] | None = None,
+) -> bytes:
     """Encode a video payload into MP4 bytes, optionally muxing audio."""
-    try:
-        from diffusers.utils import export_to_video
-    except ImportError as exc:  # pragma: no cover - optional dependency
-        raise ImportError("diffusers is required for export_to_video.") from exc
-
-    frames = _coerce_video_to_frames(video)
-    if not frames:
+    frames_np = _coerce_video_to_array(video)
+    if frames_np.size == 0:
         raise ValueError("No frames found to encode.")
+    if frames_np.ndim == 4 and frames_np.shape[-1] == 4:
+        frames_np = frames_np[..., :3]
 
-    tmp_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
-    tmp_file.close()
-    try:
-        if audio is not None:
-            from diffusers.pipelines.ltx2.export_utils import encode_video as encode_ltx2_video
+    if audio is None:
+        from diffusers.utils import export_to_video
 
-            frames_np = np.stack(frames, axis=0)
-            if frames_np.ndim == 4 and frames_np.shape[-1] == 4:
-                frames_np = frames_np[..., :3]
-            frames_np = np.clip(frames_np, 0.0, 1.0)
-            frames_u8 = (frames_np * 255).round().clip(0, 255).astype("uint8")
-            video_tensor = torch.from_numpy(frames_u8)
-            encode_ltx2_video(
-                video_tensor,
-                fps=fps,
-                audio=_coerce_audio_to_waveform(audio),
-                audio_sample_rate=audio_sample_rate,
-                output_path=tmp_file.name,
-            )
-        else:
-            export_to_video(frames, tmp_file.name, fps=fps)
-        with open(tmp_file.name, "rb") as f:
-            return f.read()
-    finally:
+        tmp_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
+        tmp_file.close()
         try:
-            os.remove(tmp_file.name)
-        except OSError:
-            pass
+            export_to_video(list(frames_np), tmp_file.name, fps=fps)
+            with open(tmp_file.name, "rb") as f:
+                return f.read()
+        finally:
+            try:
+                os.remove(tmp_file.name)
+            except OSError:
+                pass
+
+    from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
+
+    if frames_np.dtype == np.uint8:
+        frames_u8 = frames_np
+    else:
+        frames_np = np.clip(frames_np, 0.0, 1.0)
+        frames_np *= 255.0
+        frames_u8 = np.round(frames_np).astype(np.uint8)
+
+    # Ensure contiguous memory layout for faster PyAV muxing
+    frames_u8 = np.ascontiguousarray(frames_u8)
+
+    audio_np = _coerce_audio_to_numpy(audio) if audio is not None else None
+
+    return mux_video_audio_bytes(
+        frames_u8,
+        audio_np,
+        fps=float(fps),
+        audio_sample_rate=audio_sample_rate or 24000,
+        video_codec_options=video_codec_options,
+    )
 
 
-def encode_video_base64(video: Any, fps: int, audio: Any | None = None, audio_sample_rate: int | None = None) -> str:
+def encode_video_base64(
+    video: Any,
+    fps: int,
+    audio: Any | None = None,
+    audio_sample_rate: int | None = None,
+    video_codec_options: dict[str, str] | None = None,
+) -> str:
     """Encode a video (frames/array/tensor) to base64 MP4."""
-    video_bytes = _encode_video_bytes(video, fps=fps, audio=audio, audio_sample_rate=audio_sample_rate)
+    video_bytes = _encode_video_bytes(
+        video, fps=fps, audio=audio, audio_sample_rate=audio_sample_rate, video_codec_options=video_codec_options
+    )
     return base64.b64encode(video_bytes).decode("utf-8")

vllm_omni/inputs/data.py

@@ -227,6 +227,10 @@ class OmniDiffusionSamplingParams:
     frame_rate: float | None = None  # Floating-point rate used by the diffusion model when it differs from `fps`.
     height_not_provided: bool = False
     width_not_provided: bool = False
+    enable_frame_interpolation: bool = False
+    frame_interpolation_exp: int = 1
+    frame_interpolation_scale: float = 1.0
+    frame_interpolation_model_path: str | None = None
 
     # Timesteps
     timesteps: torch.Tensor | None = None

vllm_omni/outputs.py

@@ -58,6 +58,10 @@ class OmniRequestOutput:
     images: list[Image.Image] = field(default_factory=list)
     prompt: OmniPromptType | None = None
     latents: torch.Tensor | None = None
+    trajectory_timesteps: list[torch.Tensor] | None = None
+    trajectory_latents: torch.Tensor | None = None
+    trajectory_log_probs: torch.Tensor | None = None
+    trajectory_decoded: list[torch.Tensor] | list[Image.Image] | None = None
     metrics: dict[str, Any] = field(default_factory=dict)
     _multimodal_output: dict[str, Any] = field(default_factory=dict)
     _custom_output: dict[str, Any] = field(default_factory=dict)
@@ -68,6 +72,9 @@ class OmniRequestOutput:
     # memory usage info
     peak_memory_mb: float = 0.0
 
+    # error handling
+    error: str | None = None
+
     @classmethod
     def from_pipeline(
         cls,
@@ -101,6 +108,10 @@ class OmniRequestOutput:
         prompt: OmniPromptType | None = None,
         metrics: dict[str, Any] | None = None,
         latents: torch.Tensor | None = None,
+        trajectory_timesteps: list[torch.Tensor] | None = None,
+        trajectory_latents: torch.Tensor | None = None,
+        trajectory_log_probs: torch.Tensor | None = None,
+        trajectory_decoded: list[torch.Tensor] | list[Image.Image] | None = None,
         multimodal_output: dict[str, Any] | None = None,
         custom_output: dict[str, Any] | None = None,
         final_output_type: str = "image",
@@ -115,6 +126,10 @@ class OmniRequestOutput:
             prompt: The prompt used
             metrics: Generation metrics
             latents: Optional latent tensors
+            trajectory_timesteps: Optional sampled timesteps for trajectory outputs
+            trajectory_latents: Optional latent trajectory tensor
+            trajectory_log_probs: Optional per-step log probabilities
+            trajectory_decoded: Optional decoded trajectory frames
             multimodal_output: Optional multimodal output dict
             custom_output: Optional custom output dict (e.g. latent trajectories, prompt embeds)
             stage_durations: Optional stage durations (execution time of each stage) dict
@@ -129,6 +144,10 @@ class OmniRequestOutput:
             images=images,
             prompt=prompt,
             latents=latents,
+            trajectory_timesteps=trajectory_timesteps,
+            trajectory_latents=trajectory_latents,
+            trajectory_log_probs=trajectory_log_probs,
+            trajectory_decoded=trajectory_decoded,
             metrics=metrics or {},
             _multimodal_output=multimodal_output or {},
             _custom_output=custom_output or {},

vllm_omni/platforms/interface.py

@@ -1,12 +1,16 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from contextlib import nullcontext
 from enum import Enum
 from typing import Any
 
 import torch
+from vllm.logger import init_logger
 from vllm.platforms import Platform
 
+logger = init_logger(__name__)
+
 
 class OmniPlatformEnum(Enum):
     """Enum for supported Omni platforms."""
@@ -109,6 +113,23 @@ class OmniPlatform(Platform):
     def get_free_memory(cls, device: torch.device | None = None) -> int:
         raise NotImplementedError
 
+    @classmethod
+    def create_autocast_context(
+        cls,
+        *,
+        device_type: str,
+        dtype: torch.dtype,
+        enabled: bool = True,
+    ):
+        if not enabled:
+            return nullcontext()
+
+        try:
+            return torch.autocast(device_type=device_type, dtype=dtype, enabled=True)
+        except (RuntimeError, TypeError, ValueError) as exc:
+            logger.warning("autocast unavailable for device_type=%s dtype=%s: %s", device_type, dtype, exc)
+            return nullcontext()
+
     @classmethod
     def supports_cpu_offload(cls) -> bool:
         return True

vllm_omni/platforms/npu/platform.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from contextlib import nullcontext
 from typing import Any
 
 import torch
@@ -68,6 +69,9 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
 
         # Try FLASH_ATTN if mindiesd is available, otherwise fall back to SDPA
         if find_spec("mindiesd"):
+            # Configure ASCEND_CUSTOM_OPP_PATH for mindiesd custom ops upon import
+            import mindiesd  # noqa: F401
+
             logger.info("Defaulting to diffusion attention backend FLASH_ATTN")
             return DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
 
@@ -106,6 +110,24 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
         device_props = torch.npu.get_device_properties(device_id)
         return device_props.total_memory
 
+    @classmethod
+    def create_autocast_context(cls, *, device_type, dtype, enabled=True):
+        if device_type != "npu":
+            return super().create_autocast_context(
+                device_type=device_type,
+                dtype=dtype,
+                enabled=enabled,
+            )
+        if not enabled:
+            return nullcontext()
+
+        # NPU-specific fallback
+        try:
+            return torch.npu.amp.autocast(dtype=dtype)
+        except (RuntimeError, TypeError, ValueError) as exc:
+            logger.warning("autocast unavailable for device_type=%s dtype=%s: %s", device_type, dtype, exc)
+        return nullcontext()
+
     @classmethod
     def get_profiler_cls(cls) -> str:
         return "vllm_omni.platforms.npu.profiler.NPUTorchProfilerWrapper"