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

Add node to set attention mode per step and/or blocks

Browse Source
This commit is contained in:
kijai
2025-12-29 01:35:32 +02:00
parent 3730ccf603
commit 486564060f
3 changed files with 74 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes_model_loading.py
+43 -11
View File
@@ -36,6 +36,9 @@ try:
except:
PromptServer = None
attention_modes = ["sdpa", "flash_attn_2", "flash_attn_3", "sageattn", "sageattn_3", "radial_sage_attention", "sageattn_compiled",
"sageattn_ultravico", "comfy"]
#from city96's gguf nodes
def update_folder_names_and_paths(key, targets=[]):
# check for existing key
@@ -1006,6 +1009,43 @@ def add_lora_weights(patcher, lora, base_dtype, merge_loras=False):
del lora_sd
return patcher, control_lora, unianimate_sd
class WanVideoSetAttentionModeOverride:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"model": ("WANVIDEOMODEL", ),
"attention_mode": (attention_modes, {"default": "sdpa"}),
"start_step": ("INT", {"default": 1, "min": 1, "max": 10000, "step": 1, "tooltip": "Step to start applying the attention mode override"}),
"end_step": ("INT", {"default": 10000, "min": 1, "max": 10000, "step": 1, "tooltip": "Step to end applying the attention mode override"}),
"verbose": ("BOOLEAN", {"default": False, "tooltip": "Print verbose info about attention mode override during generation"}),
},
"optional": {
"blocks":("INT", {"forceInput": True} ),
}
}
RETURN_TYPES = ("WANVIDEOMODEL",)
RETURN_NAMES = ("model", )
FUNCTION = "getmodelpath"
CATEGORY = "WanVideoWrapper"
DESCRIPTION = "Override the attention mode for the model for specific step and/or block range"
def getmodelpath(self, model, attention_mode, start_step, end_step, verbose, blocks=None):
model_clone = model.clone()
attention_mode_override = {
"mode": attention_mode,
"start_step": start_step,
"end_step": end_step,
"verbose": verbose,
}
if blocks is not None:
attention_mode_override["blocks"] = blocks
model_clone.model_options['transformer_options']["attention_mode_override"] = attention_mode_override
return (model_clone,)
#region Model loading
class WanVideoModelLoader:
@classmethod
@@ -1020,17 +1060,7 @@ class WanVideoModelLoader:
"load_device": (["main_device", "offload_device"], {"default": "offload_device", "tooltip": "Initial device to load the model to, NOT recommended with the larger models unless you have 48GB+ VRAM"}),
},
"optional": {
"attention_mode": ([
"sdpa",
"flash_attn_2",
"flash_attn_3",
"sageattn",
"sageattn_3",
"radial_sage_attention",
"sageattn_compiled",
"sageattn_ultravico",
"comfy"
], {"default": "sdpa"}),
"attention_mode": (attention_modes, {"default": "sdpa"}),
"compile_args": ("WANCOMPILEARGS", ),
"block_swap_args": ("BLOCKSWAPARGS", ),
"lora": ("WANVIDLORA", {"default": None}),
@@ -2043,6 +2073,7 @@ NODE_CLASS_MAPPINGS = {
"WanVideoTorchCompileSettings": WanVideoTorchCompileSettings,
"LoadWanVideoT5TextEncoder": LoadWanVideoT5TextEncoder,
"LoadWanVideoClipTextEncoder": LoadWanVideoClipTextEncoder,
"WanVideoSetAttentionModeOverride": WanVideoSetAttentionModeOverride,
}
NODE_DISPLAY_NAME_MAPPINGS = {
@@ -2061,4 +2092,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"WanVideoTorchCompileSettings": "WanVideo Torch Compile Settings",
"LoadWanVideoT5TextEncoder": "WanVideo T5 Text Encoder Loader",
"LoadWanVideoClipTextEncoder": "WanVideo CLIP Text Encoder Loader",
"WanVideoSetAttentionModeOverride": "WanVideo Set Attention Mode Override",
}
nodes_sampler.py
+1
View File
@@ -1414,6 +1414,7 @@ class WanVideoSampler:
'is_uncond': False, # is unconditional
'current_step': idx, # current step
'current_step_percentage': current_step_percentage, # current step percentage
'attention_mode_override': transformer_options.get("attention_mode_override", None),
'last_step': len(timesteps) - 1 == idx, # is last step
'control_lora_enabled': control_lora_enabled, # control lora toggle for patch embed selection
'enhance_enabled': enhance_enabled, # enhance-a-video toggle
wanvideo/modules/model.py
+30 -15
View File
@@ -497,7 +497,7 @@ class WanSelfAttention(nn.Module):
attention_mode = self.attention_mode
if attention_mode_override is not None:
attention_mode = attention_mode_override
# Concatenate main and IP keys/values for main attention
full_k = torch.cat([k, k_ip], dim=1)
full_v = torch.cat([v, v_ip], dim=1)
@@ -1006,6 +1006,7 @@ class WanAttentionBlock(nn.Module):
longcat_num_cond_latents=0, longcat_avatar_options=None, #longcat image cond amount
x_onetoall_ref=None, onetoall_freqs=None, onetoall_ref=None, onetoall_ref_scale=1.0, #one-to-all
e_tr=None, tr_num=0, tr_start=0, #token replacement
attention_mode_override=None,
):
r"""
Args:
@@ -1150,6 +1151,10 @@ class WanAttentionBlock(nn.Module):
if enhance_enabled:
feta_scores = get_feta_scores(q, k)
if self.attention_mode == "sageattn_3" and attention_mode_override is None:
if current_step != 0 and not last_step:
attention_mode_override = "sageattn"
#self-attention
split_attn = (context is not None
and (context.shape[0] > 1 or (clip_embed is not None and clip_embed.shape[0] > 1))
@@ -1161,19 +1166,14 @@ class WanAttentionBlock(nn.Module):
y = self.self_attn.forward_split(q, k, v, seq_lens, grid_sizes, seq_chunks)
elif ref_target_masks is not None: #multi/infinite talk
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":
elif self.attention_mode == "radial_sage_attention" or attention_mode_override is not None and attention_mode_override == "radial_sage_attention":
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:
y = self.self_attn.forward(q, k, v, seq_lens)
y = self.self_attn.forward(q, k, v, seq_lens, attention_mode_override=attention_mode_override)
else:
y = self.self_attn.forward_radial(q, k, v, dense_step=False)
elif self.attention_mode == "sageattn_3":
if current_step != 0 and not last_step:
y = self.self_attn.forward(q, k, v, seq_lens, attention_mode_override="sageattn_3")
else:
y = self.self_attn.forward(q, k, v, seq_lens, attention_mode_override="sageattn")
elif x_ip is not None and self.kv_cache is None: #stand-in
# First pass: cache IP keys/values and compute attention
self.kv_cache = {"k_ip": k_ip.detach(), "v_ip": v_ip.detach()}
@@ -1184,18 +1184,18 @@ class WanAttentionBlock(nn.Module):
v_ip = self.kv_cache["v_ip"]
full_k = torch.cat([k, k_ip], dim=1)
full_v = torch.cat([v, v_ip], dim=1)
y = self.self_attn.forward(q, full_k, full_v, seq_lens)
y = self.self_attn.forward(q, full_k, full_v, seq_lens, attention_mode_override=attention_mode_override)
elif is_longcat and longcat_num_cond_latents > 0:
if longcat_num_cond_latents == 1:
num_cond_latents_thw = longcat_num_cond_latents * (N // num_latent_frames)
# process the noise tokens
x_noise = self.self_attn.forward(q[:, num_cond_latents_thw:].contiguous(), k, v, seq_lens)
x_noise = self.self_attn.forward(q[:, num_cond_latents_thw:].contiguous(), k, v, seq_lens, attention_mode_override=attention_mode_override)
# process the condition tokens
x_cond = self.self_attn.forward(
q[:, :num_cond_latents_thw].contiguous(),
k[:, :num_cond_latents_thw].contiguous(),
v[:, :num_cond_latents_thw].contiguous(),
seq_lens)
seq_lens, attention_mode_override=attention_mode_override)
# merge x_cond and x_noise
y = torch.cat([x_cond, x_noise], dim=1).contiguous()
elif longcat_num_cond_latents > 1: # video continuation
@@ -1237,13 +1237,14 @@ class WanAttentionBlock(nn.Module):
q_cond = q[:, num_ref_latents_thw:num_cond_latents_thw].contiguous()
k_cond = k[:, num_ref_latents_thw:num_cond_latents_thw].contiguous()
v_cond = v[:, num_ref_latents_thw:num_cond_latents_thw].contiguous()
x_ref = self.self_attn.forward(q_ref, k_ref, v_ref, seq_lens)
x_cond = self.self_attn.forward(q_cond, k_cond, v_cond, seq_lens)
x_ref = self.self_attn.forward(q_ref, k_ref, v_ref, seq_lens, attention_mode_override=attention_mode_override)
x_cond = self.self_attn.forward(q_cond, k_cond, v_cond, seq_lens, attention_mode_override=attention_mode_override)
# merge x_cond and x_noise
y = torch.cat([x_ref, x_cond, x_noise], dim=1).contiguous()
else:
y = self.self_attn.forward(q, k, v, seq_lens, lynx_ref_feature=lynx_ref_feature, lynx_ref_scale=lynx_ref_scale, onetoall_ref=onetoall_ref, onetoall_ref_scale=onetoall_ref_scale)
y = self.self_attn.forward(q, k, v, seq_lens, lynx_ref_feature=lynx_ref_feature, lynx_ref_scale=lynx_ref_scale,
onetoall_ref=onetoall_ref, onetoall_ref_scale=onetoall_ref_scale, attention_mode_override=attention_mode_override)
del q, k, v
@@ -2280,6 +2281,7 @@ class WanModel(torch.nn.Module):
self, x, t, context, seq_len,
is_uncond=False,
current_step_percentage=0.0, current_step=0, last_step=0, total_steps=50,
attention_mode_override=None,
clip_fea=None, y=None,
device=torch.device('cuda'),
freqs=None,
@@ -3125,8 +3127,21 @@ class WanModel(torch.nn.Module):
if lynx_ref_buffer is None and lynx_ref_feature_extractor:
lynx_ref_buffer = {}
attn_override_blocks = attention_mode = None
attention_mode_override_active = False
if attention_mode_override is not None:
attn_override_blocks = attention_mode_override.get("blocks", range(len(self.blocks)))
if attention_mode_override["start_step"] <= current_step < attention_mode_override["end_step"]:
attention_mode_override_active = True
if attention_mode_override["verbose"]:
tqdm.write(f"Applying attention mode override: {attention_mode_override['mode']} at step {current_step} on blocks: {attn_override_blocks if attn_override_blocks is not None else 'all'}")
for b, block in enumerate(self.blocks):
mm.throw_exception_if_processing_interrupted()
if attention_mode_override_active and b in attn_override_blocks:
attention_mode = attention_mode_override['mode']
else:
attention_mode = None
block_idx = f"{b:02d}"
if lynx_ref_buffer is not None and not lynx_ref_feature_extractor:
lynx_ref_feature = lynx_ref_buffer.get(block_idx, None)
@@ -3170,7 +3185,7 @@ class WanModel(torch.nn.Module):
x_onetoall_ref = onetoall_ref_block_samples[b // interval_ref]
# ---run block----#
x, x_ip, lynx_ref_feature, x_ovi = block(x, x_ip=x_ip, lynx_ref_feature=lynx_ref_feature, x_ovi=x_ovi, x_onetoall_ref=x_onetoall_ref, onetoall_freqs=onetoall_freqs, **kwargs)
x, x_ip, lynx_ref_feature, x_ovi = block(x, x_ip=x_ip, lynx_ref_feature=lynx_ref_feature, x_ovi=x_ovi, x_onetoall_ref=x_onetoall_ref, onetoall_freqs=onetoall_freqs, attention_mode_override=attention_mode, **kwargs)
# ---post block----#
if self.audio_injector is not None and s2v_audio_input is not None: