diff --git a/fastvideo/models/dits/causal_wanvideo.py b/fastvideo/models/dits/causal_wanvideo.py index a833e740f..08d86e855 100644 --- a/fastvideo/models/dits/causal_wanvideo.py +++ b/fastvideo/models/dits/causal_wanvideo.py @@ -126,26 +126,26 @@ class CausalWanSelfAttention(nn.Module): # 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 cache kv_cache_size = kv_cache["k"].shape[1] num_new_tokens = roped_query.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): + if self.local_attn_size != -1 and (current_end > kv_cache["global_end_index"]) and ( + num_new_tokens + kv_cache["local_end_index"] > kv_cache_size): # Calculate the number of new tokens added in this step # Shift existing cache content left to discard oldest tokens # Clone the source slice to avoid overlapping memory error - 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 + num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"] - kv_cache_size + num_rolled_tokens = kv_cache["local_end_index"] - 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() # Insert the new keys/values at the end - local_end_index = kv_cache["local_end_index"].item() + current_end - \ - kv_cache["global_end_index"].item() - num_evicted_tokens + local_end_index = kv_cache["local_end_index"] + current_end - \ + kv_cache["global_end_index"] - num_evicted_tokens local_start_index = local_end_index - num_new_tokens kv_cache["k"][:, local_start_index:local_end_index] = roped_key kv_cache["v"][:, local_start_index:local_end_index] = v else: # Assign new keys/values directly up to current_end - local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item() + local_end_index = kv_cache["local_end_index"] + current_end - kv_cache["global_end_index"] local_start_index = local_end_index - num_new_tokens kv_cache["k"] = kv_cache["k"].detach() kv_cache["v"] = kv_cache["v"].detach() @@ -157,8 +157,8 @@ class CausalWanSelfAttention(nn.Module): kv_cache["k"][:, max(0, local_end_index - self.max_attention_size):local_end_index], kv_cache["v"][:, max(0, local_end_index - self.max_attention_size):local_end_index] ) - kv_cache["global_end_index"].fill_(current_end) - kv_cache["local_end_index"].fill_(local_end_index) + kv_cache["global_end_index"] = current_end + kv_cache["local_end_index"] = local_end_index return x diff --git a/fastvideo/models/dits/matrixgame/action_module.py b/fastvideo/models/dits/matrixgame/action_module.py index e7a45413b..52f76632e 100644 --- a/fastvideo/models/dits/matrixgame/action_module.py +++ b/fastvideo/models/dits/matrixgame/action_module.py @@ -95,7 +95,7 @@ def _update_kv_cache_and_attend( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - kv_cache: dict[str, torch.Tensor], + kv_cache: dict[str, torch.Tensor | int], attn_layer: LocalAttention, start_frame: int, num_frame_per_block: int, @@ -139,15 +139,18 @@ def _update_kv_cache_and_attend( kv_cache_size = kv_cache["k"].shape[1] num_new_tokens = k.shape[1] if use_k_for_num_tokens else q.shape[1] + original_global_end_index = kv_cache["global_end_index"] + original_local_end_index = kv_cache["local_end_index"] + # Check if we need to evict tokens - if (current_end > kv_cache["global_end_index"].item()) and ( - num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size + if (current_end > original_global_end_index) and ( + num_new_tokens + original_local_end_index > kv_cache_size ): num_evicted_tokens = ( - num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size + num_new_tokens + original_local_end_index - kv_cache_size ) num_rolled_tokens = ( - kv_cache["local_end_index"].item() + original_local_end_index - num_evicted_tokens - sink_tokens ) @@ -171,18 +174,18 @@ def _update_kv_cache_and_attend( ) # Calculate indices with eviction adjustment local_end_index = ( - kv_cache["local_end_index"].item() + original_local_end_index + current_end - - kv_cache["global_end_index"].item() + - original_global_end_index - num_evicted_tokens ) local_start_index = local_end_index - num_new_tokens else: # Calculate indices without eviction local_end_index = ( - kv_cache["local_end_index"].item() + original_local_end_index + current_end - - kv_cache["global_end_index"].item() + - original_global_end_index ) local_start_index = local_end_index - num_new_tokens @@ -206,8 +209,8 @@ def _update_kv_cache_and_attend( attn = attn_layer(q, cached_k, cached_v) # Update indices - kv_cache["global_end_index"].fill_(current_end) - kv_cache["local_end_index"].fill_(local_end_index) + kv_cache["global_end_index"] = current_end + kv_cache["local_end_index"] = local_end_index return attn diff --git a/fastvideo/models/dits/matrixgame/causal_model.py b/fastvideo/models/dits/matrixgame/causal_model.py index 2b709d9a6..7f985eadd 100644 --- a/fastvideo/models/dits/matrixgame/causal_model.py +++ b/fastvideo/models/dits/matrixgame/causal_model.py @@ -49,7 +49,9 @@ from .model import MatrixGameCrossAttention logger = init_logger(__name__) -def causal_rope_apply(x, grid_sizes, freqs, start_frame=0): +def causal_rope_apply( + x: torch.Tensor, grid_sizes: tuple[int, int, int], freqs: torch.Tensor, start_frame: int = 0 +): n, c = x.size(2), x.size(3) // 2 # split freqs @@ -57,7 +59,7 @@ def causal_rope_apply(x, grid_sizes, freqs, start_frame=0): # loop over samples output = [] - f, h, w = grid_sizes.tolist() + f, h, w = grid_sizes for i in range(len(x)): seq_len = f * h * w @@ -176,16 +178,16 @@ class CausalMatrixGameSelfAttention(nn.Module): v: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor], block_mask: BlockMask, + grid_sizes: tuple[int, int, int], kv_cache: dict | None = None, current_start: int = 0, cache_start: int | None = None, - grid_sizes: torch.Tensor | None = None, ): if cache_start is None: cache_start = current_start # Calculate start_frame for causal mode - if kv_cache is not None and grid_sizes is not None: + if kv_cache is not None: frame_seqlen = int(grid_sizes[1] * grid_sizes[2]) start_frame = current_start // frame_seqlen else: @@ -279,17 +281,14 @@ class CausalMatrixGameSelfAttention(nn.Module): kv_cache_size = kv_cache["k"].shape[1] num_new_tokens = roped_query.shape[1] - if (current_end > kv_cache["global_end_index"].item()) and ( - num_new_tokens + kv_cache["local_end_index"].item() - > kv_cache_size + if (current_end > kv_cache["global_end_index"]) and ( + num_new_tokens + kv_cache["local_end_index"] > kv_cache_size ): num_evicted_tokens = ( - num_new_tokens - + kv_cache["local_end_index"].item() - - kv_cache_size + num_new_tokens + kv_cache["local_end_index"] - kv_cache_size ) num_rolled_tokens = ( - kv_cache["local_end_index"].item() + kv_cache["local_end_index"] - num_evicted_tokens - sink_tokens ) @@ -310,9 +309,9 @@ class CausalMatrixGameSelfAttention(nn.Module): + num_rolled_tokens, ].clone() local_end_index = ( - kv_cache["local_end_index"].item() + kv_cache["local_end_index"] + current_end - - kv_cache["global_end_index"].item() + - kv_cache["global_end_index"] - num_evicted_tokens ) local_start_index = local_end_index - num_new_tokens @@ -320,9 +319,9 @@ class CausalMatrixGameSelfAttention(nn.Module): kv_cache["v"][:, local_start_index:local_end_index] = v else: local_end_index = ( - kv_cache["local_end_index"].item() + kv_cache["local_end_index"] + current_end - - kv_cache["global_end_index"].item() + - kv_cache["global_end_index"] ) local_start_index = local_end_index - num_new_tokens @@ -339,8 +338,8 @@ class CausalMatrixGameSelfAttention(nn.Module): v_for_attn.transpose(1, 2), ).transpose(1, 2) - kv_cache["global_end_index"].fill_(current_end) - kv_cache["local_end_index"].fill_(local_end_index) + kv_cache["global_end_index"] = current_end + kv_cache["local_end_index"] = local_end_index return x @@ -460,7 +459,7 @@ class CausalMatrixGameTransformerBlock(nn.Module): temb: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor], block_mask: BlockMask, - grid_sizes: torch.Tensor, + grid_sizes: tuple[int, int, int], mouse_cond: torch.Tensor | None = None, keyboard_cond: torch.Tensor | None = None, block_mask_mouse: BlockMask | None = None, @@ -516,10 +515,10 @@ class CausalMatrixGameTransformerBlock(nn.Module): value, freqs_cis, block_mask, + grid_sizes, kv_cache, current_start, cache_start, - grid_sizes, ) attn_output = attn_output.flatten(2) attn_output, _ = self.to_out(attn_output) @@ -546,9 +545,9 @@ class CausalMatrixGameTransformerBlock(nn.Module): hidden_states = self.action_model( hidden_states, - int(grid_sizes[0]), - int(grid_sizes[1]), - int(grid_sizes[2]), + grid_sizes[0], + grid_sizes[1], + grid_sizes[2], mouse_cond, keyboard_cond, block_mask_mouse, @@ -981,10 +980,7 @@ class CausalMatrixGameWanModel(BaseDiT): freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None hidden_states = self.patch_embedding(hidden_states) - grid_sizes = torch.tensor( - [post_patch_num_frames, post_patch_height, post_patch_width], - device=hidden_states.device, - ) + grid_sizes = (post_patch_num_frames, post_patch_height, post_patch_width) hidden_states = hidden_states.flatten(2).transpose(1, 2) if timestep.dim() == 2: @@ -1200,9 +1196,9 @@ class CausalMatrixGameWanModel(BaseDiT): hidden_states.shape ) p_t, p_h, p_w = self.patch_size - post_patch_num_frames = num_frames // p_t - post_patch_height = height // p_h - post_patch_width = width // p_w + post_patch_num_frames: int = num_frames // p_t + post_patch_height: int = height // p_h + post_patch_width: int = width // p_w d = self.hidden_size // self.num_attention_heads rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] @@ -1263,10 +1259,7 @@ class CausalMatrixGameWanModel(BaseDiT): ) hidden_states = self.patch_embedding(hidden_states) - grid_sizes = torch.tensor( - [post_patch_num_frames, post_patch_height, post_patch_width], - device=hidden_states.device, - ) + grid_sizes = (post_patch_num_frames, post_patch_height, post_patch_width) hidden_states = hidden_states.flatten(2).transpose(1, 2) if timestep.dim() == 2: @@ -1393,10 +1386,10 @@ class CausalMatrixGameWanModel(BaseDiT): else: return self._forward_train(*args, **kwargs) - def unpatchify(self, x, grid_sizes): + def unpatchify(self, x: torch.Tensor, grid_sizes: tuple[int, int, int]) -> torch.Tensor: c = self.proj_out.out_features // math.prod(self.patch_size) p_t, p_h, p_w = self.patch_size - f, h, w = grid_sizes.tolist() + f, h, w = grid_sizes x = x[:, : f * h * w].view(-1, f, h, w, p_t, p_h, p_w, c) x = x.permute(0, 7, 1, 4, 2, 5, 3, 6) diff --git a/fastvideo/pipelines/stages/causal_denoising.py b/fastvideo/pipelines/stages/causal_denoising.py index 777314547..cbd44fd1e 100644 --- a/fastvideo/pipelines/stages/causal_denoising.py +++ b/fastvideo/pipelines/stages/causal_denoising.py @@ -436,9 +436,9 @@ class CausalDMDDenosingStage(DenoisingStage): dtype=dtype, device=device), "global_end_index": - torch.tensor([0], dtype=torch.long, device=device), + 0, "local_end_index": - torch.tensor([0], dtype=torch.long, device=device), + 0, }) return kv_cache1 @@ -494,4 +494,4 @@ class CausalDMDDenosingStage(DenoisingStage): result.add_check( "negative_prompt_embeds", batch.negative_prompt_embeds, lambda x: not batch.do_classifier_free_guidance or V.list_not_empty(x)) - return result \ No newline at end of file + return result diff --git a/fastvideo/pipelines/stages/matrixgame_denoising.py b/fastvideo/pipelines/stages/matrixgame_denoising.py index 3b442b068..7d46f146e 100644 --- a/fastvideo/pipelines/stages/matrixgame_denoising.py +++ b/fastvideo/pipelines/stages/matrixgame_denoising.py @@ -335,9 +335,9 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): dtype=dtype, device=device), "global_end_index": - torch.tensor([0], dtype=torch.long, device=device), + 0, "local_end_index": - torch.tensor([0], dtype=torch.long, device=device), + 0, }) return kv_cache @@ -373,9 +373,9 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): dtype=dtype, device=device), "global_end_index": - torch.tensor([0], dtype=torch.long, device=device), + 0, "local_end_index": - torch.tensor([0], dtype=torch.long, device=device), + 0, }) kv_cache_mouse.append({ "k": @@ -393,9 +393,9 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): dtype=dtype, device=device), "global_end_index": - torch.tensor([0], dtype=torch.long, device=device), + 0, "local_end_index": - torch.tensor([0], dtype=torch.long, device=device), + 0, }) return kv_cache_mouse, kv_cache_keyboard diff --git a/fastvideo/training/self_forcing_distillation_pipeline.py b/fastvideo/training/self_forcing_distillation_pipeline.py index a7da55c00..0c2f5390a 100644 --- a/fastvideo/training/self_forcing_distillation_pipeline.py +++ b/fastvideo/training/self_forcing_distillation_pipeline.py @@ -582,9 +582,9 @@ class SelfForcingDistillationPipeline(DistillationPipeline): dtype=dtype, device=device), "global_end_index": - torch.tensor([0], dtype=torch.long, device=device), + 0, "local_end_index": - torch.tensor([0], dtype=torch.long, device=device) + 0 }) # Initialize cross-attention cache @@ -616,8 +616,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline): """Reset KV cache and cross-attention cache to clean state.""" if kv_cache is not None: for cache_dict in kv_cache: - cache_dict["global_end_index"].fill_(0) - cache_dict["local_end_index"].fill_(0) + cache_dict["global_end_index"] = 0 + cache_dict["local_end_index"] = 0 cache_dict["k"].zero_() cache_dict["v"].zero_()