import math import torch import torch.distributed as dist 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 BlockMask, create_block_mask, flex_attention from wan.distributed.util import all_to_all, gather_forward from wan.modules.attention import attention from wan.modules.model import ( WAN_CROSSATTENTION_CLASSES, MLPProj, WanLayerNorm, WanRMSNorm, rope_apply, rope_params, sinusoidal_embedding_1d, ) # wan 1.3B model has a weird channel / head configurations and require max-autotune to work with flexattention # see https://github.com/pytorch/pytorch/issues/133254 # change to default for other models flex_attention = torch.compile(flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs") def causal_rope_apply(x, grid_sizes, freqs, start_frame=0): n, c = x.size(2), x.size(3) // 2 # split freqs freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) # loop over samples output = [] for i, (f, h, w) in enumerate(grid_sizes.tolist()): seq_len = f * h * w # precompute multipliers x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape(seq_len, n, -1, 2)) freqs_i = torch.cat( [ freqs[0][start_frame : start_frame + f].view(f, 1, 1, -1).expand(f, h, w, -1), freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1), ], dim=-1, ).reshape(seq_len, 1, -1) # apply rotary embedding x_i = torch.view_as_real(x_i * freqs_i).flatten(2) x_i = torch.cat([x_i, x[i, seq_len:]]) # append to collection output.append(x_i) return torch.stack(output).type_as(x) def _apply_rotary_emb(x, cos, sin, is_neox_style=False): cos = cos.unsqueeze(-2) sin = sin.unsqueeze(-2) if is_neox_style: x1, x2 = torch.chunk(x, 2, dim=-1) else: x1 = x[..., ::2] x2 = x[..., 1::2] o1 = (x1.float() * cos - x2.float() * sin).type_as(x) o2 = (x2.float() * cos + x1.float() * sin).type_as(x) if is_neox_style: return torch.cat((o1, o2), dim=-1) return torch.stack((o1, o2), dim=-1).flatten(-2) def _get_1d_rotary_pos_embed(dim, pos, theta=10000.0, dtype=torch.float64): if isinstance(pos, int): pos = torch.arange(pos, dtype=torch.float32) freqs = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=dtype) / dim)) freqs = torch.outer(pos.to(dtype), freqs) return freqs.cos(), freqs.sin() def get_causal_rotary_pos_embed( rope_sizes, hidden_size, heads_num, rope_dim_list, rope_theta=10000.0, dtype=torch.float64, start_frame=0, ): t, h, w = rope_sizes t_pos = torch.arange(t, dtype=torch.float32) + start_frame h_pos = torch.arange(h, dtype=torch.float32) w_pos = torch.arange(w, dtype=torch.float32) grid_t, grid_h, grid_w = torch.meshgrid(t_pos, h_pos, w_pos, indexing="ij") cos_t, sin_t = _get_1d_rotary_pos_embed(rope_dim_list[0], grid_t.reshape(-1), theta=rope_theta, dtype=dtype) cos_h, sin_h = _get_1d_rotary_pos_embed(rope_dim_list[1], grid_h.reshape(-1), theta=rope_theta, dtype=dtype) cos_w, sin_w = _get_1d_rotary_pos_embed(rope_dim_list[2], grid_w.reshape(-1), theta=rope_theta, dtype=dtype) return torch.cat([cos_t, cos_h, cos_w], dim=1), torch.cat([sin_t, sin_h, sin_w], dim=1) def update_kv_cache_dict(kv_cache, current_start, key, value, frame_seqlen, sink_size=0, local_attn_size=-1): num_new_tokens = key.shape[1] current_end = current_start + num_new_tokens kv_cache_size = kv_cache["k"].shape[1] sink_tokens = sink_size * frame_seqlen if ( local_attn_size != -1 and current_end > kv_cache["global_end_index"].item() and kv_cache["local_end_index"].item() + num_new_tokens > 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 with torch.no_grad(): kv_cache["k"][:, sink_tokens : sink_tokens + num_rolled_tokens] = kv_cache["k"][ :, sink_tokens + num_evicted_tokens : sink_tokens + num_evicted_tokens + num_rolled_tokens ].clone() kv_cache["v"][:, sink_tokens : sink_tokens + num_rolled_tokens] = kv_cache["v"][ :, sink_tokens + num_evicted_tokens : sink_tokens + num_evicted_tokens + num_rolled_tokens ].clone() local_end_index = ( kv_cache["local_end_index"].item() + (current_end - kv_cache["global_end_index"].item()) - num_evicted_tokens ) 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] = key kv_cache["v"][:, local_start_index:local_end_index] = value return_k = kv_cache["k"][:, max(0, local_end_index - kv_cache_size) : local_end_index] return_v = kv_cache["v"][:, max(0, local_end_index - kv_cache_size) : local_end_index] kv_cache["global_end_index"].fill_(current_end) kv_cache["local_end_index"].fill_(local_end_index) return return_k, return_v class CausalWanSelfAttention(nn.Module): def __init__(self, dim, num_heads, 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.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() self.use_sp = False def _sp_attention(self, query, key, value): if not dist.is_initialized(): raise RuntimeError("Sequence parallel attention requires an initialized torch.distributed process group.") world_size = dist.get_world_size() if world_size <= 1: raise ValueError("Sequence parallel attention requires WORLD_SIZE > 1.") if self.num_heads % world_size != 0: raise ValueError( f"Sequence parallel attention requires `num_heads` ({self.num_heads}) " f"to be divisible by WORLD_SIZE ({world_size})." ) query_len = query.shape[1] key_len = key.shape[1] value_len = value.shape[1] if query_len % world_size != 0: raise ValueError( f"Sequence parallel attention requires query length {query_len} " f"to be divisible by WORLD_SIZE ({world_size})." ) if key_len % world_size != 0: raise ValueError( f"Sequence parallel attention requires key length {key_len} " f"to be divisible by WORLD_SIZE ({world_size})." ) if value_len % world_size != 0: raise ValueError( f"Sequence parallel attention requires value length {value_len} " f"to be divisible by WORLD_SIZE ({world_size})." ) rank = dist.get_rank() query_local = torch.chunk(query, world_size, dim=1)[rank].contiguous() key_local = torch.chunk(key, world_size, dim=1)[rank].contiguous() value_local = torch.chunk(value, world_size, dim=1)[rank].contiguous() seq_lens = torch.full((query.shape[0],), key_len, dtype=torch.long, device=query.device) query_full = all_to_all(query_local, scatter_dim=2, gather_dim=1) key_full = all_to_all(key_local, scatter_dim=2, gather_dim=1) value_full = all_to_all(value_local, scatter_dim=2, gather_dim=1) output_full = attention(query_full, key_full, value_full) output_local = all_to_all(output_full, scatter_dim=1, gather_dim=2) return gather_forward(output_local, dim=1) def forward( self, x, seq_lens, grid_sizes, freqs, block_mask, sink_size, local_attn_size, max_attention_size, kv_cache=None, current_start=0, cache_start=None, ): r""" 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 # query, key, value function 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: # if it is teacher forcing training? is_tf = s == seq_lens[0].item() * 2 if is_tf: q_chunk = torch.chunk(q, 2, dim=1) k_chunk = torch.chunk(k, 2, dim=1) roped_query = [] roped_key = [] # rope should be same for clean and noisy parts 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, )[:, :, :-padded_length].transpose(2, 1) else: roped_query = rope_apply(q, grid_sizes, freqs).type_as(v) roped_key = rope_apply(k, grid_sizes, freqs).type_as(v) 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, )[:, :, :-padded_length].transpose(2, 1) else: frame_seqlen = math.prod(grid_sizes[0][1:]).item() current_start_frame = current_start // frame_seqlen if isinstance(freqs, tuple): cache_k, cache_v = update_kv_cache_dict( kv_cache, current_start, k, v, frame_seqlen=frame_seqlen, sink_size=sink_size, local_attn_size=local_attn_size, ) cos, sin = freqs if local_attn_size != -1: cache_len = cache_k.shape[1] query_len = q.shape[1] roped_cache_k = _apply_rotary_emb( cache_k, cos[:cache_len], sin[:cache_len], is_neox_style=False, ).type_as(v) roped_query = _apply_rotary_emb( q, cos[cache_len - query_len : cache_len], sin[cache_len - query_len : cache_len], is_neox_style=False, ).type_as(v) else: roped_cache_k = _apply_rotary_emb( cache_k, cos[: cache_k.shape[1]], sin[: cache_k.shape[1]], is_neox_style=False, ).type_as(v) roped_query = _apply_rotary_emb( q, cos[current_start : current_start + q.shape[1]], sin[current_start : current_start + q.shape[1]], is_neox_style=False, ).type_as(v) if self.use_sp: x = self._sp_attention(roped_query, roped_cache_k, cache_v) else: x = attention(roped_query, roped_cache_k, cache_v) else: current_end = current_start + q.shape[1] sink_tokens = sink_size * frame_seqlen kv_cache_size = kv_cache["k"].shape[1] num_new_tokens = q.shape[1] if local_attn_size == -1: cache_key = causal_rope_apply(k, grid_sizes, freqs, start_frame=current_start_frame).type_as(v) else: cache_key = k if ( 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 kv_cache["k"][:, sink_tokens : sink_tokens + num_rolled_tokens] = kv_cache["k"][ :, sink_tokens + num_evicted_tokens : sink_tokens + num_evicted_tokens + num_rolled_tokens ].clone() kv_cache["v"][:, sink_tokens : sink_tokens + num_rolled_tokens] = kv_cache["v"][ :, sink_tokens + num_evicted_tokens : sink_tokens + num_evicted_tokens + num_rolled_tokens ].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] = cache_key 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] = cache_key kv_cache["v"][:, local_start_index:local_end_index] = v cache_start_index = max(0, local_end_index - max_attention_size) cache_k = kv_cache["k"][:, cache_start_index:local_end_index] cache_v = kv_cache["v"][:, cache_start_index:local_end_index] if local_attn_size != -1: cache_num_frames = cache_k.shape[1] // frame_seqlen query_num_frames = grid_sizes[0][0].item() cache_grid_sizes = grid_sizes.clone() cache_grid_sizes[:, 0] = cache_num_frames roped_cache_k = causal_rope_apply( cache_k, cache_grid_sizes, freqs, start_frame=0, ).type_as(v) roped_query = causal_rope_apply( q, grid_sizes, freqs, start_frame=cache_num_frames - query_num_frames, ).type_as(v) else: roped_query = causal_rope_apply(q, grid_sizes, freqs, start_frame=current_start_frame).type_as(v) roped_cache_k = cache_k if self.use_sp: x = self._sp_attention(roped_query, roped_cache_k, cache_v) else: x = attention(roped_query, roped_cache_k, cache_v) kv_cache["global_end_index"].fill_(current_end) kv_cache["local_end_index"].fill_(local_end_index) # output x = x.flatten(2) x = self.o(x) return x class CausalWanAttentionBlock(nn.Module): def __init__(self, cross_attn_type, dim, ffn_dim, num_heads, qk_norm=True, cross_attn_norm=False, eps=1e-6): 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 # layers self.norm1 = WanLayerNorm(dim, eps) self.self_attn = CausalWanSelfAttention(dim, num_heads, qk_norm, eps) self.norm3 = WanLayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() self.cross_attn = WAN_CROSSATTENTION_CLASSES[cross_attn_type](dim, num_heads, (-1, -1), qk_norm, eps) self.norm2 = WanLayerNorm(dim, eps) self.ffn = nn.Sequential(nn.Linear(dim, ffn_dim), nn.GELU(approximate="tanh"), nn.Linear(ffn_dim, dim)) # modulation 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, sink_size, local_attn_size, max_attention_size, kv_cache=None, crossattn_cache=None, current_start=0, cache_start=None, ): r""" 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] # assert e.dtype == torch.float32 # with amp.autocast(dtype=torch.float32): e = (self.modulation.unsqueeze(1) + e).chunk(6, dim=2) # assert e[0].dtype == torch.float32 # self-attention 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, sink_size, local_attn_size, max_attention_size, kv_cache, current_start, cache_start, ) # with amp.autocast(dtype=torch.float32): x = x + (y.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * e[2]).flatten(1, 2) # cross-attention & ffn function 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) ) # with amp.autocast(dtype=torch.float32): 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-6): super().__init__() self.dim = dim self.out_dim = out_dim self.patch_size = patch_size self.eps = eps # layers out_dim = math.prod(patch_size) * out_dim self.norm = WanLayerNorm(dim, eps) self.head = nn.Linear(dim, out_dim) # modulation self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) def forward(self, x, e): r""" Args: x(Tensor): Shape [B, L1, C] e(Tensor): Shape [B, F, 1, C] """ # assert e.dtype == torch.float32 # with amp.autocast(dtype=torch.float32): num_frames, frame_seqlen = e.shape[1], x.shape[1] // e.shape[1] e = (self.modulation.unsqueeze(1) + e).chunk(2, dim=2) x = self.head(self.norm(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1 + e[1]) + e[0]) return x class CausalWanModel(ModelMixin, ConfigMixin): r""" 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 = ["WanAttentionBlock"] _supports_gradient_checkpointing = True @register_to_config 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, sink_size=0, local_attn_size=-1, max_attention_size=32760, qk_norm=True, cross_attn_norm=True, eps=1e-6, ): r""" Initialize the diffusion model backbone. Args: model_type (`str`, *optional*, defaults to 't2v'): Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video) 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 -1): 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 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", "i2v", "ti2v", "s2v"] self.model_type = model_type # initialization # will be updated according to config file by model later self.gradient_checkpointing = False self.block_mask = None self.num_frame_per_block = 1 self.independent_first_frame = False self.sink_size = sink_size self.local_attn_size = local_attn_size self.max_attention_size = max_attention_size 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.qk_norm = qk_norm self.cross_attn_norm = cross_attn_norm self.eps = eps # embeddings 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)) # blocks if model_type == "ti2v": cross_attn_type = "ti2v_cross_attn" elif model_type == "t2v": cross_attn_type = "t2v_cross_attn" elif model_type == "i2v": cross_attn_type = "i2v_cross_attn" else: raise ValueError(f"Invalid model type: {model_type}") self.blocks = nn.ModuleList( [ CausalWanAttentionBlock(cross_attn_type, dim, ffn_dim, num_heads, qk_norm, cross_attn_norm, eps) for _ in range(num_layers) ] ) # head self.head = CausalHead(dim, out_dim, patch_size, eps) # buffers (don't use register_buffer otherwise dtype will be changed in to()) 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.rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] if model_type == "i2v": self.img_emb = MLPProj(1280, dim) # initialize weights self.init_weights() def _set_gradient_checkpointing(self, module, value=False): self.gradient_checkpointing = value @staticmethod def _prepare_blockwise_causal_attn_mask( device: torch.device | str, num_frames: int = 21, frame_seqlen: int = 1560, num_frame_per_block=1, local_attn_size=-1, ) -> BlockMask: """ we will divide the token sequence into the following format [1 latent frame] [1 latent frame] ... [1 latent frame] We use flexattention to construct the attention mask """ total_length = num_frames * frame_seqlen # we do right padding to get to a multiple of 128 padded_length = math.ceil(total_length / 128) * 128 - total_length ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) # Block-wise causal mask will attend to all elements that are before the end of the current chunk frame_indices = torch.arange(start=0, end=total_length, step=frame_seqlen * num_frame_per_block, device=device) for tmp in frame_indices: ends[tmp : tmp + frame_seqlen * num_frame_per_block] = tmp + frame_seqlen * num_frame_per_block def attention_mask(b, h, q_idx, kv_idx): if local_attn_size == -1: return (kv_idx < ends[q_idx]) | (q_idx == kv_idx) else: return ((kv_idx < ends[q_idx]) & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))) | ( q_idx == kv_idx ) # return ((kv_idx < total_length) & (q_idx < total_length)) | (q_idx == kv_idx) # bidirectional mask block_mask = create_block_mask( attention_mask, B=None, H=None, Q_LEN=total_length + padded_length, KV_LEN=total_length + padded_length, _compile=False, device=device, ) import torch.distributed as dist if not dist.is_initialized() or dist.get_rank() == 0: print(f" cache a block wise causal mask with block size of {num_frame_per_block} frames") print(block_mask) # import imageio # import numpy as np # from torch.nn.attention.flex_attention import create_mask # mask = create_mask(attention_mask, B=None, H=None, Q_LEN=total_length + # padded_length, KV_LEN=total_length + padded_length, device=device) # import cv2 # mask = cv2.resize(mask[0, 0].cpu().float().numpy(), (1024, 1024)) # imageio.imwrite("mask_%d.jpg" % (0), np.uint8(255. * mask)) return block_mask @staticmethod def _prepare_teacher_forcing_mask( device: torch.device | str, num_frames: int = 21, frame_seqlen: int = 1560, num_frame_per_block=1 ) -> BlockMask: """ we will divide the token sequence into the following format [1 latent frame] [1 latent frame] ... [1 latent frame] We use flexattention to construct the attention mask """ # debug DEBUG = False if DEBUG: num_frames = 9 frame_seqlen = 256 total_length = num_frames * frame_seqlen * 2 # we do right padding to get to a multiple of 128 padded_length = math.ceil(total_length / 128) * 128 - total_length clean_ends = num_frames * frame_seqlen # for clean context frames, we can construct their flex attention mask based on a [start, end] interval context_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) # for noisy frames, we need two intervals to construct the flex attention mask [context_start, context_end] [noisy_start, noisy_end] noise_context_starts = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) noise_context_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) noise_noise_starts = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) noise_noise_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) # Block-wise causal mask will attend to all elements that are before the end of the current chunk attention_block_size = frame_seqlen * num_frame_per_block frame_indices = torch.arange( start=0, end=num_frames * frame_seqlen, step=attention_block_size, device=device, dtype=torch.long ) # attention for clean context frames for start in frame_indices: context_ends[start : start + attention_block_size] = start + attention_block_size noisy_image_start_list = torch.arange( num_frames * frame_seqlen, total_length, step=attention_block_size, device=device, dtype=torch.long ) noisy_image_end_list = noisy_image_start_list + attention_block_size # attention for noisy frames for block_index, (start, end) in enumerate(zip(noisy_image_start_list, noisy_image_end_list)): # attend to noisy tokens within the same block noise_noise_starts[start:end] = start noise_noise_ends[start:end] = end # attend to context tokens in previous blocks # noise_context_starts[start:end] = 0 noise_context_ends[start:end] = block_index * attention_block_size def attention_mask(b, h, q_idx, kv_idx): # first design the mask for clean frames clean_mask = (q_idx < clean_ends) & (kv_idx < context_ends[q_idx]) # then design the mask for noisy frames # noisy frames will attend to all clean preceding clean frames + itself C1 = (kv_idx < noise_noise_ends[q_idx]) & (kv_idx >= noise_noise_starts[q_idx]) C2 = (kv_idx < noise_context_ends[q_idx]) & (kv_idx >= noise_context_starts[q_idx]) noise_mask = (q_idx >= clean_ends) & (C1 | C2) eye_mask = q_idx == kv_idx return eye_mask | clean_mask | noise_mask block_mask = create_block_mask( attention_mask, B=None, H=None, Q_LEN=total_length + padded_length, KV_LEN=total_length + padded_length, _compile=False, device=device, ) if DEBUG: print(block_mask) import imageio import numpy as np from torch.nn.attention.flex_attention import create_mask mask = create_mask( attention_mask, B=None, H=None, Q_LEN=total_length + padded_length, KV_LEN=total_length + padded_length, device=device, ) import cv2 mask = cv2.resize(mask[0, 0].cpu().float().numpy(), (1024, 1024)) imageio.imwrite("mask_%d.jpg" % (0), np.uint8(255.0 * mask)) return block_mask @staticmethod def _prepare_blockwise_causal_attn_mask_i2v( device: torch.device | str, num_frames: int = 21, frame_seqlen: int = 1560, num_frame_per_block=4, local_attn_size=-1, ) -> BlockMask: """ we will divide the token sequence into the following format [1 latent frame] [N latent frame] ... [N latent frame] The first frame is separated out to support I2V generation We use flexattention to construct the attention mask """ total_length = num_frames * frame_seqlen # we do right padding to get to a multiple of 128 padded_length = math.ceil(total_length / 128) * 128 - total_length ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) # special handling for the first frame ends[:frame_seqlen] = frame_seqlen # Block-wise causal mask will attend to all elements that are before the end of the current chunk frame_indices = torch.arange( start=frame_seqlen, end=total_length, step=frame_seqlen * num_frame_per_block, device=device ) for idx, tmp in enumerate(frame_indices): ends[tmp : tmp + frame_seqlen * num_frame_per_block] = tmp + frame_seqlen * num_frame_per_block def attention_mask(b, h, q_idx, kv_idx): if local_attn_size == -1: return (kv_idx < ends[q_idx]) | (q_idx == kv_idx) else: return ((kv_idx < ends[q_idx]) & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))) | ( q_idx == kv_idx ) block_mask = create_block_mask( attention_mask, B=None, H=None, Q_LEN=total_length + padded_length, KV_LEN=total_length + padded_length, _compile=False, device=device, ) if not dist.is_initialized() or dist.get_rank() == 0: print(f" cache a block wise causal mask with block size of {num_frame_per_block} frames") print(block_mask) # import imageio # import numpy as np # from torch.nn.attention.flex_attention import create_mask # mask = create_mask(attention_mask, B=None, H=None, Q_LEN=total_length + # padded_length, KV_LEN=total_length + padded_length, device=device) # import cv2 # mask = cv2.resize(mask[0, 0].cpu().float().numpy(), (1024, 1024)) # imageio.imwrite("mask_%d.jpg" % (0), np.uint8(255. * mask)) return block_mask def _forward_inference( self, x, t, context, seq_len, clip_fea=None, y=None, kv_cache: dict = None, crossattn_cache: dict = None, current_start: int = 0, cache_start: int = 0, return_feats: bool | list[int] = False, ): r""" 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] context (List[Tensor]): List of text embeddings each with shape [L, C] seq_len (`int`): Maximum sequence length for positional encoding 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 self.model_type == "i2v": assert clip_fea is not None and y is not None # params device = self.patch_embedding.weight.device p_t, p_h, p_w = self.patch_size post_patch_num_frames = x[0].shape[1] // p_t post_patch_height = x[0].shape[2] // p_h post_patch_width = x[0].shape[3] // p_w frame_seqlen = post_patch_height * post_patch_width current_start_frame = current_start // frame_seqlen if kv_cache is not None: max_cache_len = kv_cache[0]["k"].shape[1] if self.local_attn_size != -1: total_frames = (max_cache_len // frame_seqlen) + post_patch_num_frames else: total_frames = current_start_frame + post_patch_num_frames freqs = get_causal_rotary_pos_embed( (total_frames, post_patch_height, post_patch_width), self.dim, self.num_heads, self.rope_dim_list, rope_theta=10000.0, dtype=torch.float64, start_frame=0, ) freqs = tuple(u.to(device) for u in freqs) else: if self.freqs.device != device: self.freqs = self.freqs.to(device) freqs = self.freqs if y is not None: x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)] # embeddings 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) assert seq_lens.max() <= seq_len x = torch.cat(x) """ torch.cat([ torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1) for u in x ]) """ # time embeddings # with amp.autocast(dtype=torch.float32): 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) # assert e.dtype == torch.float32 and e0.dtype == torch.float32 # context 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]) ) if clip_fea is not None: context_clip = self.img_emb(clip_fea) # bs x 257 x dim context = torch.concat([context_clip, context], dim=1) # arguments kwargs = { "e": e0, "seq_lens": seq_lens, "grid_sizes": grid_sizes, "freqs": freqs, "context": context, "context_lens": context_lens, "block_mask": self.block_mask, "sink_size": self.sink_size, "local_attn_size": self.local_attn_size, "max_attention_size": self.max_attention_size, } def create_custom_forward(module): def custom_forward(*inputs, **kwargs): return module(*inputs, **kwargs) return custom_forward feats: list[torch.Tensor] = [] feature_indices = return_feats if isinstance(return_feats, list) else [] 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) if isinstance(return_feats, bool) and return_feats: feats.append(x) elif block_index in feature_indices: feats.append(x) # if not dist.is_initialized() or dist.get_rank() == 0: # print( # f"seq_lens={seq_lens}&{seq_len}, grid_sizes={grid_sizes}," # f"freqs={self.freqs.shape}, context_shape={context.shape}," # f"sink_size={self.sink_size}, local_attn_size={self.local_attn_size}," # f"max_attention_size={self.max_attention_size}" # ) # head x = self.head(x, e.unflatten(dim=0, sizes=t.shape).unsqueeze(2)) # unpatchify x = self.unpatchify(x, grid_sizes) output = torch.stack(x) if return_feats: return output, feats return output def _forward_train( self, x, t, context, seq_len, clean_x=None, aug_t=None, clip_fea=None, y=None, ): r""" 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] context (List[Tensor]): List of text embeddings each with shape [L, C] seq_len (`int`): Maximum sequence length for positional encoding 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 self.model_type == "i2v": assert clip_fea is not None and y is not None # params device = self.patch_embedding.weight.device if self.freqs.device != device: self.freqs = self.freqs.to(device) # Construct blockwise causal attn mask if self.block_mask is None: if clean_x is not None: if self.independent_first_frame: raise NotImplementedError() else: 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, ) else: if self.independent_first_frame: self.block_mask = self._prepare_blockwise_causal_attn_mask_i2v( 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, ) else: self.block_mask = self._prepare_blockwise_causal_attn_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, ) if y is not None: x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)] # embeddings 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) assert seq_lens.max() <= seq_len 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]) # time embeddings # with amp.autocast(dtype=torch.float32): 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) # assert e.dtype == torch.float32 and e0.dtype == torch.float32 # context 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]) ) if clip_fea is not None: context_clip = self.img_emb(clip_fea) # bs x 257 x dim context = torch.concat([context_clip, context], dim=1) if clean_x is not None: clean_x = [self.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) assert seq_lens_clean.max() <= seq_len 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) if aug_t is None: 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) # arguments 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, "sink_size": self.sink_size, "local_attn_size": self.local_attn_size, "max_attention_size": self.max_attention_size, } 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) # if not dist.is_initialized() or dist.get_rank() == 0: # print( # f"seq_lens={seq_lens}&{seq_len}, grid_sizes={grid_sizes}," # f"freqs={self.freqs.shape}, context_shape={context.shape}," # f"sink_size={self.sink_size}, local_attn_size={self.local_attn_size}," # f"max_attention_size={self.max_attention_size}" # ) if clean_x is not None: x = x[:, x.shape[1] // 2 :] # head x = self.head(x, e.unflatten(dim=0, sizes=t.shape).unsqueeze(2)) # unpatchify 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): r""" 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()): 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)]) out.append(u) return out def init_weights(self): r""" Initialize model parameters using Xavier initialization. """ # basic init 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) # init embeddings 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) # init output layer nn.init.zeros_(self.head.head.weight)