diff --git a/nodes.py b/nodes.py index ef5cd71..3cc4b53 100644 --- a/nodes.py +++ b/nodes.py @@ -19,7 +19,7 @@ from .wanvideo.utils.scheduling_flow_match_lcm import FlowMatchLCMScheduler from .multitalk.multitalk import timestep_transform, add_noise -from .enhance_a_video.globals import enable_enhance, disable_enhance, set_enhance_weight, set_num_frames +from .enhance_a_video.globals import set_enhance_weight, set_num_frames from .taehv import TAEHV from einops import rearrange @@ -2308,6 +2308,12 @@ class WanVideoSampler: rope_params(1024, 2 * (d // 6)) ], dim=1) + transformer.rope_func = rope_function + for block in transformer.blocks: + block.rope_func = rope_function + if transformer.vace_layers is not None: + for block in transformer.vace_blocks: + block.rope_func = rope_function if not isinstance(cfg, list): cfg = [cfg] * (steps +1) @@ -2388,10 +2394,10 @@ class WanVideoSampler: set_num_frames(context_frames) else: set_num_frames(latent_video_length) - enable_enhance() + enhance_enabled = True else: feta_args = None - disable_enhance() + enhance_enabled = False # Initialize Cache if enabled transformer.enable_teacache = transformer.enable_magcache = False @@ -2628,6 +2634,7 @@ class WanVideoSampler: 't': timestep, 'current_step': idx, 'control_lora_enabled': control_lora_enabled, + 'enhance_enabled': enhance_enabled, 'camera_embed': camera_embed, 'unianim_data': unianim_data, 'fun_ref': fun_ref_input if fun_ref_image is not None else None, @@ -2642,7 +2649,6 @@ class WanVideoSampler: "nag_context": text_embeds.get("nag_prompt_embeds", None), "multitalk_audio": multitalk_audio_input if multitalk_audio_embedding is not None else None, "ref_target_masks": ref_target_masks if multitalk_audio_embedding is not None else None, - "rope_func": rope_function } batch_size = 1 @@ -2877,10 +2883,9 @@ class WanVideoSampler: latent_model_input = torch.cat([latent_model_input[:, shift_idx:]] + [latent_model_input[:, :shift_idx]], dim=1) #enhance-a-video + enhance_enabled = False if feta_args is not None and feta_start_percent <= current_step_percentage <= feta_end_percent: - enable_enhance() - else: - disable_enhance() + enhance_enabled = True #flow-edit if flowedit_args is not None: diff --git a/nodes_model_loading.py b/nodes_model_loading.py index 47449db..4125471 100644 --- a/nodes_model_loading.py +++ b/nodes_model_loading.py @@ -668,6 +668,7 @@ class WanVideoModelLoader: "num_heads": num_heads, "num_layers": num_layers, "attention_mode": attention_mode, + "rope_func": "comfy", "main_device": device, "offload_device": offload_device, "teacache_coefficients": teacache_coefficients_map[model_variant], diff --git a/utils.py b/utils.py index 0e9ccc2..ea4bc50 100644 --- a/utils.py +++ b/utils.py @@ -33,6 +33,10 @@ def get_module_memory_mb(module): memory += param.nelement() * param.element_size() return memory / (1024 * 1024) # Convert to MB +def get_tensor_memory(tensor): + memory_bytes = tensor.element_size() * tensor.nelement() + return f"{memory_bytes / (1024 * 1024):.2f} MB" + def apply_lora(model, device_to, transformer_load_device, params_to_keep=None, dtype=None, base_dtype=None, state_dict=None, low_mem_load=False): to_load = [] for n, m in model.model.named_modules(): diff --git a/wanvideo/modules/model.py b/wanvideo/modules/model.py index 69497a9..6c760e2 100644 --- a/wanvideo/modules/model.py +++ b/wanvideo/modules/model.py @@ -7,7 +7,6 @@ from diffusers.configuration_utils import ConfigMixin, register_to_config from diffusers.models.modeling_utils import ModelMixin from einops import repeat, rearrange from ...enhance_a_video.enhance import get_feta_scores -from ...enhance_a_video.globals import is_enhance_enabled try: from torch.nn.attention.flex_attention import create_block_mask, flex_attention, BlockMask @@ -234,8 +233,15 @@ class WanSelfAttention(nn.Module): self.o = nn.Linear(in_features, out_features) self.norm_q = WanRMSNorm(out_features, eps=eps) if qk_norm else nn.Identity() self.norm_k = WanRMSNorm(out_features, eps=eps) if qk_norm else nn.Identity() + + def qkv_fn(self, x): + b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim + q = self.norm_q(self.q(x)).view(b, s, n, d) + k = self.norm_k(self.k(x)).view(b, s, n, d) + v = self.v(x).view(b, s, n, d) + return q, k, v - def forward(self, x, seq_lens, grid_sizes, freqs, rope_func = "default", block_mask=None, ref_target_masks=None): + def forward(self, q, k, v, seq_lens, grid_sizes, freqs, block_mask=None): r""" Args: x(Tensor): Shape [B, L, num_heads, C / num_heads] @@ -243,42 +249,23 @@ class WanSelfAttention(nn.Module): grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] """ - b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim - - # query, key, value - q = self.norm_q(self.q(x)).view(b, s, n, d) - k = self.norm_k(self.k(x)).view(b, s, n, d) - v = self.v(x).view(b, s, n, d) - - if is_enhance_enabled(): - feta_scores = get_feta_scores(q, k) if self.attention_mode == 'flex_attention': - if rope_func == "comfy": - roped_query, roped_key = apply_rope_comfy(q, k, freqs) - else: - roped_query = rope_apply(q, grid_sizes, freqs).type_as(v) - roped_key = rope_apply(k, grid_sizes, freqs).type_as(v) - padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1] padded_roped_query = torch.cat( - [roped_query, - torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]], - device=q.device, dtype=v.dtype)], - dim=1 - ) + [q, torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]], + device=q.device, dtype=v.dtype)], dim=1 + ) padded_roped_key = torch.cat( - [roped_key, torch.zeros([k.shape[0], padded_length, k.shape[2], k.shape[3]], - device=k.device, dtype=v.dtype)], - dim=1 - ) + [k, torch.zeros([k.shape[0], padded_length, k.shape[2], k.shape[3]], + device=k.device, dtype=v.dtype)], dim=1 + ) padded_v = torch.cat( [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]], - device=v.device, dtype=v.dtype)], - dim=1 - ) + device=v.device, dtype=v.dtype)], dim=1 + ) x = flex_attention( query=padded_roped_query.transpose(2, 1), @@ -286,39 +273,21 @@ class WanSelfAttention(nn.Module): value=padded_v.transpose(2, 1), block_mask=block_mask )[:, :, :-padded_length].transpose(2, 1) - - else: - if rope_func == "comfy": - q, k = apply_rope_comfy(q, k, freqs) - elif rope_func == "comfy_chunked": - q, k = apply_rope_comfy_chunked(q, k, freqs) - else: - q=rope_apply(q, grid_sizes, freqs) - k=rope_apply(k, grid_sizes, freqs) - + else: x = attention( - q=q, - k=k, - v=v, + q, k, v, k_lens=seq_lens, window_size=self.window_size, - attention_mode=self.attention_mode) + attention_mode=self.attention_mode + ) # output x = x.flatten(2) x = self.o(x) - if is_enhance_enabled(): - x *= feta_scores - - #multitalk - x_ref_attn_map = None - if ref_target_masks is not None: - x_ref_attn_map = get_attn_map_with_target(q.type_as(x), k.type_as(x), grid_sizes[0], - ref_target_masks=ref_target_masks) - return x, x_ref_attn_map + return x - def forward_split(self, x, seq_lens, grid_sizes, freqs, seq_chunks=1,current_step=0, video_attention_split_steps = [], rope_func = "default"): + def forward_split(self, q, k, v, seq_lens, grid_sizes, freqs, seq_chunks=1,current_step=0, video_attention_split_steps = []): r""" Args: x(Tensor): Shape [B, L, num_heads, C / num_heads] @@ -326,25 +295,6 @@ class WanSelfAttention(nn.Module): grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] """ - b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim - - # query, key, value function - def qkv_fn(x): - q = self.norm_q(self.q(x)).view(b, s, n, d) - k = self.norm_k(self.k(x)).view(b, s, n, d) - v = self.v(x).view(b, s, n, d) - return q, k, v - - q, k, v = qkv_fn(x) - - if rope_func == "comfy": - q, k = apply_rope_comfy(q, k, freqs) - else: - q=rope_apply(q, grid_sizes, freqs) - k=rope_apply(k, grid_sizes, freqs) - - if is_enhance_enabled(): - feta_scores = get_feta_scores(q, k) # Split by frames if multiple prompts are provided if seq_chunks > 1 and current_step in video_attention_split_steps: @@ -407,9 +357,6 @@ class WanSelfAttention(nn.Module): x = x.flatten(2) x = self.o(x) - if is_enhance_enabled(): - x *= feta_scores - return x def normalized_attention_guidance(self, b, n, d, q, context, nag_context=None, nag_params={}): @@ -568,7 +515,9 @@ class WanAttentionBlock(nn.Module): qk_norm=True, cross_attn_norm=False, eps=1e-6, - attention_mode='sdpa'): + attention_mode='sdpa', + rope_func="comfy" + ): super().__init__() self.dim = out_features self.ffn_dim = ffn_dim @@ -578,6 +527,7 @@ class WanAttentionBlock(nn.Module): self.cross_attn_norm = cross_attn_norm self.eps = eps self.attention_mode = attention_mode + self.rope_func = rope_func # layers self.norm1 = WanLayerNorm(out_features, eps) @@ -602,7 +552,7 @@ class WanAttentionBlock(nn.Module): # modulation self.modulation = nn.Parameter(torch.randn(1, 6, out_features) / in_features**0.5) - @torch.compiler.disable() + #@torch.compiler.disable() def get_mod(self, e): if e.dim() == 3: modulation = self.modulation # 1, 6, dim @@ -628,13 +578,13 @@ class WanAttentionBlock(nn.Module): context_lens, current_step, video_attention_split_steps=[], - rope_func = "default", clip_embed=None, camera_embed=None, audio_proj=None, audio_context_lens=None, audio_scale=1.0, num_latent_frames=21, + enhance_enabled=False, block_mask=None, nag_params={}, nag_context=None, @@ -665,28 +615,52 @@ class WanAttentionBlock(nn.Module): input_x += camera_embed # self-attention + x_ref_attn_map = None + + #query, key, value + q, k, v = self.self_attn.qkv_fn(input_x) + del input_x + + # FETA + if enhance_enabled: + feta_scores = get_feta_scores(q, k) + + #RoPE + if self.rope_func == "comfy": + q, k = apply_rope_comfy(q, k, freqs) + elif self.rope_func == "comfy_chunked": + q, k = apply_rope_comfy_chunked(q, k, freqs) + else: + q=rope_apply(q, grid_sizes, freqs) + k=rope_apply(k, grid_sizes, freqs) + + #self-attention if context is not None and (context.shape[0] > 1 or (clip_embed is not None and clip_embed.shape[0] > 1)) and x.shape[0] == 1: y = self.self_attn.forward_split( - input_x, - seq_lens, grid_sizes, - freqs, rope_func=rope_func, + q, k, v, + seq_lens, grid_sizes, freqs, seq_chunks=max(context.shape[0], clip_embed.shape[0] if clip_embed is not None else 0), current_step=current_step, video_attention_split_steps=video_attention_split_steps ) else: - y, x_ref_attn_map = self.self_attn.forward( - input_x, + y = self.self_attn.forward( + q, k, v, seq_lens, grid_sizes, - freqs, rope_func=rope_func, - block_mask=block_mask, - ref_target_masks=ref_target_masks, + freqs, block_mask=block_mask ) + + #multitalk mask + if ref_target_masks is not None: + x_ref_attn_map = get_attn_map_with_target(q.type_as(x), k.type_as(x), grid_sizes[0], ref_target_masks=ref_target_masks) + + # FETA + if enhance_enabled: + y.mul_(feta_scores) + #ReCamMaster if camera_embed is not None: - y = self.projector(y) - - del input_x + y = self.projector(y) x.add_(y.mul_(e[2])) del y @@ -707,7 +681,6 @@ class WanAttentionBlock(nn.Module): x.add_(y.mul_(e[5])) del y - del e return x #@torch.compiler.disable() @@ -796,9 +769,11 @@ class VaceWanAttentionBlock(WanAttentionBlock): qk_norm=True, cross_attn_norm=False, eps=1e-6, - block_id=0 + block_id=0, + attention_mode='sdpa', + rope_func="comfy" ): - super().__init__(cross_attn_type, in_features, out_features, ffn_dim, ffn2_dim, num_heads, window_size, qk_norm, cross_attn_norm, eps) + super().__init__(cross_attn_type, in_features, out_features, ffn_dim, ffn2_dim, num_heads, window_size, qk_norm, cross_attn_norm, eps, attention_mode, rope_func) self.block_id = block_id if block_id == 0: self.before_proj = nn.Linear(in_features, out_features) @@ -832,9 +807,10 @@ class BaseWanAttentionBlock(WanAttentionBlock): cross_attn_norm=False, eps=1e-6, block_id=None, - attention_mode='sdpa' + attention_mode='sdpa', + rope_func="comfy" ): - super().__init__(cross_attn_type, in_features, out_features, ffn_dim, ffn2_dim, num_heads, window_size, qk_norm, cross_attn_norm, eps, attention_mode) + super().__init__(cross_attn_type, in_features, out_features, ffn_dim, ffn2_dim, num_heads, window_size, qk_norm, cross_attn_norm, eps, attention_mode, rope_func) self.block_id = block_id def forward(self, x, vace_hints=None, vace_context_scale=[1.0], **kwargs): @@ -940,6 +916,7 @@ class WanModel(ModelMixin, ConfigMixin): cross_attn_norm=True, eps=1e-6, attention_mode='sdpa', + rope_func='comfy', main_device=torch.device('cuda'), offload_device=torch.device('cpu'), teacache_coefficients=[], @@ -1010,6 +987,7 @@ class WanModel(ModelMixin, ConfigMixin): self.cross_attn_norm = cross_attn_norm self.eps = eps self.attention_mode = attention_mode + self.rope_func = rope_func self.main_device = main_device self.offload_device = offload_device @@ -1072,7 +1050,7 @@ class WanModel(ModelMixin, ConfigMixin): # vace blocks self.vace_blocks = nn.ModuleList([ VaceWanAttentionBlock('t2v_cross_attn', self.in_features, self.out_features, self.ffn_dim, self.ffn2_dim,self.num_heads, self.window_size, self.qk_norm, - self.cross_attn_norm, self.eps, block_id=i) + self.cross_attn_norm, self.eps, block_id=i, attention_mode=self.attention_mode, rope_func=self.rope_func) for i in self.vace_layers ]) @@ -1083,7 +1061,7 @@ class WanModel(ModelMixin, ConfigMixin): self.blocks = nn.ModuleList([ BaseWanAttentionBlock('t2v_cross_attn', self.in_features, self.out_features, ffn_dim, self.ffn2_dim, num_heads, window_size, qk_norm, cross_attn_norm, eps, - attention_mode=self.attention_mode, + attention_mode=self.attention_mode, rope_func=self.rope_func, block_id=self.vace_layers_mapping[i] if i in self.vace_layers else None) for i in range(num_layers) ]) @@ -1099,7 +1077,7 @@ class WanModel(ModelMixin, ConfigMixin): self.blocks = nn.ModuleList([ WanAttentionBlock(cross_attn_type, self.in_features, self.out_features, ffn_dim, ffn2_dim, num_heads, window_size, qk_norm, cross_attn_norm, eps, - attention_mode=self.attention_mode) + attention_mode=self.attention_mode, rope_func=self.rope_func) for _ in range(num_layers) ]) @@ -1280,6 +1258,7 @@ class WanModel(ModelMixin, ConfigMixin): y=None, device=torch.device('cuda'), freqs=None, + enhance_enabled=False, pred_id=None, control_lora_enabled=False, vace_data=None, @@ -1298,8 +1277,7 @@ class WanModel(ModelMixin, ConfigMixin): nag_params={}, nag_context=None, multitalk_audio=None, - ref_target_masks=None, - rope_func="comfy", + ref_target_masks=None ): r""" Forward pass through the diffusion model @@ -1642,13 +1620,13 @@ class WanModel(ModelMixin, ConfigMixin): context=context, context_lens=context_lens, clip_embed=clip_embed, - rope_func=rope_func, current_step=current_step, video_attention_split_steps=self.video_attention_split_steps, camera_embed=camera_embed, audio_proj=audio_proj, audio_context_lens=audio_context_lens, num_latent_frames = F, + enhance_enabled=enhance_enabled, audio_scale=audio_scale, block_mask=self.block_mask, nag_params=nag_params, @@ -1837,7 +1815,3 @@ def relative_l1_distance(last_tensor, current_tensor): norm = torch.abs(last_tensor).mean() relative_l1_distance = l1_distance / norm return relative_l1_distance.to(torch.float32).to(current_tensor.device) - -def get_tensor_memory(tensor): - memory_bytes = tensor.element_size() * tensor.nelement() - return f"{memory_bytes / (1024 * 1024):.2f} MB" \ No newline at end of file