From 0dba876979aa08d38e6d381a7904584fa0dbe17d Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Sun, 20 Jul 2025 13:22:56 +0300 Subject: [PATCH] Correct radial attn mask size when using context windows --- nodes.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/nodes.py b/nodes.py index 0139322..9ae020e 100644 --- a/nodes.py +++ b/nodes.py @@ -1672,7 +1672,7 @@ class WanVideoSampler: for i, block in enumerate(transformer.blocks): block.self_attn.mask_map = block.dense_attention_mode = block.dense_timesteps = block.self_attn.decay_factor = None block.dense_block = True if i < dense_blocks else False - block.self_attn.mask_map = MaskMap(video_token_num=seq_len, num_frame=latent_video_length, block_size=block_size) + block.self_attn.mask_map = MaskMap(video_token_num=seq_len, num_frame=latent_video_length if context_options is None else context_frames, block_size=block_size) block.dense_attention_mode = dense_attention_mode block.dense_timesteps = dense_timesteps block.self_attn.decay_factor = decay_factor @@ -1680,7 +1680,7 @@ class WanVideoSampler: for i, block in enumerate(transformer.vace_blocks): block.self_attn.mask_map = block.dense_attention_mode = block.dense_timesteps = block.self_attn.decay_factor = None block.dense_block = True if i < dense_vace_blocks else False - block.self_attn.mask_map = MaskMap(video_token_num=seq_len, num_frame=latent_video_length, block_size=block_size) + block.self_attn.mask_map = MaskMap(video_token_num=seq_len, num_frame=latent_video_length if context_options is None else context_frames, block_size=block_size) block.dense_attention_mode = dense_attention_mode block.dense_timesteps = dense_timesteps block.self_attn.decay_factor = decay_factor