From 40b9ab7369fbe713aa29ddd26376571753e18e34 Mon Sep 17 00:00:00 2001 From: KimKyungPyo Date: Thu, 26 Mar 2026 16:36:16 +0900 Subject: [PATCH 1/3] fix(save-images): handle tensor dimensions and flatten image lists before saving --- comfy_extras/nodes_dataset.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/comfy_extras/nodes_dataset.py b/comfy_extras/nodes_dataset.py index 98ed25d7e..f7193f1be 100644 --- a/comfy_extras/nodes_dataset.py +++ b/comfy_extras/nodes_dataset.py @@ -167,6 +167,12 @@ def save_images_to_folder(image_list, output_dir, prefix="image"): os.makedirs(output_dir, exist_ok=True) saved_files = [] + if isinstance(image_list, torch.Tensor): + if image_list.dim() == 4: + image_list = [image_list[i] for i in range(image_list.shape[0])] + else: + image_list = [image_list] + for idx, img_tensor in enumerate(image_list): # Handle different tensor shapes if isinstance(img_tensor, torch.Tensor): @@ -188,6 +194,9 @@ 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: + img_array = img_array[0] + img = Image.fromarray(img_array) else: raise ValueError(f"Expected torch.Tensor, got {type(img_tensor)}") @@ -278,6 +287,16 @@ class SaveImageTextDataSetToFolderNode(io.ComfyNode): output_dir = os.path.join(folder_paths.get_output_directory(), folder_name) saved_files = save_images_to_folder(images, output_dir, filename_prefix) + flat_images = [] + for img in images: + if isinstance(img, torch.Tensor) and img.dim() == 4: + for i in range(img.shape[0]): + flat_images.append(img[i]) + else: + flat_images.append(img) + images = flat_images + + # Save captions for idx, (filename, caption) in enumerate(zip(saved_files, texts)): caption_filename = filename.replace(".png", ".txt") From 2a141d9927e7af2cedfac500d23cf2a436594248 Mon Sep 17 00:00:00 2001 From: KyungPyoKim Date: Thu, 26 Mar 2026 17:02:45 +0900 Subject: [PATCH 2/3] 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)}") From 3e4586ede7047a5f50e35cde7928db27b3f25ac2 Mon Sep 17 00:00:00 2001 From: KimKyungPyo Date: Thu, 26 Mar 2026 17:46:46 +0900 Subject: [PATCH 3/3] Fix: properly handle batched tensors to prevent data loss --- comfy_extras/nodes_dataset.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/comfy_extras/nodes_dataset.py b/comfy_extras/nodes_dataset.py index d0dab055e..d05c8c257 100644 --- a/comfy_extras/nodes_dataset.py +++ b/comfy_extras/nodes_dataset.py @@ -168,10 +168,15 @@ def save_images_to_folder(image_list, output_dir, prefix="image"): saved_files = [] if isinstance(image_list, torch.Tensor): - if image_list.dim() == 4: - image_list = [image_list[i] for i in range(image_list.shape[0])] + image_list = [image_list] + + normalized_images = [] + for img in image_list: + if isinstance(img, torch.Tensor) and img.dim() == 4: + normalized_images.extend([img[i] for i in range(img.shape[0])]) else: - image_list = [image_list] + normalized_images.append(img) + image_list = normalized_images for idx, img_tensor in enumerate(image_list): # Handle different tensor shapes