From a79fe2d5fbf9898be0f6070b5c1bd52d8f34202a Mon Sep 17 00:00:00 2001
From: ZHOU He <461607353@qq.com>
Date: Sun, 16 Nov 2025 01:31:35 +0800
Subject: [PATCH] add new node.

---
 node.py | 109 ++++++++++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 107 insertions(+), 2 deletions(-)

diff --git a/node.py b/node.py
index eeaa127..3261914 100644
--- a/node.py
+++ b/node.py
@@ -632,11 +632,11 @@ class PreRun:
             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
-            accum_sim = None
-            text_len = fg_condition_mask.shape[-1]
+            accum_sim = None            
             fg_condition_mask = fg_condition_mask.to(device)
             def masked_attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask: Tensor) -> Tensor:
                 """
@@ -787,3 +787,108 @@ class CrossSimMaskExtract:
             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": ("SUBJECT",),
+                "cross_sim": ("CROSSSIM",),
+                "fg_condition_mask": ("MASK",),
+                "kv_start_step":  ("INT",{"default":1,"min":0}),
+                "kv_end_step":  ("INT",{"default":5,"min":0}),
+            },
+            "optional": {
+                "latent":("LATENT",)
+            }
+        }
+    @classmethod
+    def go(cls, model, subject, cross_sim, fg_condition_mask, kv_start_step, kv_end_step, latent):
+        
+        index = len(subject) - 1
+        def prepare_callback(model, max_steps, x0_output_dict=None):
+            nonlocal subject, fg_condition_mask
+            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
+                # 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]
+                    # todo 
+                
+                # Apply rotary positional encoding for Flux models
+                if pe is not None:
+                    q, k = flux_math.apply_rope(q, k, pe)
+                
+                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)")
+
+                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
+                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,)
\ No newline at end of file