Spaces:
Running on Zero
Running on Zero
| import logging | |
| import math | |
| import os | |
| import torch | |
| import torch.nn as nn | |
| from diffusers.configuration_utils import ConfigMixin, register_to_config | |
| from diffusers.models.modeling_utils import ModelMixin | |
| from torch.nn.attention.flex_attention import flex_attention | |
| from src.utils.causal_mask import CausalMaskMixin, visualize_causal_mask | |
| from src.utils.causal_rope import causal_rope_apply | |
| from src.utils.distributed import is_main_process | |
| from src.wan.modules.model import WAN_CROSSATTENTION_CLASSES | |
| from src.wan22.modules.model import ( | |
| WanLayerNorm, | |
| WanRMSNorm, | |
| rope_apply, | |
| rope_params, | |
| sinusoidal_embedding_1d, | |
| ) | |
| from .attention import attention | |
| flex_attention = torch.compile(flex_attention, dynamic=False, mode="default") | |
| class CausalWanSelfAttention(nn.Module): | |
| def __init__( | |
| self, | |
| dim: int, | |
| num_heads: int, | |
| local_attn_size: int = -1, | |
| train_frame_size: int = -1, | |
| sink_size: int = 0, | |
| qk_norm: bool = True, | |
| eps: float = 1e-06, | |
| ): | |
| assert dim % num_heads == 0 | |
| super().__init__() | |
| self.dim = dim | |
| self.num_heads = num_heads | |
| self.head_dim = dim // num_heads | |
| self.local_attn_size = local_attn_size | |
| self.sink_size = sink_size | |
| self.qk_norm = qk_norm | |
| self.eps = eps | |
| self.train_frame_size = train_frame_size | |
| self.inference_frame_size = local_attn_size if local_attn_size != -1 else train_frame_size | |
| self.q = nn.Linear(dim, dim) | |
| self.k = nn.Linear(dim, dim) | |
| self.v = nn.Linear(dim, dim) | |
| self.o = nn.Linear(dim, dim) | |
| self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() | |
| self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() | |
| def forward( | |
| self, | |
| x, | |
| seq_lens, | |
| grid_sizes, | |
| freqs, | |
| block_mask, | |
| kv_cache=None, | |
| current_start=0, | |
| cache_start=None, | |
| ): | |
| """ | |
| Args: | |
| x(Tensor): Shape [B, L, num_heads, C / num_heads] | |
| seq_lens(Tensor): Shape [B] | |
| grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) | |
| freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] | |
| block_mask (BlockMask) | |
| """ | |
| b, s, n, d = (*x.shape[:2], self.num_heads, self.head_dim) | |
| if cache_start is None: | |
| cache_start = current_start | |
| def qkv_fn(x): | |
| q = self.norm_q(self.q(x)).view(b, s, n, d) | |
| k = self.norm_k(self.k(x)).view(b, s, n, d) | |
| v = self.v(x).view(b, s, n, d) | |
| return (q, k, v) | |
| q, k, v = qkv_fn(x) | |
| if kv_cache is None: | |
| is_tf = s == seq_lens[0].item() * 2 | |
| if not is_tf: | |
| raise ValueError( | |
| "Uncached attention requires a teacher-forcing clean/noisy sequence." | |
| ) | |
| q_chunk = torch.chunk(q, 2, dim=1) | |
| k_chunk = torch.chunk(k, 2, dim=1) | |
| roped_query = [] | |
| roped_key = [] | |
| for ii in range(2): | |
| rq = rope_apply(q_chunk[ii], grid_sizes, freqs).type_as(v) | |
| rk = rope_apply(k_chunk[ii], grid_sizes, freqs).type_as(v) | |
| roped_query.append(rq) | |
| roped_key.append(rk) | |
| roped_query = torch.cat(roped_query, dim=1) | |
| roped_key = torch.cat(roped_key, dim=1) | |
| padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1] | |
| padded_roped_query = torch.cat( | |
| [ | |
| roped_query, | |
| torch.zeros( | |
| [q.shape[0], padded_length, q.shape[2], q.shape[3]], | |
| device=q.device, | |
| dtype=v.dtype, | |
| ), | |
| ], | |
| dim=1, | |
| ) | |
| padded_roped_key = torch.cat( | |
| [ | |
| roped_key, | |
| torch.zeros( | |
| [k.shape[0], padded_length, k.shape[2], k.shape[3]], | |
| device=k.device, | |
| dtype=v.dtype, | |
| ), | |
| ], | |
| dim=1, | |
| ) | |
| padded_v = torch.cat( | |
| [ | |
| v, | |
| torch.zeros( | |
| [v.shape[0], padded_length, v.shape[2], v.shape[3]], | |
| device=v.device, | |
| dtype=v.dtype, | |
| ), | |
| ], | |
| dim=1, | |
| ) | |
| x = flex_attention( | |
| query=padded_roped_query.transpose(2, 1), | |
| key=padded_roped_key.transpose(2, 1), | |
| value=padded_v.transpose(2, 1), | |
| block_mask=block_mask, | |
| ) | |
| if padded_length > 0: | |
| x = x[:, :, :-padded_length] | |
| x = x.transpose(2, 1) | |
| else: | |
| frame_seqlen = math.prod(grid_sizes[0][1:]).item() | |
| max_attention_size = self.inference_frame_size * frame_seqlen | |
| current_end = current_start + q.shape[1] | |
| sink_tokens = self.sink_size * frame_seqlen | |
| kv_cache_size = kv_cache["k"].shape[1] | |
| num_new_tokens = q.shape[1] | |
| if ( | |
| self.local_attn_size != -1 | |
| and current_end > kv_cache["global_end_index"].item() | |
| and (num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size) | |
| ): | |
| num_evicted_tokens = ( | |
| num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size | |
| ) | |
| num_rolled_tokens = ( | |
| kv_cache["local_end_index"].item() - num_evicted_tokens - sink_tokens | |
| ) | |
| rolling_dst = slice(sink_tokens, sink_tokens + num_rolled_tokens) | |
| rolling_src = slice( | |
| sink_tokens + num_evicted_tokens, | |
| sink_tokens + num_evicted_tokens + num_rolled_tokens, | |
| ) | |
| kv_cache["k"][:, rolling_dst] = kv_cache["k"][:, rolling_src].clone() | |
| kv_cache["v"][:, rolling_dst] = kv_cache["v"][:, rolling_src].clone() | |
| local_end_index = ( | |
| kv_cache["local_end_index"].item() | |
| + current_end | |
| - kv_cache["global_end_index"].item() | |
| - num_evicted_tokens | |
| ) | |
| local_start_index = local_end_index - num_new_tokens | |
| kv_cache["k"][:, local_start_index:local_end_index] = k | |
| kv_cache["v"][:, local_start_index:local_end_index] = v | |
| else: | |
| local_end_index = ( | |
| kv_cache["local_end_index"].item() | |
| + current_end | |
| - kv_cache["global_end_index"].item() | |
| ) | |
| local_start_index = local_end_index - num_new_tokens | |
| kv_cache["k"][:, local_start_index:local_end_index] = k | |
| kv_cache["v"][:, local_start_index:local_end_index] = v | |
| local_start_frame_index = local_start_index // frame_seqlen | |
| local_end_frame_index = local_end_index // frame_seqlen | |
| assert local_start_frame_index + grid_sizes[0][0].item() == local_end_frame_index, ( | |
| f"local_start_frame_index: {local_start_frame_index}, local_end_frame_index: {local_end_frame_index}, grid_sizes[0][0]: {grid_sizes[0][0].item()}" | |
| ) | |
| query_relative_indices = torch.arange( | |
| local_start_frame_index, local_end_frame_index, device=q.device | |
| ) | |
| cache_relative_indices = torch.arange(0, local_end_frame_index, device=k.device) | |
| if self.train_frame_size == -1: | |
| query_shift = 0 | |
| cache_shift = 0 | |
| else: | |
| max_train_frame_index = self.train_frame_size - 1 | |
| query_shift = max(query_relative_indices.max().item() - max_train_frame_index, 0) | |
| cache_shift = max(cache_relative_indices.max().item() - max_train_frame_index, 0) | |
| query_relative_indices = torch.clamp(query_relative_indices - query_shift, min=0) | |
| cache_relative_indices = torch.clamp(cache_relative_indices - cache_shift, min=0) | |
| if ( | |
| self.local_attn_size != -1 | |
| and current_end > kv_cache["global_end_index"].item() | |
| and (num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size) | |
| ): | |
| tail_start = max(sink_tokens, local_end_index - max_attention_size + sink_tokens) | |
| tail_end = local_end_index | |
| sink_end = min(kv_cache["local_end_index"].item(), sink_tokens) | |
| if sink_end > 0: | |
| k = torch.cat( | |
| [kv_cache["k"][:, :sink_end], kv_cache["k"][:, tail_start:tail_end]], dim=1 | |
| ) | |
| v = torch.cat( | |
| [kv_cache["v"][:, :sink_end], kv_cache["v"][:, tail_start:tail_end]], dim=1 | |
| ) | |
| else: | |
| k = kv_cache["k"][:, tail_start:tail_end] | |
| v = kv_cache["v"][:, tail_start:tail_end] | |
| else: | |
| k = kv_cache["k"][:, max(0, local_end_index - max_attention_size) : local_end_index] | |
| v = kv_cache["v"][:, max(0, local_end_index - max_attention_size) : local_end_index] | |
| roped_query = causal_rope_apply( | |
| q, grid_sizes, freqs, relative_frame_indices=query_relative_indices | |
| ).type_as(v) | |
| roped_key = causal_rope_apply( | |
| k, grid_sizes, freqs, relative_frame_indices=cache_relative_indices | |
| ).type_as(v) | |
| x = attention(roped_query, roped_key, v) | |
| kv_cache["global_end_index"].fill_(current_end) | |
| kv_cache["local_end_index"].fill_(local_end_index) | |
| x = x.flatten(2) | |
| x = self.o(x) | |
| return x | |
| class CausalWanAttentionBlock(nn.Module): | |
| def __init__( | |
| self, | |
| model_type: str, | |
| dim: int, | |
| ffn_dim: int, | |
| num_heads: int, | |
| local_attn_size: int = -1, | |
| sink_size: int = 0, | |
| qk_norm: bool = True, | |
| cross_attn_norm: bool = False, | |
| eps: float = 1e-06, | |
| train_frame_size: int = 21, | |
| ): | |
| super().__init__() | |
| self.dim = dim | |
| self.ffn_dim = ffn_dim | |
| self.num_heads = num_heads | |
| self.qk_norm = qk_norm | |
| self.cross_attn_norm = cross_attn_norm | |
| self.eps = eps | |
| self.local_attn_size = local_attn_size | |
| self.sink_size = sink_size | |
| self.norm1 = WanLayerNorm(dim, eps) | |
| self.self_attn = CausalWanSelfAttention( | |
| dim=dim, | |
| num_heads=num_heads, | |
| local_attn_size=local_attn_size, | |
| sink_size=sink_size, | |
| qk_norm=qk_norm, | |
| eps=eps, | |
| train_frame_size=train_frame_size, | |
| ) | |
| self.norm3 = ( | |
| WanLayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() | |
| ) | |
| WanCrossAttentionClass = WAN_CROSSATTENTION_CLASSES["t2v_cross_attn"] | |
| self.cross_attn = WanCrossAttentionClass( | |
| dim=dim, num_heads=num_heads, window_size=(-1, -1), qk_norm=qk_norm, eps=eps | |
| ) | |
| self.norm2 = WanLayerNorm(dim, eps) | |
| self.ffn = nn.Sequential( | |
| nn.Linear(dim, ffn_dim), nn.GELU(approximate="tanh"), nn.Linear(ffn_dim, dim) | |
| ) | |
| self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) | |
| def forward( | |
| self, | |
| x, | |
| e, | |
| seq_lens, | |
| grid_sizes, | |
| freqs, | |
| context, | |
| context_lens, | |
| block_mask, | |
| kv_cache=None, | |
| crossattn_cache=None, | |
| current_start=0, | |
| cache_start=None, | |
| ): | |
| """ | |
| Args: | |
| x(Tensor): Shape [B, L, C] | |
| e(Tensor): Shape [B, F, 6, C] | |
| seq_lens(Tensor): Shape [B], length of each sequence in batch | |
| grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) | |
| freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] | |
| """ | |
| num_frames, frame_seqlen = (e.shape[1], x.shape[1] // e.shape[1]) | |
| e = (self.modulation.unsqueeze(1) + e).chunk(6, dim=2) | |
| y = self.self_attn( | |
| ( | |
| self.norm1(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1 + e[1]) + e[0] | |
| ).flatten(1, 2), | |
| seq_lens, | |
| grid_sizes, | |
| freqs, | |
| block_mask, | |
| kv_cache, | |
| current_start, | |
| cache_start, | |
| ) | |
| x = x + (y.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * e[2]).flatten(1, 2) | |
| def cross_attn_ffn(x, context, context_lens, e, crossattn_cache=None): | |
| x = x + self.cross_attn( | |
| self.norm3(x), context, context_lens, crossattn_cache=crossattn_cache | |
| ) | |
| y = self.ffn( | |
| ( | |
| self.norm2(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1 + e[4]) | |
| + e[3] | |
| ).flatten(1, 2) | |
| ) | |
| x = x + (y.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * e[5]).flatten(1, 2) | |
| return x | |
| x = cross_attn_ffn(x, context, context_lens, e, crossattn_cache) | |
| return x | |
| class CausalHead(nn.Module): | |
| def __init__(self, dim, out_dim, patch_size, eps=1e-06): | |
| super().__init__() | |
| self.dim = dim | |
| self.out_dim = out_dim | |
| self.patch_size = patch_size | |
| self.eps = eps | |
| out_dim = math.prod(patch_size) * out_dim | |
| self.norm = WanLayerNorm(dim, eps) | |
| self.head = nn.Linear(dim, out_dim) | |
| self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) | |
| def forward(self, x, e): | |
| """ | |
| Args: | |
| x(Tensor): Shape [B, L1, C] | |
| e(Tensor): Shape [B, F, 1, C] | |
| """ | |
| num_frames, frame_seqlen = (e.shape[1], x.shape[1] // e.shape[1]) | |
| shift, scale = (self.modulation.unsqueeze(1) + e).chunk(2, dim=2) | |
| x = self.norm(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) | |
| x = self.head(x * (1 + scale) + shift) | |
| return x | |
| class CausalWanModel(ModelMixin, ConfigMixin, CausalMaskMixin): | |
| """ | |
| Wan diffusion backbone supporting both text-to-video and image-to-video. | |
| """ | |
| ignore_for_config = ["patch_size", "cross_attn_norm", "qk_norm", "text_dim"] | |
| _no_split_modules = ["CausalWanAttentionBlock"] | |
| _supports_gradient_checkpointing = True | |
| def __init__( | |
| self, | |
| model_type="t2v", | |
| patch_size=(1, 2, 2), | |
| text_len=512, | |
| in_dim=16, | |
| dim=2048, | |
| ffn_dim=8192, | |
| freq_dim=256, | |
| text_dim=4096, | |
| out_dim=16, | |
| num_heads=16, | |
| num_layers=32, | |
| local_attn_size=21, | |
| sink_size=0, | |
| train_frame_size=21, | |
| qk_norm=True, | |
| cross_attn_norm=True, | |
| eps=1e-06, | |
| num_frame_per_block: int = 1, | |
| independent_first_frame: bool = False, | |
| ): | |
| """ | |
| Initialize the diffusion model backbone. | |
| Args: | |
| model_type (`str`, *optional*, defaults to 't2v'): | |
| Selected backbone family: 't2v' or 'ti2v'; sampling requires supplied C0. | |
| patch_size (`tuple`, *optional*, defaults to (1, 2, 2)): | |
| 3D patch dimensions for video embedding (t_patch, h_patch, w_patch) | |
| text_len (`int`, *optional*, defaults to 512): | |
| Fixed length for text embeddings | |
| in_dim (`int`, *optional*, defaults to 16): | |
| Input video channels (C_in) | |
| dim (`int`, *optional*, defaults to 2048): | |
| Hidden dimension of the transformer | |
| ffn_dim (`int`, *optional*, defaults to 8192): | |
| Intermediate dimension in feed-forward network | |
| freq_dim (`int`, *optional*, defaults to 256): | |
| Dimension for sinusoidal time embeddings | |
| text_dim (`int`, *optional*, defaults to 4096): | |
| Input dimension for text embeddings | |
| out_dim (`int`, *optional*, defaults to 16): | |
| Output video channels (C_out) | |
| num_heads (`int`, *optional*, defaults to 16): | |
| Number of attention heads | |
| num_layers (`int`, *optional*, defaults to 32): | |
| Number of transformer blocks | |
| local_attn_size (`int`, *optional*, defaults to 21): | |
| Window size for temporal local attention (-1 indicates global attention) | |
| sink_size (`int`, *optional*, defaults to 0): | |
| Size of the attention sink, we keep the first `sink_size` frames unchanged when rolling the KV cache | |
| train_frame_size (`int`, *optional*, defaults to 21): | |
| Number of frames in training sequences, used for block-relative RoPE in inference | |
| qk_norm (`bool`, *optional*, defaults to True): | |
| Enable query/key normalization | |
| cross_attn_norm (`bool`, *optional*, defaults to False): | |
| Enable cross-attention normalization | |
| eps (`float`, *optional*, defaults to 1e-6): | |
| Epsilon value for normalization layers | |
| """ | |
| super().__init__() | |
| assert model_type in ["t2v", "ti2v"] | |
| self.model_type = model_type | |
| self.patch_size = patch_size | |
| self.text_len = text_len | |
| self.in_dim = in_dim | |
| self.dim = dim | |
| self.ffn_dim = ffn_dim | |
| self.freq_dim = freq_dim | |
| self.text_dim = text_dim | |
| self.out_dim = out_dim | |
| self.num_heads = num_heads | |
| self.num_layers = num_layers | |
| self.local_attn_size = local_attn_size | |
| self.train_frame_size = train_frame_size | |
| self.sink_size = sink_size | |
| self.qk_norm = qk_norm | |
| self.cross_attn_norm = cross_attn_norm | |
| self.eps = eps | |
| self.num_frame_per_block = num_frame_per_block | |
| self.independent_first_frame = independent_first_frame | |
| self.patch_embedding = nn.Conv3d(in_dim, dim, kernel_size=patch_size, stride=patch_size) | |
| self.text_embedding = nn.Sequential( | |
| nn.Linear(text_dim, dim), nn.GELU(approximate="tanh"), nn.Linear(dim, dim) | |
| ) | |
| self.time_embedding = nn.Sequential( | |
| nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim) | |
| ) | |
| self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6)) | |
| block_cls = CausalWanAttentionBlock | |
| self.blocks = nn.ModuleList( | |
| [ | |
| block_cls( | |
| model_type=model_type, | |
| dim=dim, | |
| ffn_dim=ffn_dim, | |
| num_heads=num_heads, | |
| local_attn_size=local_attn_size, | |
| sink_size=sink_size, | |
| train_frame_size=train_frame_size, | |
| qk_norm=qk_norm, | |
| cross_attn_norm=cross_attn_norm, | |
| eps=eps, | |
| ) | |
| for _ in range(num_layers) | |
| ] | |
| ) | |
| self.head = CausalHead(dim, out_dim, patch_size, eps) | |
| assert dim % num_heads == 0 and dim // num_heads % 2 == 0 | |
| d = dim // num_heads | |
| self.freqs = torch.cat( | |
| [ | |
| rope_params(1024, d - 4 * (d // 6)), | |
| rope_params(1024, 2 * (d // 6)), | |
| rope_params(1024, 2 * (d // 6)), | |
| ], | |
| dim=1, | |
| ) | |
| self.init_weights() | |
| self.gradient_checkpointing = False | |
| self.block_mask = None | |
| def _forward_inference( | |
| self, | |
| x, | |
| t, | |
| context, | |
| kv_cache: dict = None, | |
| crossattn_cache: dict = None, | |
| current_start: int = 0, | |
| cache_start: int = 0, | |
| ): | |
| """ | |
| Run the diffusion model with kv caching. | |
| See Algorithm 2 of CausVid paper https://arxiv.org/abs/2412.07772 for details. | |
| This function will be run for num_frame times. | |
| Process the latent frames one by one (1560 tokens each) | |
| Args: | |
| x (List[Tensor]): | |
| List of input video tensors, each with shape [C_in, F, H, W] | |
| t (Tensor): | |
| Diffusion timesteps tensor of shape [B, F] | |
| context (List[Tensor]): | |
| List of text embeddings each with shape [L, C] | |
| clip_fea (Tensor, *optional*): | |
| CLIP image features for image-to-video mode | |
| y (List[Tensor], *optional*): | |
| Conditional video inputs for image-to-video mode, same shape as x | |
| Returns: | |
| List[Tensor]: | |
| List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8] | |
| """ | |
| device = self.patch_embedding.weight.device | |
| if self.freqs.device != device: | |
| self.freqs = self.freqs.to(device) | |
| patch_embedding = self.patch_embedding | |
| x = [patch_embedding(u.unsqueeze(0)) for u in x] | |
| grid_sizes = torch.stack([torch.tensor(u.shape[2:], dtype=torch.long) for u in x]) | |
| x = [u.flatten(2).transpose(1, 2) for u in x] | |
| seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long) | |
| x = torch.cat(x) | |
| e = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, t.flatten()).type_as(x)) | |
| e0 = self.time_projection(e).unflatten(1, (6, self.dim)).unflatten(dim=0, sizes=t.shape) | |
| context_lens = None | |
| context = self.text_embedding( | |
| torch.stack( | |
| [torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in context] | |
| ) | |
| ) | |
| kwargs = { | |
| "e": e0, | |
| "seq_lens": seq_lens, | |
| "grid_sizes": grid_sizes, | |
| "freqs": self.freqs, | |
| "context": context, | |
| "context_lens": context_lens, | |
| "block_mask": self.block_mask, | |
| } | |
| def create_custom_forward(module): | |
| def custom_forward(*inputs, **kwargs): | |
| return module(*inputs, **kwargs) | |
| return custom_forward | |
| for block_index, block in enumerate(self.blocks): | |
| if torch.is_grad_enabled() and self.gradient_checkpointing: | |
| kwargs.update( | |
| { | |
| "kv_cache": kv_cache[block_index], | |
| "current_start": current_start, | |
| "cache_start": cache_start, | |
| } | |
| ) | |
| x = torch.utils.checkpoint.checkpoint( | |
| create_custom_forward(block), x, **kwargs, use_reentrant=False | |
| ) | |
| else: | |
| kwargs.update( | |
| { | |
| "kv_cache": kv_cache[block_index], | |
| "crossattn_cache": crossattn_cache[block_index], | |
| "current_start": current_start, | |
| "cache_start": cache_start, | |
| } | |
| ) | |
| x = block(x, **kwargs) | |
| x = self.head(x, e.unflatten(dim=0, sizes=t.shape).unsqueeze(2)) | |
| x = self.unpatchify(x, grid_sizes) | |
| output = torch.stack(x) | |
| return output | |
| def _forward_train(self, x, t, context, clean_x=None): | |
| """ | |
| Forward pass through the diffusion model | |
| Args: | |
| x (List[Tensor]): | |
| List of input video tensors, each with shape [C_in, F, H, W] | |
| t (Tensor): | |
| Diffusion timesteps tensor of shape [B, F] | |
| context (List[Tensor]): | |
| List of text embeddings each with shape [L, C] | |
| clip_fea (Tensor, *optional*): | |
| CLIP image features for image-to-video mode | |
| y (List[Tensor], *optional*): | |
| Conditional video inputs for image-to-video mode, same shape as x | |
| Returns: | |
| List[Tensor]: | |
| List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8] | |
| """ | |
| if clean_x is None: | |
| raise ValueError("Teacher-forcing forward requires clean_x.") | |
| device = self.patch_embedding.weight.device | |
| if self.freqs.device != device: | |
| self.freqs = self.freqs.to(device) | |
| if self.block_mask is None: | |
| self.block_mask = self._prepare_teacher_forcing_mask( | |
| device, | |
| num_frames=x.shape[2], | |
| frame_seqlen=x.shape[-2] * x.shape[-1] // (self.patch_size[1] * self.patch_size[2]), | |
| num_frame_per_block=self.num_frame_per_block, | |
| local_attn_size=self.local_attn_size, | |
| sink_size=self.sink_size, | |
| independent_first_frame=self.independent_first_frame, | |
| ) | |
| _vis_flag = os.environ.get("WAN_VIS_CAUSAL_MASK", "0") | |
| if _vis_flag and _vis_flag != "0": | |
| if is_main_process(): | |
| if not hasattr(self, "_causal_mask_vis_done"): | |
| self._causal_mask_vis_done = set() | |
| mode = "TF" | |
| vis_key = (mode, int(x.shape[2]), bool(self.independent_first_frame)) | |
| if vis_key not in self._causal_mask_vis_done: | |
| out_dir = os.environ.get( | |
| "WAN_VIS_CAUSAL_MASK_DIR", os.path.join("outputs", "causal_mask_vis") | |
| ) | |
| try: | |
| saved_path = visualize_causal_mask( | |
| mode=mode, | |
| num_frames=int(x.shape[2]), | |
| num_frame_per_block=self.num_frame_per_block, | |
| local_attn_size=self.local_attn_size, | |
| sink_size=self.sink_size, | |
| independent_first_frame=bool(self.independent_first_frame), | |
| out_dir=out_dir, | |
| ) | |
| logging.info( | |
| "[causal_mask_vis] saved %s mask (num_frames=%d, num_frame_per_block=%d, local_attn_size=%d, sink_size=%d, independent_first_frame=%s) -> %s", | |
| mode, | |
| int(x.shape[2]), | |
| self.num_frame_per_block, | |
| self.local_attn_size, | |
| self.sink_size, | |
| bool(self.independent_first_frame), | |
| saved_path, | |
| ) | |
| except Exception as _e: | |
| logging.warning("[causal_mask_vis] failed: %s", _e) | |
| self._causal_mask_vis_done.add(vis_key) | |
| x = [self.patch_embedding(u.unsqueeze(0)) for u in x] | |
| grid_sizes = torch.stack([torch.tensor(u.shape[2:], dtype=torch.long) for u in x]) | |
| x = [u.flatten(2).transpose(1, 2) for u in x] | |
| seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long) | |
| x = torch.cat( | |
| [torch.cat([u, u.new_zeros(1, seq_lens[0] - u.size(1), u.size(2))], dim=1) for u in x] | |
| ) | |
| e = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, t.flatten()).type_as(x)) | |
| e0 = self.time_projection(e).unflatten(1, (6, self.dim)).unflatten(dim=0, sizes=t.shape) | |
| context_lens = None | |
| context = self.text_embedding( | |
| torch.stack( | |
| [torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in context] | |
| ) | |
| ) | |
| clean_patch_embedding = self.patch_embedding | |
| clean_x = [clean_patch_embedding(u.unsqueeze(0)) for u in clean_x] | |
| clean_x = [u.flatten(2).transpose(1, 2) for u in clean_x] | |
| seq_lens_clean = torch.tensor([u.size(1) for u in clean_x], dtype=torch.long) | |
| clean_x = torch.cat( | |
| [ | |
| torch.cat([u, u.new_zeros(1, seq_lens_clean[0] - u.size(1), u.size(2))], dim=1) | |
| for u in clean_x | |
| ] | |
| ) | |
| x = torch.cat([clean_x, x], dim=1) | |
| aug_t = torch.zeros_like(t) | |
| e_clean = self.time_embedding( | |
| sinusoidal_embedding_1d(self.freq_dim, aug_t.flatten()).type_as(x) | |
| ) | |
| e0_clean = ( | |
| self.time_projection(e_clean) | |
| .unflatten(1, (6, self.dim)) | |
| .unflatten(dim=0, sizes=t.shape) | |
| ) | |
| e0 = torch.cat([e0_clean, e0], dim=1) | |
| kwargs = { | |
| "e": e0, | |
| "seq_lens": seq_lens, | |
| "grid_sizes": grid_sizes, | |
| "freqs": self.freqs, | |
| "context": context, | |
| "context_lens": context_lens, | |
| "block_mask": self.block_mask, | |
| } | |
| def create_custom_forward(module): | |
| def custom_forward(*inputs, **kwargs): | |
| return module(*inputs, **kwargs) | |
| return custom_forward | |
| for block in self.blocks: | |
| if torch.is_grad_enabled() and self.gradient_checkpointing: | |
| x = torch.utils.checkpoint.checkpoint( | |
| create_custom_forward(block), x, **kwargs, use_reentrant=False | |
| ) | |
| else: | |
| x = block(x, **kwargs) | |
| x = x[:, x.shape[1] // 2 :] | |
| x = self.head(x, e.unflatten(dim=0, sizes=t.shape).unsqueeze(2)) | |
| x = self.unpatchify(x, grid_sizes) | |
| return torch.stack(x) | |
| def forward(self, *args, **kwargs): | |
| if kwargs.get("kv_cache", None) is not None: | |
| return self._forward_inference(*args, **kwargs) | |
| else: | |
| return self._forward_train(*args, **kwargs) | |
| def unpatchify(self, x, grid_sizes): | |
| """ | |
| Reconstruct video tensors from patch embeddings. | |
| Args: | |
| x (List[Tensor]): | |
| List of patchified features, each with shape [L, C_out * prod(patch_size)] | |
| grid_sizes (Tensor): | |
| Original spatial-temporal grid dimensions before patching, | |
| shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches) | |
| Returns: | |
| List[Tensor]: | |
| Reconstructed video tensors with shape [C_out, F, H / 8, W / 8] | |
| """ | |
| c = self.out_dim | |
| out = [] | |
| for u, v in zip(x, grid_sizes.tolist(), strict=False): | |
| u = u[: math.prod(v)].view(*v, *self.patch_size, c) | |
| u = torch.einsum("fhwpqrc->cfphqwr", u) | |
| u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size, strict=False)]) | |
| out.append(u) | |
| return out | |
| def init_weights(self): | |
| """ | |
| Initialize model parameters using Xavier initialization. | |
| """ | |
| for m in self.modules(): | |
| if isinstance(m, nn.Linear): | |
| nn.init.xavier_uniform_(m.weight) | |
| if m.bias is not None: | |
| nn.init.zeros_(m.bias) | |
| nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1)) | |
| for m in self.text_embedding.modules(): | |
| if isinstance(m, nn.Linear): | |
| nn.init.normal_(m.weight, std=0.02) | |
| for m in self.time_embedding.modules(): | |
| if isinstance(m, nn.Linear): | |
| nn.init.normal_(m.weight, std=0.02) | |
| nn.init.zeros_(self.head.head.weight) | |