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

Add cross sim logic.

Browse Source
This commit is contained in:
ZHOU He
2025-11-16 00:03:50 +08:00
parent e43bb5d4d5
commit a0499dd2b7
2 changed files with 83 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+3 -1
View File
@@ -1,4 +1,4 @@
from .node import BatchedMaskedReferenceGen,ReferenceConditionCombine,MasksToPatches,ConditioningMatchMask,GenSubject,PreviewSubjectMask,PreRun
from .node import BatchedMaskedReferenceGen,ReferenceConditionCombine,MasksToPatches,ConditioningMatchMask,GenSubject,PreviewSubjectMask,PreRun,CrossSimMaskExtract
NODE_CLASS_MAPPINGS = {
"BatchedMaskedReferenceGen": BatchedMaskedReferenceGen,
@@ -8,6 +8,7 @@ NODE_CLASS_MAPPINGS = {
"GenSubject": GenSubject,
"PreviewSubjectMask" : PreviewSubjectMask,
"PreRun": PreRun,
"CrossSimMaskExtract": CrossSimMaskExtract
}
NODE_DISPLAY_NAME_MAPPINGS = {
"ReferenceConditionCombine": "Reference Condition Combine",
@@ -17,6 +18,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"GenSubject": "Gen Subject",
"PreviewSubjectMask": "Preview Subject Mask",
"PreRun": "PreRun",
"CrossSimMaskExtract": "CrossSim Mask Extract"
}
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
node.py
+80 -28
View File
@@ -13,6 +13,7 @@ from functools import partial
from einops import rearrange
from nodes import CLIPTextEncode
import latent_preview
import copy
orig_attention = flux_math.attention
orig_prepare_callback = latent_preview.prepare_callback
@@ -458,9 +459,9 @@ def get_shared_fg_mask(id_fg_mask, curr_fg_mask, curr2id_argmax_indices, curr2id
curr_valid_mask = curr2id_valid_mask.flatten()
id_fg_mask = id_fg_mask.flatten().to(curr2id_argmax_indices.device)
curr_fg_mask = curr_fg_mask.flatten().to(curr2id_argmax_indices.device)
rearrange_id_fg_mask = id_fg_mask[curr2id_argmax_indices[0]]
rearrange_id_fg_mask = id_fg_mask[curr2id_argmax_indices]
share_fg_mask = curr_fg_mask & rearrange_id_fg_mask & curr_valid_mask
id_share_fg_indices = curr2id_argmax_indices[0][share_fg_mask]
id_share_fg_indices = curr2id_argmax_indices[share_fg_mask]
curr_share_fg_indices = torch.nonzero(share_fg_mask).squeeze()
return id_share_fg_indices, curr_share_fg_indices
@@ -594,7 +595,8 @@ class GenSubject:
#if step == total_steps - 1:
# latent_preview.prepare_callback = orig_prepare_callback
return callback_change_sampled_step
model = copy.copy(model)
model.model_options = copy.copy(model.model_options)
model.model_options["prepare_callback"] = prepare_callback
hook_prepare_callback()
return (model, subject,)
@@ -611,14 +613,20 @@ class PreRun:
"model": ("MODEL",),
"subject": ("SUBJECT",),
"fg_condition_mask": ("MASK",),
"step": ("INT",),
"cross_sim_step": ("INT",),
"sim_thredhold": ("FLOAT",{"default":0.5,"min":0, "max":2}),
},
"optional":{
"latent": ("LATENT",),
}
}
@classmethod
def go(cls, model, subject, fg_condition_mask, step):
cross_sim = {"cross_sim": None}
def go(cls, model, subject, fg_condition_mask, cross_sim_step,sim_thredhold,latent):
cross_sim = {}
def prepare_callback(model, max_steps, x0_output_dict=None):
nonlocal fg_condition_mask, cross_sim_step
sampled_step = -1
sampling_step = -1
layer = 0
@@ -628,7 +636,6 @@ class PreRun:
attn_dtype = torch.bfloat16 if model_management.should_use_bf16(device=device) else torch.float16
fg_bg_mask = None
accum_sim = None
accum_count = 0
text_len = fg_condition_mask.shape[-1]
fg_condition_mask = fg_condition_mask.to(device)
def masked_attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask: Tensor) -> Tensor:
@@ -645,7 +652,7 @@ class PreRun:
Returns:
Output tensor of shape [B, L, H*D]
"""
nonlocal sampling_step, sampled_step, layer, subject, text_len, accum_sim, accum_count, fg_condition_mask
nonlocal sampling_step, sampled_step, layer, subject, text_len, accum_sim, fg_bg_mask
if sampling_step != sampled_step + 1:
layer = 0
sampling_step = sampled_step + 1
@@ -653,10 +660,10 @@ class PreRun:
if pe is not None:
q, k = flux_math.apply_rope(q, k, pe)
if sampling_step == step:
if sampling_step == cross_sim_step:
accum_mask = []
for i in enumerate(len(fg_condition_mask)):
accum_mask = get_fg_attn_weights(text_len, text_len, q, k, fg_condition_mask[i])
for i in range(fg_condition_mask.shape[0]):
accum_mask.append(get_fg_attn_weights(text_len, text_len, q, k, fg_condition_mask[i:i+1]))
accum_mask = torch.cat(accum_mask)
if fg_bg_mask is None:
fg_bg_mask = accum_mask
@@ -668,16 +675,15 @@ class PreRun:
# Reshape back to the expected output format
x = rearrange(x, "B H L D -> B L (H D)")
# Store attention result x along with k and v
if sampling_step == step:
if sampling_step == cross_sim_step:
all_cross_sim = []
for s in subject:
all_cross_sim.append(torch.matmul(F.normalize(x[:, text_len:, :], dim=-1), s[sampling_step]["layers"][layer]["x"]))
all_cross_sim.append(torch.matmul(F.normalize(x[:, text_len:, :], dim=-1), s[sampling_step]["layers"][layer]["x"].to(device = device, dtype = attn_dtype)))
this_cross_sim = torch.cat(all_cross_sim)
if accum_sim is None:
accum_sim = this_cross_sim
else:
accum_sim += this_cross_sim
accum_count += 1
layer += 1
return x
@@ -701,37 +707,83 @@ class PreRun:
else: # chroma
chroma_layers.attention = overrider
origin_callback = orig_prepare_callback(model, max_steps, x0_output_dict)
def callback_change_sampled_step(step, x0, x, total_steps):
nonlocal sampling_step, sampled_step, layer, model, cross_sim, accum_sim
if sampling_step == step:
def callback_change_sampled_step(step, x0, x, total_steps):
nonlocal sampling_step, sampled_step, layer, model, cross_sim, accum_sim, cross_sim_step,fg_bg_mask
sampled_step = step
if sampling_step == cross_sim_step:
B,L,H,W = x0.shape
w = W // model.model.diffusion_model.patch_size
h = H // model.model.diffusion_model.patch_size
bitmask = fg_bg_mask[:,0] >= fg_bg_mask[:,1]
bitmask = rearrange(bitmask, "b (h w) -> b h w", h=h, w=w)
fg_bg_mask = None
value = accum_sim / accum_count
value = accum_sim / layer
max_sim, argmax_indices = torch.max(value, dim=-1)
all_id_fg_inds = []
all_curr_fg_inds = []
for i in enumerate(len(fg_condition_mask)):
for i in range(len(fg_condition_mask)):
id_fg_inds, curr_fg_inds = get_shared_fg_mask(
subject[i][sampling_step]["fg_mask"].to(device),
subject[i][sampling_step]["fg_mask"][0].to(device),
bitmask[i],
argmax_indices[i],
max_sim[i]>0.5
max_sim[i]>sim_thredhold
)
all_id_fg_inds.append(id_fg_inds)
all_curr_fg_inds.append(curr_fg_inds)
cross_sim["id_fg_inds"] = torch.stack(curr_fg_inds)
cross_sim["curr_fg_inds"] = torch.stack(all_curr_fg_inds)
cross_sim["max_sim"] = max_sim
cross_sim["argmax_indices"] = argmax_indices
cross_sim["debug_mask"] = value
cross_sim["fg_mask"] = bitmask
cross_sim["id_fg_inds"] = torch.stack(all_id_fg_inds).cpu()
cross_sim["curr_fg_inds"] = torch.stack(all_curr_fg_inds).cpu()
cross_sim["max_sim"] = max_sim.cpu()
cross_sim["argmax_indices"] = argmax_indices.cpu()
cross_sim["debug_mask"] = value.to(dtype=torch.float32).cpu()
cross_sim["fg_mask"] = bitmask.cpu()
return callback_change_sampled_step
model = copy.copy(model)
model.model_options = copy.copy(model.model_options)
model.model_options["prepare_callback"] = prepare_callback
hook_prepare_callback()
return (model, cross_sim,)
return (model, cross_sim,)
class CrossSimMaskExtract:
RETURN_TYPES = ("MASK", "MASK")
RETURN_NAMES = ("cross_mask", "fg_mask")
FUNCTION = "extract"
CATEGORY = "chara_consist"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"cross_sim": ("CROSSSIM",),
"latent": ("LATENT",),
}
}
def extract(self, cross_sim, latent):
# Extract fg_mask directly from cross_sim
fg_mask = cross_sim.get("fg_mask", None)
# Create cross_mask from max_sim or debug_mask
# Using debug_mask as it appears to contain the cross attention similarity data
cross_mask = cross_sim.get("debug_mask", None)
# If debug_mask is not available, try to create from max_sim
#if cross_mask is None and "max_sim" in cross_sim:
# # max_sim is likely a tensor with similarity values
# max_sim = cross_sim["max_sim"]
# # Normalize to 0-1 range for mask
# if max_sim is not None:
# cross_mask = (max_sim - max_sim.min()) / (max_sim.max() - max_sim.min() + 1e-8)
# Handle case where neither is available
#if fg_mask is None:
# # Create an empty mask if fg_mask is not available
# fg_mask = torch.zeros(1, 1, 1, dtype=torch.float32)
if cross_mask is None:
# Create an empty mask if cross_mask is not available
cross_mask = torch.zeros(1, 1, 1, dtype=torch.float32)
return (cross_mask, fg_mask)