diff --git a/comfy_extras/nodes_custom_sampler.py b/comfy_extras/nodes_custom_sampler.py
index 4e8510ca9..36f70ce82 100644
--- a/comfy_extras/nodes_custom_sampler.py
+++ b/comfy_extras/nodes_custom_sampler.py
@@ -753,7 +753,7 @@ class SamplerCustom(io.ComfyNode):
             noise_mask = latent["noise_mask"]
 
         x0_output = {}
-        callback = latent_preview.prepare_callback(model, sigmas.shape[-1] - 1, x0_output, shape=latent_image.shape)
+        callback = latent_preview.prepare_callback(model, sigmas.shape[-1] - 1, x0_output, shape=latent_image.shape if latent_image.is_nested else None)
 
         disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
         samples = comfy.sample.sample_custom(model, noise, cfg, sampler, sigmas, positive, negative, latent_image, noise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=noise_seed)
@@ -944,7 +944,7 @@ class SamplerCustomAdvanced(io.ComfyNode):
             noise_mask = latent["noise_mask"]
 
         x0_output = {}
-        callback = latent_preview.prepare_callback(guider.model_patcher, sigmas.shape[-1] - 1, x0_output, shape=latent_image.shape)
+        callback = latent_preview.prepare_callback(guider.model_patcher, sigmas.shape[-1] - 1, x0_output, shape=latent_image.shape if latent_image.is_nested else None)
 
         disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
         samples = guider.sample(noise.generate_noise(latent), latent_image, sampler, sigmas, denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=noise.seed)
diff --git a/latent_preview.py b/latent_preview.py
index 7997b036d..807931e03 100644
--- a/latent_preview.py
+++ b/latent_preview.py
@@ -122,7 +122,7 @@ def prepare_callback(model, steps, x0_output_dict=None, shape=None):
         if x0_output_dict is not None:
             x0_output_dict["x0"] = x0
 
-        if x0.ndim == 3:
+        if shape is not None:
             cut = math.prod(shape[1:])
             x0 = x0[:, :, :cut].reshape([x0.shape[0]] + list(shape)[1:])
 
diff --git a/nodes.py b/nodes.py
index 4fda13849..7e9676e74 100644
--- a/nodes.py
+++ b/nodes.py
@@ -1504,7 +1504,7 @@ def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive,
     if "noise_mask" in latent:
         noise_mask = latent["noise_mask"]
 
-    callback = latent_preview.prepare_callback(model, steps, shape=latent_image.shape)
+    callback = latent_preview.prepare_callback(model, steps, shape=latent_image.shape if latent_image.is_nested else None)
     disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
     samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,
                                   denoise=denoise, disable_noise=disable_noise, start_step=start_step, last_step=last_step,