Spaces:
Running on Zero
Running on Zero
| # Modified for Looking-Ahead: selected TF/TI2V components; unrelated workflows removed. | |
| # Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. | |
| import torch | |
| import torch.nn as nn | |
| from .attention import attention | |
| class WanRMSNorm(nn.Module): | |
| def __init__(self, dim, eps=1e-5): | |
| super().__init__() | |
| self.dim = dim | |
| self.eps = eps | |
| self.weight = nn.Parameter(torch.ones(dim)) | |
| def forward(self, x): | |
| r""" | |
| Args: | |
| x(Tensor): Shape [B, L, C] | |
| """ | |
| return self._norm(x.float()).type_as(x) * self.weight | |
| def _norm(self, x): | |
| return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) | |
| def reset_parameters(self): | |
| nn.init.ones_(self.weight) | |
| class WanSelfAttention(nn.Module): | |
| def __init__(self, | |
| dim, | |
| num_heads, | |
| window_size=(-1, -1), | |
| qk_norm=True, | |
| eps=1e-6): | |
| assert dim % num_heads == 0 | |
| super().__init__() | |
| self.dim = dim | |
| self.num_heads = num_heads | |
| self.head_dim = dim // num_heads | |
| self.window_size = window_size | |
| self.qk_norm = qk_norm | |
| self.eps = eps | |
| # layers | |
| 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() | |
| class WanT2VCrossAttention(WanSelfAttention): | |
| def forward(self, x, context, context_lens, crossattn_cache=None): | |
| r""" | |
| Args: | |
| x(Tensor): Shape [B, L1, C] | |
| context(Tensor): Shape [B, L2, C] | |
| context_lens(Tensor): Shape [B] | |
| crossattn_cache (List[dict], *optional*): Contains the cached key and value tensors for context embedding. | |
| """ | |
| b, n, d = x.size(0), self.num_heads, self.head_dim | |
| # compute query, key, value | |
| q = self.norm_q(self.q(x)).view(b, -1, n, d) | |
| if crossattn_cache is not None: | |
| if not crossattn_cache["is_init"]: | |
| crossattn_cache["is_init"] = True | |
| k = self.norm_k(self.k(context)).view(b, -1, n, d) | |
| v = self.v(context).view(b, -1, n, d) | |
| crossattn_cache["k"] = k | |
| crossattn_cache["v"] = v | |
| else: | |
| k = crossattn_cache["k"] | |
| v = crossattn_cache["v"] | |
| else: | |
| k = self.norm_k(self.k(context)).view(b, -1, n, d) | |
| v = self.v(context).view(b, -1, n, d) | |
| # compute attention | |
| x = attention(q, k, v, k_lens=context_lens) | |
| # output | |
| x = x.flatten(2) | |
| x = self.o(x) | |
| return x | |
| WAN_CROSSATTENTION_CLASSES = {"t2v_cross_attn": WanT2VCrossAttention} | |