nodes/kijai-ComfyUI-WanVideoWrapper
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix and rename args to be clearer

Browse Source
This commit is contained in:
kijai
2025-07-18 01:28:01 +03:00
parent 712eecf64b
commit a7166fc1a4
2 changed files with 15 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes.py
+13 -11
View File
@@ -280,8 +280,8 @@ class WanVideoSetRadialAttention:
"sageattn",
"sparse_sage_attention",
], {"default": "sageattn", "tooltip": "The attention mode for dense attention"}),
"dense_block": ("INT", {"default": 1, "min": 0, "max": 8, "step": 1, "tooltip": "Number of blocks to apply normal attention to"}),
"dense_timestep": ("INT", {"default": 10, "min": 0, "max": 100, "step": 1, "tooltip": "The step to start applying sparse attention"}),
"dense_blocks": ("INT", {"default": 1, "min": 0, "max": 8, "step": 1, "tooltip": "Number of blocks to apply normal attention to"}),
"dense_timesteps": ("INT", {"default": 10, "min": 0, "max": 100, "step": 1, "tooltip": "The step to start applying sparse attention"}),
"decay_factor": ("FLOAT", {"default": 0.2, "min": 0, "max": 1, "step": 0.01, "tooltip": "Controls how quickly the attention window shrinks as the distance between frames increases in the sparse attention mask."}),
}
}
@@ -291,7 +291,7 @@ class WanVideoSetRadialAttention:
FUNCTION = "loadmodel"
CATEGORY = "WanVideoWrapper"
def loadmodel(self, model, dense_attention_mode, dense_block, dense_timestep, decay_factor):
def loadmodel(self, model, dense_attention_mode, dense_blocks, dense_timesteps, decay_factor):
if "radial" not in model.model.diffusion_model.attention_mode:
raise Exception("Enable radial attention first in the model loader.")
@@ -300,8 +300,8 @@ class WanVideoSetRadialAttention:
patcher.model_options['transformer_options'] = {}
patcher.model_options["transformer_options"]["dense_attention_mode"] = dense_attention_mode
patcher.model_options["transformer_options"]["dense_block"] = dense_block
patcher.model_options["transformer_options"]["dense_timestep"] = dense_timestep
patcher.model_options["transformer_options"]["dense_blocks"] = dense_blocks
patcher.model_options["transformer_options"]["dense_timesteps"] = dense_timesteps
patcher.model_options["transformer_options"]["decay_factor"] = decay_factor
return (patcher,)
@@ -2537,22 +2537,24 @@ class WanVideoSampler:
# Radial attention setup
if transformer.attention_mode == "radial_sage_attention":
if latent.shape[2] % 16 != 0 or latent.shape[3] % 16 != 0:
raise Exception(f"Radial attention mode only supports image size divisible by 128.")
if transformer_options is not None:
dense_timestep = transformer_options.get("dense_timestep", 10)
dense_block = transformer_options.get("dense_block", 1)
dense_timesteps = transformer_options.get("dense_timesteps", 10)
dense_blocks = transformer_options.get("dense_blocks", 1)
decay_factor = transformer_options.get("decay_factor", 0.2)
dense_attention_mode = transformer_options.get("dense_attention_mode", "sageattn")
from .wanvideo.radial_attention.attn_mask import MaskMap
for i, block in enumerate(transformer.blocks):
block.self_attn.mask_map = block.dense_attention_mode = block.dense_timestep = block.self_attn.decay_factor = None
block.dense_block = True if i < dense_block else False
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.dense_attention_mode = dense_attention_mode
block.dense_timestep = dense_timestep
block.dense_timesteps = dense_timesteps
block.self_attn.decay_factor = decay_factor
log.info(f"Radial attention mode enabled. dense_attention_mode: {dense_attention_mode}, dense_timestep: {dense_timestep}, dense_block: {dense_block}, decay_factor: {decay_factor}")
log.info(f"Radial attention mode enabled. dense_attention_mode: {dense_attention_mode}, dense_timesteps: {dense_timesteps}, dense_blocks: {dense_blocks}, decay_factor: {decay_factor}")
self.cache_state = [None, None]
if phantom_latents is not None:
wanvideo/modules/model.py
+2 -2
View File
@@ -600,7 +600,7 @@ class WanAttentionBlock(nn.Module):
self.attention_mode = attention_mode
self.rope_func = rope_func
#radial attn
self.dense_timestep = 10
self.dense_timesteps = 10
self.dense_block = False
self.dense_attention_mode = "sageattn"
@@ -738,7 +738,7 @@ class WanAttentionBlock(nn.Module):
elif ref_target_masks is not None:
y, x_ref_attn_map = self.self_attn.forward_multitalk(q, k, v, seq_lens, grid_sizes, ref_target_masks)
elif self.attention_mode == "radial_sage_attention":
if self.dense_block and self.dense_timestep is not None and current_step < self.dense_timestep:
if self.dense_block or self.dense_timesteps is not None and current_step < self.dense_timesteps:
if self.dense_attention_mode == "sparse_sage_attn":
y = self.self_attn.forward_radial(q, k, v, dense_step=True)
else: