diff --git a/__init__.py b/__init__.py index 73e29f9..d762b87 100644 --- a/__init__.py +++ b/__init__.py @@ -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'] \ No newline at end of file diff --git a/node.py b/node.py index 83bde24..2490734 100644 --- a/node.py +++ b/node.py @@ -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,) \ No newline at end of file + return (model, cross_sim,) \ No newline at end of file