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 5ba7330e3c into 3f77450ef1

Browse Source
This commit is contained in:
Charles Chan 2026-03-29 10:57:52 +08:00 committed by GitHub
commit b65992c444
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 24 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
comfy_extras/nodes_dataset.py
View File

@ -153,7 +153,7 @@ class LoadImageTextDataSetFromFolderNode(io.ComfyNode):
return io.NodeOutput(output_tensor, captions) return io.NodeOutput(output_tensor, captions)
def save_images_to_folder(image_list, output_dir, prefix="image"): def save_images_to_folder(image_list, output_dir, prefix="image", overwrite=True):
"""Utility function to save a list of image tensors to disk. """Utility function to save a list of image tensors to disk.
Args: Args:
@ -193,7 +193,11 @@ def save_images_to_folder(image_list, output_dir, prefix="image"):
raise ValueError(f"Expected torch.Tensor, got {type(img_tensor)}") raise ValueError(f"Expected torch.Tensor, got {type(img_tensor)}")
# Save image # Save image
filename = f"{prefix}_{idx:05d}.png" if overwrite:
filename = f"{prefix}_{idx:05d}.png"
else:
_, _, counter, _, resolved_prefix = folder_paths.get_save_image_path(prefix, output_dir)
filename = f"{resolved_prefix}_{counter:05}_{idx:05d}.png"
filepath = os.path.join(output_dir, filename) filepath = os.path.join(output_dir, filename)
img.save(filepath) img.save(filepath)
saved_files.append(filename) saved_files.append(filename)
@ -224,18 +228,25 @@ class SaveImageDataSetToFolderNode(io.ComfyNode):
tooltip="Prefix for saved image filenames.", tooltip="Prefix for saved image filenames.",
advanced=True, advanced=True,
), ),
io.Combo.Input(
"mode",
default="overwrite",
options=["overwrite", "increment"],
tooltip="Whether to overwrite existing files or increment filenames to avoid overwriting."
),
], ],
outputs=[], outputs=[],
) )
@classmethod @classmethod
def execute(cls, images, folder_name, filename_prefix): def execute(cls, images, folder_name, filename_prefix, mode):
# Extract scalar values # Extract scalar values
folder_name = folder_name[0] folder_name = folder_name[0]
filename_prefix = filename_prefix[0] filename_prefix = filename_prefix[0]
mode = mode[0]
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, mode=='overwrite')
logging.info(f"Saved {len(saved_files)} images to {output_dir}.") logging.info(f"Saved {len(saved_files)} images to {output_dir}.")
return io.NodeOutput() return io.NodeOutput()
@ -265,18 +276,25 @@ class SaveImageTextDataSetToFolderNode(io.ComfyNode):
tooltip="Prefix for saved image filenames.", tooltip="Prefix for saved image filenames.",
advanced=True, advanced=True,
), ),
io.Combo.Input(
"mode",
default="overwrite",
options=["overwrite", "increment"],
tooltip="Whether to overwrite existing files or increment filenames to avoid overwriting."
),
], ],
outputs=[], outputs=[],
) )
@classmethod @classmethod
def execute(cls, images, texts, folder_name, filename_prefix): def execute(cls, images, texts, folder_name, filename_prefix, mode):
# Extract scalar values # Extract scalar values
folder_name = folder_name[0] folder_name = folder_name[0]
filename_prefix = filename_prefix[0] filename_prefix = filename_prefix[0]
mode = mode[0]
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, mode=='overwrite')
# Save captions # Save captions
for idx, (filename, caption) in enumerate(zip(saved_files, texts)): for idx, (filename, caption) in enumerate(zip(saved_files, texts)):