# SPDX-License-Identifier: Apache-2.0 # Copyright 2025 The Wan Team and The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import contextvars import functools import importlib.util import os import sys from contextlib import contextmanager from types import ModuleType import torch import torch.nn as nn import torch.nn.functional as F from einops import rearrange import fastvideo.envs as envs from fastvideo.distributed.parallel_state import GroupCoordinator, get_sp_group, model_parallel_is_initialized from fastvideo.layers.activation import get_act_fn from fastvideo.logger import init_logger from fastvideo.models.vaes.common import (DiagonalGaussianDistribution, ParallelTiledVAE) from fastvideo.models.wan.vae_config import WanVAEConfig logger = init_logger(__name__) CACHE_T = 2 # fastvideo-kernel/wan22_vae_decode decodes the Wan2.2-TI2V-5B VAE's latents [1, 48, T, 44, 78] (704 x 1248) on # sm_100a (B200 / GB200) and sm_103a (B300). _DECODE_KERNEL_LATENT = (48, 44, 78) _DECODE_KERNEL_VIDEO = (704, 1248) _DECODE_KERNEL_CAPABILITIES = ((10, 0), (10, 3)) # The group sizes it splits a decode over (decode_kernel_parallel), each checked to decode the same video as one GPU. _DECODE_KERNEL_WORLDS = (2, 4, 8) @functools.cache def _wan22_vae_decode() -> ModuleType: """The wrapper of fastvideo-kernel/wan22_vae_decode (``FASTVIDEO_WAN22_VAE_DECODE_DIR``, by default this checkout's), loaded once per process as ``wan22_vae_decode``, as its own tests import it. Its extension builds on first use (into ``WAN22_VAE_DECODE_BUILD_DIR``, see the wrapper).""" module = sys.modules.get("wan22_vae_decode") if module is None: root = envs.FASTVIDEO_WAN22_VAE_DECODE_DIR.get() or os.path.join( os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))), "fastvideo-kernel", "wan22_vae_decode") path = os.path.join(root, "wan22_vae_decode.py") if not os.path.isfile(path): raise RuntimeError(f"decode_kernel needs fastvideo-kernel/wan22_vae_decode; {path} does not exist " "(FASTVIDEO_WAN22_VAE_DECODE_DIR).") spec = importlib.util.spec_from_file_location("wan22_vae_decode", path) assert spec is not None and spec.loader is not None module = importlib.util.module_from_spec(spec) sys.modules["wan22_vae_decode"] = module try: spec.loader.exec_module(module) except BaseException: del sys.modules["wan22_vae_decode"] raise return module is_first_frame = contextvars.ContextVar("is_first_frame", default=False) feat_cache = contextvars.ContextVar("feat_cache", default=None) feat_idx = contextvars.ContextVar("feat_idx", default=0) first_chunk = contextvars.ContextVar("first_chunk", default=None) use_light_vae = contextvars.ContextVar("use_light_vae", default=False) @contextmanager def forward_context(first_frame_arg=False, feat_cache_arg=None, feat_idx_arg=None, first_chunk_arg=None, use_light_vae_arg=False): is_first_frame_token = is_first_frame.set(first_frame_arg) feat_cache_token = feat_cache.set(feat_cache_arg) feat_idx_token = feat_idx.set(feat_idx_arg) first_chunk_token = first_chunk.set(first_chunk_arg) use_light_vae_token = use_light_vae.set(use_light_vae_arg) try: yield finally: is_first_frame.reset(is_first_frame_token) feat_cache.reset(feat_cache_token) feat_idx.reset(feat_idx_token) first_chunk.reset(first_chunk_token) use_light_vae.reset(use_light_vae_token) class AvgDown3D(nn.Module): def __init__( self, in_channels, out_channels, factor_t, factor_s=1, ): super().__init__() self.in_channels = in_channels self.out_channels = out_channels self.factor_t = factor_t self.factor_s = factor_s self.factor = self.factor_t * self.factor_s * self.factor_s assert in_channels * self.factor % out_channels == 0 self.group_size = in_channels * self.factor // out_channels def forward(self, x: torch.Tensor) -> torch.Tensor: pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t pad = (0, 0, 0, 0, pad_t, 0) x = F.pad(x, pad) B, C, T, H, W = x.shape x = x.view( B, C, T // self.factor_t, self.factor_t, H // self.factor_s, self.factor_s, W // self.factor_s, self.factor_s, ) x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() x = x.view( B, C * self.factor, T // self.factor_t, H // self.factor_s, W // self.factor_s, ) x = x.view( B, self.out_channels, self.group_size, T // self.factor_t, H // self.factor_s, W // self.factor_s, ) x = x.mean(dim=2) return x class DupUp3D(nn.Module): def __init__( self, in_channels: int, out_channels: int, factor_t, factor_s=1, ): super().__init__() self.in_channels = in_channels self.out_channels = out_channels self.factor_t = factor_t self.factor_s = factor_s self.factor = self.factor_t * self.factor_s * self.factor_s assert out_channels * self.factor % in_channels == 0 self.repeats = out_channels * self.factor // in_channels def forward(self, x: torch.Tensor) -> torch.Tensor: x = x.repeat_interleave(self.repeats, dim=1) x = x.view( x.size(0), self.out_channels, self.factor_t, self.factor_s, self.factor_s, x.size(2), x.size(3), x.size(4), ) x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() x = x.view( x.size(0), self.out_channels, x.size(2) * self.factor_t, x.size(4) * self.factor_s, x.size(6) * self.factor_s, ) _first_chunk = first_chunk.get() if _first_chunk: x = x[:, :, self.factor_t - 1:, :, :] return x class WanCausalConv3d(nn.Conv3d): r""" A custom 3D causal convolution layer with feature caching support. This layer extends the standard Conv3D layer by ensuring causality in the time dimension and handling feature caching for efficient inference. Args: in_channels (int): Number of channels in the input image out_channels (int): Number of channels produced by the convolution kernel_size (int or tuple): Size of the convolving kernel stride (int or tuple, optional): Stride of the convolution. Default: 1 padding (int or tuple, optional): Zero-padding added to all three sides of the input. Default: 0 """ def __init__( self, in_channels: int, out_channels: int, kernel_size: int | tuple[int, int, int], stride: int | tuple[int, int, int] = 1, padding: int | tuple[int, int, int] = 0, ) -> None: super().__init__( in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, ) self.padding: tuple[int, int, int] # Set up causal padding self._padding: tuple[int, ...] = (self.padding[2], self.padding[2], self.padding[1], self.padding[1], 2 * self.padding[0], 0) self.padding = (0, 0, 0) def forward(self, x, cache_x=None): from fastvideo.platforms import current_platform padding = list(self._padding) if cache_x is not None and self._padding[4] > 0: cache_x = cache_x.to(x.device) x = torch.cat([cache_x, x], dim=2) padding[4] -= cache_x.shape[2] x = F.pad(x, padding) x = x.to( self.weight.dtype) if current_platform.is_mps() else x # casting needed for mps since amp isn't supported return super().forward(x) class WanRMS_norm(nn.Module): r""" A custom RMS normalization layer. Args: dim (int): The number of dimensions to normalize over. channel_first (bool, optional): Whether the input tensor has channels as the first dimension. Default is True. images (bool, optional): Whether the input represents image data. Default is True. bias (bool, optional): Whether to include a learnable bias term. Default is False. """ def __init__(self, dim: int, channel_first: bool = True, images: bool = True, bias: bool = False) -> None: super().__init__() broadcastable_dims = (1, 1, 1) if not images else (1, 1) shape = (dim, *broadcastable_dims) if channel_first else (dim, ) self.channel_first = channel_first self.scale = dim**0.5 self.gamma = nn.Parameter(torch.ones(shape)) self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0 def forward(self, x): return F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias class WanUpsample(nn.Upsample): r""" Perform upsampling while ensuring the output tensor has the same data type as the input. Args: x (torch.Tensor): Input tensor to be upsampled. Returns: torch.Tensor: Upsampled tensor with the same data type as the input. """ def forward(self, x): return super().forward(x.float()).type_as(x) class WanResample(nn.Module): r""" A custom resampling module for 2D and 3D data. Args: dim (int): The number of input/output channels. mode (str): The resampling mode. Must be one of: - 'none': No resampling (identity operation). - 'upsample2d': 2D upsampling with nearest-exact interpolation and convolution. - 'upsample3d': 3D upsampling with nearest-exact interpolation, convolution, and causal 3D convolution. - 'downsample2d': 2D downsampling with zero-padding and convolution. - 'downsample3d': 3D downsampling with zero-padding, convolution, and causal 3D convolution. """ def __init__(self, dim: int, mode: str, upsample_out_dim: int = None) -> None: super().__init__() self.dim = dim self.mode = mode # default to dim //2 if upsample_out_dim is None: upsample_out_dim = dim // 2 # layers if mode == "upsample2d": self.resample = nn.Sequential( WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), nn.Conv2d(dim, upsample_out_dim, 3, padding=1), ) elif mode == "upsample3d": self.resample = nn.Sequential( WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), nn.Conv2d(dim, upsample_out_dim, 3, padding=1), ) self.time_conv = WanCausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) elif mode == "downsample2d": self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) elif mode == "downsample3d": self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) self.time_conv = WanCausalConv3d(dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0)) else: self.resample = nn.Identity() def forward(self, x): b, c, t, h, w = x.size() first_frame = is_first_frame.get() _first_chunk = first_chunk.get() _use_light_vae = use_light_vae.get() if first_frame: assert t == 1 _feat_cache = feat_cache.get() _feat_idx = feat_idx.get() if self.mode == "upsample3d": if _feat_cache is not None: idx = _feat_idx if _use_light_vae: if _feat_cache[idx] is None: _feat_cache[idx] = "Rep" cache_x = x[:, :, -CACHE_T:, :, :].clone() # "Rep" sentinel on first chunk, cache tensor after; isinstance avoids `tensor == str` version surprises. if isinstance(_feat_cache[idx], str): x = self.time_conv(x) else: if cache_x.shape[2] < 2: cache_x = torch.cat([ _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(x.device), cache_x, ], dim=2) x = self.time_conv(x, _feat_cache[idx]) _feat_cache[idx] = cache_x _feat_idx += 1 x = x.reshape(b, 2, c, t, h, w) x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3) x = x.reshape(b, c, t * 2, h, w) if _first_chunk: x = x[:, :, 1:, :, :] else: if _feat_cache[idx] is None: _feat_cache[idx] = "Rep" _feat_idx += 1 else: cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and _feat_cache[idx] is not None and _feat_cache[idx] != "Rep": cache_x = torch.cat( [_feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) if cache_x.shape[2] < 2 and _feat_cache[idx] is not None and _feat_cache[idx] == "Rep": cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2) x = self.time_conv(x) if _feat_cache[idx] == "Rep" else self.time_conv(x, _feat_cache[idx]) _feat_cache[idx] = cache_x _feat_idx += 1 x = x.reshape(b, 2, c, t, h, w) x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3) x = x.reshape(b, c, t * 2, h, w) feat_cache.set(_feat_cache) feat_idx.set(_feat_idx) elif not first_frame and hasattr(self, "time_conv"): x = self.time_conv(x) x = x.reshape(b, 2, c, t, h, w) x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3) x = x.reshape(b, c, t * 2, h, w) t = x.shape[2] x = x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w) x = self.resample(x) x = x.view(b, t, x.size(1), x.size(2), x.size(3)).permute(0, 2, 1, 3, 4) _feat_cache = feat_cache.get() _feat_idx = feat_idx.get() if self.mode == "downsample3d": if _feat_cache is not None: idx = _feat_idx if _feat_cache[idx] is None: _feat_cache[idx] = x.clone() _feat_idx += 1 else: cache_x = x[:, :, -1:, :, :].clone() x = self.time_conv(torch.cat([_feat_cache[idx][:, :, -1:, :, :], x], 2)) _feat_cache[idx] = cache_x _feat_idx += 1 feat_cache.set(_feat_cache) feat_idx.set(_feat_idx) elif not first_frame and hasattr(self, "time_conv"): x = self.time_conv(x) return x class WanResidualBlock(nn.Module): r""" A custom residual block module. Args: in_dim (int): Number of input channels. out_dim (int): Number of output channels. dropout (float, optional): Dropout rate for the dropout layer. Default is 0.0. non_linearity (str, optional): Type of non-linearity to use. Default is "silu". """ def __init__( self, in_dim: int, out_dim: int, dropout: float = 0.0, non_linearity: str = "silu", ) -> None: super().__init__() self.in_dim = in_dim self.out_dim = out_dim self.nonlinearity = get_act_fn(non_linearity) # layers self.norm1 = WanRMS_norm(in_dim, images=False) self.conv1 = WanCausalConv3d(in_dim, out_dim, 3, padding=1) self.norm2 = WanRMS_norm(out_dim, images=False) self.dropout = nn.Dropout(dropout) self.conv2 = WanCausalConv3d(out_dim, out_dim, 3, padding=1) self.conv_shortcut = WanCausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity() def forward(self, x): # Apply shortcut connection h = self.conv_shortcut(x) # First normalization and activation x = self.norm1(x) x = self.nonlinearity(x) _feat_cache = feat_cache.get() _feat_idx = feat_idx.get() if _feat_cache is not None: idx = _feat_idx cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and _feat_cache[idx] is not None: cache_x = torch.cat([_feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) x = self.conv1(x, _feat_cache[idx]) _feat_cache[idx] = cache_x _feat_idx += 1 feat_cache.set(_feat_cache) feat_idx.set(_feat_idx) else: x = self.conv1(x) # Second normalization and activation x = self.norm2(x) x = self.nonlinearity(x) # Dropout x = self.dropout(x) _feat_cache = feat_cache.get() _feat_idx = feat_idx.get() if _feat_cache is not None: idx = _feat_idx cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and _feat_cache[idx] is not None: cache_x = torch.cat([_feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) x = self.conv2(x, _feat_cache[idx]) _feat_cache[idx] = cache_x _feat_idx += 1 feat_cache.set(_feat_cache) feat_idx.set(_feat_idx) else: x = self.conv2(x) # Add residual connection return x + h class WanAttentionBlock(nn.Module): r""" Causal self-attention with a single head. Args: dim (int): The number of channels in the input tensor. """ def __init__(self, dim) -> None: super().__init__() self.dim = dim # layers self.norm = WanRMS_norm(dim) self.to_qkv = nn.Conv2d(dim, dim * 3, 1) self.proj = nn.Conv2d(dim, dim, 1) def forward(self, x): identity = x batch_size, channels, time, height, width = x.size() x = x.permute(0, 2, 1, 3, 4).reshape(batch_size * time, channels, height, width) x = self.norm(x) # compute query, key, value qkv = self.to_qkv(x) qkv = qkv.reshape(batch_size * time, 1, channels * 3, -1) qkv = qkv.permute(0, 1, 3, 2).contiguous() q, k, v = qkv.chunk(3, dim=-1) # apply attention x = F.scaled_dot_product_attention(q, k, v) x = x.squeeze(1).permute(0, 2, 1).reshape(batch_size * time, channels, height, width) # output projection x = self.proj(x) # Reshape back: [(b*t), c, h, w] -> [b, c, t, h, w] x = x.view(batch_size, time, channels, height, width) x = x.permute(0, 2, 1, 3, 4) return x + identity class WanMidBlock(nn.Module): """ Middle block for WanVAE encoder and decoder. Args: dim (int): Number of input/output channels. dropout (float): Dropout rate. non_linearity (str): Type of non-linearity to use. """ def __init__(self, dim: int, dropout: float = 0.0, non_linearity: str = "silu", num_layers: int = 1): super().__init__() self.dim = dim # Create the components resnets = [WanResidualBlock(dim, dim, dropout, non_linearity)] attentions = [] for _ in range(num_layers): attentions.append(WanAttentionBlock(dim)) resnets.append(WanResidualBlock(dim, dim, dropout, non_linearity)) self.attentions = nn.ModuleList(attentions) self.resnets = nn.ModuleList(resnets) self.gradient_checkpointing = False def forward(self, x): # First residual block x = self.resnets[0](x) # Process through attention and residual blocks for attn, resnet in zip(self.attentions, self.resnets[1:], strict=True): if attn is not None: x = attn(x) x = resnet(x) return x class WanResidualDownBlock(nn.Module): def __init__(self, in_dim, out_dim, dropout, num_res_blocks, temperal_downsample=False, down_flag=False): super().__init__() # Shortcut path with downsample self.avg_shortcut = AvgDown3D( in_dim, out_dim, factor_t=2 if temperal_downsample else 1, factor_s=2 if down_flag else 1, ) # Main path with residual blocks and downsample resnets = [] for _ in range(num_res_blocks): resnets.append(WanResidualBlock(in_dim, out_dim, dropout)) in_dim = out_dim self.resnets = nn.ModuleList(resnets) # Add the final downsample block if down_flag: mode = "downsample3d" if temperal_downsample else "downsample2d" self.downsampler = WanResample(out_dim, mode=mode) else: self.downsampler = None def forward(self, x): x_copy = x.clone() for resnet in self.resnets: x = resnet(x) if self.downsampler is not None: x = self.downsampler(x) return x + self.avg_shortcut(x_copy) class WanEncoder3d(nn.Module): r""" A 3D encoder module. Args: dim (int): The base number of channels in the first layer. z_dim (int): The dimensionality of the latent space. dim_mult (list of int): Multipliers for the number of channels in each block. num_res_blocks (int): Number of residual blocks in each block. attn_scales (list of float): Scales at which to apply attention mechanisms. temperal_downsample (list of bool): Whether to downsample temporally in each block. dropout (float): Dropout rate for the dropout layers. non_linearity (str): Type of non-linearity to use. """ def __init__( self, in_channels: int = 3, dim=128, z_dim=4, dim_mult=(1, 2, 4, 4), num_res_blocks=2, attn_scales=(), temperal_downsample=(True, True, False), dropout=0.0, non_linearity: str = "silu", is_residual: bool = False, # wan 2.2 vae use a residual downblock ): super().__init__() self.dim = dim self.z_dim = z_dim dim_mult = list(dim_mult) self.dim_mult = dim_mult self.num_res_blocks = num_res_blocks self.attn_scales = list(attn_scales) self.temperal_downsample = list(temperal_downsample) self.nonlinearity = get_act_fn(non_linearity) # dimensions dims = [dim * u for u in [1] + dim_mult] scale = 1.0 # init block self.conv_in = WanCausalConv3d(in_channels, dims[0], 3, padding=1) # downsample blocks self.down_blocks = nn.ModuleList([]) for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:], strict=True)): # residual (+attention) blocks if is_residual: self.down_blocks.append( WanResidualDownBlock( in_dim, out_dim, dropout, num_res_blocks, temperal_downsample=temperal_downsample[i] if i != len(dim_mult) - 1 else False, down_flag=i != len(dim_mult) - 1, )) else: for _ in range(num_res_blocks): self.down_blocks.append(WanResidualBlock(in_dim, out_dim, dropout)) if scale in attn_scales: self.down_blocks.append(WanAttentionBlock(out_dim)) in_dim = out_dim # downsample block if i != len(dim_mult) - 1: mode = "downsample3d" if temperal_downsample[i] else "downsample2d" self.down_blocks.append(WanResample(out_dim, mode=mode)) scale /= 2.0 # middle blocks self.mid_block = WanMidBlock(out_dim, dropout, non_linearity, num_layers=1) # output blocks self.norm_out = WanRMS_norm(out_dim, images=False) self.conv_out = WanCausalConv3d(out_dim, z_dim, 3, padding=1) self.gradient_checkpointing = False def forward(self, x): _feat_cache = feat_cache.get() _feat_idx = feat_idx.get() if _feat_cache is not None: idx = _feat_idx cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and _feat_cache[idx] is not None: # cache last frame of last two chunk cache_x = torch.cat([_feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) x = self.conv_in(x, _feat_cache[idx]) _feat_cache[idx] = cache_x _feat_idx += 1 feat_cache.set(_feat_cache) feat_idx.set(_feat_idx) else: x = self.conv_in(x) ## downsamples for layer in self.down_blocks: x = layer(x) ## middle x = self.mid_block(x) ## head x = self.norm_out(x) x = self.nonlinearity(x) _feat_cache = feat_cache.get() _feat_idx = feat_idx.get() if _feat_cache is not None: idx = _feat_idx cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and _feat_cache[idx] is not None: # cache last frame of last two chunk cache_x = torch.cat([_feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) x = self.conv_out(x, _feat_cache[idx]) _feat_cache[idx] = cache_x _feat_idx += 1 feat_cache.set(_feat_cache) feat_idx.set(_feat_idx) else: x = self.conv_out(x) return x # adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/autoencoders/autoencoder_kl_wan.py class WanResidualUpBlock(nn.Module): """ A block that handles upsampling for the WanVAE decoder. Args: in_dim (int): Input dimension out_dim (int): Output dimension num_res_blocks (int): Number of residual blocks dropout (float): Dropout rate temperal_upsample (bool): Whether to upsample on temporal dimension up_flag (bool): Whether to upsample or not non_linearity (str): Type of non-linearity to use """ def __init__( self, in_dim: int, out_dim: int, num_res_blocks: int, dropout: float = 0.0, temperal_upsample: bool = False, up_flag: bool = False, non_linearity: str = "silu", ): super().__init__() self.in_dim = in_dim self.out_dim = out_dim if up_flag: self.avg_shortcut = DupUp3D( in_dim, out_dim, factor_t=2 if temperal_upsample else 1, factor_s=2, ) else: self.avg_shortcut = None # create residual blocks resnets = [] current_dim = in_dim for _ in range(num_res_blocks + 1): resnets.append(WanResidualBlock(current_dim, out_dim, dropout, non_linearity)) current_dim = out_dim self.resnets = nn.ModuleList(resnets) # Add upsampling layer if needed if up_flag: upsample_mode = "upsample3d" if temperal_upsample else "upsample2d" self.upsampler = WanResample(out_dim, mode=upsample_mode, upsample_out_dim=out_dim) else: self.upsampler = None self.gradient_checkpointing = False def forward(self, x): """ Forward pass through the upsampling block. Args: x (torch.Tensor): Input tensor feat_cache (list, optional): Feature cache for causal convolutions feat_idx (list, optional): Feature index for cache management Returns: torch.Tensor: Output tensor """ if self.avg_shortcut is not None: x_copy = x.clone() for resnet in self.resnets: x = resnet(x) if self.upsampler is not None: x = self.upsampler(x) if self.avg_shortcut is not None: x = x + self.avg_shortcut(x_copy) return x class WanUpBlock(nn.Module): """ A block that handles upsampling for the WanVAE decoder. Args: in_dim (int): Input dimension out_dim (int): Output dimension num_res_blocks (int): Number of residual blocks dropout (float): Dropout rate upsample_mode (str, optional): Mode for upsampling ('upsample2d' or 'upsample3d') non_linearity (str): Type of non-linearity to use """ def __init__( self, in_dim: int, out_dim: int, num_res_blocks: int, dropout: float = 0.0, upsample_mode: str | None = None, non_linearity: str = "silu", ): super().__init__() self.in_dim = in_dim self.out_dim = out_dim # Create layers list resnets = [] # Add residual blocks and attention if needed current_dim = in_dim for _ in range(num_res_blocks + 1): resnets.append(WanResidualBlock(current_dim, out_dim, dropout, non_linearity)) current_dim = out_dim self.resnets = nn.ModuleList(resnets) # Add upsampling layer if needed self.upsamplers = None if upsample_mode is not None: self.upsamplers = nn.ModuleList([WanResample(out_dim, mode=upsample_mode)]) self.gradient_checkpointing = False def forward(self, x): """ Forward pass through the upsampling block. Args: x (torch.Tensor): Input tensor feat_cache (list, optional): Feature cache for causal convolutions feat_idx (list, optional): Feature index for cache management Returns: torch.Tensor: Output tensor """ for resnet in self.resnets: x = resnet(x) if self.upsamplers is not None: x = self.upsamplers[0](x) return x class WanDecoder3d(nn.Module): r""" A 3D decoder module. Args: dim (int): The base number of channels in the first layer. z_dim (int): The dimensionality of the latent space. dim_mult (list of int): Multipliers for the number of channels in each block. num_res_blocks (int): Number of residual blocks in each block. attn_scales (list of float): Scales at which to apply attention mechanisms. temperal_upsample (list of bool): Whether to upsample temporally in each block. dropout (float): Dropout rate for the dropout layers. non_linearity (str): Type of non-linearity to use. """ def __init__( self, dim=128, z_dim=4, dim_mult=(1, 2, 4, 4), num_res_blocks=2, attn_scales=(), temperal_upsample=(False, True, True), dropout=0.0, non_linearity: str = "silu", out_channels: int = 3, is_residual: bool = False, ): super().__init__() self.dim = dim self.z_dim = z_dim dim_mult = list(dim_mult) self.dim_mult = dim_mult self.num_res_blocks = num_res_blocks self.attn_scales = list(attn_scales) self.temperal_upsample = list(temperal_upsample) self.nonlinearity = get_act_fn(non_linearity) # dimensions dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] # init block self.conv_in = WanCausalConv3d(z_dim, dims[0], 3, padding=1) # middle blocks self.mid_block = WanMidBlock(dims[0], dropout, non_linearity, num_layers=1) # upsample blocks self.up_blocks = nn.ModuleList([]) for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:], strict=True)): # residual (+attention) blocks if i > 0 and not is_residual: # wan vae 2.1 in_dim = in_dim // 2 # determine if we need upsampling up_flag = i != len(dim_mult) - 1 # determine upsampling mode, if not upsampling, set to None upsample_mode = None if up_flag and temperal_upsample[i]: upsample_mode = "upsample3d" elif up_flag: upsample_mode = "upsample2d" # Create and add the upsampling block if is_residual: up_block = WanResidualUpBlock( in_dim=in_dim, out_dim=out_dim, num_res_blocks=num_res_blocks, dropout=dropout, temperal_upsample=temperal_upsample[i] if up_flag else False, up_flag=up_flag, non_linearity=non_linearity, ) else: up_block = WanUpBlock( in_dim=in_dim, out_dim=out_dim, num_res_blocks=num_res_blocks, dropout=dropout, upsample_mode=upsample_mode, non_linearity=non_linearity, ) self.up_blocks.append(up_block) # output blocks self.norm_out = WanRMS_norm(out_dim, images=False) self.conv_out = WanCausalConv3d(out_dim, out_channels, 3, padding=1) self.gradient_checkpointing = False def forward(self, x): ## conv1 _feat_cache = feat_cache.get() _feat_idx = feat_idx.get() if _feat_cache is not None: idx = _feat_idx cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and _feat_cache[idx] is not None: # cache last frame of last two chunk cache_x = torch.cat([_feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) x = self.conv_in(x, _feat_cache[idx]) _feat_cache[idx] = cache_x _feat_idx += 1 feat_cache.set(_feat_cache) feat_idx.set(_feat_idx) else: x = self.conv_in(x) ## middle x = self.mid_block(x) ## upsamples for up_block in self.up_blocks: x = up_block(x) ## head x = self.norm_out(x) x = self.nonlinearity(x) _feat_cache = feat_cache.get() _feat_idx = feat_idx.get() if _feat_cache is not None: idx = _feat_idx cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and _feat_cache[idx] is not None: # cache last frame of last two chunk cache_x = torch.cat([_feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) x = self.conv_out(x, _feat_cache[idx]) _feat_cache[idx] = cache_x _feat_idx += 1 feat_cache.set(_feat_cache) feat_idx.set(_feat_idx) else: x = self.conv_out(x) return x def patchify(x, patch_size): if patch_size == 1: return x if x.dim() == 4: x = rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size) elif x.dim() == 5: x = rearrange( x, "b c f (h q) (w r) -> b (c r q) f h w", q=patch_size, r=patch_size, ) else: raise ValueError(f"Invalid input shape: {x.shape}") return x def unpatchify(x, patch_size): if patch_size == 1: return x if x.dim() == 4: x = rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size) elif x.dim() == 5: x = rearrange( x, "b (c r q) f h w -> b c f (h q) (w r)", q=patch_size, r=patch_size, ) return x def _is_wan_vae_codec(name: str, submodule: nn.Module) -> bool: """Match the Wan VAE's encoder/decoder for in-place torch.compile. Compiling the codec submodules' ``forward`` in place (rather than replacing the whole VAE module) ensures the compiled version is the one held by the already-constructed DecodingStage. Without this, ``enable_torch_compile_vae`` falls back to replacing ``self.modules["vae"]`` and is a no-op for Wan. """ return (name == "encoder" and isinstance(submodule, WanEncoder3d)) or (name == "decoder" and isinstance(submodule, WanDecoder3d)) class AutoencoderKLWan(nn.Module, ParallelTiledVAE): r""" A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. Introduced in [Wan 2.1]. """ _supports_gradient_checkpointing = False _compile_conditions = [_is_wan_vae_codec] def __init__( self, config: WanVAEConfig, ) -> None: nn.Module.__init__(self) ParallelTiledVAE.__init__(self, config) self.z_dim = config.z_dim self.temperal_downsample = list(config.temperal_downsample) self.temperal_upsample = list(config.temperal_downsample)[::-1] decoder_base_dim = config.base_dim if config.decoder_base_dim is None else config.decoder_base_dim self.latents_mean = list(config.latents_mean) self.latents_std = list(config.latents_std) self.shift_factor = config.shift_factor if config.load_encoder: self.encoder = WanEncoder3d( in_channels=config.in_channels, dim=config.base_dim, z_dim=self.z_dim * 2, dim_mult=config.dim_mult, num_res_blocks=config.num_res_blocks, attn_scales=config.attn_scales, temperal_downsample=self.temperal_downsample, dropout=config.dropout, is_residual=config.is_residual, ) self.quant_conv = WanCausalConv3d(self.z_dim * 2, self.z_dim * 2, 1) self.post_quant_conv = WanCausalConv3d(self.z_dim, self.z_dim, 1) if config.load_decoder: self.decoder = WanDecoder3d( dim=decoder_base_dim, z_dim=self.z_dim, dim_mult=config.dim_mult, num_res_blocks=config.num_res_blocks, attn_scales=config.attn_scales, temperal_upsample=self.temperal_upsample, dropout=config.dropout, out_channels=config.out_channels, is_residual=config.is_residual, ) self.use_feature_cache = config.use_feature_cache # The kernel decoders of decode_kernel, one per device and process group (their weights prepared once; None # where the group's ranks could not all build one), and why the kernel cannot run. self._kernel_decoders: dict[tuple[torch.device, tuple[int, ...] | None], object | None] = {} self._kernel_unavailable: str | None = None def clear_cache(self) -> None: def _count_conv3d(model) -> int: count = 0 for m in model.modules(): if isinstance(m, WanCausalConv3d): count += 1 return count if self.config.load_decoder: self._conv_num = _count_conv3d(self.decoder) self._conv_idx = 0 self._feat_map = [None] * self._conv_num # cache encode if self.config.load_encoder: self._enc_conv_num = _count_conv3d(self.encoder) self._enc_conv_idx = 0 self._enc_feat_map = [None] * self._enc_conv_num def optimize_memory_format(self) -> None: """Optional CUDA-only weight layout optimization used by MG3 LightVAE.""" def _convert(module: nn.Module) -> None: for child in module.children(): if isinstance(child, nn.Conv3d): child.weight.data = child.weight.data.to(memory_format=torch.channels_last_3d) else: _convert(child) _convert(self) def encode(self, x: torch.Tensor) -> torch.Tensor: if self.use_feature_cache: self.clear_cache() if self.config.patch_size is not None: x = patchify(x, patch_size=self.config.patch_size) with forward_context(feat_cache_arg=self._enc_feat_map, feat_idx_arg=self._enc_conv_idx, use_light_vae_arg=self.config.use_light_vae): t = x.shape[2] iter_ = 1 + (t - 1) // 4 for i in range(iter_): feat_idx.set(0) if i == 0: out = self.encoder(x[:, :, :1, :, :]) else: out_ = self.encoder(x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :]) out = torch.cat([out, out_], 2) enc = self.quant_conv(out) mu, logvar = enc[:, :self.z_dim, :, :, :], enc[:, self.z_dim:, :, :, :] enc = torch.cat([mu, logvar], dim=1) enc = DiagonalGaussianDistribution(enc) self.clear_cache() else: for block in self.encoder.down_blocks: if isinstance(block, WanResample) and block.mode == "downsample3d": _padding = list(block.time_conv._padding) _padding[4] = 2 block.time_conv._padding = tuple(_padding) enc = ParallelTiledVAE.encode(self, x) return enc def _encode(self, x: torch.Tensor, first_frame=False) -> torch.Tensor: with forward_context(first_frame_arg=first_frame, use_light_vae_arg=self.config.use_light_vae): out = self.encoder(x) enc = self.quant_conv(out) mu, logvar = enc[:, :self.z_dim, :, :, :], enc[:, self.z_dim:, :, :, :] enc = torch.cat([mu, logvar], dim=1) return enc def tiled_encode(self, x: torch.Tensor) -> torch.Tensor: first_frame = x[:, :, 0, :, :].unsqueeze(2) first_frame = self._encode(first_frame, first_frame=True) enc = ParallelTiledVAE.tiled_encode(self, x) enc = enc[:, :, 1:] enc = torch.cat([first_frame, enc], dim=2) return enc def spatial_tiled_encode(self, x: torch.Tensor) -> torch.Tensor: first_frame = x[:, :, 0, :, :].unsqueeze(2) first_frame = self._encode(first_frame, first_frame=True) enc = ParallelTiledVAE.spatial_tiled_encode(self, x) enc = enc[:, :, 1:] enc = torch.cat([first_frame, enc], dim=2) return enc def decode(self, z: torch.Tensor) -> torch.Tensor: if getattr(self.config, "decode_kernel", False): video = self._kernel_decode(z) if video is not None: return video if self.use_feature_cache: self.clear_cache() iter_ = z.shape[2] x = self.post_quant_conv(z) with forward_context(feat_cache_arg=self._feat_map, feat_idx_arg=self._conv_idx, use_light_vae_arg=self.config.use_light_vae): for i in range(iter_): feat_idx.set(0) if i == 0: first_chunk.set(True) out = self.decoder(x[:, :, i:i + 1, :, :]) else: first_chunk.set(False) out_ = self.decoder(x[:, :, i:i + 1, :, :]) out = torch.cat([out, out_], 2) if self.config.patch_size is not None: out = unpatchify(out, patch_size=self.config.patch_size) out = out.float() out = torch.clamp(out, min=-1.0, max=1.0) self.clear_cache() else: out = ParallelTiledVAE.decode(self, z) return out def _kernel_decode(self, z: torch.Tensor) -> torch.Tensor | None: """``decode`` by fastvideo-kernel's Wan2.2 decoder (``decode_kernel``): bf16 with fp32 accumulation, within the error of a bf16 decode of the fp32 one, or None where it does not apply (other devices, geometries, batches or VAEs), and the decode below runs as before. Every rank of a multi-GPU run returns the whole video, as the decode below does: decoded on its GPU alone, or with ``decode_kernel_parallel`` split over the sequence-parallel group and broadcast from the group's first rank.""" # Only what is the same on every rank decides here, so that the ranks of a group take the split together; a # rank whose own build failed still joins the group's agreement (see _kernel_decoder). reason = None if not z.is_cuda: reason = "the latents are not on a GPU" elif torch.cuda.get_device_capability(z.device) not in _DECODE_KERNEL_CAPABILITIES: major, minor = torch.cuda.get_device_capability(z.device) reason = f"the GPU is sm_{major}{minor}, not sm_100 or sm_103" elif z.dim() != 5 or z.shape[0] != 1 or (z.shape[1], *z.shape[3:]) != _DECODE_KERNEL_LATENT: reason = f"latents {tuple(z.shape)} are not [1, 48, T, 44, 78] (a 704 x 1248 video, batch 1)" if reason is not None: logger.warning_once(f"decode_kernel: {reason}; decoding with AutoencoderKLWan.") return None group = self._kernel_group(z) if group is not None and (decoder := self._kernel_decoder(z.device, group)) is not None: logger.info_once(f"decode_kernel: decoding with fastvideo-kernel's Wan2.2 decoder over the " f"{group.world_size} GPUs of the sequence-parallel group.") video = decoder.decode(z) # the video on the group's first rank, None on the others if video is None: video = torch.empty((1, 3, 4 * z.shape[2] - 3, *_DECODE_KERNEL_VIDEO), dtype=torch.bfloat16, device=z.device) return group.broadcast(video, src=0).float() decoder = self._kernel_decoder(z.device, None) if decoder is None: logger.warning_once(f"decode_kernel: {self._kernel_unavailable}; decoding with AutoencoderKLWan.") return None logger.info_once("decode_kernel: decoding with fastvideo-kernel's Wan2.2 decoder on this GPU.") return decoder.decode(z).float() def _kernel_group(self, z: torch.Tensor) -> GroupCoordinator | None: """The sequence-parallel group ``decode_kernel_parallel`` splits the decode of ``z`` over, or None (decode alone): the switch is off, there is no group of a size the kernel splits over, or the video has fewer latent frames than the group has GPUs, where the split's fixed cost (about 50 ms over 4 B300s, 100 ms over 8) exceeds the decode.""" if not getattr(self.config, "decode_kernel_parallel", False) or not model_parallel_is_initialized(): return None group = get_sp_group() return group if group.world_size in _DECODE_KERNEL_WORLDS and z.shape[2] >= group.world_size else None def _kernel_decoder(self, device: torch.device, group: GroupCoordinator | None) -> object | None: """The kernel's decoder on ``device``, alone (``group`` None) or as this rank of ``group``, built on first use (its extension built or loaded, its weights prepared), or None where it cannot be. Over a group every rank builds its own without communicating, then the ranks agree: the split decode needs every one of them, so a rank that could not build makes the whole group decode alone.""" key = (device, None if group is None else tuple(group.ranks)) if key in self._kernel_decoders: return self._kernel_decoders[key] decoder = None if self._kernel_unavailable is None: try: decoder = _wan22_vae_decode().WanVAEDecoder(self.state_dict(), device=device, pg=None if group is None else group.device_group, local=group is None) except (KeyError, ValueError) as error: # not the Wan2.2-TI2V-5B VAE's decoder self._kernel_unavailable = f"this VAE's decoder does not fit the kernel ({error})" except Exception as error: # the kernel failed to load or build; generation goes on without it self._kernel_unavailable = f"the kernel did not load ({type(error).__name__}: {error})" if group is not None: ready = torch.tensor([int(decoder is not None)], dtype=torch.int32) torch.distributed.all_reduce(ready, op=torch.distributed.ReduceOp.MIN, group=group.cpu_group) if not ready.item(): if decoder is not None: logger.warning("decode_kernel_parallel: a rank of the sequence-parallel group could not build " "the kernel's decoder; every rank decodes alone.") decoder = None self._kernel_decoders[key] = decoder return decoder def _decode(self, z: torch.Tensor, first_frame=False) -> torch.Tensor: x = self.post_quant_conv(z) with forward_context(first_frame_arg=first_frame, use_light_vae_arg=self.config.use_light_vae): out = self.decoder(x) out = torch.clamp(out, min=-1.0, max=1.0) return out def tiled_decode(self, z: torch.Tensor) -> torch.Tensor: self.blend_num_frames *= 2 dec = ParallelTiledVAE.tiled_decode(self, z) start_frame_idx = self.temporal_compression_ratio - 1 dec = dec[:, :, start_frame_idx:] return dec def spatial_tiled_decode(self, z: torch.Tensor) -> torch.Tensor: dec = ParallelTiledVAE.spatial_tiled_decode(self, z) start_frame_idx = self.temporal_compression_ratio - 1 dec = dec[:, :, start_frame_idx:] return dec def parallel_tiled_decode(self, z: torch.FloatTensor) -> torch.FloatTensor: self.blend_num_frames *= 2 dec = ParallelTiledVAE.parallel_tiled_decode(self, z) start_frame_idx = self.temporal_compression_ratio - 1 dec = dec[:, :, start_frame_idx:] return dec def get_streaming_cache(self) -> list[torch.Tensor | None]: def _count_conv3d(model) -> int: count = 0 for m in model.modules(): if isinstance(m, WanCausalConv3d): count += 1 return count conv_num = _count_conv3d(self.decoder) return [None] * conv_num def streaming_decode( self, z: torch.Tensor, cache: list[torch.Tensor | None], is_first_chunk: bool = False, ) -> tuple[torch.Tensor, list[torch.Tensor | None]]: """ Args: z (`torch.Tensor`): Latent tensor of shape [B, C, T, H, W]. cache (`list[torch.Tensor | None]`): The VAE cache. is_first_chunk (`bool`): Whether this is the first chunk in the sequence. Returns: A tuple of (decoded_frames, updated_cache). """ iter_ = z.shape[2] x = self.post_quant_conv(z) with forward_context(feat_cache_arg=cache, feat_idx_arg=0, use_light_vae_arg=self.config.use_light_vae): outputs = [] for i in range(iter_): feat_idx.set(0) first_chunk.set(is_first_chunk and i == 0) decoded_chunk = self.decoder(x[:, :, i:i + 1, :, :]) outputs.append(decoded_chunk) out = torch.cat(outputs, dim=2) if self.config.patch_size is not None: out = unpatchify(out, patch_size=self.config.patch_size) out = out.float() out = torch.clamp(out, min=-1.0, max=1.0) return out, cache def forward( self, sample: torch.Tensor, sample_posterior: bool = False, generator: torch.Generator | None = None, ) -> torch.Tensor: """ Args: sample (`torch.Tensor`): Input sample. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`DecoderOutput`] instead of a plain tuple. """ x = sample posterior = self.encode(x).latent_dist z = posterior.sample(generator=generator) if sample_posterior else posterior.mode() dec = self.decode(z) return dec # Entry point for model registry EntryClass = AutoencoderKLWan