Update nodes.py

This commit is contained in:
kijai
2025-07-18 14:37:11 +03:00
parent a1cc320c36
commit ed11e2bb3e
+5 -2
View File
@@ -281,6 +281,7 @@ class WanVideoSetRadialAttention:
"sparse_sage_attention",
], {"default": "sageattn", "tooltip": "The attention mode for dense attention"}),
"dense_blocks": ("INT", {"default": 1, "min": 0, "max": 40, "step": 1, "tooltip": "Number of blocks to apply normal attention to"}),
"dense_vace_blocks": ("INT", {"default": 15, "min": 0, "max": 40, "step": 1, "tooltip": "Number of vace 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."}),
}
@@ -292,7 +293,7 @@ class WanVideoSetRadialAttention:
CATEGORY = "WanVideoWrapper"
DESCRIPTION = "Sets radial attention parameters, dense attention refers to normal attention"
def loadmodel(self, model, dense_attention_mode, dense_blocks, dense_timesteps, decay_factor):
def loadmodel(self, model, dense_attention_mode, dense_blocks, dense_vace_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.")
@@ -302,6 +303,7 @@ class WanVideoSetRadialAttention:
patcher.model_options["transformer_options"]["dense_attention_mode"] = dense_attention_mode
patcher.model_options["transformer_options"]["dense_blocks"] = dense_blocks
patcher.model_options["transformer_options"]["dense_vace_blocks"] = dense_vace_blocks
patcher.model_options["transformer_options"]["dense_timesteps"] = dense_timesteps
patcher.model_options["transformer_options"]["decay_factor"] = decay_factor
@@ -2543,6 +2545,7 @@ class WanVideoSampler:
dense_timesteps = transformer_options.get("dense_timesteps", None)
dense_blocks = transformer_options.get("dense_blocks", None)
dense_vace_blocks = transformer_options.get("dense_vace_blocks", None)
decay_factor = transformer_options.get("decay_factor", None)
dense_attention_mode = transformer_options.get("dense_attention_mode", None)
if dense_timesteps is None:
@@ -2559,7 +2562,7 @@ class WanVideoSampler:
if transformer.vace_layers is not None:
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_blocks else False
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.dense_attention_mode = dense_attention_mode
block.dense_timesteps = dense_timesteps