From 40b9ab7369fbe713aa29ddd26376571753e18e34 Mon Sep 17 00:00:00 2001 From: KimKyungPyo Date: Thu, 26 Mar 2026 16:36:16 +0900 Subject: [PATCH] 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")