diff --git a/nodes_model_loading.py b/nodes_model_loading.py index 83abf6f..01977f4 100644 --- a/nodes_model_loading.py +++ b/nodes_model_loading.py @@ -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", } diff --git a/nodes_sampler.py b/nodes_sampler.py index d054b90..6072608 100644 --- a/nodes_sampler.py +++ b/nodes_sampler.py @@ -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 diff --git a/wanvideo/modules/model.py b/wanvideo/modules/model.py index 96df41c..6a31cb9 100644 --- a/wanvideo/modules/model.py +++ b/wanvideo/modules/model.py @@ -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: