From 2a141d9927e7af2cedfac500d23cf2a436594248 Mon Sep 17 00:00:00 2001 From: KyungPyoKim Date: Thu, 26 Mar 2026 17:02:45 +0900 Subject: [PATCH] Update comfy_extras/nodes_dataset.py Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> --- comfy_extras/nodes_dataset.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/comfy_extras/nodes_dataset.py b/comfy_extras/nodes_dataset.py index f7193f1be..d0dab055e 100644 --- a/comfy_extras/nodes_dataset.py +++ b/comfy_extras/nodes_dataset.py @@ -194,9 +194,12 @@ def save_images_to_folder(image_list, output_dir, prefix="image"): img_array = np.clip(img_array * 255.0, 0, 255).astype(np.uint8) # Convert to PIL Image - while img_array.ndim > 3: + while img_array.ndim > 3 and img_array.shape[0] == 1: img_array = img_array[0] - + if img_array.ndim > 3: + raise ValueError( + f"Unsupported image tensor shape after normalization: {tuple(img_array.shape)}" + ) img = Image.fromarray(img_array) else: raise ValueError(f"Expected torch.Tensor, got {type(img_tensor)}")