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

PreRun

Browse Source
This commit is contained in:
ZHOU He
2025-11-15 19:39:42 +08:00
parent fdab9f0902
commit e43bb5d4d5
2 changed files with 68 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+4 -2
View File
@@ -1,4 +1,4 @@
from .node import BatchedMaskedReferenceGen,ReferenceConditionCombine,MasksToPatches,ConditioningMatchMask,GenSubject,PreviewSubjectMask
from .node import BatchedMaskedReferenceGen,ReferenceConditionCombine,MasksToPatches,ConditioningMatchMask,GenSubject,PreviewSubjectMask,PreRun
NODE_CLASS_MAPPINGS = {
"BatchedMaskedReferenceGen": BatchedMaskedReferenceGen,
@@ -6,7 +6,8 @@ NODE_CLASS_MAPPINGS = {
"ReferenceConditionCombine": ReferenceConditionCombine,
"ConditioningMatchMask": ConditioningMatchMask,
"GenSubject": GenSubject,
"PreviewSubjectMask" : PreviewSubjectMask
"PreviewSubjectMask" : PreviewSubjectMask,
"PreRun": PreRun,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"ReferenceConditionCombine": "Reference Condition Combine",
@@ -15,6 +16,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"BatchedMaskedReferenceGen": "Batched Masked Reference Gen",
"GenSubject": "Gen Subject",
"PreviewSubjectMask": "Preview Subject Mask",
"PreRun": "PreRun",
}
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
node.py
+64 -37
View File
@@ -464,6 +464,14 @@ def get_shared_fg_mask(id_fg_mask, curr_fg_mask, curr2id_argmax_indices, curr2id
curr_share_fg_indices = torch.nonzero(share_fg_mask).squeeze()
return id_share_fg_indices, curr_share_fg_indices
def get_fg_attn_weights(text_len, image_start, q, k, fg_condition_mask):
attn_weights = torch.matmul(q[:, :, image_start:], k[:, :, :text_len].transpose(-2, -1)).to(torch.float16)
attn_weights = attn_weights.softmax(-1).sum(1, keepdim=True)
fg_condition_mask_gpu = fg_condition_mask
bg_condition_mask_gpu = 1.0 - fg_condition_mask_gpu
masks = torch.nn.functional.normalize(torch.stack([fg_condition_mask_gpu, bg_condition_mask_gpu], dim=1), p=1, dim=-1)
attn_weights = (attn_weights * masks).sum(-1)
return attn_weights
class GenSubject:
RETURN_TYPES = ("MODEL","SUBJECT",)
@@ -488,7 +496,7 @@ class GenSubject:
subject = [{}]
index = len(subject) - 1
def prepare_callback(model, max_steps, x0_output_dict=None):
nonlocal subject
nonlocal subject, fg_condition_mask
subject[index] = {}
sampled_step = -1
sampling_step = -1
@@ -499,16 +507,7 @@ class GenSubject:
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
fg_bg_mask = None
def get_fg_attn_weights(text_len, image_start, q, k):
nonlocal fg_condition_mask
attn_weights = torch.matmul(q[:, :, image_start:], k[:, :, :text_len].transpose(-2, -1)).to(torch.float16)
attn_weights = attn_weights.softmax(-1).sum(1, keepdim=True)
fg_condition_mask_gpu = fg_condition_mask.to(device)
bg_condition_mask_gpu = 1.0 - fg_condition_mask_gpu
masks = torch.nn.functional.normalize(torch.stack([fg_condition_mask_gpu, bg_condition_mask_gpu], dim=1), p=1, dim=-1)
attn_weights = (attn_weights * masks).sum(-1)
return attn_weights
fg_condition_mask = fg_condition_mask.to(device)
def masked_attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask: Tensor) -> Tensor:
"""
Implementation using PyTorch's scaled_dot_product_attention with custom masking.
@@ -539,9 +538,9 @@ class GenSubject:
if sampling_step >= mask_start_step and sampling_step < mask_end_step:
if fg_bg_mask is None:
fg_bg_mask = get_fg_attn_weights(text_len, text_len, q, k)
fg_bg_mask = get_fg_attn_weights(text_len, text_len, q, k, fg_condition_mask)
else:
fg_bg_mask += get_fg_attn_weights(text_len, text_len, q, k)
fg_bg_mask += get_fg_attn_weights(text_len, text_len, q, k, fg_condition_mask)
x = F.scaled_dot_product_attention(
q, k, v, attn_mask=mask
@@ -552,7 +551,7 @@ class GenSubject:
if sampling_step >= mask_start_step and sampling_step < mask_end_step:
if not layer in subject[index][sampling_step]["layers"]:
subject[index][sampling_step]["layers"][layer] = {}
subject[index][sampling_step]["layers"][layer]["x"] = x[:, text_len:,:].cpu()
subject[index][sampling_step]["layers"][layer]["x"] = F.normalize(x[:, text_len:,:], dim=-1).transpose(-2, -1).cpu()
layer += 1
return x
@@ -611,22 +610,27 @@ class PreRun:
"required": {
"model": ("MODEL",),
"subject": ("SUBJECT",),
"fg_condition_mask": ("MASK",),
"step": ("INT",),
}
}
@classmethod
def go(cls, model, subject, step):
def go(cls, model, subject, fg_condition_mask, step):
cross_sim = {"cross_sim": None}
def prepare_callback(model, max_steps, x0_output_dict=None):
sampled_step = -1
sampling_step = -1
layer = 0
model_type = "flux"
image_patches = 0
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
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:
"""
Implementation using PyTorch's scaled_dot_product_attention with custom masking.
@@ -641,14 +645,23 @@ class PreRun:
Returns:
Output tensor of shape [B, L, H*D]
"""
nonlocal sampling_step, sampled_step, layer, subject, image_patches, cross_sim
nonlocal sampling_step, sampled_step, layer, subject, text_len, accum_sim, accum_count, fg_condition_mask
if sampling_step != sampled_step + 1:
layer = 0
sampling_step = sampled_step + 1
if pe is not None:
q, k = flux_math.apply_rope(q, k, pe)
if sampling_step == 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])
accum_mask = torch.cat(accum_mask)
if fg_bg_mask is None:
fg_bg_mask = accum_mask
else:
fg_bg_mask += accum_mask
x = F.scaled_dot_product_attention(
q, k, v, attn_mask=mask
)
@@ -656,7 +669,15 @@ class PreRun:
x = rearrange(x, "B H L D -> B L (H D)")
# Store attention result x along with k and v
if sampling_step == step:
cross_sim
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"]))
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
@@ -681,30 +702,36 @@ class PreRun:
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 max_steps, image_patches
nonlocal sampling_step, sampled_step, layer, fg_bg_mask, model
B,L,H,W = x0.shape
w = W // model.model.diffusion_model.patch_size
h = H // model.model.diffusion_model.patch_size
image_patches = w * h
sampled_step = step
if sampling_step >= mask_start_step and sampling_step < mask_end_step:
nonlocal sampling_step, sampled_step, layer, model, cross_sim, accum_sim
if sampling_step == 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
# Store fg_mask in a dictionary with step id as key
if sampling_step not in subject[index]:
subject[index][sampling_step] = {"layers": {}}
subject[index][sampling_step]["fg_mask"] = bitmask.cpu()
origin_callback(step, x0, x, total_steps)
#if step == total_steps - 1:
# latent_preview.prepare_callback = orig_prepare_callback
value = accum_sim / accum_count
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)):
id_fg_inds, curr_fg_inds = get_shared_fg_mask(
subject[i][sampling_step]["fg_mask"].to(device),
bitmask[i],
argmax_indices[i],
max_sim[i]>0.5
)
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
return callback_change_sampled_step
model.model_options["prepare_callback"] = prepare_callback
hook_prepare_callback()
return (model, subject,)
return (model, cross_sim,)