nodes/thatname-Comfyui_CharaConsist
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add support for Qwen models.

Browse Source
This commit is contained in:
ZHOU He
2025-11-18 20:29:44 +08:00
parent a31e893fd9
commit 3c8e8d3753
1 changed files with 59 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
node.py
+59 -55
View File
@@ -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