fix(save-images): handle tensor dimensions and flatten image lists before saving

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")