import gc import torch import torch.nn.functional as F import numpy as np from torch import Tensor from comfy.ldm.modules import attention as comfy_attention from comfy.ldm.flux import math as flux_math from comfy.ldm.flux import layers as flux_layers from comfy.ldm.chroma import layers as chroma_layers from comfy import model_management from typing import List, Dict, Optional from functools import partial from einops import rearrange from nodes import CLIPTextEncode import latent_preview import copy 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: """ Convert pixel masks to patch-level masks for attention. Args: masks: [N, H, W] tensor - subject regions in images Returns: patch_masks: [N, h*w] tensor where h=H_latent*8/16, w=W_latent*8/16 h, w: patch dimensions """ N = masks.shape[0] H, W = masks.shape[1], masks.shape[2] # Calculate patch dimensions (assuming 16x16 pixels per patch, but 1x1 in latent space) h = H // 16 w = W // 16 patch_masks = torch.zeros(N + pad, h, w, device=masks.device, dtype=masks.dtype) for i in range(N): # Convert to patch level using nearest-neighbor interpolation mask_resized = F.interpolate( masks[i:i+1, None, :, :], size=(h, w), mode='nearest-exact' ) # Flatten to patch sequence patch_masks[i] = mask_resized[0][0] for i in range(pad): patch_masks[i + N] = torch.ones((h, w), device=masks.device, dtype=masks.dtype) return patch_masks, h, w def match_condition(reference_conditions: List, target_condition: List) -> torch.Tensor: """ Match reference conditions to target condition using dot product similarity. Args: reference_conditions: List of N conditioning tensors (one per reference subject) target_condition: Conditioning tensor for target images (shared) Returns: conditions_mask: 2D tensor [num_references, max_embedding_seq_length] with 1s where reference embeddings match target subsequences """ # Extract embeddings ref_embeds = [cond[0][0] for cond in reference_conditions] # text embeddings target_embed = target_condition[0][0] # [1, seq_len, embed_dim] num_references = len(ref_embeds) target_seq_len = target_embed.shape[1] # Initialize mask tensor conditions_mask = torch.zeros(num_references, target_seq_len, dtype=target_embed.dtype, device=target_embed.device) # Process each reference embedding for i, ref_embed in enumerate(ref_embeds): ref_seq_len = ref_embed.shape[1] # If reference is longer than target, skip (shouldn't happen per spec) if ref_seq_len > target_seq_len: continue # Get reference embedding (remove batch dimension) ref_vec = ref_embed[0] # [ref_seq_len, embed_dim] # Use unfold to get all subsequences of target with length ref_seq_len # target_embed[0] is [target_seq_len, embed_dim] # unfolded will be [target_seq_len - ref_seq_len + 1, ref_seq_len, embed_dim] unfolded = target_embed[0].unfold(0, ref_seq_len, 1).permute(0,2,1) # [num_subseq, ref_seq_len, embed_dim] # Normalize for cosine similarity ref_normalized = F.normalize(ref_vec, p=2, dim=-1) # [ref_seq_len, embed_dim] target_normalized = F.normalize(unfolded, p=2, dim=-1) # [num_subseq, ref_seq_len, embed_dim] # Compute cosine similarity for all subsequences at once # Result: [num_subseq] similarities = torch.mean(torch.sum(ref_normalized.unsqueeze(0) * target_normalized, dim=-1), dim=-1) # Find the best matching subsequence best_match_idx = torch.argmax(similarities) # Mask the best matching subsequence start_idx = best_match_idx.item() conditions_mask[i, start_idx:start_idx + ref_seq_len] = 1.0 return conditions_mask def prepare_batched_conditions(reference_conditions: List, target_condition: List, num_targets) -> tuple: """ Prepare the batched conditions for reference and target images. Args: reference_conditions: List of N conditioning tensors (one per reference subject) target_condition: Conditioning tensor for target images (shared) Returns: extended_condition: Batched conditioning tensor [B, num_reference + 1, max_seq, 4096] conditions_mask: Mask tensor [B, num_references + 1, max_seq] pooled_output: CLIP pooled output (None for Chroma, tensor for Flux) """ # Extract text embeddings from conditions ref_embeds = [] ref_pooled = [] seq_lens = [] for cond in reference_conditions: if cond is None: continue ref_embeds.append(cond[0][0]) # text embedding ref_pooled.append(cond[0][1].get('pooled_output', None)) # Target condition target_embed = target_condition[0][0] target_pooled = target_condition[0][1].get('pooled_output', None) # Calculate max sequence length all_embeds = [target_embed] + ref_embeds max_seq_len = max(embed.shape[1] for embed in all_embeds) # Get batch size and embedding dimension embed_dim = target_embed.shape[2] # Should be 4096 # Create padded tensors with shape [B, num_reference + 1, max_seq, 4096] num_references = len(ref_embeds) padded_embeds = torch.zeros(num_references + num_targets, max_seq_len, embed_dim, dtype=target_embed.dtype, device=target_embed.device) conditions_mask = torch.zeros(num_references + num_targets, max_seq_len, dtype=target_embed.dtype, device=target_embed.device) # Fill in target embeds (first position) target_seq_len = target_embed.shape[1] # Fill in reference embeds for i, ref_embed in enumerate(ref_embeds): ref_seq_len = ref_embed.shape[1] padded_embeds[i, :ref_seq_len, :] = ref_embed[0] conditions_mask[i, :ref_seq_len] = 1.0 seq_lens.append(ref_seq_len) for i in range(num_targets): padded_embeds[i + num_references, :target_seq_len, :] = target_embed[0] conditions_mask[i + num_references, :target_seq_len] = 1.0 seq_lens.append(target_seq_len) # Handle pooled output (only used by Flux) if target_pooled is not None: # For Flux, use the target pooled output pooled_output = target_pooled else: pooled_output = None extended_condition = [[ padded_embeds, {'pooled_output': pooled_output} ]] return extended_condition, conditions_mask, max_seq_len, seq_lens class ReferenceConditionCombine: RETURN_TYPES = ("CONDITIONING", "MASK") RETURN_NAMES = ("condition", "condition_mask") FUNCTION = "go" CATEGORY = "chara_consist" @classmethod def INPUT_TYPES(cls): return { "required": { "num_targets": ("INT",), }, "optional": { "reference_condition1": ("CONDITIONING",), "reference_condition2": ("CONDITIONING",), "reference_condition3": ("CONDITIONING",), "reference_condition4": ("CONDITIONING",), "reference_condition5": ("CONDITIONING",), "target_condition": ("CONDITIONING",), } } @classmethod def go(cls, *, num_targets, reference_condition1=None, reference_condition2=None, reference_condition3=None, reference_condition4=None, reference_condition5=None, target_condition=None): reference_conditions = [reference_condition1, reference_condition2, reference_condition3, reference_condition4, reference_condition5] extended_condition, conditions_mask, max_seq_len, seq_lens = prepare_batched_conditions(reference_conditions, target_condition, num_targets) return extended_condition, conditions_mask.unsqueeze(0) class BatchedMaskedReferenceGen: RETURN_TYPES = ("MODEL",) RETURN_NAMES = ("model",) FUNCTION = "go" CATEGORY = "chara_consist" @classmethod def INPUT_TYPES(cls): return { "required": { "model": ("MODEL",), "source_patch_masks": ("MASK",), "target_patch_masks": ("MASK",), "conditions_mask": ("MASK",), "conditions_ref_mask": ("MASK",), }, } @classmethod def go(cls, *, model: object, source_patch_masks: torch.Tensor, target_patch_masks: torch.Tensor, conditions_mask: torch.Tensor, conditions_ref_mask=torch.Tensor): print(f'BatchedMaskedReferenceGen - Processing with model: {model}') if conditions_mask.dim() == 3: conditions_mask = conditions_mask[0] if conditions_ref_mask.dim() == 3: conditions_ref_mask = conditions_ref_mask[0] # Detect model type model_type = "flux" if model.model.__class__.__name__ == "Chroma": model_type = "chroma" print(f'Detected model type: {model_type}') num_references = conditions_ref_mask.shape[0] num_targets = conditions_mask.shape[0] - num_references 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 w = source_patch_masks.shape[-1] h = source_patch_masks.shape[1] patch_count = h * w source_patch_masks = source_patch_masks.flatten(start_dim=1) target_patch_masks = target_patch_masks.flatten(start_dim=1) reference_seq_mask = torch.log(torch.cat([conditions_mask[:num_references], torch.ones((num_references, patch_count), dtype=attn_dtype)], dim=1)) source_seq_masks = torch.log(torch.cat([conditions_mask[:num_references+1], source_patch_masks], dim=1)) target_seq_masks = torch.log(torch.cat([conditions_mask[num_references].unsqueeze(0).repeat(num_references+1,1), target_patch_masks], dim=1)) def masked_attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, reference_seq_mask, target_seq_masks, source_seq_masks) -> Tensor: """ Implementation using PyTorch's scaled_dot_product_attention with custom masking. Args: q: Query tensor of shape [B, H, L, D] k: Key tensor of shape [B, H, L, D] v: Value tensor of shape [B, H, L, D] pe: Positional encoding tensor attn_mask: Optional attention mask of shape [B, H, L, L] Returns: Output tensor of shape [B, L, H*D] """ target_count = q.shape[0] - num_references # Apply rotary positional encoding for Flux models if pe is not None: q, k = flux_math.apply_rope(q, k, pe) # Get actual sequence length from QKV tensors actual_seq_len = q.shape[2] #if attn_mask is None or attn_mask.shape[-2] != actual_seq_len: # this_source = torch.cat([torch.zeros([num_references, text_len]), source_patch_masks], dim=-1) # this_target = torch.cat([torch.zeros([num_references, text_len]), source_patch_masks], dim=-1) # attn_mask = torch.stack(torch.min(this_source.unsqueeze(1), this_target.unsqueeze(2)), dim=-1).to(device) reference_seq_mask_ = reference_seq_mask.to(device, dtype=attn_dtype) reference_attn_mask = torch.min(reference_seq_mask_.unsqueeze(1), reference_seq_mask_.unsqueeze(2)) target_attn_mask = rearrange(torch.min(source_seq_masks.to(device, dtype=attn_dtype).unsqueeze(1), target_seq_masks.to(device, dtype=attn_dtype).unsqueeze(2)), "B H W -> 1 W (B H)") h = torch.zeros(q.shape, device=device) # PyTorch's scaled_dot_product_attention h[:num_references] = F.scaled_dot_product_attention( q[:num_references], k[:num_references], v[:num_references], attn_mask=reference_attn_mask ) B, H, L, D = k.shape k = k.permute(1, 0, 2, 3).reshape(1, H, B * L, D) B, H, L, D = v.shape v = v.permute(1, 0, 2, 3).reshape(1, H, B * L, D) for i in range(target_count): target_idx = i + num_references if i > 0: k[:, :, num_references*actual_seq_len : num_references*actual_seq_len + actual_seq_len] = k[:, :, target_idx*actual_seq_len : target_idx*actual_seq_len + actual_seq_len] v[:, :, num_references*actual_seq_len : num_references*actual_seq_len + actual_seq_len] = v[:, :, target_idx*actual_seq_len : target_idx*actual_seq_len + actual_seq_len] h[target_idx:target_idx + 1] = F.scaled_dot_product_attention( q[target_idx: target_idx + 1], k[:, :, :num_references*actual_seq_len + actual_seq_len], v[:, :, :num_references*actual_seq_len + actual_seq_len], attn_mask=target_attn_mask) # Reshape back to the expected output format h = rearrange(h, "B H L D -> B L (H D)") return h # Create the attention override function def batched_masked_attention_override(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask, transformer_options, reference_seq_mask, target_seq_masks, source_seq_masks) -> Tensor: if q.shape[0] != num_references + num_targets or reference_seq_mask is None or reference_seq_mask.shape[-1] != q.shape[2]: return orig_attention(q, k, v, pe) else: return masked_attention(q, k, v, pe, reference_seq_mask, target_seq_masks, source_seq_masks) # Apply the attention override based on model type overrider = partial(batched_masked_attention_override, reference_seq_mask=reference_seq_mask, target_seq_masks=target_seq_masks, source_seq_masks=source_seq_masks) if model_type == "flux": flux_math.attention = overrider else: # chroma chroma_layers.attention = overrider # Clean up #del source_patch_masks, target_patch_masks gc.collect() torch.cuda.empty_cache() print('BatchedMaskedReferenceGen processing completed') return (model,) class ConditioningMatchMask: RETURN_TYPES = ("MASK",) RETURN_NAMES = ("mask",) FUNCTION = "match" CATEGORY = "chara_consist" @classmethod def INPUT_TYPES(cls): return { "required": { }, "optional": { "reference_condition1": ("CONDITIONING",), "reference_condition2": ("CONDITIONING",), "reference_condition3": ("CONDITIONING",), "reference_condition4": ("CONDITIONING",), "reference_condition5": ("CONDITIONING",), "target_condition": ("CONDITIONING",), } } def match(self, reference_condition1=None, reference_condition2=None, reference_condition3=None, reference_condition4=None, reference_condition5=None, target_condition=None): # Collect reference conditions reference_conditions = [] ref_conditions = [reference_condition1, reference_condition2, reference_condition3, reference_condition4, reference_condition5] for cond in ref_conditions: if cond is not None: reference_conditions.append(cond) if not reference_conditions: raise ValueError("At least one reference condition must be provided") if target_condition is None: raise ValueError("Target condition must be provided") # Call the match_condition function conditions_mask = match_condition(reference_conditions, target_condition) return (conditions_mask.unsqueeze(0),) class MasksToPatches: RETURN_TYPES = ("MASK", "INT", "INT") RETURN_NAMES = ("patch_masks", "h", "w") FUNCTION = "convert" CATEGORY = "chara_consist" @classmethod def INPUT_TYPES(cls): return { "required": { "masks": ("MASK",), "target_count": ("INT",{"default":1,"min":0}), } } def convert(self, masks, target_count): # Call the convert_masks_to_patches function patch_masks, h, w = convert_masks_to_patches(masks, target_count) print(f'Patch dimensions: {h}x{w}, Total patches: {h*w}') return (patch_masks, h, w) class PreviewSubjectMask: RETURN_TYPES = ("MASK","LATENT",) RETURN_NAMES = ("mask","latent",) FUNCTION = "go" CATEGORY = "chara_consist" @classmethod def INPUT_TYPES(cls): return { "required": { "latent": ("LATENT",), "subject": ("ATTN",), }, } @classmethod def go(cls, latent, subject): # Iterate through subject[0] and gather all fg_masks fg_masks = [] for s in subject: # Now s is a dict where keys are step ids for step_id, step_data in s.items(): if "fg_mask" in step_data: fg_masks.append(step_data["fg_mask"]) # Stack all fg_masks into a single tensor if fg_masks: # Stack masks to create a tensor with shape [B, H, W] combined_mask = torch.cat(fg_masks, dim=0) else: # If no fg_masks found, return empty tensor combined_mask = torch.empty(0) return (combined_mask, latent, ) class MaskToPatchMask: RETURN_TYPES = ("MASK",) FUNCTION = "set_mask" CATEGORY = "chara_consist" @classmethod def INPUT_TYPES(cls): return { "required": { "mask": ("MASK",), "patch_size": ("INT",{"default": 16}), "threshold": ("FLOAT",{"default": 1.0}), } } @classmethod def set_mask(cls, mask: torch.Tensor, patch_size, threshold): # Get the dimensions of the input mask B, H, W = mask.shape # Calculate patch dimensions (divide by 16) h_patches = H // patch_size w_patches = W // patch_size # Reshape mask to [B, H//16, 16, W//16, 16] for processing 16x16 blocks mask_reshaped = mask.view(B, h_patches, patch_size, w_patches, patch_size) # Check if any pixel in each 16x16 block is non-zero # Take max over the 16x16 blocks, then convert to boolean patch_mask = mask_reshaped.mean(dim=2).mean(dim=3) >= threshold # Convert to boolean tensor 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(): 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) else: return orig_prepare_callback(model, max_steps, x0_output_dict) latent_preview.prepare_callback = prepare_callback def get_shared_fg_mask(id_fg_mask, curr_fg_mask, curr2id_argmax_indices, curr2id_valid_mask): 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] share_fg_mask = curr_fg_mask & rearrange_id_fg_mask & curr_valid_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 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 ExtractAttn: RETURN_TYPES = ("MODEL","ATTN",) FUNCTION = "go" CATEGORY = "chara_consist" @classmethod def INPUT_TYPES(cls): return { "required": { "model": ("MODEL",), "kv_start_step": ("INT",{"default":1,"min":1}), "kv_end_step": ("INT",{"default":9,"min":1}), "x_start_step": ("INT",{"default":1,"min":1}), "x_end_step": ("INT",{"default":9,"min":1}), "mask_start_step": ("INT",{"default":8,"min":1}), "mask_end_step": ("INT",{"default":9,"min":1}), }, "optional":{ "fg_condition_mask": ("MASK",), } } @classmethod def go(cls, model, kv_start_step, kv_end_step, x_start_step, x_end_step, mask_start_step, mask_end_step, fg_condition_mask = None): attn_data = [{}] index = len(attn_data) - 1 def prepare_callback(model, max_steps, x0_output_dict=None): nonlocal attn_data, fg_condition_mask attn_data[index] = {} sampled_step = -1 sampling_step = -1 layer = 0 w = 0 h = 0 #text_len = fg_condition_mask.shape[-1] 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 if fg_condition_mask is not None: 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. Args: q: Query tensor of shape [B, H, L, D] k: Key tensor of shape [B, H, L, D] v: Value tensor of shape [B, H, L, D] pe: Positional encoding tensor attn_mask: Optional attention mask of shape [B, H, L, L] Returns: Output tensor of shape [B, L, H*D] """ nonlocal sampling_step, sampled_step, layer, attn_data, w, h, fg_bg_mask text_len = q.shape[2] - w * h if sampling_step != sampled_step + 1: layer = 0 sampling_step = sampled_step + 1 # Initialize step in attn_data dict if not exists if sampling_step not in attn_data[index]: attn_data[index][sampling_step] = {"layers": {}} # Store k, v, and x in a dict with layer id as key if sampling_step >= kv_start_step and sampling_step < kv_end_step: attn_data[index][sampling_step]["layers"][layer] = {"k": k[:, :, text_len:,:].cpu(), "v": v[:, :, text_len:,:].cpu()} # Apply rotary positional encoding for Flux models if pe is not None: 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: 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_condition_mask) x = F.scaled_dot_product_attention( q, k, v, attn_mask=mask ) # 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 >= x_start_step and sampling_step < x_end_step: if not layer in attn_data[index][sampling_step]["layers"]: attn_data[index][sampling_step]["layers"][layer] = {} attn_data[index][sampling_step]["layers"][layer]["x"] = x[:, text_len:,:].cpu() layer += 1 return x 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 nonlocal sampling_step, sampled_step, layer, fg_bg_mask, model, w, h w = x0.shape[-1] // model.model.diffusion_model.patch_size h = x0.shape[-2] // model.model.diffusion_model.patch_size sampled_step = step if fg_condition_mask is not None and sampling_step >= mask_start_step and sampling_step < mask_end_step: 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 attn_data[index]: attn_data[index][sampling_step] = {"layers": {}} attn_data[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 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, attn_data,) class GetCrossSim: RETURN_TYPES = ("CROSSSIM",) RETURN_NAMES = ("cross_sim",) FUNCTION = "compute_cross_sim" CATEGORY = "chara_consist" @classmethod def INPUT_TYPES(cls): return { "required": { "subject_attn": ("ATTN",), "curr_attn": ("ATTN",), "id_mask": ("MASK",), "curr_mask": ("MASK",), "cross_sim_step": ("INT",), "sim_threshold": ("FLOAT", {"default":0.5,"min":0,"max":2}), } } @classmethod def compute_cross_sim(cls, subject_attn, curr_attn, id_mask, curr_mask, cross_sim_step, sim_threshold): cross_sim = {} 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 # Validate inputs if not subject_attn or len(subject_attn) == 0: raise ValueError("subject_attn cannot be empty") if not curr_attn or len(curr_attn) == 0: raise ValueError("curr_attn cannot be empty") # Use the first (and only) current image attention data curr_attn_data = curr_attn[0] # Check if the specified step exists in current image attention data if cross_sim_step not in curr_attn_data: raise ValueError(f"cross_sim_step {cross_sim_step} not found in curr_attn") # Get current attention layers for the specified step curr_layers = curr_attn_data[cross_sim_step].get("layers", {}) if not curr_layers: raise ValueError(f"No layers found for step {cross_sim_step} in curr_attn") # Initialize accumulators all_cross_sim = [] all_max_sim = [] all_argmax_indices = [] all_id_fg_inds = [] all_curr_fg_inds = [] 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 # Iterate through each ID subject for id_idx, id_attn_data in enumerate(subject_attn): # Check if the specified step exists in this ID subject if cross_sim_step not in id_attn_data: continue # Skip this subject if step doesn't exist # Get ID subject layers for the specified step id_layers = id_attn_data[cross_sim_step].get("layers", {}) if not id_layers: continue # Skip this subject if no layers id_fg_mask = id_mask[id_idx].to(device).flatten() curr_fg_mask = curr_mask[id_idx].to(device).flatten() # todo multiple images in 1 batch # Accumulate cross similarity across all layers layer_sim = None for layer_idx in curr_layers: if layer_idx not in id_layers: continue # Skip layer if not present in both subjects # Get attention vectors for this layer curr_x = F.normalize(curr_layers[layer_idx]["x"].to(device)) # [B, D, L] format id_x = F.normalize(id_layers[layer_idx]["x"].to(device)) # [B, D, L] format # Compute cross similarity: current @ ID^T # curr_x is [B, D, L], id_x is [B, D, L] # We want: curr_x^T @ id_x -> [L, L] similarity = torch.matmul(curr_x, id_x.transpose(-2, -1)) # [B, L, L] if layer_sim is None: layer_sim = similarity else: layer_sim += similarity if layer_sim is not None: # Average across layers avg_sim = layer_sim / len(curr_layers) max_sim, argmax_indices = torch.max(avg_sim, dim=-1) id_fg_inds, curr_fg_inds = get_shared_fg_mask( # todo multiple batch gen. id_fg_mask, curr_fg_mask, argmax_indices[0], max_sim[0]>sim_threshold ) all_id_fg_inds.append(id_fg_inds.cpu()) all_curr_fg_inds.append(curr_fg_inds.cpu()) all_max_sim.append(max_sim) all_argmax_indices.append(argmax_indices) all_cross_sim.append(avg_sim) cross_sim["id_fg_inds"] = all_id_fg_inds cross_sim["curr_fg_inds"] = all_curr_fg_inds cross_sim["max_sim"] = torch.cat(all_max_sim).cpu() cross_sim["argmax_indices"] = torch.cat(all_argmax_indices).cpu() cross_sim["debug_mask"] = torch.cat(all_cross_sim).cpu() cross_sim["id_mask"] = id_mask cross_sim["curr_mask"] = curr_mask return (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) class GenConsistent: RETURN_TYPES = ("MODEL",) FUNCTION = "go" CATEGORY = "chara_consist" @classmethod def INPUT_TYPES(cls): return { "required": { "model": ("MODEL",), "subject": ("ATTN",), "cross_sim": ("CROSSSIM",), "kv_start_step": ("INT",{"default":1,"min":1}), "kv_end_step": ("INT",{"default":5,"min":1}), "x_start_step": ("INT",{"default":1,"min":1}), "x_end_step": ("INT",{"default":5,"min":1}), "ada_tome_alpha": ("FLOAT",{"default":0.0,"min":0,"max":1,"step":0.1}), }, } @classmethod def go(cls, model, subject, cross_sim, kv_start_step, kv_end_step, x_start_step, x_end_step, ada_tome_alpha): def prepare_callback(model, max_steps, x0_output_dict=None): nonlocal subject sampled_step = -1 sampling_step = -1 layer = 0 w = 0 h = 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 def masked_attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, _: Tensor) -> Tensor: """ Implementation using PyTorch's scaled_dot_product_attention with custom masking. Args: q: Query tensor of shape [B, H, L, D] k: Key tensor of shape [B, H, L, D] v: Value tensor of shape [B, H, L, D] pe: Positional encoding tensor attn_mask: Optional attention mask of shape [B, H, L, L] Returns: Output tensor of shape [B, L, H*D] """ nonlocal sampling_step, sampled_step, layer, subject, w, h, fg_bg_mask if sampling_step != sampled_step + 1: layer = 0 sampling_step = sampled_step + 1 text_len = q.shape[2] - w*h original_seq_len = q.shape[2] # Apply rotary positional encoding for Flux models if pe is not None: 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] # Store k, v, and x in a dict with layer id as key if sampling_step >= kv_start_step and sampling_step < kv_end_step: # For all other steps, we may apply the key-value expansion algorithm # Get cross-sim data for this layer id_fg_inds = cross_sim.get("id_fg_inds", None) curr_fg_inds = cross_sim.get("curr_fg_inds", None) curr_mask = cross_sim.get("curr_mask", None) # 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 (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 "v" in subject[subj_idx][sampling_step]["layers"][layer]): ori_saved_key = subject[subj_idx][sampling_step]["layers"][layer]["k"].to(k.device) ori_saved_value = subject[subj_idx][sampling_step]["layers"][layer]["v"].to(v.device) # Ensure we have valid indices within tensor bounds # Handle batch dimension properly - cross_sim tensors have [B, ...] form where B = len(subject) subj_id_fg_inds = id_fg_inds[subj_idx] subj_curr_fg_inds = curr_fg_inds[subj_idx] valid_id_inds = subj_id_fg_inds[subj_id_fg_inds < ori_saved_key.shape[2]] valid_curr_inds = subj_curr_fg_inds[subj_curr_fg_inds < k.shape[2]] if len(valid_id_inds) > 0 and len(valid_curr_inds) > 0: # Apply key-value expansion based on cross-sim data # Extract the saved key-value pairs using valid indices saved_key_fg = ori_saved_key[:, :, valid_id_inds] saved_value_fg = ori_saved_value[:, :, valid_id_inds] # Apply rotary positional encoding to saved keys (fix missing apply_rope) if pe is not None: # 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 = 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) v = torch.cat([v, saved_value_fg], dim=2) x = F.scaled_dot_product_attention( q, k, v, attn_mask=None ) # Reshape back to the expected output format x = rearrange(x, "B H L D -> B L (H D)") # Apply ada_tome blending if within specified step range if sampling_step >= x_start_step and sampling_step < x_end_step and ada_tome_alpha > 0: # Get cross-similarity data id_fg_inds = cross_sim.get("id_fg_inds", None)#.to(device) curr_fg_inds = cross_sim.get("curr_fg_inds", None)#.to(device) max_sim = cross_sim.get("max_sim", None).to(device) if (id_fg_inds is not None and curr_fg_inds is not None and max_sim is not None): # Process all subjects for blending for subj_idx in range(len(subject)): # Check if we have valid subject data for this step and layer if (sampling_step in subject[subj_idx] and layer in subject[subj_idx][sampling_step]["layers"] and "x" in subject[subj_idx][sampling_step]["layers"][layer]): # Get saved hidden states from subject saved_x = subject[subj_idx][sampling_step]["layers"][layer]["x"].to(x.device, dtype=attn_dtype) # Ensure we have valid indices within tensor bounds if subj_idx < len(id_fg_inds) and subj_idx < len(curr_fg_inds) and subj_idx < max_sim.shape[0]: subj_id_fg_inds = id_fg_inds[subj_idx] subj_curr_fg_inds = curr_fg_inds[subj_idx] subj_max_sim = max_sim[subj_idx] # Validate indices 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_current_x = vision_hidden_states[:, valid_curr_inds, :] # [B, valid_curr_inds, D] # 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 # Weight by similarity scores sim_weight = subj_max_sim.flatten()[valid_curr_inds] if len(valid_curr_inds) <= subj_max_sim.numel() else torch.ones(len(valid_curr_inds), dtype=attn_dtype, device=x.device) alpha_tensor = alpha_tensor * sim_weight alpha_tensor = alpha_tensor.view(1, -1, 1) # Blend: (1 - alpha) * current + alpha * saved blended_states = (1 - alpha_tensor) * matched_current_x + alpha_tensor * matched_saved_x # Update current vision hidden states with blended values vision_hidden_states[:, valid_curr_inds, :] = blended_states # Reconstruct full hidden states with text + blended vision x = torch.cat([x[:, :text_len, :], vision_hidden_states], dim=-2) layer += 1 return x 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 nonlocal sampling_step, sampled_step, layer, fg_bg_mask, model, w, h w = x0.shape[-1] // model.model.diffusion_model.patch_size h = x0.shape[-2] // model.model.diffusion_model.patch_size sampled_step = step origin_callback(step, x0, x, total_steps) 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,)