diff --git a/node.py b/node.py
index f54e761..5d34e85 100644
--- a/node.py
+++ b/node.py
@@ -549,33 +549,30 @@ class ExtractAttn:
         return {
             "required": {
                 "model": ("MODEL",),
-                "latent": ("LATENT",),
-                "kv_start_step":  ("INT",{"default":1,"min":0}),
-                "kv_end_step":  ("INT",{"default":9,"min":0}),
-                "x_start_step":  ("INT",{"default":1,"min":0}),
-                "x_end_step":  ("INT",{"default":9,"min":0}),
-                "mask_start_step": ("INT",{"default":8,"min":0}),
-                "mask_end_step": ("INT",{"default":9,"min":0}),
+                "kv_start_step":  ("INT",{"default":1,"min":1}),
+                "kv_end_step":  ("INT",{"default":9,"min":1}),
+                "x_start_step":  ("INT",{"default":1,"min":1}),
+                "x_end_step":  ("INT",{"default":9,"min":1}),
+                "mask_start_step": ("INT",{"default":8,"min":1}),
+                "mask_end_step": ("INT",{"default":9,"min":1}),
             },
             "optional":{
                 "fg_condition_mask": ("MASK",),
             }
         }
     @classmethod
-    def go(cls, model, latent, kv_start_step, kv_end_step, x_start_step, x_end_step, mask_start_step, mask_end_step, fg_condition_mask = None):
+    def go(cls, model, kv_start_step, kv_end_step, x_start_step, x_end_step, mask_start_step, mask_end_step, fg_condition_mask = None):
         
         attn_data = [{}]
         index = len(attn_data) - 1
         def prepare_callback(model, max_steps, x0_output_dict=None):
-            nonlocal attn_data, fg_condition_mask, latent
+            nonlocal attn_data, fg_condition_mask
             attn_data[index] = {}
             sampled_step = -1
             sampling_step = -1
             layer = 0
-            
-            w = latent["samples"].shape[-1] // model.model.diffusion_model.patch_size
-            h = latent["samples"].shape[-2] // model.model.diffusion_model.patch_size
-
+            w = 0
+            h = 0
             #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
@@ -606,7 +603,7 @@ class ExtractAttn:
                     attn_data[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:
-                    attn_data[index][sampling_step]["layers"][layer] = {"k": k.cpu(), "v": v.cpu()}
+                    attn_data[index][sampling_step]["layers"][layer] = {"k": k[:, :, text_len:,:].cpu(), "v": v[:, :, text_len:,:].cpu()}
                 # Apply rotary positional encoding for Flux models
                 if pe is not None:
                     q = orig_apply_rope1(q, pe)
@@ -635,12 +632,12 @@ class ExtractAttn:
             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
+                nonlocal sampling_step, sampled_step, layer, fg_bg_mask, model, w, h
+                B,L,H,W = x0.shape
+                w = W // model.model.diffusion_model.patch_size
+                h = H // model.model.diffusion_model.patch_size
                 sampled_step = step
                 if fg_condition_mask is not None and sampling_step >= mask_start_step and sampling_step < mask_end_step:
-                    B,L,H,W = x0.shape
-                    w = W // model.model.diffusion_model.patch_size
-                    h = H // model.model.diffusion_model.patch_size
                     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
@@ -823,25 +820,23 @@ class GenConsistent:
                 "model": ("MODEL",),
                 "subject": ("ATTN",),
                 "cross_sim": ("CROSSSIM",),
-                "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}),
+                "kv_start_step":  ("INT",{"default":1,"min":1}),
+                "kv_end_step":  ("INT",{"default":5,"min":1}),
+                "x_start_step":  ("INT",{"default":1,"min":1}),
+                "x_end_step":  ("INT",{"default":5,"min":1}),
                 "ada_tome_alpha": ("FLOAT",{"default":0.0,"min":0,"max":1,"step":0.1}),
-                "latent":("LATENT",),
             },
         }
     @classmethod
-    def go(cls, model, subject, cross_sim, kv_start_step, kv_end_step, x_start_step, x_end_step, ada_tome_alpha, latent):
+    def go(cls, model, subject, cross_sim, kv_start_step, kv_end_step, x_start_step, x_end_step, ada_tome_alpha):
         
         def prepare_callback(model, max_steps, x0_output_dict=None):
-            nonlocal subject, latent
+            nonlocal subject
             sampled_step = -1
             sampling_step = -1
             layer = 0
-            
-            w = latent["samples"].shape[-1] // model.model.diffusion_model.patch_size
-            h = latent["samples"].shape[-2] // model.model.diffusion_model.patch_size
+            w = 0
+            h = 0
             
             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
@@ -989,7 +984,10 @@ class GenConsistent:
             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
+                nonlocal sampling_step, sampled_step, layer, fg_bg_mask, model, w, h
+                B,L,H,W = x0.shape
+                w = W // model.model.diffusion_model.patch_size
+                h = H // model.model.diffusion_model.patch_size
                 sampled_step = step
                 origin_callback(step, x0, x, total_steps)
             return callback_change_sampled_step