diff --git a/node.py b/node.py
index 5d34e85..64c465a 100644
--- a/node.py
+++ b/node.py
@@ -633,9 +633,9 @@ class ExtractAttn:
             def callback_change_sampled_step(step, x0, x, total_steps):
                 nonlocal max_steps
                 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
+                
+                w = x0.shape[-1] // model.model.diffusion_model.patch_size
+                h = x0.shape[-2] // 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:
                     bitmask = fg_bg_mask[:,0] >= fg_bg_mask[:,1]
@@ -985,9 +985,9 @@ class GenConsistent:
             def callback_change_sampled_step(step, x0, x, total_steps):
                 nonlocal max_steps
                 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
+                
+                w = x0.shape[-1] // model.model.diffusion_model.patch_size
+                h = x0.shape[-2] // model.model.diffusion_model.patch_size
                 sampled_step = step
                 origin_callback(step, x0, x, total_steps)
             return callback_change_sampled_step