From d93ed94c26e16a7053b36c8abaee4104685418eb Mon Sep 17 00:00:00 2001 From: ZHOU He <461607353@qq.com> Date: Fri, 14 Nov 2025 00:30:10 +0800 Subject: [PATCH] Add subject gen. --- __init__.py | 6 +- node.py | 218 ++++++++++++++++++++++++---------------------------- 2 files changed, 107 insertions(+), 117 deletions(-) diff --git a/__init__.py b/__init__.py index b4c9954..73e29f9 100644 --- a/__init__.py +++ b/__init__.py @@ -1,16 +1,20 @@ -from .node import BatchedMaskedReferenceGen,ReferenceConditionCombine,MasksToPatches,ConditioningMatchMask +from .node import BatchedMaskedReferenceGen,ReferenceConditionCombine,MasksToPatches,ConditioningMatchMask,GenSubject,PreviewSubjectMask NODE_CLASS_MAPPINGS = { "BatchedMaskedReferenceGen": BatchedMaskedReferenceGen, "MasksToPatches": MasksToPatches, "ReferenceConditionCombine": ReferenceConditionCombine, "ConditioningMatchMask": ConditioningMatchMask, + "GenSubject": GenSubject, + "PreviewSubjectMask" : PreviewSubjectMask } NODE_DISPLAY_NAME_MAPPINGS = { "ReferenceConditionCombine": "Reference Condition Combine", "MasksToPatches": "Masks to Patches", "ConditioningMatchMask": "Conditioning Match Mask", "BatchedMaskedReferenceGen": "Batched Masked Reference Gen", + "GenSubject": "Gen Subject", + "PreviewSubjectMask": "Preview Subject Mask", } __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/node.py b/node.py index 28cfda2..1858ce3 100644 --- a/node.py +++ b/node.py @@ -400,10 +400,9 @@ class MasksToPatches: print(f'Patch dimensions: {h}x{w}, Total patches: {h*w}') return (patch_masks, h, w) - -class GenID: - RETURN_TYPES = ("MODEL",) - RETURN_NAMES = ("model",) +class PreviewSubjectMask: + RETURN_TYPES = ("MASK","LATENT",) + RETURN_NAMES = ("mask","latent",) FUNCTION = "go" CATEGORY = "chara_consist" @@ -411,129 +410,116 @@ class GenID: def INPUT_TYPES(cls): return { "required": { - "model": ("MODEL",), - "latent": ("LATENTS",), - "target_patch_masks": ("MASK",), - "conditions_mask": ("MASK",), - "conditions_ref_mask": ("MASK",), + "latent": ("LATENT",), + "subject": ("SUBJECT",), }, } @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}') + def go(cls, latent, subject): + return (subject[0]["front_mask"], latent, ) + +class GenSubject: + RETURN_TYPES = ("LATENT","SUBJECT",) + RETURN_NAMES = ("latent","subject",) + FUNCTION = "go" + CATEGORY = "chara_consist" - 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)) - sampling_step = -1 + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "latent": ("LATENT",), + "kv_start_step": ("INT",{"default":1,"min":0}), + "kv_end_step": ("INT",{"default":5,"min":0}), + "save_mask_step": ("INT",{"default":5,"min":0}), + }, + "optional": { + "fg_condition_mask": ("MASK",), + } + } + @classmethod + def go(cls, *, latent: torch.Tensor, kv_start_step, kv_end_step, save_mask_step, fg_condition_mask = None): + B,L,H,W = latent["samples"].shape + subject = [{"k": [], "v":[]}] def prepare_callback(model, steps, x0_output_dict=None): - origin_callback = orig_prepare_callback(model, steps, x0_output_dict) - def callback(step, x0, x, total_steps): - sampling_step = step - origin_callback(step, x0, x, total_steps) - return callback - - 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] + sampled_step = -1 + sampling_step = -1 + layer = 0 + model_type = "flux" + w = W // model.model.diffusion_model.patch_size + h = H // model.model.diffusion_model.patch_size + def get_fg_attn_weights(text_len, image_start, q, k): + attn_weights = torch.matmul(q[:, :, image_start:], k[:, :, :text_len].transpose(-2, -1)).to(torch.float16) + attn_weights = attn_weights.softmax(-1).sum(1) + if fg_condition_mask is not None: + fg_condition_mask = rearrange(fg_condition_mask, "s -> 1 1 s") + attn_weights = (attn_weights * fg_condition_mask).mean(1) + mask = rearrange(attn_weights, "b (h w) -> b h w", h=h, w=w) + return mask + 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. - 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) + 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] + """ + if sampling_step != sampled_step + 1: + layer = 0 + sampling_step = sampled_step + 1 + if sampling_step >= kv_start_step and sampling_step < kv_end_step: + subject[0]["k"] = k + subject[0]["v"] = v + # Apply rotary positional encoding for Flux models + if pe is not None: + q, k = flux_math.apply_rope(q, k, pe) - 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)") - - x = torch.zeros(q.shape, device=device) - # PyTorch's scaled_dot_product_attention - x[:num_references] = F.scaled_dot_product_attention( - q[:num_references], k[:num_references], v[:num_references], - attn_mask=reference_attn_mask - ) + # Get actual sequence length from QKV tensors + full_len = q.shape[2] + text_len = full_len - h * w # because the latent size is equal to patch count for flux and chroma. - 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 + if sampling_step == save_mask_step: + subject[0]["front_mask"] = get_fg_attn_weights(text_len, text_len, q, k) - # 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: + h = F.scaled_dot_product_attention( + q, k, v, attn_mask=mask + ) + # Reshape back to the expected output format + h = rearrange(h, "B H L D -> B L (H D)") + layer += 1 + return h + + # Create the attention override function + def attention_override(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask: Tensor, transformer_options) -> Tensor: latent_preview.prepare_callback = prepare_callback - 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,) + 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, steps, x0_output_dict) + def callback_change_sampled_step(step, x0, x, total_steps): + sampled_step = step + origin_callback(step, x0, x, total_steps) + return callback_change_sampled_step + latent_preview.prepare_callback = prepare_callback + return (latent,)