wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-03-29 21:13:33 +08:00
Code Issues Actions 9 Packages Projects Releases Wiki Activity

Merge 3e4586ede7 into fc1fdf3389

Browse Source
This commit is contained in:
KyungPyoKim 2026-03-28 22:22:53 -04:00 committed by GitHub
commit df0297535e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 27 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
comfy_extras/nodes_dataset.py
View File

@ -167,6 +167,17 @@ 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):
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:
normalized_images.append(img)
image_list = normalized_images
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 +199,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) 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 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) 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 +295,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")