From a7663bf0940c555da56464ef41f641a7795e5d95 Mon Sep 17 00:00:00 2001
From: ZHOU He <461607353@qq.com>
Date: Sun, 16 Nov 2025 01:42:27 +0800
Subject: [PATCH] Add expand kv, but have not impl multi-subject and rope yet

---
 node.py | 37 ++++++++++++++++++++++++++++++++-----
 1 file changed, 32 insertions(+), 5 deletions(-)

diff --git a/node.py b/node.py
index 3261914..25b4957 100644
--- a/node.py
+++ b/node.py
@@ -843,14 +843,41 @@ class GenConsistent:
                 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)
+                # 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)
+                    
+                    # Check if we have valid cross-sim data and subject data to work with
+                    if (id_fg_inds is not None and curr_fg_inds is not None and 
+                        len(id_fg_inds) > 0 and len(curr_fg_inds) > 0 and
+                        sampling_step in subject[index] and 
+                        layer in subject[index][sampling_step]["layers"] and
+                        "k" in subject[index][sampling_step]["layers"][layer] and 
+                        "v" in subject[index][sampling_step]["layers"][layer]):
+                        
+                        ori_saved_key = subject[index][sampling_step]["layers"][layer]["k"].to(k.device)
+                        ori_saved_value = subject[index][sampling_step]["layers"][layer]["v"].to(v.device)
+                        
+                        # Ensure we have valid indices within tensor bounds
+                        valid_id_inds = id_fg_inds[id_fg_inds < ori_saved_key.shape[2]]
+                        valid_curr_inds = curr_fg_inds[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]
+                            
+                            # 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=mask
@@ -891,4 +918,4 @@ class GenConsistent:
         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
+        return (model, subject,)