From d1ea42fa9fbfa530a81d80127cd10c60fc421314 Mon Sep 17 00:00:00 2001
From: ZHOU He <461607353@qq.com>
Date: Sat, 15 Nov 2025 12:07:19 +0800
Subject: [PATCH] Save mask done.

---
 node.py | 48 +++++++++++++++++++++++++++++++++---------------
 1 file changed, 33 insertions(+), 15 deletions(-)

diff --git a/node.py b/node.py
index 826d98e..839b1e0 100644
--- a/node.py
+++ b/node.py
@@ -419,8 +419,9 @@ class PreviewSubjectMask:
         # Iterate through subject[0] and gather all fg_masks
         fg_masks = []
         for s in subject:
-            for item in s["fg_mask"]:
-                fg_masks.append(item)
+            for item in s:
+                if "fg_mask" in item:
+                    fg_masks.append(item["fg_mask"])
         
         # Stack all fg_masks into a single tensor
         if fg_masks:
@@ -441,6 +442,28 @@ def hook_prepare_callback():
             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[0]]
+    share_fg_mask = curr_fg_mask & rearrange_id_fg_mask & curr_valid_mask
+    id_share_fg_indices = curr2id_argmax_indices[0][share_fg_mask]
+    curr_share_fg_indices = torch.nonzero(share_fg_mask).squeeze()
+    return id_share_fg_indices, curr_share_fg_indices
+
+
 class GenSubject:
     RETURN_TYPES = ("MODEL","LATENT","SUBJECT",)
     FUNCTION = "go"
@@ -462,10 +485,10 @@ class GenSubject:
     @classmethod
     def go(cls, model, latent: torch.Tensor, fg_condition_mask, kv_start_step, kv_end_step, mask_start_step, mask_end_step):
         B,L,H,W = latent["samples"].shape
-        subject = [{"data":[], "fg_mask":[]}]
+        subject = []
         def prepare_callback(model, max_steps, x0_output_dict=None):
-            subject[0] = {"data":[], "fg_mask":[]}
-            nonlocal W, H
+            nonlocal subject, W, H
+            subject.append([])
             sampled_step = -1
             sampling_step = -1
             layer = 0
@@ -504,13 +527,10 @@ class GenSubject:
                 if sampling_step != sampled_step + 1:
                     layer = 0
                     sampling_step = sampled_step + 1
-
-                current_data = {"step": sampling_step, "layer": layer}
-                subject[0]["data"].append(current_data)
+                if len(subject[-1]) == 0 or subject[-1][-1]["step"] != sampling_step:
+                    subject[-1].append({"step":sampling_step, "layers": []})
                 if sampling_step >= kv_start_step and sampling_step < kv_end_step:
-                    print("save kv at step ", sampling_step, " layer ", layer)
-                    current_data["k"] = k.cpu()
-                    current_data["v"] = v.cpu()
+                    subject[-1][-1]["layers"].append({"layer": 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)
@@ -519,8 +539,6 @@ class GenSubject:
                 full_len = q.shape[2]
                 text_len = full_len - h * w # because the latent size is equal to patch count for flux and chroma.
                 
-
-                
                 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)
@@ -537,7 +555,7 @@ class GenSubject:
 
             # Create the attention override function
             def attention_override(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask: Tensor, transformer_options) -> Tensor:
-                if q.shape[2] != text_len + (w * h):
+                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)
@@ -562,7 +580,7 @@ class GenSubject:
                 if sampling_step >= mask_start_step and sampling_step < mask_end_step:
                     bitmask = fg_bg_mask[:,0] >= fg_bg_mask[:,1]
                     fg_bg_mask = None
-                    subject[0]["fg_mask"].append(bitmask.cpu())
+                    subject[-1][-1]["fg_mask"] = bitmask.cpu()
                 origin_callback(step, x0, x, total_steps)
                 #if step == total_steps - 1:
                 #    latent_preview.prepare_callback = orig_prepare_callback