Fix tensor in flux.

This commit is contained in:
ZHOU He
2025-12-06 12:32:14 +08:00
parent 9aa12a5cbd
commit c5b77e50e9
+6 -6
View File
@@ -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