diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 9e572ea6b..c8a0e622c 100644 --- a/docs/models/supported_models.md +++ b/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 diff --git a/examples/offline_inference/wan22/text_to_video.py b/examples/offline_inference/wan22/text_to_video.py new file mode 100644 index 000000000..6a93f4b83 --- /dev/null +++ b/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() diff --git a/tests/single_stage/test_video_diffusion_model.py b/tests/single_stage/test_video_diffusion_model.py new file mode 100644 index 000000000..b7d5f7229 --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index d3af974d2..5c8c0bb21 100644 --- a/vllm_omni/diffusion/data.py +++ b/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" diff --git a/vllm_omni/diffusion/models/wan2_2/__init__.py b/vllm_omni/diffusion/models/wan2_2/__init__.py new file mode 100644 index 000000000..9978f9e2a --- /dev/null +++ b/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"] diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py new file mode 100644 index 000000000..04483f58e --- /dev/null +++ b/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.") diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py new file mode 100644 index 000000000..d65c6839e --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index d3e93c8b9..f47333209 100644 --- a/vllm_omni/diffusion/registry.py +++ b/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 = {