ALSv/self-forcing
0
1from wan.modules.attention import attention2from wan.modules.model import (3 WanRMSNorm,4 rope_apply,5 WanLayerNorm,6 WAN_CROSSATTENTION_CLASSES,7 rope_params,8 MLPProj,9 sinusoidal_embedding_1d10)11from torch.nn.attention.flex_attention import create_block_mask, flex_attention12from diffusers.configuration_utils import ConfigMixin, register_to_config13from torch.nn.attention.flex_attention import BlockMask14from diffusers.models.modeling_utils import ModelMixin15import torch.nn as nn16import torch17import math18import torch.distributed as dist19 20# wan 1.3B model has a weird channel / head configurations and require max-autotune to work with flexattention21# see https://github.com/pytorch/pytorch/issues/13325422# change to default for other models23flex_attention = torch.compile(24 flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs")25 26 27def causal_rope_apply(x, grid_sizes, freqs, start_frame=0):28 n, c = x.size(2), x.size(3) // 229 30 # split freqs31 freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)32 33 # loop over samples34 output = []35 36 for i, (f, h, w) in enumerate(grid_sizes.tolist()):37 seq_len = f * h * w38 39 # precompute multipliers40 x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape(41 seq_len, n, -1, 2))42 freqs_i = torch.cat([43 freqs[0][start_frame:start_frame + f].view(f, 1, 1, -1).expand(f, h, w, -1),44 freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),45 freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)46 ],47 dim=-1).reshape(seq_len, 1, -1)48 49 # apply rotary embedding50 x_i = torch.view_as_real(x_i * freqs_i).flatten(2)51 x_i = torch.cat([x_i, x[i, seq_len:]])52 53 # append to collection54 output.append(x_i)55 return torch.stack(output).type_as(x)56 57 58class CausalWanSelfAttention(nn.Module):59 60 def __init__(self,61 dim,62 num_heads,63 local_attn_size=-1,64 sink_size=0,65 qk_norm=True,66 eps=1e-6):67 assert dim % num_heads == 068 super().__init__()69 self.dim = dim70 self.num_heads = num_heads71 self.head_dim = dim // num_heads72 self.local_attn_size = local_attn_size73 self.sink_size = sink_size74 self.qk_norm = qk_norm75 self.eps = eps76 self.max_attention_size = 32760 if local_attn_size == -1 else local_attn_size * 156077 78 # layers79 self.q = nn.Linear(dim, dim)80 self.k = nn.Linear(dim, dim)81 self.v = nn.Linear(dim, dim)82 self.o = nn.Linear(dim, dim)83 self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()84 self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()85 86 def forward(87 self,88 x,89 seq_lens,90 grid_sizes,91 freqs,92 block_mask,93 kv_cache=None,94 current_start=0,95 cache_start=None96 ):97 r"""98 Args:99 x(Tensor): Shape [B, L, num_heads, C / num_heads]100 seq_lens(Tensor): Shape [B]101 grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)102 freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]103 block_mask (BlockMask)104 """105 b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim106 if cache_start is None:107 cache_start = current_start108 109 # query, key, value function110 def qkv_fn(x):111 q = self.norm_q(self.q(x)).view(b, s, n, d)112 k = self.norm_k(self.k(x)).view(b, s, n, d)113 v = self.v(x).view(b, s, n, d)114 return q, k, v115 116 q, k, v = qkv_fn(x)117 118 if kv_cache is None:119 # if it is teacher forcing training?120 is_tf = (s == seq_lens[0].item() * 2)121 if is_tf:122 q_chunk = torch.chunk(q, 2, dim=1)123 k_chunk = torch.chunk(k, 2, dim=1)124 roped_query = []125 roped_key = []126 # rope should be same for clean and noisy parts127 for ii in range(2):128 rq = rope_apply(q_chunk[ii], grid_sizes, freqs).type_as(v)129 rk = rope_apply(k_chunk[ii], grid_sizes, freqs).type_as(v)130 roped_query.append(rq)131 roped_key.append(rk)132 133 roped_query = torch.cat(roped_query, dim=1)134 roped_key = torch.cat(roped_key, dim=1)135 136 padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1]137 padded_roped_query = torch.cat(138 [roped_query,139 torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]],140 device=q.device, dtype=v.dtype)],141 dim=1142 )143 144 padded_roped_key = torch.cat(145 [roped_key, torch.zeros([k.shape[0], padded_length, k.shape[2], k.shape[3]],146 device=k.device, dtype=v.dtype)],147 dim=1148 )149 150 padded_v = torch.cat(151 [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]],152 device=v.device, dtype=v.dtype)],153 dim=1154 )155 156 x = flex_attention(157 query=padded_roped_query.transpose(2, 1),158 key=padded_roped_key.transpose(2, 1),159 value=padded_v.transpose(2, 1),160 block_mask=block_mask161 )[:, :, :-padded_length].transpose(2, 1)162 163 else:164 roped_query = rope_apply(q, grid_sizes, freqs).type_as(v)165 roped_key = rope_apply(k, grid_sizes, freqs).type_as(v)166 167 padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1]168 padded_roped_query = torch.cat(169 [roped_query,170 torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]],171 device=q.device, dtype=v.dtype)],172 dim=1173 )174 175 padded_roped_key = torch.cat(176 [roped_key, torch.zeros([k.shape[0], padded_length, k.shape[2], k.shape[3]],177 device=k.device, dtype=v.dtype)],178 dim=1179 )180 181 padded_v = torch.cat(182 [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]],183 device=v.device, dtype=v.dtype)],184 dim=1185 )186 187 x = flex_attention(188 query=padded_roped_query.transpose(2, 1),189 key=padded_roped_key.transpose(2, 1),190 value=padded_v.transpose(2, 1),191 block_mask=block_mask192 )[:, :, :-padded_length].transpose(2, 1)193 else:194 frame_seqlen = math.prod(grid_sizes[0][1:]).item()195 current_start_frame = current_start // frame_seqlen196 roped_query = causal_rope_apply(197 q, grid_sizes, freqs, start_frame=current_start_frame).type_as(v)198 roped_key = causal_rope_apply(199 k, grid_sizes, freqs, start_frame=current_start_frame).type_as(v)200 201 current_end = current_start + roped_query.shape[1]202 sink_tokens = self.sink_size * frame_seqlen203 # If we are using local attention and the current KV cache size is larger than the local attention size, we need to truncate the KV cache204 kv_cache_size = kv_cache["k"].shape[1]205 num_new_tokens = roped_query.shape[1]206 if self.local_attn_size != -1 and (current_end > kv_cache["global_end_index"].item()) and (207 num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size):208 # Calculate the number of new tokens added in this step209 # Shift existing cache content left to discard oldest tokens210 # Clone the source slice to avoid overlapping memory error211 num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size212 num_rolled_tokens = kv_cache["local_end_index"].item() - num_evicted_tokens - sink_tokens213 kv_cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \214 kv_cache["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()215 kv_cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \216 kv_cache["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()217 # Insert the new keys/values at the end218 local_end_index = kv_cache["local_end_index"].item() + current_end - \219 kv_cache["global_end_index"].item() - num_evicted_tokens220 local_start_index = local_end_index - num_new_tokens221 kv_cache["k"][:, local_start_index:local_end_index] = roped_key222 kv_cache["v"][:, local_start_index:local_end_index] = v223 else:224 # Assign new keys/values directly up to current_end225 local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item()226 local_start_index = local_end_index - num_new_tokens227 kv_cache["k"][:, local_start_index:local_end_index] = roped_key228 kv_cache["v"][:, local_start_index:local_end_index] = v229 x = attention(230 roped_query,231 kv_cache["k"][:, max(0, local_end_index - self.max_attention_size):local_end_index],232 kv_cache["v"][:, max(0, local_end_index - self.max_attention_size):local_end_index]233 )234 kv_cache["global_end_index"].fill_(current_end)235 kv_cache["local_end_index"].fill_(local_end_index)236 237 # output238 x = x.flatten(2)239 x = x.to(self.o.weight.dtype)240 x = self.o(x)241 return x242 243 244class CausalWanAttentionBlock(nn.Module):245 246 def __init__(self,247 cross_attn_type,248 dim,249 ffn_dim,250 num_heads,251 local_attn_size=-1,252 sink_size=0,253 qk_norm=True,254 cross_attn_norm=False,255 eps=1e-6):256 super().__init__()257 self.dim = dim258 self.ffn_dim = ffn_dim259 self.num_heads = num_heads260 self.local_attn_size = local_attn_size261 self.qk_norm = qk_norm262 self.cross_attn_norm = cross_attn_norm263 self.eps = eps264 265 # layers266 self.norm1 = WanLayerNorm(dim, eps)267 self.self_attn = CausalWanSelfAttention(dim, num_heads, local_attn_size, sink_size, qk_norm, eps)268 self.norm3 = WanLayerNorm(269 dim, eps,270 elementwise_affine=True) if cross_attn_norm else nn.Identity()271 self.cross_attn = WAN_CROSSATTENTION_CLASSES[cross_attn_type](dim,272 num_heads,273 (-1, -1),274 qk_norm,275 eps)276 self.norm2 = WanLayerNorm(dim, eps)277 self.ffn = nn.Sequential(278 nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'),279 nn.Linear(ffn_dim, dim))280 281 # modulation282 self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)283 284 def forward(285 self,286 x,287 e,288 seq_lens,289 grid_sizes,290 freqs,291 context,292 context_lens,293 block_mask,294 kv_cache=None,295 crossattn_cache=None,296 current_start=0,297 cache_start=None298 ):299 r"""300 Args:301 x(Tensor): Shape [B, L, C]302 e(Tensor): Shape [B, F, 6, C]303 seq_lens(Tensor): Shape [B], length of each sequence in batch304 grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)305 freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]306 """307 num_frames, frame_seqlen = e.shape[1], x.shape[1] // e.shape[1]308 # assert e.dtype == torch.float32309 # with amp.autocast(dtype=torch.float32):310 e = (self.modulation.unsqueeze(1) + e).chunk(6, dim=2)311 # assert e[0].dtype == torch.float32312 313 # self-attention314 y = self.self_attn(315 (self.norm1(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1 + e[1]) + e[0]).flatten(1, 2),316 seq_lens, grid_sizes,317 freqs, block_mask, kv_cache, current_start, cache_start)318 319 # with amp.autocast(dtype=torch.float32):320 x = x + (y.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * e[2]).flatten(1, 2)321 322 # cross-attention & ffn function323 def cross_attn_ffn(x, context, context_lens, e, crossattn_cache=None):324 x = x + self.cross_attn(self.norm3(x), context,325 context_lens, crossattn_cache=crossattn_cache)326 y = self.ffn(327 (self.norm2(x).unflatten(dim=1, sizes=(num_frames,328 frame_seqlen)) * (1 + e[4]) + e[3]).flatten(1, 2)329 )330 # with amp.autocast(dtype=torch.float32):331 x = x + (y.unflatten(dim=1, sizes=(num_frames,332 frame_seqlen)) * e[5]).flatten(1, 2)333 return x334 335 x = cross_attn_ffn(x, context, context_lens, e, crossattn_cache)336 return x337 338 339class CausalHead(nn.Module):340 341 def __init__(self, dim, out_dim, patch_size, eps=1e-6):342 super().__init__()343 self.dim = dim344 self.out_dim = out_dim345 self.patch_size = patch_size346 self.eps = eps347 348 # layers349 out_dim = math.prod(patch_size) * out_dim350 self.norm = WanLayerNorm(dim, eps)351 self.head = nn.Linear(dim, out_dim)352 353 # modulation354 self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)355 356 def forward(self, x, e):357 r"""358 Args:359 x(Tensor): Shape [B, L1, C]360 e(Tensor): Shape [B, F, 1, C]361 """362 # assert e.dtype == torch.float32363 # with amp.autocast(dtype=torch.float32):364 num_frames, frame_seqlen = e.shape[1], x.shape[1] // e.shape[1]365 e = (self.modulation.unsqueeze(1) + e).chunk(2, dim=2)366 x = (self.head(self.norm(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1 + e[1]) + e[0]))367 return x368 369 370class CausalWanModel(ModelMixin, ConfigMixin):371 r"""372 Wan diffusion backbone supporting both text-to-video and image-to-video.373 """374 375 ignore_for_config = [376 'patch_size', 'cross_attn_norm', 'qk_norm', 'text_dim'377 ]378 _no_split_modules = ['WanAttentionBlock']379 _supports_gradient_checkpointing = True380 381 @register_to_config382 def __init__(self,383 model_type='t2v',384 patch_size=(1, 2, 2),385 text_len=512,386 in_dim=16,387 dim=2048,388 ffn_dim=8192,389 freq_dim=256,390 text_dim=4096,391 out_dim=16,392 num_heads=16,393 num_layers=32,394 local_attn_size=-1,395 sink_size=0,396 qk_norm=True,397 cross_attn_norm=True,398 eps=1e-6):399 r"""400 Initialize the diffusion model backbone.401 402 Args:403 model_type (`str`, *optional*, defaults to 't2v'):404 Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video)405 patch_size (`tuple`, *optional*, defaults to (1, 2, 2)):406 3D patch dimensions for video embedding (t_patch, h_patch, w_patch)407 text_len (`int`, *optional*, defaults to 512):408 Fixed length for text embeddings409 in_dim (`int`, *optional*, defaults to 16):410 Input video channels (C_in)411 dim (`int`, *optional*, defaults to 2048):412 Hidden dimension of the transformer413 ffn_dim (`int`, *optional*, defaults to 8192):414 Intermediate dimension in feed-forward network415 freq_dim (`int`, *optional*, defaults to 256):416 Dimension for sinusoidal time embeddings417 text_dim (`int`, *optional*, defaults to 4096):418 Input dimension for text embeddings419 out_dim (`int`, *optional*, defaults to 16):420 Output video channels (C_out)421 num_heads (`int`, *optional*, defaults to 16):422 Number of attention heads423 num_layers (`int`, *optional*, defaults to 32):424 Number of transformer blocks425 local_attn_size (`int`, *optional*, defaults to -1):426 Window size for temporal local attention (-1 indicates global attention)427 sink_size (`int`, *optional*, defaults to 0):428 Size of the attention sink, we keep the first `sink_size` frames unchanged when rolling the KV cache429 qk_norm (`bool`, *optional*, defaults to True):430 Enable query/key normalization431 cross_attn_norm (`bool`, *optional*, defaults to False):432 Enable cross-attention normalization433 eps (`float`, *optional*, defaults to 1e-6):434 Epsilon value for normalization layers435 """436 437 super().__init__()438 439 assert model_type in ['t2v', 'i2v']440 self.model_type = model_type441 442 self.patch_size = patch_size443 self.text_len = text_len444 self.in_dim = in_dim445 self.dim = dim446 self.ffn_dim = ffn_dim447 self.freq_dim = freq_dim448 self.text_dim = text_dim449 self.out_dim = out_dim450 self.num_heads = num_heads451 self.num_layers = num_layers452 self.local_attn_size = local_attn_size453 self.qk_norm = qk_norm454 self.cross_attn_norm = cross_attn_norm455 self.eps = eps456 457 # embeddings458 self.patch_embedding = nn.Conv3d(459 in_dim, dim, kernel_size=patch_size, stride=patch_size)460 self.text_embedding = nn.Sequential(461 nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'),462 nn.Linear(dim, dim))463 464 self.time_embedding = nn.Sequential(465 nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))466 self.time_projection = nn.Sequential(467 nn.SiLU(), nn.Linear(dim, dim * 6))468 469 # blocks470 cross_attn_type = 't2v_cross_attn' if model_type == 't2v' else 'i2v_cross_attn'471 self.blocks = nn.ModuleList([472 CausalWanAttentionBlock(cross_attn_type, dim, ffn_dim, num_heads,473 local_attn_size, sink_size, qk_norm, cross_attn_norm, eps)474 for _ in range(num_layers)475 ])476 477 # head478 self.head = CausalHead(dim, out_dim, patch_size, eps)479 480 # buffers (don't use register_buffer otherwise dtype will be changed in to())481 assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0482 d = dim // num_heads483 self.freqs = torch.cat([484 rope_params(1024, d - 4 * (d // 6)),485 rope_params(1024, 2 * (d // 6)),486 rope_params(1024, 2 * (d // 6))487 ],488 dim=1)489 490 if model_type == 'i2v':491 self.img_emb = MLPProj(1280, dim)492 493 # initialize weights494 self.init_weights()495 496 self.gradient_checkpointing = False497 498 self.block_mask = None499 500 self.num_frame_per_block = 1501 self.independent_first_frame = False502 503 def _set_gradient_checkpointing(self, module, value=False):504 self.gradient_checkpointing = value505 506 @staticmethod507 def _prepare_blockwise_causal_attn_mask(508 device: torch.device | str, num_frames: int = 21,509 frame_seqlen: int = 1560, num_frame_per_block=1, local_attn_size=-1510 ) -> BlockMask:511 """512 we will divide the token sequence into the following format513 [1 latent frame] [1 latent frame] ... [1 latent frame]514 We use flexattention to construct the attention mask515 """516 total_length = num_frames * frame_seqlen517 518 # we do right padding to get to a multiple of 128519 padded_length = math.ceil(total_length / 128) * 128 - total_length520 521 ends = torch.zeros(total_length + padded_length,522 device=device, dtype=torch.long)523 524 # Block-wise causal mask will attend to all elements that are before the end of the current chunk525 frame_indices = torch.arange(526 start=0,527 end=total_length,528 step=frame_seqlen * num_frame_per_block,529 device=device530 )531 532 for tmp in frame_indices:533 ends[tmp:tmp + frame_seqlen * num_frame_per_block] = tmp + \534 frame_seqlen * num_frame_per_block535 536 def attention_mask(b, h, q_idx, kv_idx):537 if local_attn_size == -1:538 return (kv_idx < ends[q_idx]) | (q_idx == kv_idx)539 else:540 return ((kv_idx < ends[q_idx]) & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))) | (q_idx == kv_idx)541 # return ((kv_idx < total_length) & (q_idx < total_length)) | (q_idx == kv_idx) # bidirectional mask542 543 block_mask = create_block_mask(attention_mask, B=None, H=None, Q_LEN=total_length + padded_length,544 KV_LEN=total_length + padded_length, _compile=False, device=device)545 546 import torch.distributed as dist547 if not dist.is_initialized() or dist.get_rank() == 0:548 print(549 f" cache a block wise causal mask with block size of {num_frame_per_block} frames")550 print(block_mask)551 552 # import imageio553 # import numpy as np554 # from torch.nn.attention.flex_attention import create_mask555 556 # mask = create_mask(attention_mask, B=None, H=None, Q_LEN=total_length +557 # padded_length, KV_LEN=total_length + padded_length, device=device)558 # import cv2559 # mask = cv2.resize(mask[0, 0].cpu().float().numpy(), (1024, 1024))560 # imageio.imwrite("mask_%d.jpg" % (0), np.uint8(255. * mask))561 562 return block_mask563 564 @staticmethod565 def _prepare_teacher_forcing_mask(566 device: torch.device | str, num_frames: int = 21,567 frame_seqlen: int = 1560, num_frame_per_block=1568 ) -> BlockMask:569 """570 we will divide the token sequence into the following format571 [1 latent frame] [1 latent frame] ... [1 latent frame]572 We use flexattention to construct the attention mask573 """574 # debug575 DEBUG = False576 if DEBUG:577 num_frames = 9578 frame_seqlen = 256579 580 total_length = num_frames * frame_seqlen * 2581 582 # we do right padding to get to a multiple of 128583 padded_length = math.ceil(total_length / 128) * 128 - total_length584 585 clean_ends = num_frames * frame_seqlen586 # for clean context frames, we can construct their flex attention mask based on a [start, end] interval587 context_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)588 # for noisy frames, we need two intervals to construct the flex attention mask [context_start, context_end] [noisy_start, noisy_end]589 noise_context_starts = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)590 noise_context_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)591 noise_noise_starts = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)592 noise_noise_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)593 594 # Block-wise causal mask will attend to all elements that are before the end of the current chunk595 attention_block_size = frame_seqlen * num_frame_per_block596 frame_indices = torch.arange(597 start=0,598 end=num_frames * frame_seqlen,599 step=attention_block_size,600 device=device, dtype=torch.long601 )602 603 # attention for clean context frames604 for start in frame_indices:605 context_ends[start:start + attention_block_size] = start + attention_block_size606 607 noisy_image_start_list = torch.arange(608 num_frames * frame_seqlen, total_length,609 step=attention_block_size,610 device=device, dtype=torch.long611 )612 noisy_image_end_list = noisy_image_start_list + attention_block_size613 614 # attention for noisy frames615 for block_index, (start, end) in enumerate(zip(noisy_image_start_list, noisy_image_end_list)):616 # attend to noisy tokens within the same block617 noise_noise_starts[start:end] = start618 noise_noise_ends[start:end] = end619 # attend to context tokens in previous blocks620 # noise_context_starts[start:end] = 0621 noise_context_ends[start:end] = block_index * attention_block_size622 623 def attention_mask(b, h, q_idx, kv_idx):624 # first design the mask for clean frames625 clean_mask = (q_idx < clean_ends) & (kv_idx < context_ends[q_idx])626 # then design the mask for noisy frames627 # noisy frames will attend to all clean preceeding clean frames + itself628 C1 = (kv_idx < noise_noise_ends[q_idx]) & (kv_idx >= noise_noise_starts[q_idx])629 C2 = (kv_idx < noise_context_ends[q_idx]) & (kv_idx >= noise_context_starts[q_idx])630 noise_mask = (q_idx >= clean_ends) & (C1 | C2)631 632 eye_mask = q_idx == kv_idx633 return eye_mask | clean_mask | noise_mask634 635 block_mask = create_block_mask(attention_mask, B=None, H=None, Q_LEN=total_length + padded_length,636 KV_LEN=total_length + padded_length, _compile=False, device=device)637 638 if DEBUG:639 print(block_mask)640 import imageio641 import numpy as np642 from torch.nn.attention.flex_attention import create_mask643 644 mask = create_mask(attention_mask, B=None, H=None, Q_LEN=total_length +645 padded_length, KV_LEN=total_length + padded_length, device=device)646 import cv2647 mask = cv2.resize(mask[0, 0].cpu().float().numpy(), (1024, 1024))648 imageio.imwrite("mask_%d.jpg" % (0), np.uint8(255. * mask))649 650 return block_mask651 652 @staticmethod653 def _prepare_blockwise_causal_attn_mask_i2v(654 device: torch.device | str, num_frames: int = 21,655 frame_seqlen: int = 1560, num_frame_per_block=4, local_attn_size=-1656 ) -> BlockMask:657 """658 we will divide the token sequence into the following format659 [1 latent frame] [N latent frame] ... [N latent frame]660 The first frame is separated out to support I2V generation661 We use flexattention to construct the attention mask662 """663 total_length = num_frames * frame_seqlen664 665 # we do right padding to get to a multiple of 128666 padded_length = math.ceil(total_length / 128) * 128 - total_length667 668 ends = torch.zeros(total_length + padded_length,669 device=device, dtype=torch.long)670 671 # special handling for the first frame672 ends[:frame_seqlen] = frame_seqlen673 674 # Block-wise causal mask will attend to all elements that are before the end of the current chunk675 frame_indices = torch.arange(676 start=frame_seqlen,677 end=total_length,678 step=frame_seqlen * num_frame_per_block,679 device=device680 )681 682 for idx, tmp in enumerate(frame_indices):683 ends[tmp:tmp + frame_seqlen * num_frame_per_block] = tmp + \684 frame_seqlen * num_frame_per_block685 686 def attention_mask(b, h, q_idx, kv_idx):687 if local_attn_size == -1:688 return (kv_idx < ends[q_idx]) | (q_idx == kv_idx)689 else:690 return ((kv_idx < ends[q_idx]) & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))) | \691 (q_idx == kv_idx)692 693 block_mask = create_block_mask(attention_mask, B=None, H=None, Q_LEN=total_length + padded_length,694 KV_LEN=total_length + padded_length, _compile=False, device=device)695 696 if not dist.is_initialized() or dist.get_rank() == 0:697 print(698 f" cache a block wise causal mask with block size of {num_frame_per_block} frames")699 print(block_mask)700 701 # import imageio702 # import numpy as np703 # from torch.nn.attention.flex_attention import create_mask704 705 # mask = create_mask(attention_mask, B=None, H=None, Q_LEN=total_length +706 # padded_length, KV_LEN=total_length + padded_length, device=device)707 # import cv2708 # mask = cv2.resize(mask[0, 0].cpu().float().numpy(), (1024, 1024))709 # imageio.imwrite("mask_%d.jpg" % (0), np.uint8(255. * mask))710 711 return block_mask712 713 def _forward_inference(714 self,715 x,716 t,717 context,718 seq_len,719 clip_fea=None,720 y=None,721 kv_cache: dict = None,722 crossattn_cache: dict = None,723 current_start: int = 0,724 cache_start: int = 0725 ):726 r"""727 Run the diffusion model with kv caching.728 See Algorithm 2 of CausVid paper https://arxiv.org/abs/2412.07772 for details.729 This function will be run for num_frame times.730 Process the latent frames one by one (1560 tokens each)731 732 Args:733 x (List[Tensor]):734 List of input video tensors, each with shape [C_in, F, H, W]735 t (Tensor):736 Diffusion timesteps tensor of shape [B]737 context (List[Tensor]):738 List of text embeddings each with shape [L, C]739 seq_len (`int`):740 Maximum sequence length for positional encoding741 clip_fea (Tensor, *optional*):742 CLIP image features for image-to-video mode743 y (List[Tensor], *optional*):744 Conditional video inputs for image-to-video mode, same shape as x745 746 Returns:747 List[Tensor]:748 List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8]749 """750 751 if self.model_type == 'i2v':752 assert clip_fea is not None and y is not None753 # params754 device = self.patch_embedding.weight.device755 if self.freqs.device != device:756 self.freqs = self.freqs.to(device)757 758 if y is not None:759 x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]760 761 # embeddings762 x = [self.patch_embedding(u.unsqueeze(0)) for u in x]763 grid_sizes = torch.stack(764 [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])765 x = [u.flatten(2).transpose(1, 2) for u in x]766 seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)767 assert seq_lens.max() <= seq_len768 x = torch.cat(x)769 """770 torch.cat([771 torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))],772 dim=1) for u in x773 ])774 """775 776 # time embeddings777 # with amp.autocast(dtype=torch.float32):778 e = self.time_embedding(779 sinusoidal_embedding_1d(self.freq_dim, t.flatten()).type_as(x))780 e0 = self.time_projection(e).unflatten(781 1, (6, self.dim)).unflatten(dim=0, sizes=t.shape)782 # assert e.dtype == torch.float32 and e0.dtype == torch.float32783 784 # context785 context_lens = None786 context = self.text_embedding(787 torch.stack([788 torch.cat(789 [u, u.new_zeros(self.text_len - u.size(0), u.size(1))])790 for u in context791 ]))792 793 if clip_fea is not None:794 context_clip = self.img_emb(clip_fea) # bs x 257 x dim795 context = torch.concat([context_clip, context], dim=1)796 797 # arguments798 kwargs = dict(799 e=e0,800 seq_lens=seq_lens,801 grid_sizes=grid_sizes,802 freqs=self.freqs,803 context=context,804 context_lens=context_lens,805 block_mask=self.block_mask806 )807 808 def create_custom_forward(module):809 def custom_forward(*inputs, **kwargs):810 return module(*inputs, **kwargs)811 return custom_forward812 813 for block_index, block in enumerate(self.blocks):814 if torch.is_grad_enabled() and self.gradient_checkpointing:815 kwargs.update(816 {817 "kv_cache": kv_cache[block_index],818 "current_start": current_start,819 "cache_start": cache_start820 }821 )822 x = torch.utils.checkpoint.checkpoint(823 create_custom_forward(block),824 x, **kwargs,825 use_reentrant=False,826 )827 else:828 kwargs.update(829 {830 "kv_cache": kv_cache[block_index],831 "crossattn_cache": crossattn_cache[block_index],832 "current_start": current_start,833 "cache_start": cache_start834 }835 )836 x = block(x, **kwargs)837 838 # head839 x = self.head(x, e.unflatten(dim=0, sizes=t.shape).unsqueeze(2))840 # unpatchify841 x = self.unpatchify(x, grid_sizes)842 return torch.stack(x)843 844 def _forward_train(845 self,846 x,847 t,848 context,849 seq_len,850 clean_x=None,851 aug_t=None,852 clip_fea=None,853 y=None,854 ):855 r"""856 Forward pass through the diffusion model857 858 Args:859 x (List[Tensor]):860 List of input video tensors, each with shape [C_in, F, H, W]861 t (Tensor):862 Diffusion timesteps tensor of shape [B]863 context (List[Tensor]):864 List of text embeddings each with shape [L, C]865 seq_len (`int`):866 Maximum sequence length for positional encoding867 clip_fea (Tensor, *optional*):868 CLIP image features for image-to-video mode869 y (List[Tensor], *optional*):870 Conditional video inputs for image-to-video mode, same shape as x871 872 Returns:873 List[Tensor]:874 List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8]875 """876 if self.model_type == 'i2v':877 assert clip_fea is not None and y is not None878 # params879 device = self.patch_embedding.weight.device880 if self.freqs.device != device:881 self.freqs = self.freqs.to(device)882 883 # Construct blockwise causal attn mask884 if self.block_mask is None:885 if clean_x is not None:886 if self.independent_first_frame:887 raise NotImplementedError()888 else:889 self.block_mask = self._prepare_teacher_forcing_mask(890 device, num_frames=x.shape[2],891 frame_seqlen=x.shape[-2] * x.shape[-1] // (self.patch_size[1] * self.patch_size[2]),892 num_frame_per_block=self.num_frame_per_block893 )894 else:895 if self.independent_first_frame:896 self.block_mask = self._prepare_blockwise_causal_attn_mask_i2v(897 device, num_frames=x.shape[2],898 frame_seqlen=x.shape[-2] * x.shape[-1] // (self.patch_size[1] * self.patch_size[2]),899 num_frame_per_block=self.num_frame_per_block,900 local_attn_size=self.local_attn_size901 )902 else:903 self.block_mask = self._prepare_blockwise_causal_attn_mask(904 device, num_frames=x.shape[2],905 frame_seqlen=x.shape[-2] * x.shape[-1] // (self.patch_size[1] * self.patch_size[2]),906 num_frame_per_block=self.num_frame_per_block,907 local_attn_size=self.local_attn_size908 )909 910 if y is not None:911 x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]912 913 # embeddings914 x = [self.patch_embedding(u.unsqueeze(0)) for u in x]915 916 grid_sizes = torch.stack(917 [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])918 x = [u.flatten(2).transpose(1, 2) for u in x]919 920 seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)921 assert seq_lens.max() <= seq_len922 x = torch.cat([923 torch.cat([u, u.new_zeros(1, seq_lens[0] - u.size(1), u.size(2))],924 dim=1) for u in x925 ])926 927 # time embeddings928 # with amp.autocast(dtype=torch.float32):929 e = self.time_embedding(930 sinusoidal_embedding_1d(self.freq_dim, t.flatten()).type_as(x))931 e0 = self.time_projection(e).unflatten(932 1, (6, self.dim)).unflatten(dim=0, sizes=t.shape)933 # assert e.dtype == torch.float32 and e0.dtype == torch.float32934 935 # context936 context_lens = None937 context = self.text_embedding(938 torch.stack([939 torch.cat(940 [u, u.new_zeros(self.text_len - u.size(0), u.size(1))])941 for u in context942 ]))943 944 if clip_fea is not None:945 context_clip = self.img_emb(clip_fea) # bs x 257 x dim946 context = torch.concat([context_clip, context], dim=1)947 948 if clean_x is not None:949 clean_x = [self.patch_embedding(u.unsqueeze(0)) for u in clean_x]950 clean_x = [u.flatten(2).transpose(1, 2) for u in clean_x]951 952 seq_lens_clean = torch.tensor([u.size(1) for u in clean_x], dtype=torch.long)953 assert seq_lens_clean.max() <= seq_len954 clean_x = torch.cat([955 torch.cat([u, u.new_zeros(1, seq_lens_clean[0] - u.size(1), u.size(2))], dim=1) for u in clean_x956 ])957 958 x = torch.cat([clean_x, x], dim=1)959 if aug_t is None:960 aug_t = torch.zeros_like(t)961 e_clean = self.time_embedding(962 sinusoidal_embedding_1d(self.freq_dim, aug_t.flatten()).type_as(x))963 e0_clean = self.time_projection(e_clean).unflatten(964 1, (6, self.dim)).unflatten(dim=0, sizes=t.shape)965 e0 = torch.cat([e0_clean, e0], dim=1)966 967 # arguments968 kwargs = dict(969 e=e0,970 seq_lens=seq_lens,971 grid_sizes=grid_sizes,972 freqs=self.freqs,973 context=context,974 context_lens=context_lens,975 block_mask=self.block_mask)976 977 def create_custom_forward(module):978 def custom_forward(*inputs, **kwargs):979 return module(*inputs, **kwargs)980 return custom_forward981 982 for block in self.blocks:983 if torch.is_grad_enabled() and self.gradient_checkpointing:984 x = torch.utils.checkpoint.checkpoint(985 create_custom_forward(block),986 x, **kwargs,987 use_reentrant=False,988 )989 else:990 x = block(x, **kwargs)991 992 if clean_x is not None:993 x = x[:, x.shape[1] // 2:]994 995 # head996 x = self.head(x, e.unflatten(dim=0, sizes=t.shape).unsqueeze(2))997 998 # unpatchify999 x = self.unpatchify(x, grid_sizes)1000 return torch.stack(x)1001 1002 def forward(1003 self,1004 *args,1005 **kwargs1006 ):1007 if kwargs.get('kv_cache', None) is not None:1008 return self._forward_inference(*args, **kwargs)1009 else:1010 return self._forward_train(*args, **kwargs)1011 1012 def unpatchify(self, x, grid_sizes):1013 r"""1014 Reconstruct video tensors from patch embeddings.1015 1016 Args:1017 x (List[Tensor]):1018 List of patchified features, each with shape [L, C_out * prod(patch_size)]1019 grid_sizes (Tensor):1020 Original spatial-temporal grid dimensions before patching,1021 shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches)1022 1023 Returns:1024 List[Tensor]:1025 Reconstructed video tensors with shape [C_out, F, H / 8, W / 8]1026 """1027 1028 c = self.out_dim1029 out = []1030 for u, v in zip(x, grid_sizes.tolist()):1031 u = u[:math.prod(v)].view(*v, *self.patch_size, c)1032 u = torch.einsum('fhwpqrc->cfphqwr', u)1033 u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)])1034 out.append(u)1035 return out1036 1037 def init_weights(self):1038 r"""1039 Initialize model parameters using Xavier initialization.1040 """1041 1042 # basic init1043 for m in self.modules():1044 if isinstance(m, nn.Linear):1045 nn.init.xavier_uniform_(m.weight)1046 if m.bias is not None:1047 nn.init.zeros_(m.bias)1048 1049 # init embeddings1050 nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1))1051 for m in self.text_embedding.modules():1052 if isinstance(m, nn.Linear):1053 nn.init.normal_(m.weight, std=.02)1054 for m in self.time_embedding.modules():1055 if isinstance(m, nn.Linear):1056 nn.init.normal_(m.weight, std=.02)1057 1058 # init output layer1059 nn.init.zeros_(self.head.head.weight)1060