[Model] Add Wan2.2 text-to-video support (#202)

Signed-off-by: linyueqian <linyueqian@outlook.com>

docs/models/supported_models.md

@@ -22,6 +22,7 @@ th {
 | `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |
 | `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` |
 |`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |
+| `WanPipeline` | Wan2.2 | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` |
 
 
 ## List of Supported Models for NPU

examples/offline_inference/wan22/text_to_video.py

@@ -0,0 +1,105 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.utils.platform_utils import detect_device_type, is_npu
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Generate a video with Wan2.2 T2V.")
+    parser.add_argument(
+        "--model",
+        default="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+        help="Diffusers Wan2.2 model ID or local path.",
+    )
+    parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.")
+    parser.add_argument("--negative_prompt", default="", help="Negative prompt.")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed.")
+    parser.add_argument("--guidance_scale", type=float, default=4.0, help="CFG scale (applied to low/high).")
+    parser.add_argument("--guidance_scale_high", type=float, default=None, help="Optional separate CFG for high-noise.")
+    parser.add_argument("--height", type=int, default=720, help="Video height.")
+    parser.add_argument("--width", type=int, default=1280, help="Video width.")
+    parser.add_argument("--num_frames", type=int, default=81, help="Number of frames (Wan default is 81).")
+    parser.add_argument("--num_inference_steps", type=int, default=40, help="Sampling steps.")
+    parser.add_argument("--boundary_ratio", type=float, default=0.875, help="Boundary split ratio for low/high DiT.")
+    parser.add_argument(
+        "--flow_shift", type=float, default=5.0, help="Scheduler flow_shift (5.0 for 720p, 12.0 for 480p)."
+    )
+    parser.add_argument("--output", type=str, default="wan22_output.mp4", help="Path to save the video (mp4).")
+    parser.add_argument("--fps", type=int, default=24, help="Frames per second for the output video.")
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    device = detect_device_type()
+    generator = torch.Generator(device=device).manual_seed(args.seed)
+
+    # Enable VAE memory optimizations on NPU
+    vae_use_slicing = is_npu()
+    vae_use_tiling = is_npu()
+
+    omni = Omni(
+        model=args.model,
+        vae_use_slicing=vae_use_slicing,
+        vae_use_tiling=vae_use_tiling,
+        boundary_ratio=args.boundary_ratio,
+        flow_shift=args.flow_shift,
+    )
+
+    frames = omni.generate(
+        args.prompt,
+        negative_prompt=args.negative_prompt,
+        height=args.height,
+        width=args.width,
+        generator=generator,
+        guidance_scale=args.guidance_scale,
+        guidance_scale_2=args.guidance_scale_high,
+        num_inference_steps=args.num_inference_steps,
+        num_frames=args.num_frames,
+    )
+
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    try:
+        from diffusers.utils import export_to_video
+    except ImportError:
+        raise ImportError("diffusers is required for export_to_video.")
+
+    # frames may be np.ndarray (preferred) or torch.Tensor
+    # export_to_video expects a list of frames with values in [0, 1]
+    if isinstance(frames, torch.Tensor):
+        video_tensor = frames.detach().cpu()
+        if video_tensor.dim() == 5:
+            # [B, C, F, H, W] or [B, F, H, W, C]
+            if video_tensor.shape[1] in (3, 4):
+                video_tensor = video_tensor[0].permute(1, 2, 3, 0)
+            else:
+                video_tensor = video_tensor[0]
+        elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4):
+            video_tensor = video_tensor.permute(1, 2, 3, 0)
+        # If float, assume [-1,1] and normalize to [0,1]
+        if video_tensor.is_floating_point():
+            video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5
+        video_array = video_tensor.float().numpy()
+    else:
+        video_array = frames
+        if hasattr(video_array, "shape") and video_array.ndim == 5:
+            video_array = video_array[0]
+
+    # Convert 4D array (frames, H, W, C) to list of frames for export_to_video
+    if isinstance(video_array, np.ndarray) and video_array.ndim == 4:
+        video_array = list(video_array)
+
+    export_to_video(video_array, str(output_path), fps=args.fps)
+    print(f"Saved generated video to {output_path}")
+
+
+if __name__ == "__main__":
+    main()

tests/single_stage/test_video_diffusion_model.py

@@ -0,0 +1,48 @@
+import os
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+
+# ruff: noqa: E402
+REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from vllm_omni import Omni
+
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+
+models = ["Wan-AI/Wan2.2-T2V-A14B-Diffusers"]
+
+
+@pytest.mark.parametrize("model_name", models)
+def test_video_diffusion_model(model_name: str):
+    m = Omni(
+        model=model_name,
+        boundary_ratio=0.875,
+        flow_shift=5.0,
+    )
+    # Use minimal settings for testing
+    # num_frames must satisfy: num_frames % vae_scale_factor_temporal == 1
+    # For Wan2.2, vae_scale_factor_temporal=4, so valid values are 5, 9, 13, 17, ...
+    height = 480
+    width = 640
+    num_frames = 5
+    frames = m.generate(
+        "A cat sitting on a table",
+        height=height,
+        width=width,
+        num_frames=num_frames,
+        num_inference_steps=2,
+        guidance_scale=1.0,
+        generator=torch.Generator("cuda").manual_seed(42),
+    )
+
+    assert frames is not None
+    assert hasattr(frames, "shape")
+    # frames shape: (batch, num_frames, height, width, channels)
+    assert frames.shape[1] == num_frames
+    assert frames.shape[2] == height
+    assert frames.shape[3] == width

vllm_omni/diffusion/data.py

@@ -162,6 +162,8 @@ class OmniDiffusionConfig:
 
     # MoE parameters used by Wan2.2
     boundary_ratio: float | None = None
+    # Scheduler flow_shift for Wan2.2 (12.0 for 480p, 5.0 for 720p)
+    flow_shift: float | None = None
 
     # Logging
     log_level: str = "info"

vllm_omni/diffusion/models/wan2_2/__init__.py

@@ -0,0 +1,4 @@
+from .pipeline_wan2_2 import Wan22Pipeline, get_wan22_post_process_func
+from .wan2_2_transformer import WanTransformer3DModel
+
+__all__ = ["Wan22Pipeline", "get_wan22_post_process_func", "WanTransformer3DModel"]

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py

@@ -0,0 +1,426 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import os
+from collections.abc import Iterable
+
+import torch
+from diffusers import AutoencoderKLWan, FlowMatchEulerDiscreteScheduler
+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.utils import get_local_device
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+
+
+def get_wan22_post_process_func(
+    od_config: OmniDiffusionConfig,
+):
+    from diffusers.video_processor import VideoProcessor
+
+    video_processor = VideoProcessor(vae_scale_factor=8)
+
+    def post_process_func(
+        video: torch.Tensor,
+        output_type: str = "np",
+    ):
+        if output_type == "latent":
+            return video
+        return video_processor.postprocess_video(video, output_type=output_type)
+
+    return post_process_func
+
+
+class Wan22Pipeline(nn.Module):
+    def __init__(
+        self,
+        *,
+        od_config: OmniDiffusionConfig,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.od_config = od_config
+
+        self.device = get_local_device()
+        dtype = getattr(od_config, "dtype", torch.bfloat16)
+
+        model = od_config.model
+        local_files_only = os.path.exists(model)
+
+        # Set up weights sources for transformer(s)
+        self.weights_sources = [
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=od_config.model,
+                subfolder="transformer",
+                revision=None,
+                prefix="transformer.",
+                fall_back_to_pt=True,
+            ),
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=od_config.model,
+                subfolder="transformer_2",
+                revision=None,
+                prefix="transformer_2.",
+                fall_back_to_pt=True,
+            ),
+        ]
+
+        self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
+        self.text_encoder = UMT5EncoderModel.from_pretrained(
+            model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only
+        ).to(self.device)
+        self.vae = AutoencoderKLWan.from_pretrained(
+            model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only
+        ).to(self.device)
+
+        # Initialize transformers without from_pretrained (weights loaded via load_weights)
+        self.transformer = WanTransformer3DModel()
+        self.transformer_2 = WanTransformer3DModel()
+
+        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+            model, subfolder="scheduler", local_files_only=local_files_only
+        )
+        # Apply flow_shift if specified (12.0 for 480p, 5.0 for 720p recommended for Wan2.2)
+        if od_config.flow_shift is not None:
+            self.scheduler.config.flow_shift = od_config.flow_shift
+
+        self.vae_scale_factor_temporal = self.vae.config.scale_factor_temporal if getattr(self, "vae", None) else 4
+        self.vae_scale_factor_spatial = self.vae.config.scale_factor_spatial if getattr(self, "vae", None) else 8
+        self.boundary_ratio = od_config.boundary_ratio
+
+        self._guidance_scale = None
+        self._guidance_scale_2 = None
+        self._num_timesteps = None
+        self._current_timestep = None
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def do_classifier_free_guidance(self):
+        return self._guidance_scale is not None and self._guidance_scale > 1.0
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def current_timestep(self):
+        return self._current_timestep
+
+    @torch.no_grad()
+    def forward(
+        self,
+        req: OmniDiffusionRequest,
+        prompt: str | None = None,
+        negative_prompt: str | None = None,
+        height: int | None = None,
+        width: int | None = None,
+        num_inference_steps: int | None = None,
+        guidance_scale: float | tuple[float, float] = 4.0,
+        frame_num: int | None = None,
+        output_type: str | None = "np",
+        generator: torch.Generator | None = None,
+        prompt_embeds: torch.Tensor | None = None,
+        negative_prompt_embeds: torch.Tensor | None = None,
+        attention_kwargs: dict | None = None,
+        **kwargs,
+    ) -> DiffusionOutput:
+        prompt = req.prompt if req.prompt is not None else prompt
+        negative_prompt = req.negative_prompt if req.negative_prompt is not None else negative_prompt
+        if prompt is None and prompt_embeds is None:
+            raise ValueError("Prompt or prompt_embeds is required for Wan2.2 generation.")
+
+        height = req.height or height or 720
+        width = req.width or width or 1280
+        num_frames = req.num_frames if req.num_frames else frame_num or 81
+        num_steps = req.num_inference_steps or num_inference_steps or 40
+
+        guidance_low = guidance_scale if isinstance(guidance_scale, (int, float)) else guidance_scale[0]
+        guidance_high = (
+            req.guidance_scale_2
+            if req.guidance_scale_2 is not None
+            else (
+                guidance_scale[1]
+                if isinstance(guidance_scale, (list, tuple)) and len(guidance_scale) > 1
+                else guidance_low
+            )
+        )
+
+        # record guidance for properties
+        self._guidance_scale = guidance_low
+        self._guidance_scale_2 = guidance_high
+
+        # validate shapes
+        self.check_inputs(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            height=height,
+            width=width,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            guidance_scale_2=guidance_high if self.boundary_ratio is not None else None,
+        )
+
+        if num_frames % self.vae_scale_factor_temporal != 1:
+            num_frames = num_frames // self.vae_scale_factor_temporal * self.vae_scale_factor_temporal + 1
+        num_frames = max(num_frames, 1)
+
+        device = self.device
+        dtype = self.transformer.dtype
+
+        # Seed / generator
+        if generator is None:
+            generator = req.generator
+        if generator is None and req.seed is not None:
+            generator = torch.Generator(device=device).manual_seed(req.seed)
+
+        # Encode prompts
+        if prompt_embeds is None:
+            prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0,
+                num_videos_per_prompt=req.num_outputs_per_prompt or 1,
+                max_sequence_length=req.max_sequence_length or 512,
+                device=device,
+                dtype=dtype,
+            )
+        else:
+            prompt_embeds = prompt_embeds.to(device=device, dtype=dtype)
+            if negative_prompt_embeds is not None:
+                negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype)
+            elif guidance_low > 1.0 or guidance_high > 1.0:
+                raise ValueError(
+                    "negative_prompt_embeds must be provided when prompt_embeds are given and guidance > 1."
+                )
+
+        # Timesteps
+        self.scheduler.set_timesteps(num_steps, device=device)
+        timesteps = self.scheduler.timesteps
+        self._num_timesteps = len(timesteps)
+        boundary_timestep = None
+        if self.boundary_ratio is not None:
+            boundary_timestep = self.boundary_ratio * self.scheduler.config.num_train_timesteps
+
+        # Latents
+        num_channels_latents = self.transformer.config.in_channels
+        latents = self.prepare_latents(
+            batch_size=prompt_embeds.shape[0],
+            num_channels_latents=num_channels_latents,
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            dtype=torch.float32,
+            device=device,
+            generator=generator,
+            latents=req.latents,
+        )
+
+        if attention_kwargs is None:
+            attention_kwargs = {}
+
+        # Denoising
+        for t in timesteps:
+            self._current_timestep = t
+            current_model = self.transformer
+            current_guidance_scale = guidance_low
+            if boundary_timestep is not None and t < boundary_timestep:
+                current_model = self.transformer_2
+                current_guidance_scale = guidance_high
+
+            latent_model_input = latents.to(dtype)
+            timestep = t.expand(latents.shape[0])
+
+            noise_pred = current_model(
+                hidden_states=latent_model_input,
+                timestep=timestep,
+                encoder_hidden_states=prompt_embeds,
+                attention_kwargs=attention_kwargs,
+                return_dict=False,
+            )[0]
+
+            if current_guidance_scale > 1.0 and negative_prompt_embeds is not None:
+                noise_uncond = current_model(
+                    hidden_states=latent_model_input,
+                    timestep=timestep,
+                    encoder_hidden_states=negative_prompt_embeds,
+                    attention_kwargs=attention_kwargs,
+                    return_dict=False,
+                )[0]
+                noise_pred = noise_uncond + current_guidance_scale * (noise_pred - noise_uncond)
+
+            latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+        self._current_timestep = None
+
+        # Decode
+        if output_type == "latent":
+            output = latents
+        else:
+            latents = latents.to(self.vae.dtype)
+            latents_mean = (
+                torch.tensor(self.vae.config.latents_mean)
+                .view(1, self.vae.config.z_dim, 1, 1, 1)
+                .to(latents.device, latents.dtype)
+            )
+            latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
+                latents.device, latents.dtype
+            )
+            latents = latents / latents_std + latents_mean
+            output = self.vae.decode(latents, return_dict=False)[0]
+
+        return DiffusionOutput(output=output)
+
+    def encode_prompt(
+        self,
+        prompt: str | list[str],
+        negative_prompt: str | list[str] | None = None,
+        do_classifier_free_guidance: bool = True,
+        num_videos_per_prompt: int = 1,
+        max_sequence_length: int = 512,
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
+    ):
+        device = device or self.device
+        dtype = dtype or self.text_encoder.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        prompt_clean = [self._prompt_clean(p) for p in prompt]
+        batch_size = len(prompt_clean)
+
+        text_inputs = self.tokenizer(
+            prompt_clean,
+            padding="max_length",
+            max_length=max_sequence_length,
+            truncation=True,
+            add_special_tokens=True,
+            return_attention_mask=True,
+            return_tensors="pt",
+        )
+        ids, mask = text_inputs.input_ids, text_inputs.attention_mask
+        seq_lens = mask.gt(0).sum(dim=1).long()
+
+        prompt_embeds = self.text_encoder(ids.to(device), mask.to(device)).last_hidden_state
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+        prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)]
+        prompt_embeds = torch.stack(
+            [torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0
+        )
+
+        _, seq_len, _ = prompt_embeds.shape
+        prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
+
+        negative_prompt_embeds = None
+        if do_classifier_free_guidance:
+            negative_prompt = negative_prompt or ""
+            negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+            neg_text_inputs = self.tokenizer(
+                [self._prompt_clean(p) for p in negative_prompt],
+                padding="max_length",
+                max_length=max_sequence_length,
+                truncation=True,
+                add_special_tokens=True,
+                return_attention_mask=True,
+                return_tensors="pt",
+            )
+            ids_neg, mask_neg = neg_text_inputs.input_ids, neg_text_inputs.attention_mask
+            seq_lens_neg = mask_neg.gt(0).sum(dim=1).long()
+            negative_prompt_embeds = self.text_encoder(ids_neg.to(device), mask_neg.to(device)).last_hidden_state
+            negative_prompt_embeds = negative_prompt_embeds.to(dtype=dtype, device=device)
+            negative_prompt_embeds = [u[:v] for u, v in zip(negative_prompt_embeds, seq_lens_neg)]
+            negative_prompt_embeds = torch.stack(
+                [
+                    torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))])
+                    for u in negative_prompt_embeds
+                ],
+                dim=0,
+            )
+            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_videos_per_prompt, 1)
+            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
+
+        return prompt_embeds, negative_prompt_embeds
+
+    @staticmethod
+    def _prompt_clean(text: str) -> str:
+        return " ".join(text.strip().split())
+
+    def prepare_latents(
+        self,
+        batch_size: int,
+        num_channels_latents: int,
+        height: int,
+        width: int,
+        num_frames: int,
+        dtype: torch.dtype | None,
+        device: torch.device | None,
+        generator: torch.Generator | list[torch.Generator] | None,
+        latents: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if latents is not None:
+            return latents.to(device=device, dtype=dtype)
+
+        num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
+        shape = (
+            batch_size,
+            num_channels_latents,
+            num_latent_frames,
+            int(height) // self.vae_scale_factor_spatial,
+            int(width) // self.vae_scale_factor_spatial,
+        )
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(f"Generator list length {len(generator)} does not match batch size {batch_size}.")
+        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        return latents
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Load weights using AutoWeightsLoader for vLLM integration."""
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights)
+
+    def check_inputs(
+        self,
+        prompt,
+        negative_prompt,
+        height,
+        width,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        guidance_scale_2=None,
+    ):
+        if height % 16 != 0 or width % 16 != 0:
+            raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.")
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif negative_prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and "
+                f"`negative_prompt_embeds`: {negative_prompt_embeds}. "
+                "Please make sure to only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+        elif negative_prompt is not None and (
+            not isinstance(negative_prompt, str) and not isinstance(negative_prompt, list)
+        ):
+            raise ValueError(f"`negative_prompt` has to be of type `str` or `list` but is {type(negative_prompt)}")
+
+        if self.boundary_ratio is None and guidance_scale_2 is not None:
+            raise ValueError("`guidance_scale_2` is only supported when `boundary_ratio` is set.")

vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py

@@ -0,0 +1,704 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import math
+from collections.abc import Iterable
+from typing import Any, Optional, Union
+
+import torch
+import torch.nn as nn
+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.logger import init_logger
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import QKVParallelLinear, ReplicatedLinear
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+from vllm_omni.diffusion.attention.layer import Attention
+
+logger = init_logger(__name__)
+
+
+def apply_rotary_emb_wan(
+    hidden_states: torch.Tensor,
+    freqs_cos: torch.Tensor,
+    freqs_sin: torch.Tensor,
+) -> torch.Tensor:
+    """
+    Apply rotary embeddings to input tensors using the given frequency tensors.
+
+    Args:
+        hidden_states: Input tensor of shape [B, S, H, D]
+        freqs_cos: Cosine frequencies
+        freqs_sin: Sine frequencies
+
+    Returns:
+        Tensor with rotary embeddings applied
+    """
+    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)
+
+
+class WanRotaryPosEmbed(nn.Module):
+    """
+    Rotary position embeddings for 3D video data (temporal + spatial dimensions).
+    """
+
+    def __init__(
+        self,
+        attention_head_dim: int,
+        patch_size: tuple[int, int, int],
+        max_seq_len: int,
+        theta: float = 10000.0,
+    ):
+        super().__init__()
+
+        self.attention_head_dim = attention_head_dim
+        self.patch_size = patch_size
+        self.max_seq_len = max_seq_len
+
+        # Split dimensions for temporal, height, width
+        h_dim = w_dim = 2 * (attention_head_dim // 6)
+        t_dim = attention_head_dim - h_dim - w_dim
+        freqs_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64
+
+        freqs_cos = []
+        freqs_sin = []
+
+        for dim in [t_dim, h_dim, w_dim]:
+            freq_cos, freq_sin = self._get_1d_rotary_pos_embed(dim, max_seq_len, theta, freqs_dtype)
+            freqs_cos.append(freq_cos)
+            freqs_sin.append(freq_sin)
+
+        self.register_buffer("freqs_cos", torch.cat(freqs_cos, dim=1), persistent=False)
+        self.register_buffer("freqs_sin", torch.cat(freqs_sin, dim=1), persistent=False)
+
+    @staticmethod
+    def _get_1d_rotary_pos_embed(
+        dim: int,
+        max_seq_len: int,
+        theta: float,
+        freqs_dtype: torch.dtype,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Generate 1D rotary position embeddings."""
+        freqs = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=freqs_dtype) / dim))
+        t = torch.arange(max_seq_len, dtype=freqs_dtype)
+        freqs = torch.outer(t, freqs)
+        # Repeat interleave for real representation
+        freqs_cos = freqs.cos().repeat_interleave(2, dim=-1)
+        freqs_sin = freqs.sin().repeat_interleave(2, dim=-1)
+        return freqs_cos.float(), freqs_sin.float()
+
+    def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+        p_t, p_h, p_w = self.patch_size
+        ppf, pph, ppw = num_frames // p_t, height // p_h, width // p_w
+
+        split_sizes = [
+            self.attention_head_dim - 2 * (self.attention_head_dim // 3),
+            self.attention_head_dim // 3,
+            self.attention_head_dim // 3,
+        ]
+
+        freqs_cos = self.freqs_cos.split(split_sizes, dim=1)
+        freqs_sin = self.freqs_sin.split(split_sizes, dim=1)
+
+        freqs_cos_f = freqs_cos[0][:ppf].view(ppf, 1, 1, -1).expand(ppf, pph, ppw, -1)
+        freqs_cos_h = freqs_cos[1][:pph].view(1, pph, 1, -1).expand(ppf, pph, ppw, -1)
+        freqs_cos_w = freqs_cos[2][:ppw].view(1, 1, ppw, -1).expand(ppf, pph, ppw, -1)
+
+        freqs_sin_f = freqs_sin[0][:ppf].view(ppf, 1, 1, -1).expand(ppf, pph, ppw, -1)
+        freqs_sin_h = freqs_sin[1][:pph].view(1, pph, 1, -1).expand(ppf, pph, ppw, -1)
+        freqs_sin_w = freqs_sin[2][:ppw].view(1, 1, ppw, -1).expand(ppf, pph, ppw, -1)
+
+        freqs_cos = torch.cat([freqs_cos_f, freqs_cos_h, freqs_cos_w], dim=-1).reshape(1, ppf * pph * ppw, 1, -1)
+        freqs_sin = torch.cat([freqs_sin_f, freqs_sin_h, freqs_sin_w], dim=-1).reshape(1, ppf * pph * ppw, 1, -1)
+
+        return freqs_cos, freqs_sin
+
+
+class WanImageEmbedding(nn.Module):
+    """Image embedding module for I2V tasks."""
+
+    def __init__(self, in_features: int, out_features: int, pos_embed_seq_len: Optional[int] = None):
+        super().__init__()
+
+        self.norm1 = FP32LayerNorm(in_features)
+        self.ff = FeedForward(in_features, out_features, mult=1, activation_fn="gelu")
+        self.norm2 = FP32LayerNorm(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:
+            self.pos_embed = None
+
+    def forward(self, encoder_hidden_states_image: torch.Tensor) -> torch.Tensor:
+        if self.pos_embed is not None:
+            batch_size, seq_len, embed_dim = encoder_hidden_states_image.shape
+            encoder_hidden_states_image = encoder_hidden_states_image.view(-1, 2 * seq_len, embed_dim)
+            encoder_hidden_states_image = encoder_hidden_states_image + self.pos_embed
+
+        hidden_states = self.norm1(encoder_hidden_states_image)
+        hidden_states = self.ff(hidden_states)
+        hidden_states = self.norm2(hidden_states)
+        return hidden_states
+
+
+class WanTimeTextImageEmbedding(nn.Module):
+    """Combined time, text, and image condition embeddings."""
+
+    def __init__(
+        self,
+        dim: int,
+        time_freq_dim: int,
+        time_proj_dim: int,
+        text_embed_dim: int,
+        image_embed_dim: Optional[int] = None,
+        pos_embed_seq_len: Optional[int] = None,
+    ):
+        super().__init__()
+
+        self.timesteps_proj = Timesteps(num_channels=time_freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0)
+        self.time_embedder = TimestepEmbedding(in_channels=time_freq_dim, time_embed_dim=dim)
+        self.act_fn = nn.SiLU()
+        self.time_proj = nn.Linear(dim, time_proj_dim)
+        self.text_embedder = PixArtAlphaTextProjection(text_embed_dim, dim, act_fn="gelu_tanh")
+
+        self.image_embedder = None
+        if image_embed_dim is not None:
+            self.image_embedder = WanImageEmbedding(image_embed_dim, dim, pos_embed_seq_len=pos_embed_seq_len)
+
+    def forward(
+        self,
+        timestep: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        encoder_hidden_states_image: Optional[torch.Tensor] = None,
+        timestep_seq_len: Optional[int] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+        timestep = self.timesteps_proj(timestep)
+        if timestep_seq_len is not None:
+            timestep = timestep.unflatten(0, (-1, timestep_seq_len))
+
+        time_embedder_dtype = next(iter(self.time_embedder.parameters())).dtype
+        if timestep.dtype != time_embedder_dtype and time_embedder_dtype != torch.int8:
+            timestep = timestep.to(time_embedder_dtype)
+        temb = self.time_embedder(timestep).type_as(encoder_hidden_states)
+        timestep_proj = self.time_proj(self.act_fn(temb))
+
+        encoder_hidden_states = self.text_embedder(encoder_hidden_states)
+        if encoder_hidden_states_image is not None:
+            encoder_hidden_states_image = self.image_embedder(encoder_hidden_states_image)
+
+        return temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image
+
+
+class WanSelfAttention(nn.Module):
+    """
+    Optimized self-attention module using vLLM layers.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        head_dim: int,
+        eps: float = 1e-5,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.inner_dim = num_heads * head_dim
+
+        # Fused QKV projection using vLLM's optimized layer
+        self.to_qkv = QKVParallelLinear(
+            hidden_size=dim,
+            head_size=head_dim,
+            total_num_heads=num_heads,
+            bias=True,
+            disable_tp=True,
+        )
+
+        # QK normalization using vLLM's RMSNorm
+        self.norm_q = RMSNorm(self.inner_dim, eps=eps)
+        self.norm_k = RMSNorm(self.inner_dim, eps=eps)
+
+        # Output projection
+        self.to_out = nn.ModuleList(
+            [
+                ReplicatedLinear(self.inner_dim, dim, bias=True),
+                nn.Dropout(dropout),
+            ]
+        )
+
+        # Unified attention layer
+        self.attn = Attention(
+            num_heads=num_heads,
+            head_size=head_dim,
+            softmax_scale=1.0 / (head_dim**0.5),
+            causal=False,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        rotary_emb: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+    ) -> torch.Tensor:
+        # Fused QKV projection
+        qkv, _ = self.to_qkv(hidden_states)
+        query, key, value = qkv.chunk(3, dim=-1)
+
+        # Apply QK normalization
+        query = self.norm_q(query)
+        key = self.norm_k(key)
+
+        # Reshape for multi-head attention
+        query = query.unflatten(2, (self.num_heads, -1))
+        key = key.unflatten(2, (self.num_heads, -1))
+        value = value.unflatten(2, (self.num_heads, -1))
+
+        # Apply rotary embeddings
+        if rotary_emb is not None:
+            freqs_cos, freqs_sin = rotary_emb
+            query = apply_rotary_emb_wan(query, freqs_cos, freqs_sin)
+            key = apply_rotary_emb_wan(key, freqs_cos, freqs_sin)
+
+        # Compute attention using unified attention layer
+        hidden_states = self.attn(query, key, value)
+        hidden_states = hidden_states.flatten(2, 3)
+        hidden_states = hidden_states.type_as(query)
+
+        # Output projection
+        hidden_states, _ = self.to_out[0](hidden_states)
+        hidden_states = self.to_out[1](hidden_states)
+
+        return hidden_states
+
+
+class WanCrossAttention(nn.Module):
+    """
+    Optimized cross-attention module using vLLM layers.
+    Handles both text cross-attention and optional image cross-attention (I2V).
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        head_dim: int,
+        eps: float = 1e-5,
+        dropout: float = 0.0,
+        added_kv_proj_dim: Optional[int] = None,
+    ):
+        super().__init__()
+
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.inner_dim = num_heads * head_dim
+        self.kv_inner_dim = head_dim * num_heads  # For cross-attention, K/V come from encoder
+
+        # Query projection
+        self.to_q = ReplicatedLinear(dim, self.inner_dim, bias=True)
+
+        # Separate K and V projections for cross-attention
+        self.to_k = ReplicatedLinear(dim, self.kv_inner_dim, bias=True)
+        self.to_v = ReplicatedLinear(dim, self.kv_inner_dim, bias=True)
+
+        # QK normalization
+        self.norm_q = RMSNorm(self.inner_dim, eps=eps)
+        self.norm_k = RMSNorm(self.kv_inner_dim, eps=eps)
+
+        # Optional added KV projections for I2V (image embeddings)
+        self.added_kv_proj_dim = added_kv_proj_dim
+        if added_kv_proj_dim is not None:
+            self.add_k_proj = ReplicatedLinear(added_kv_proj_dim, self.inner_dim, bias=True)
+            self.add_v_proj = ReplicatedLinear(added_kv_proj_dim, self.inner_dim, bias=True)
+            self.norm_added_k = RMSNorm(self.inner_dim, eps=eps)
+        else:
+            self.add_k_proj = None
+            self.add_v_proj = None
+            self.norm_added_k = None
+
+        # Output projection
+        self.to_out = nn.ModuleList(
+            [
+                ReplicatedLinear(self.inner_dim, dim, bias=True),
+                nn.Dropout(dropout),
+            ]
+        )
+
+        # Unified attention layer
+        self.attn = Attention(
+            num_heads=num_heads,
+            head_size=head_dim,
+            softmax_scale=1.0 / (head_dim**0.5),
+            causal=False,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        # Handle I2V case where encoder_hidden_states contains both image and text
+        encoder_hidden_states_img = None
+        if self.add_k_proj is not None:
+            # 512 is the context length of the text encoder, hardcoded for now
+            image_context_length = encoder_hidden_states.shape[1] - 512
+            encoder_hidden_states_img = encoder_hidden_states[:, :image_context_length]
+            encoder_hidden_states = encoder_hidden_states[:, image_context_length:]
+
+        # Query projection
+        query, _ = self.to_q(hidden_states)
+        query = self.norm_q(query)
+
+        # KV projection from encoder
+        key, _ = self.to_k(encoder_hidden_states)
+        value, _ = self.to_v(encoder_hidden_states)
+        key = self.norm_k(key)
+
+        # Reshape for multi-head attention
+        query = query.unflatten(2, (self.num_heads, -1))
+        key = key.unflatten(2, (self.num_heads, -1))
+        value = value.unflatten(2, (self.num_heads, -1))
+
+        # I2V: Additional attention with image embeddings
+        hidden_states_img = None
+        if encoder_hidden_states_img is not None:
+            key_img, _ = self.add_k_proj(encoder_hidden_states_img)
+            value_img, _ = self.add_v_proj(encoder_hidden_states_img)
+            key_img = self.norm_added_k(key_img)
+
+            key_img = key_img.unflatten(2, (self.num_heads, -1))
+            value_img = value_img.unflatten(2, (self.num_heads, -1))
+
+            hidden_states_img = self.attn(query, key_img, value_img)
+            hidden_states_img = hidden_states_img.flatten(2, 3)
+            hidden_states_img = hidden_states_img.type_as(query)
+
+        # Main cross-attention using unified attention layer
+        hidden_states = self.attn(query, key, value)
+        hidden_states = hidden_states.flatten(2, 3)
+        hidden_states = hidden_states.type_as(query)
+
+        # Add image attention output if present
+        if hidden_states_img is not None:
+            hidden_states = hidden_states + hidden_states_img
+
+        # Output projection
+        hidden_states, _ = self.to_out[0](hidden_states)
+        hidden_states = self.to_out[1](hidden_states)
+
+        return hidden_states
+
+
+class WanTransformerBlock(nn.Module):
+    """
+    Transformer block for Wan model with self-attention, cross-attention, and FFN.
+    Uses scale-shift modulation from timestep embeddings.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        ffn_dim: int,
+        num_heads: int,
+        eps: float = 1e-6,
+        added_kv_proj_dim: Optional[int] = None,
+        cross_attn_norm: bool = False,
+    ):
+        super().__init__()
+
+        head_dim = dim // num_heads
+
+        # 1. Self-attention
+        self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
+        self.attn1 = WanSelfAttention(
+            dim=dim,
+            num_heads=num_heads,
+            head_dim=head_dim,
+            eps=eps,
+        )
+
+        # 2. Cross-attention
+        self.attn2 = WanCrossAttention(
+            dim=dim,
+            num_heads=num_heads,
+            head_dim=head_dim,
+            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()
+
+        # 3. Feed-forward
+        self.ffn = FeedForward(dim, inner_dim=ffn_dim, activation_fn="gelu-approximate")
+        self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False)
+
+        # Scale-shift table for modulation
+        self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        rotary_emb: tuple[torch.Tensor, torch.Tensor],
+    ) -> torch.Tensor:
+        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()
+            ).chunk(6, dim=2)
+            shift_msa = shift_msa.squeeze(2)
+            scale_msa = scale_msa.squeeze(2)
+            gate_msa = gate_msa.squeeze(2)
+            c_shift_msa = c_shift_msa.squeeze(2)
+            c_scale_msa = c_scale_msa.squeeze(2)
+            c_gate_msa = c_gate_msa.squeeze(2)
+        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()
+            ).chunk(6, dim=1)
+
+        # 1. Self-attention
+        norm_hidden_states = (self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa).type_as(hidden_states)
+        attn_output = self.attn1(norm_hidden_states, rotary_emb)
+        hidden_states = (hidden_states.float() + attn_output * gate_msa).type_as(hidden_states)
+
+        # 2. Cross-attention
+        norm_hidden_states = self.norm2(hidden_states.float()).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
+        )
+        ff_output = self.ffn(norm_hidden_states)
+        hidden_states = (hidden_states.float() + ff_output.float() * c_gate_msa).type_as(hidden_states)
+
+        return hidden_states
+
+
+class WanTransformer3DModel(nn.Module):
+    """
+    Optimized Wan Transformer model for video generation using vLLM layers.
+
+    This is an optimized version of the diffusers WanTransformer3DModel that uses
+    vLLM's efficient QKVParallelLinear and RMSNorm implementations.
+
+    Args:
+        patch_size: 3D patch dimensions for video embedding (t_patch, h_patch, w_patch)
+        num_attention_heads: Number of attention heads
+        attention_head_dim: Dimension of each attention head
+        in_channels: Number of input channels
+        out_channels: Number of output channels
+        text_dim: Input dimension for text embeddings
+        freq_dim: Dimension for sinusoidal time embeddings
+        ffn_dim: Intermediate dimension in feed-forward network
+        num_layers: Number of transformer blocks
+        cross_attn_norm: Enable cross-attention normalization
+        eps: Epsilon value for normalization layers
+        image_dim: Optional image embedding dimension for I2V
+        added_kv_proj_dim: Optional added KV projection dimension for I2V
+        rope_max_seq_len: Maximum sequence length for rotary embeddings
+        pos_embed_seq_len: Optional position embedding sequence length
+    """
+
+    def __init__(
+        self,
+        patch_size: tuple[int, int, int] = (1, 2, 2),
+        num_attention_heads: int = 40,
+        attention_head_dim: int = 128,
+        in_channels: int = 16,
+        out_channels: int = 16,
+        text_dim: int = 4096,
+        freq_dim: int = 256,
+        ffn_dim: int = 13824,
+        num_layers: int = 40,
+        cross_attn_norm: bool = True,
+        eps: float = 1e-6,
+        image_dim: Optional[int] = None,
+        added_kv_proj_dim: Optional[int] = None,
+        rope_max_seq_len: int = 1024,
+        pos_embed_seq_len: Optional[int] = None,
+    ):
+        super().__init__()
+
+        # Store config for compatibility
+        self.config = type(
+            "Config",
+            (),
+            {
+                "patch_size": patch_size,
+                "num_attention_heads": num_attention_heads,
+                "attention_head_dim": attention_head_dim,
+                "in_channels": in_channels,
+                "out_channels": out_channels,
+                "text_dim": text_dim,
+                "freq_dim": freq_dim,
+                "ffn_dim": ffn_dim,
+                "num_layers": num_layers,
+                "cross_attn_norm": cross_attn_norm,
+                "eps": eps,
+                "image_dim": image_dim,
+                "added_kv_proj_dim": added_kv_proj_dim,
+                "rope_max_seq_len": rope_max_seq_len,
+                "pos_embed_seq_len": pos_embed_seq_len,
+            },
+        )()
+
+        inner_dim = num_attention_heads * attention_head_dim
+        out_channels = out_channels or in_channels
+
+        # 1. Patch & position embedding
+        self.rope = WanRotaryPosEmbed(attention_head_dim, patch_size, rope_max_seq_len)
+        self.patch_embedding = nn.Conv3d(in_channels, inner_dim, kernel_size=patch_size, stride=patch_size)
+
+        # 2. Condition embeddings
+        self.condition_embedder = WanTimeTextImageEmbedding(
+            dim=inner_dim,
+            time_freq_dim=freq_dim,
+            time_proj_dim=inner_dim * 6,
+            text_embed_dim=text_dim,
+            image_embed_dim=image_dim,
+            pos_embed_seq_len=pos_embed_seq_len,
+        )
+
+        # 3. Transformer blocks
+        self.blocks = nn.ModuleList(
+            [
+                WanTransformerBlock(inner_dim, ffn_dim, num_attention_heads, eps, added_kv_proj_dim, cross_attn_norm)
+                for _ in range(num_layers)
+            ]
+        )
+
+        # 4. Output norm & projection
+        self.norm_out = FP32LayerNorm(inner_dim, eps, elementwise_affine=False)
+        self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size))
+        self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5)
+
+    @property
+    def dtype(self) -> torch.dtype:
+        """Return the dtype of the model parameters."""
+        return next(self.parameters()).dtype
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.LongTensor,
+        encoder_hidden_states: torch.Tensor,
+        encoder_hidden_states_image: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+        attention_kwargs: Optional[dict[str, Any]] = None,
+    ) -> Union[torch.Tensor, Transformer2DModelOutput]:
+        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+        p_t, p_h, p_w = self.config.patch_size
+        post_patch_num_frames = num_frames // p_t
+        post_patch_height = height // p_h
+        post_patch_width = width // p_w
+
+        rotary_emb = self.rope(hidden_states)
+
+        hidden_states = self.patch_embedding(hidden_states)
+        hidden_states = hidden_states.flatten(2).transpose(1, 2)
+
+        # Handle timestep shape
+        if timestep.ndim == 2:
+            ts_seq_len = timestep.shape[1]
+            timestep = timestep.flatten()
+        else:
+            ts_seq_len = None
+
+        temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
+            timestep, encoder_hidden_states, encoder_hidden_states_image, timestep_seq_len=ts_seq_len
+        )
+        if ts_seq_len is not None:
+            timestep_proj = timestep_proj.unflatten(2, (6, -1))
+        else:
+            timestep_proj = timestep_proj.unflatten(1, (6, -1))
+
+        if encoder_hidden_states_image is not None:
+            encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1)
+
+        # Transformer blocks
+        for block in self.blocks:
+            hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb)
+
+        # Output norm, projection & unpatchify
+        if temb.ndim == 3:
+            shift, scale = (self.scale_shift_table.unsqueeze(0).to(temb.device) + temb.unsqueeze(2)).chunk(2, dim=2)
+            shift = shift.squeeze(2)
+            scale = scale.squeeze(2)
+        else:
+            shift, scale = (self.scale_shift_table.to(temb.device) + temb.unsqueeze(1)).chunk(2, dim=1)
+
+        shift = shift.to(hidden_states.device)
+        scale = scale.to(hidden_states.device)
+
+        hidden_states = (self.norm_out(hidden_states.float()) * (1 + scale) + shift).type_as(hidden_states)
+        hidden_states = self.proj_out(hidden_states)
+
+        hidden_states = hidden_states.reshape(
+            batch_size, post_patch_num_frames, post_patch_height, post_patch_width, p_t, p_h, p_w, -1
+        )
+        hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
+        output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
+
+        if not return_dict:
+            return (output,)
+
+        return Transformer2DModelOutput(sample=output)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """
+        Load weights from a pretrained model, handling the mapping from
+        separate Q/K/V projections to fused QKV projections for self-attention.
+
+        Diffusers weight names:
+        - blocks.N.attn1.to_q/to_k/to_v -> fused to blocks.N.attn1.to_qkv (self-attention)
+        - blocks.N.attn2.to_q/to_k/to_v -> kept separate (cross-attention)
+        - blocks.N.attn1.norm_q/norm_k -> QK normalization for self-attention
+
+        Returns:
+            Set of parameter names that were successfully loaded.
+        """
+        # Stacked params mapping for self-attention QKV fusion
+        # Format: (param_name, shard_name, shard_id)
+        # Note: Only fuse attn1 (self-attention), NOT attn2 (cross-attention)
+        stacked_params_mapping = [
+            # self-attention QKV fusion (attn1 only)
+            (".attn1.to_qkv", ".attn1.to_q", "q"),
+            (".attn1.to_qkv", ".attn1.to_k", "k"),
+            (".attn1.to_qkv", ".attn1.to_v", "v"),
+        ]
+
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+
+        for name, loaded_weight in weights:
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+
+        return loaded_params

vllm_omni/diffusion/registry.py

@@ -24,6 +24,11 @@ _DIFFUSION_MODELS = {
         "pipeline_z_image",
         "ZImagePipeline",
     ),
+    "WanPipeline": (
+        "wan2_2",
+        "pipeline_wan2_2",
+        "Wan22Pipeline",
+    ),
 }
 
 
@@ -61,6 +66,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "QwenImagePipeline": "get_qwen_image_post_process_func",
     "QwenImageEditPipeline": "get_qwen_image_edit_post_process_func",
     "ZImagePipeline": "get_post_process_func",
+    "WanPipeline": "get_wan22_post_process_func",
 }
 
 _DIFFUSION_PRE_PROCESS_FUNCS = {