| from typing import Dict, Union |
|
|
| import torch |
| import torch.nn as nn |
|
|
| from ..wanvideo.modules.model import WanLayerNorm, WanSelfAttention, EmbedND_RifleX, sinusoidal_embedding_1d, apply_rotary_emb_split, apply_rope_comfy1 |
|
|
| class WanAttentionBlock(nn.Module): |
| def __init__(self, in_features, out_features, ffn_dim, ffn2_dim, num_heads, qk_norm=True, cross_attn_norm=False, eps=1e-6, attention_mode="sdpa", rope_func="comfy", rms_norm_function="default"): |
| super().__init__() |
| self.dim = out_features |
| self.ffn_dim = ffn_dim |
| self.num_heads = num_heads |
| self.head_dim = out_features // num_heads |
| self.qk_norm = qk_norm |
| self.cross_attn_norm = cross_attn_norm |
| self.eps = eps |
| self.attention_mode = attention_mode |
| self.rope_func = rope_func |
|
|
| |
| self.norm1 = WanLayerNorm(self.dim, eps) |
| self.self_attn = WanSelfAttention(in_features, out_features, num_heads, qk_norm, eps, self.attention_mode, rms_norm_function=rms_norm_function, head_norm=False) |
| self.norm2 = WanLayerNorm(self.dim, eps) |
| self.ffn = nn.Sequential(nn.Linear(in_features, ffn_dim), nn.GELU(approximate='tanh'), nn.Linear(ffn2_dim, out_features)) |
|
|
| self.modulation = nn.Parameter(torch.randn(1, 6, out_features) / in_features**0.5) |
|
|
|
|
| def get_mod(self, e, modulation): |
| if e.dim() == 3: |
| if e.shape[-1] == 512: |
| e = self.modulation(e) |
| return e.unsqueeze(2).chunk(6, dim=-1) |
| return (modulation + e).chunk(6, dim=1) |
| elif e.dim() == 4: |
| e_mod = modulation.unsqueeze(2) + e |
| return [ei.squeeze(1) for ei in e_mod.unbind(dim=1)] |
|
|
|
|
| def modulate(self, norm_x, shift_msa, scale_msa): |
| return torch.addcmul(shift_msa, norm_x, 1 + scale_msa) |
|
|
| def ffn_chunked(self, mod_x, num_chunks=4): |
| seq_len = mod_x.shape[1] |
| if seq_len <= 8192 or num_chunks <= 1: |
| return self.ffn(mod_x) |
| return torch.cat([self.ffn(chunk.contiguous()) for chunk in mod_x.chunk(num_chunks, dim=1)], dim=1) |
|
|
| |
| def forward(self, x, e, seq_lens, freqs, split_rope=True, e_tr=None, tr_start=0, tr_num=0): |
|
|
| use_token_replace = False |
| if e_tr is not None and tr_num > 0: |
| tr_shift_msa, tr_scale_msa, tr_gate_msa, tr_shift_mlp, tr_scale_mlp, tr_gate_mlp = self.get_mod(e_tr.to(x.device), self.modulation) |
| use_token_replace = True |
| tr_start = tr_start or 0 |
| tr_end = tr_start + (tr_num or 0) |
|
|
| shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.get_mod(e.to(x.device), self.modulation) |
| del e |
| input_dtype = x.dtype |
|
|
| if use_token_replace: |
| norm_x = self.norm1(x.to(shift_msa.dtype)) |
| input_x = torch.cat([ |
| torch.addcmul(shift_msa, norm_x[:, :tr_start], 1 + scale_msa), |
| torch.addcmul(tr_shift_msa, norm_x[:, tr_start:tr_end], 1 + tr_scale_msa), |
| torch.addcmul(shift_msa, norm_x[:, tr_end:], 1 + scale_msa) |
| ], dim=1).to(input_dtype) |
| else: |
| input_x = self.modulate(self.norm1(x.to(shift_msa.dtype)), shift_msa, scale_msa).to(input_dtype) |
| del shift_msa, scale_msa |
|
|
| b, s, n, d = *x.shape[:2], self.self_attn.num_heads, self.self_attn.head_dim |
| h_dim = w_dim = 2 * (self.head_dim // 6) |
| t_dim = self.head_dim - h_dim - w_dim |
|
|
| q = self.self_attn.norm_q(self.self_attn.q(input_x)).to(self.self_attn.norm_q.weight.dtype).view(b, s, n, d) |
| if split_rope: |
| q = apply_rotary_emb_split(q, freqs, t_dim) |
| else: |
| q = apply_rope_comfy1(q, freqs) |
|
|
| k = self.self_attn.norm_k(self.self_attn.k(input_x).to(self.self_attn.norm_k.weight.dtype)).to(input_x.dtype).view(b, s, n, d) |
| if split_rope: |
| k = apply_rotary_emb_split(k, freqs, t_dim) |
| else: |
| k = apply_rope_comfy1(k, freqs) |
|
|
| v = self.self_attn.v(input_x).view(b, s, n, d) |
| del input_x |
|
|
| y = self.self_attn.forward(q, k, v, seq_lens) |
| del q, k, v |
| if use_token_replace: |
| x = x + torch.cat([ |
| y[:, :tr_start] * gate_msa, |
| y[:, tr_start:tr_end] * tr_gate_msa, |
| y[:, tr_end:] * gate_msa |
| ], dim=1).to(input_dtype) |
| else: |
| x = x.addcmul(y, gate_msa) |
| del y, gate_msa |
|
|
| |
| if use_token_replace: |
| norm2_x = self.norm2(x.to(shift_mlp.dtype)) |
| mod_x = torch.cat([ |
| torch.addcmul(shift_mlp, norm2_x[:, :tr_start], 1 + scale_mlp), |
| torch.addcmul(tr_shift_mlp, norm2_x[:, tr_start:tr_end], 1 + tr_scale_mlp), |
| torch.addcmul(shift_mlp, norm2_x[:, tr_end:], 1 + scale_mlp) |
| ], dim=1) |
| else: |
| mod_x = torch.addcmul(shift_mlp, self.norm2(x.to(shift_mlp.dtype)), 1 + scale_mlp) |
| del shift_mlp, scale_mlp |
| x_ffn = self.ffn_chunked(mod_x.to(input_dtype), num_chunks=1) |
| del mod_x |
|
|
| |
| if use_token_replace: |
| x = x + torch.cat([ |
| x_ffn[:, :tr_start] * gate_mlp, |
| x_ffn[:, tr_start:tr_end] * tr_gate_mlp, |
| x_ffn[:, tr_end:] * gate_mlp |
| ], dim=1).to(input_dtype) |
| else: |
| x = x.addcmul(x_ffn.to(gate_mlp.dtype), gate_mlp).to(input_dtype) |
| del gate_mlp |
|
|
| return x |
|
|
|
|
| class WanRefextractor(nn.Module): |
| def __init__(self, patch_size=(1, 2, 2), in_dim=16, dim=5120, in_features=5120, out_features=5120, ffn_dim=8192, ffn2_dim=8192, |
| freq_dim=256, num_heads=16, num_layers=32, eps=1e-6, |
| qk_norm=True, cross_attn_norm=True, |
| attention_mode='sdpa', rope_func='comfy', rms_norm_function='default', |
| main_device=torch.device('cuda'), offload_device=torch.device('cpu'), dtype=torch.float16): |
| super().__init__() |
| self.patch_size = patch_size |
| self.freq_dim = freq_dim |
| self.dim = dim |
| self.main_device = main_device |
| self.base_dtype = dtype |
| self.attention_mode = attention_mode |
| self.patch_embedding = nn.Conv3d(in_dim, dim, kernel_size=patch_size, stride=patch_size) |
| 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)) |
|
|
| self.blocks = nn.ModuleList([ |
| WanAttentionBlock(in_features, out_features, ffn_dim, ffn2_dim, num_heads, |
| qk_norm, cross_attn_norm, eps, attention_mode="sdpa", rope_func=rope_func, rms_norm_function=rms_norm_function) |
| for i in range(num_layers) |
| ]) |
|
|
| self.ref_blocks = nn.ModuleList([]) |
| for _ in range(len(self.blocks)+1): |
| self.ref_blocks.append(nn.Linear(in_features, out_features)) |
|
|
| d = dim // num_heads |
| self.rope_embedder = EmbedND_RifleX(d,10000.0, [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)], num_frames=1, k=0) |
|
|
| def rope_encode_comfy(self, t, h, w, freq_offset=0, t_start=0, steps_t=None, steps_h=None, steps_w=None, ntk_alphas=[1,1,1], device=None, dtype=None): |
| patch_size = self.patch_size |
| t_len = ((t + (patch_size[0] // 2)) // patch_size[0]) |
| h_len = ((h + (patch_size[1] // 2)) // patch_size[1]) |
| w_len = ((w + (patch_size[2] // 2)) // patch_size[2]) |
|
|
| if steps_t is None: |
| steps_t = t_len |
| if steps_h is None: |
| steps_h = h_len |
| if steps_w is None: |
| steps_w = w_len |
|
|
| img_ids = torch.zeros((steps_t, steps_h, steps_w, 3), device=device, dtype=dtype) |
| img_ids[:, :, :, 0] = img_ids[:, :, :, 0] + torch.linspace(t_start+freq_offset, t_start + (t_len - 1), steps=steps_t, device=device, dtype=dtype).reshape(-1, 1, 1) |
| img_ids[:, :, :, 1] = img_ids[:, :, :, 1] + torch.linspace(freq_offset, h_len - 1, steps=steps_h, device=device, dtype=dtype).reshape(1, -1, 1) |
| img_ids[:, :, :, 2] = img_ids[:, :, :, 2] + torch.linspace(freq_offset, w_len - 1, steps=steps_w, device=device, dtype=dtype).reshape(1, 1, -1) |
| img_ids = img_ids.reshape(1, -1, img_ids.shape[-1]) |
|
|
| freqs = self.rope_embedder(img_ids, ntk_alphas).movedim(1, 2) |
| return freqs |
|
|
| def forward( |
| self, |
| x: torch.Tensor, |
| timestep: torch.LongTensor, |
| ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: |
| B, C, F, H, W = x.shape |
|
|
| freqs = self.rope_encode_comfy(F, H, W, device=x.device, dtype=x.dtype) |
|
|
| self.patch_embedding.to(self.main_device) |
| x = self.patch_embedding(x.float()).to(x.dtype).flatten(2).transpose(1, 2).to(self.base_dtype) |
|
|
| seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.int32) |
|
|
| time_embed_dtype = self.time_embedding[0].weight.dtype |
| if time_embed_dtype not in [torch.float16, torch.bfloat16, torch.float32]: |
| time_embed_dtype = self.base_dtype |
| e = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, timestep.flatten()).to(time_embed_dtype)) |
| e0 = self.time_projection(e).unflatten(1, (6, self.dim)).to(self.base_dtype) |
| del e |
|
|
| |
| block_samples = () |
| for block in self.blocks: |
| block_samples = block_samples + (x, ) |
| x = block(x, e0, seq_lens, freqs) |
|
|
| block_samples = block_samples + (x, ) |
|
|
| ref_block_samples = () |
| for block_sample, ref_block in zip(block_samples, self.ref_blocks): |
| block_sample = ref_block(block_sample) |
| ref_block_samples = ref_block_samples + (block_sample, ) |
|
|
| return ref_block_samples, freqs |
|
|
