From 3c8e8d375364731fb4e4d87620d1c95487337430 Mon Sep 17 00:00:00 2001 From: ZHOU He <461607353@qq.com> Date: Tue, 18 Nov 2025 20:29:44 +0800 Subject: [PATCH] Add support for Qwen models. --- node.py | 114 +++++++++++++++++++++++++++++--------------------------- 1 file changed, 59 insertions(+), 55 deletions(-) diff --git a/node.py b/node.py index b8bf63a..f54e761 100644 --- a/node.py +++ b/node.py @@ -14,7 +14,10 @@ from einops import rearrange from nodes import CLIPTextEncode import latent_preview import copy -orig_attention = flux_math.attention +import comfy.ldm.qwen_image.model as qwen_image_model +#import comfy.ldm.wan.model as wan_model +orig_attention = comfy_attention.optimized_attention +orig_apply_rope1 = flux_math.apply_rope1 orig_prepare_callback = latent_preview.prepare_callback def convert_masks_to_patches(masks: torch.Tensor, pad) -> tuple: @@ -470,10 +473,46 @@ class MaskToPatchMask: patch_mask = patch_mask.bool() return (patch_mask,) +def hook_math(hooked_attention): + save_freq = None + def hook_apply_rope1(x, freqs_cis): + nonlocal save_freq + save_freq = freqs_cis + return x + def hook_optimized_attention(q, k, v, heads, mask=None, skip_reshape = False, skip_output_reshape=False, transformer_options={}, **kwargs) -> Tensor: + nonlocal save_freq + if "cond_or_uncond" in transformer_options and transformer_options["cond_or_uncond"][-1]: # negative prompt, ignore + q = orig_apply_rope1(q, save_freq) + k = orig_apply_rope1(k, save_freq) + return orig_attention(q, k, v, heads, mask=mask, skip_reshape=skip_reshape, skip_output_reshape=skip_output_reshape, transformer_options=transformer_options, **kwargs) + if not skip_reshape: # not B H L D but B L (H D), we must reshape to B H L D + B, L, HD = q.shape + D = HD // heads + q = q.view(B, L, heads, D).permute(0, 2, 1, 3) + B, L, HD = k.shape + D = HD // heads + k = k.view(B, L, heads, D).permute(0, 2, 1, 3) + B, L, HD = v.shape + D = HD // heads + v = v.view(B, L, heads, D).permute(0, 2, 1, 3) + output = hooked_attention(q, k, v, save_freq, mask, **kwargs) # no reshape for hooked attention + if skip_output_reshape: # because we reshaped output to B L (H D) before, we need to reshape back + output = rearrange(output, "B L (H D) -> B H L D", H=heads) + return output + + flux_math.apply_rope1 = hook_apply_rope1 + flux_math.optimized_attention = hook_optimized_attention + qwen_image_model.apply_rope1 = hook_apply_rope1 + qwen_image_model.optimized_attention_masked = hook_optimized_attention + +def unhook_math(): + flux_math.apply_rope1 = orig_apply_rope1 + flux_math.optimized_attention = orig_attention + qwen_image_model.apply_rope1 = orig_apply_rope1 + qwen_image_model.optimized_attention_masked = orig_attention def hook_prepare_callback(): - flux_layers.attention = orig_attention - chroma_layers.attention = orig_attention + unhook_math() def prepare_callback(model, max_steps, x0_output_dict=None): if "prepare_callback" in model.model_options: return model.model_options["prepare_callback"](model, max_steps, x0_output_dict) @@ -533,11 +572,9 @@ class ExtractAttn: sampled_step = -1 sampling_step = -1 layer = 0 - model_type = "flux" - B,L,H,W = latent["samples"].shape - w = W // model.model.diffusion_model.patch_size - h = H // model.model.diffusion_model.patch_size + w = latent["samples"].shape[-1] // model.model.diffusion_model.patch_size + h = latent["samples"].shape[-2] // model.model.diffusion_model.patch_size #text_len = fg_condition_mask.shape[-1] device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') @@ -572,7 +609,8 @@ class ExtractAttn: attn_data[index][sampling_step]["layers"][layer] = {"k": k.cpu(), "v": v.cpu()} # Apply rotary positional encoding for Flux models if pe is not None: - q, k = flux_math.apply_rope(q, k, pe) + q = orig_apply_rope1(q, pe) + k = orig_apply_rope1(k, pe) if fg_condition_mask is not None and sampling_step >= mask_start_step and sampling_step < mask_end_step: if fg_bg_mask is None: @@ -592,23 +630,8 @@ class ExtractAttn: attn_data[index][sampling_step]["layers"][layer]["x"] = x[:, text_len:,:].cpu() layer += 1 return x - - # Create the attention override function - def attention_override(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask: Tensor, transformer_options) -> Tensor: - if transformer_options["cond_or_uncond"][-1]:#q.shape[2] != text_len + (w * h) - return orig_attention(q, k, v, pe) - else: - return masked_attention(q, k, v, pe, mask) - - # Apply the attention override based on model type - overrider = attention_override - - if model.model.__class__.__name__ == "Chroma": - model_type = "chroma" - if model_type == "flux": - flux_layers.attention = overrider - else: # chroma - chroma_layers.attention = overrider + + hook_math(masked_attention) origin_callback = orig_prepare_callback(model, max_steps, x0_output_dict) def callback_change_sampled_step(step, x0, x, total_steps): nonlocal max_steps @@ -816,11 +839,9 @@ class GenConsistent: sampled_step = -1 sampling_step = -1 layer = 0 - model_type = "flux" - B,L,H,W = latent["samples"].shape - w = W // model.model.diffusion_model.patch_size - h = H // model.model.diffusion_model.patch_size + w = latent["samples"].shape[-1] // model.model.diffusion_model.patch_size + h = latent["samples"].shape[-2] // model.model.diffusion_model.patch_size device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') attn_dtype = torch.bfloat16 if model_management.should_use_bf16(device=device) else torch.float16 @@ -848,7 +869,8 @@ class GenConsistent: original_seq_len = q.shape[2] # Apply rotary positional encoding for Flux models if pe is not None: - q, k = flux_math.apply_rope(q, k, pe) + q = orig_apply_rope1(q, pe) + k = orig_apply_rope1(k, pe) # Store original sequence length before K,V concatenation original_k_len = k.shape[2] @@ -864,8 +886,8 @@ class GenConsistent: # Process all subjects, not just the last one (index = len(subject) - 1) for subj_idx in range(len(subject)): # Check if we have valid cross-sim data and subject data to work with - if (len(id_fg_inds[subj_idx]) > 0 and - len(curr_fg_inds[subj_idx]) > 0 and + if (id_fg_inds[subj_idx].shape[0] > 0 and + curr_fg_inds[subj_idx].shape[0] > 0 and sampling_step in subject[subj_idx] and layer in subject[subj_idx][sampling_step]["layers"] and "k" in subject[subj_idx][sampling_step]["layers"][layer] and @@ -893,7 +915,7 @@ class GenConsistent: # Extract positional encoding for the saved key positions # This mimics the approach in attention_processor_characonsist.py image_rotary_emb_fg = pe[:,:,valid_curr_inds] - saved_key_fg = flux_math.apply_rope1(saved_key_fg, image_rotary_emb_fg) + saved_key_fg = orig_apply_rope1(saved_key_fg, image_rotary_emb_fg) # Concatenate the expanded key-value pairs k = torch.cat([k, saved_key_fg], dim=2) @@ -931,20 +953,17 @@ class GenConsistent: # Validate indices - valid_id_inds = subj_id_fg_inds[subj_id_fg_inds < saved_x.shape[-1]] - valid_curr_inds = subj_curr_fg_inds[subj_curr_fg_inds < x.shape[-1]] + valid_id_inds = subj_id_fg_inds[subj_id_fg_inds < saved_x.shape[-2]] + valid_curr_inds = subj_curr_fg_inds[subj_curr_fg_inds < x.shape[-2] - text_len] if len(valid_id_inds) > 0 and len(valid_curr_inds) > 0: # Extract vision tokens (after text_len) for blending vision_hidden_states = x[:, text_len:, :] # Get matched hidden states from saved subject and current generation - matched_saved_x = saved_x[:, :, valid_id_inds] # [B, D, valid_id_inds] + matched_saved_x = saved_x[:, valid_id_inds, :] # [B, D, valid_id_inds] matched_current_x = vision_hidden_states[:, valid_curr_inds, :] # [B, valid_curr_inds, D] - # Transpose saved_x to match current_x format [B, valid_id_inds, D] - matched_saved_x = matched_saved_x.transpose(-2, -1) - # Apply ada_tome blending with similarity-weighted alpha alpha_tensor = torch.ones(len(valid_curr_inds), dtype=attn_dtype, device=x.device) alpha_tensor = alpha_tensor * ada_tome_alpha @@ -966,22 +985,7 @@ class GenConsistent: layer += 1 return x - # Create the attention override function - def attention_override(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask: Tensor, transformer_options) -> Tensor: - if transformer_options["cond_or_uncond"][-1]:#q.shape[2] != text_len + (w * h) - return orig_attention(q, k, v, pe) - else: - return masked_attention(q, k, v, pe, mask) - - # Apply the attention override based on model type - overrider = attention_override - - if model.model.__class__.__name__ == "Chroma": - model_type = "chroma" - if model_type == "flux": - flux_layers.attention = overrider - else: # chroma - chroma_layers.attention = overrider + hook_math(masked_attention) origin_callback = orig_prepare_callback(model, max_steps, x0_output_dict) def callback_change_sampled_step(step, x0, x, total_steps): nonlocal max_steps