790 lines
34 KiB
Python
790 lines
34 KiB
Python
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
|
|
orig_attention = flux_math.attention
|
|
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": ("SUBJECT",),
|
|
},
|
|
}
|
|
@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 the 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, )
|
|
def hook_prepare_callback():
|
|
flux_layers.attention = orig_attention
|
|
chroma_layers.attention = orig_attention
|
|
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_interpolate_weight(weight, start_step, decay_step, end_step):
|
|
steps = np.arange(0, end_step - decay_step)
|
|
decay_weights = weight * 0.5 * (1 + np.cos(np.pi * steps / (end_step - decay_step)))
|
|
decay_weights = decay_weights.tolist()
|
|
constant_weights = [weight] * (decay_step - start_step)
|
|
weight_list = constant_weights + decay_weights
|
|
weight_dict = dict()
|
|
for ind, interpolate_step in enumerate(range(start_step, end_step)):
|
|
weight_dict[interpolate_step] = weight_list[ind]
|
|
return weight_dict
|
|
|
|
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 GenSubject:
|
|
RETURN_TYPES = ("MODEL","SUBJECT",)
|
|
FUNCTION = "go"
|
|
CATEGORY = "chara_consist"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"fg_condition_mask": ("MASK",),
|
|
"kv_start_step": ("INT",{"default":1,"min":0}),
|
|
"kv_end_step": ("INT",{"default":5,"min":0}),
|
|
"mask_start_step": ("INT",{"default":7,"min":0}),
|
|
"mask_end_step": ("INT",{"default":8,"min":0}),
|
|
}
|
|
}
|
|
@classmethod
|
|
def go(cls, model, fg_condition_mask, kv_start_step, kv_end_step, mask_start_step, mask_end_step):
|
|
|
|
subject = [{}]
|
|
index = len(subject) - 1
|
|
def prepare_callback(model, max_steps, x0_output_dict=None):
|
|
nonlocal subject, fg_condition_mask
|
|
subject[index] = {}
|
|
sampled_step = -1
|
|
sampling_step = -1
|
|
layer = 0
|
|
model_type = "flux"
|
|
|
|
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
|
|
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, subject, text_len, fg_bg_mask
|
|
if sampling_step != sampled_step + 1:
|
|
layer = 0
|
|
sampling_step = sampled_step + 1
|
|
# Initialize step in subject dict if not exists
|
|
if sampling_step not in subject[index]:
|
|
subject[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:
|
|
subject[index][sampling_step]["layers"][layer] = {"k": k.cpu(), "v": v.cpu()}
|
|
# Apply rotary positional encoding for Flux models
|
|
if pe is not None:
|
|
q, k = flux_math.apply_rope(q, k, pe)
|
|
|
|
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_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 >= 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"] = F.normalize(x[:, text_len:,:], dim=-1).transpose(-2, -1).cpu()
|
|
layer += 1
|
|
return x
|
|
|
|
# Create the attention override function
|
|
def attention_override(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask: Tensor, transformer_options) -> Tensor:
|
|
if transformer_options["cond_or_uncond"][-1]:#q.shape[2] != text_len + (w * h)
|
|
return orig_attention(q, k, v, pe)
|
|
else:
|
|
return masked_attention(q, k, v, pe, mask)
|
|
|
|
# Apply the attention override based on model type
|
|
overrider = attention_override
|
|
|
|
if model.model.__class__.__name__ == "Chroma":
|
|
model_type = "chroma"
|
|
|
|
print(f'GenSubject Detected model type: {model_type}')
|
|
|
|
if model_type == "flux":
|
|
flux_layers.attention = overrider
|
|
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 max_steps
|
|
nonlocal sampling_step, sampled_step, layer, fg_bg_mask, model
|
|
sampled_step = step
|
|
if sampling_step >= mask_start_step and sampling_step < mask_end_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
|
|
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,)
|
|
|
|
class PreRun:
|
|
RETURN_TYPES = ("MODEL","CROSSSIM",)
|
|
FUNCTION = "go"
|
|
CATEGORY = "chara_consist"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"subject": ("SUBJECT",),
|
|
"fg_condition_mask": ("MASK",),
|
|
"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, 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
|
|
model_type = "flux"
|
|
|
|
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
|
|
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.
|
|
|
|
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, text_len, accum_sim, fg_bg_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 == cross_sim_step:
|
|
accum_mask = []
|
|
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
|
|
else:
|
|
fg_bg_mask += accum_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 == 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"].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
|
|
layer += 1
|
|
return x
|
|
|
|
# Create the attention override function
|
|
def attention_override(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask: Tensor, transformer_options) -> Tensor:
|
|
if transformer_options["cond_or_uncond"][-1]:#q.shape[2] != text_len + (w * h)
|
|
return orig_attention(q, k, v, pe)
|
|
else:
|
|
return masked_attention(q, k, v, pe, mask)
|
|
|
|
# Apply the attention override based on model type
|
|
overrider = attention_override
|
|
|
|
if model.model.__class__.__name__ == "Chroma":
|
|
model_type = "chroma"
|
|
|
|
print(f'GenSubject Detected model type: {model_type}')
|
|
|
|
if model_type == "flux":
|
|
flux_layers.attention = overrider
|
|
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, 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 / layer
|
|
max_sim, argmax_indices = torch.max(value, dim=-1)
|
|
all_id_fg_inds = []
|
|
all_curr_fg_inds = []
|
|
for i in range(len(fg_condition_mask)):
|
|
id_fg_inds, curr_fg_inds = get_shared_fg_mask(
|
|
subject[i][sampling_step]["fg_mask"][0].to(device),
|
|
bitmask[i],
|
|
argmax_indices[i],
|
|
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(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,)
|
|
|
|
|
|
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)
|