Change to dict.

node.py

@@ -419,9 +419,10 @@ class PreviewSubjectMask:
         # Iterate through subject[0] and gather all fg_masks
         fg_masks = []
         for s in subject:
-            for item in s:
-                if "fg_mask" in item:
-                    fg_masks.append(item["fg_mask"])
+            # 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:
@@ -477,18 +478,20 @@ class GenSubject:
                 "fg_condition_mask": ("MASK",),
                 "kv_start_step":  ("INT",{"default":1,"min":0}),
                 "kv_end_step":  ("INT",{"default":5,"min":0}),
+                "x_start_step":  ("INT",{"default":1,"min":0}),
+                "x_end_step":  ("INT",{"default":5,"min":0}),
                 "mask_start_step": ("INT",{"default":5,"min":0}),
                 "mask_end_step": ("INT",{"default":6,"min":0}),
             }
         }
     @classmethod
-    def go(cls, model, fg_condition_mask, kv_start_step, kv_end_step, mask_start_step, mask_end_step):
+    def go(cls, model, fg_condition_mask, kv_start_step, kv_end_step, x_start_step, x_end_step, mask_start_step, mask_end_step):
         
-        subject = [[]]
+        subject = [{}]
         index = len(subject) - 1
         def prepare_callback(model, max_steps, x0_output_dict=None):
             nonlocal subject
-            subject[index] = []
+            subject[index] = {}
             sampled_step = -1
             sampling_step = -1
             layer = 0
@@ -526,10 +529,12 @@ class GenSubject:
                 if sampling_step != sampled_step + 1:
                     layer = 0
                     sampling_step = sampled_step + 1
-                if len(subject[index]) == 0 or subject[index][-1]["step"] != sampling_step:
-                    subject[index].append({"step":sampling_step, "layers": []})
+                # 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][-1]["layers"].append({"layer": layer,"k": k.cpu(), "v": v.cpu()})
+                    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)
@@ -545,6 +550,9 @@ class GenSubject:
                 )
                 # 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 >= x_start_step and sampling_step < x_end_step:
+                    subject[index][sampling_step]["layers"][layer]["x"] = x.cpu()
                 layer += 1
                 return x
 
@@ -579,7 +587,10 @@ class GenSubject:
                     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
-                    subject[index][-1]["fg_mask"] = bitmask.cpu()
+                    # 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