wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-04-14 20:42:31 +08:00
Code Issues Actions 15 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
KimKyungPyo 2026-03-26 16:36:16 +09:00
parent 2a1f402601
commit 40b9ab7369
1 changed files with 19 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
comfy_extras/nodes_dataset.py
View File

@ -167,6 +167,12 @@ def save_images_to_folder(image_list, output_dir, prefix="image"):
os.makedirs(output_dir, exist_ok=True) os.makedirs(output_dir, exist_ok=True)
saved_files = [] 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): for idx, img_tensor in enumerate(image_list):
# Handle different tensor shapes # Handle different tensor shapes
if isinstance(img_tensor, torch.Tensor): 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) img_array = np.clip(img_array * 255.0, 0, 255).astype(np.uint8)
# Convert to PIL Image # Convert to PIL Image
while img_array.ndim > 3:
img_array = img_array[0]
img = Image.fromarray(img_array) img = Image.fromarray(img_array)
else: else:
raise ValueError(f"Expected torch.Tensor, got {type(img_tensor)}") 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) output_dir = os.path.join(folder_paths.get_output_directory(), folder_name)
saved_files = save_images_to_folder(images, output_dir, filename_prefix) 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 # Save captions
for idx, (filename, caption) in enumerate(zip(saved_files, texts)): for idx, (filename, caption) in enumerate(zip(saved_files, texts)):
caption_filename = filename.replace(".png", ".txt") caption_filename = filename.replace(".png", ".txt")