diff --git a/__init__.py b/__init__.py index d762b87..a729c7e 100644 --- a/__init__.py +++ b/__init__.py @@ -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'] \ No newline at end of file diff --git a/node.py b/node.py index 2490734..eeaa127 100644 --- a/node.py +++ b/node.py @@ -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,) \ No newline at end of file + 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)