From c5b77e50e93d9ecb493a9872e9ce75ed1b565181 Mon Sep 17 00:00:00 2001
From: ZHOU He <461607353@qq.com>
Date: Sat, 6 Dec 2025 12:32:14 +0800
Subject: [PATCH] Fix tensor in flux.

---
 node.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

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