Solo YoloandClaude Opus 4.6
8116f97fab
Eight custom nodes for loading, iterating, and saving image/caption datasets in ComfyUI workflows. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
463 lines
15 KiB
Python
463 lines
15 KiB
Python
import os
|
|
import logging
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
|
|
import folder_paths
|
|
|
|
|
|
class LoadCaptionDataset:
|
|
"""Load caption .txt files from a subfolder of ComfyUI/input/.
|
|
|
|
Outputs parallel lists of captions, save prefixes (folder/stem), and
|
|
the corresponding source images (if they exist alongside the captions).
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"folder": (sorted(folder_paths.get_input_subfolders()),),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING", "IMAGE")
|
|
RETURN_NAMES = ("captions", "save_prefixes", "images")
|
|
OUTPUT_IS_LIST = (True, True, True)
|
|
FUNCTION = "load"
|
|
CATEGORY = "dataset"
|
|
|
|
def load(self, folder):
|
|
input_dir = folder_paths.get_input_directory()
|
|
folder_path = os.path.join(input_dir, folder)
|
|
|
|
if not os.path.isdir(folder_path):
|
|
raise ValueError(f"Folder not found: {folder_path}")
|
|
|
|
valid_image_exts = [".png", ".jpg", ".jpeg", ".webp"]
|
|
|
|
txt_files = sorted(
|
|
f for f in os.listdir(folder_path) if f.lower().endswith(".txt")
|
|
)
|
|
|
|
if not txt_files:
|
|
raise ValueError(f"No .txt caption files found in {folder_path}")
|
|
|
|
captions = []
|
|
save_prefixes = []
|
|
images = []
|
|
|
|
for txt_file in txt_files:
|
|
stem = os.path.splitext(txt_file)[0]
|
|
txt_path = os.path.join(folder_path, txt_file)
|
|
|
|
with open(txt_path, "r", encoding="utf-8") as f:
|
|
caption = f.read().strip()
|
|
|
|
captions.append(caption)
|
|
save_prefixes.append(f"{folder}/{stem}")
|
|
|
|
# Try to load the matching image
|
|
img_tensor = None
|
|
for ext in valid_image_exts:
|
|
img_path = os.path.join(folder_path, stem + ext)
|
|
if os.path.exists(img_path):
|
|
img = Image.open(img_path).convert("RGB")
|
|
img_array = np.array(img).astype(np.float32) / 255.0
|
|
img_tensor = torch.from_numpy(img_array)[None,]
|
|
break
|
|
|
|
if img_tensor is None:
|
|
img_tensor = torch.zeros(1, 64, 64, 3)
|
|
|
|
images.append(img_tensor)
|
|
|
|
logging.info(
|
|
f"LoadCaptionDataset: loaded {len(captions)} captions from {folder}"
|
|
)
|
|
return (captions, save_prefixes, images)
|
|
|
|
|
|
class SaveImageExact:
|
|
"""Save an image to the output directory using the exact save_prefix as the
|
|
filename (no counter suffix). save_prefix may contain a subdirectory
|
|
component, e.g. 'test_sample/0a4597f94d8fe8752302c274b5d8a166'.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.output_dir = folder_paths.get_output_directory()
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"save_prefix": ("STRING", {"default": "output/image"}),
|
|
},
|
|
"hidden": {
|
|
"prompt": "PROMPT",
|
|
"extra_pnginfo": "EXTRA_PNGINFO",
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "save"
|
|
CATEGORY = "dataset"
|
|
|
|
def save(self, images, save_prefix, prompt=None, extra_pnginfo=None):
|
|
# Sanitize: prevent directory traversal
|
|
save_prefix = save_prefix.replace("..", "").strip("/")
|
|
|
|
filepath = os.path.join(self.output_dir, save_prefix + ".png")
|
|
dirpath = os.path.dirname(filepath)
|
|
os.makedirs(dirpath, exist_ok=True)
|
|
|
|
# Take the first image in the batch
|
|
i = 255.0 * images[0].cpu().numpy()
|
|
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
|
|
|
|
metadata = None
|
|
try:
|
|
from comfy.cli_args import args
|
|
if not args.disable_metadata:
|
|
from PIL.PngImagePlugin import PngInfo
|
|
metadata = PngInfo()
|
|
if prompt is not None:
|
|
import json
|
|
metadata.add_text("prompt", json.dumps(prompt))
|
|
if extra_pnginfo is not None:
|
|
import json
|
|
for x in extra_pnginfo:
|
|
metadata.add_text(x, json.dumps(extra_pnginfo[x]))
|
|
except Exception:
|
|
pass
|
|
|
|
img.save(filepath, pnginfo=metadata, compress_level=4)
|
|
|
|
subfolder = os.path.dirname(save_prefix)
|
|
filename = os.path.basename(save_prefix) + ".png"
|
|
|
|
logging.info(f"SaveImageExact: saved {filepath}")
|
|
return {"ui": {"images": [{"filename": filename, "subfolder": subfolder, "type": "output"}]}}
|
|
|
|
|
|
class CountFilesInFolder:
|
|
"""Count image or text files in a subfolder of ComfyUI/input/."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"folder": (sorted(folder_paths.get_input_subfolders()),),
|
|
"file_type": (["images", "text"],),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("INT",)
|
|
RETURN_NAMES = ("count",)
|
|
FUNCTION = "count"
|
|
CATEGORY = "dataset"
|
|
|
|
def count(self, folder, file_type):
|
|
input_dir = folder_paths.get_input_directory()
|
|
folder_path = os.path.join(input_dir, folder)
|
|
|
|
if not os.path.isdir(folder_path):
|
|
raise ValueError(f"Folder not found: {folder_path}")
|
|
|
|
if file_type == "images":
|
|
exts = {".png", ".jpg", ".jpeg", ".webp"}
|
|
else:
|
|
exts = {".txt"}
|
|
|
|
count = sum(
|
|
1 for f in os.listdir(folder_path)
|
|
if os.path.splitext(f.lower())[1] in exts
|
|
)
|
|
logging.info(f"CountFilesInFolder: {count} {file_type} files in {folder}")
|
|
return (count,)
|
|
|
|
|
|
class SequentialIndex:
|
|
"""Index node for sequential processing with control_after_generate.
|
|
|
|
Set the control to 'increment' and queue N jobs to iterate through items.
|
|
Connect max_count from CountFilesInFolder to guard against overrun.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"index": ("INT", {"default": 0, "min": 0, "max": 999999, "control_after_generate": True}),
|
|
},
|
|
"optional": {
|
|
"max_count": ("INT", {"default": 0, "min": 0, "max": 999999, "forceInput": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("INT",)
|
|
RETURN_NAMES = ("index",)
|
|
FUNCTION = "execute"
|
|
CATEGORY = "dataset"
|
|
|
|
def execute(self, index, max_count=0):
|
|
if max_count > 0 and index >= max_count:
|
|
raise Exception(
|
|
f"Index {index} is out of range (max_count={max_count}). "
|
|
f"Valid indices: 0 to {max_count - 1}."
|
|
)
|
|
|
|
logging.info(
|
|
f"SequentialIndex: index {index} of {max_count if max_count > 0 else '?'}"
|
|
)
|
|
return (index,)
|
|
|
|
|
|
class LoadImageAtIndex:
|
|
"""Load a single image from a folder at a given index.
|
|
|
|
Images are sorted alphabetically by filename. Outputs the image tensor
|
|
and the filename stem (without extension). Set control_after_generate to
|
|
'increment' and queue N jobs to iterate through a folder.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"folder": (sorted(folder_paths.get_input_subfolders()),),
|
|
"index": ("INT", {"default": 0, "min": 0, "max": 999999, "control_after_generate": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING")
|
|
RETURN_NAMES = ("image", "filename_stem")
|
|
FUNCTION = "load"
|
|
CATEGORY = "dataset"
|
|
|
|
def load(self, folder, index):
|
|
input_dir = folder_paths.get_input_directory()
|
|
folder_path = os.path.join(input_dir, folder)
|
|
|
|
if not os.path.isdir(folder_path):
|
|
raise ValueError(f"Folder not found: {folder_path}")
|
|
|
|
valid_exts = {".png", ".jpg", ".jpeg", ".webp"}
|
|
image_files = sorted(
|
|
f for f in os.listdir(folder_path)
|
|
if os.path.splitext(f.lower())[1] in valid_exts
|
|
)
|
|
|
|
if not image_files:
|
|
raise ValueError(f"No image files found in {folder_path}")
|
|
|
|
if index >= len(image_files):
|
|
raise ValueError(
|
|
f"Index {index} out of range (folder has {len(image_files)} images)"
|
|
)
|
|
|
|
filename = image_files[index]
|
|
stem = os.path.splitext(filename)[0]
|
|
img_path = os.path.join(folder_path, filename)
|
|
|
|
img = Image.open(img_path).convert("RGB")
|
|
img_array = np.array(img).astype(np.float32) / 255.0
|
|
img_tensor = torch.from_numpy(img_array)[None,]
|
|
|
|
logging.info(f"LoadImageAtIndex: loaded {filename} (index {index})")
|
|
return (img_tensor, stem)
|
|
|
|
|
|
class LoadCaptionAtIndex:
|
|
"""Load a single caption (.txt) from a folder at a given index.
|
|
|
|
Caption files are sorted alphabetically. Outputs the caption text
|
|
and a save_prefix in the format 'folder/stem'. Set control_after_generate
|
|
to 'increment' and queue N jobs to iterate through a folder.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"folder": (sorted(folder_paths.get_input_subfolders()),),
|
|
"index": ("INT", {"default": 0, "min": 0, "max": 999999, "control_after_generate": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING")
|
|
RETURN_NAMES = ("caption", "save_prefix")
|
|
FUNCTION = "load"
|
|
CATEGORY = "dataset"
|
|
|
|
def load(self, folder, index):
|
|
input_dir = folder_paths.get_input_directory()
|
|
folder_path = os.path.join(input_dir, folder)
|
|
|
|
if not os.path.isdir(folder_path):
|
|
raise ValueError(f"Folder not found: {folder_path}")
|
|
|
|
txt_files = sorted(
|
|
f for f in os.listdir(folder_path) if f.lower().endswith(".txt")
|
|
)
|
|
|
|
if not txt_files:
|
|
raise ValueError(f"No .txt files found in {folder_path}")
|
|
|
|
if index >= len(txt_files):
|
|
raise ValueError(
|
|
f"Index {index} out of range (folder has {len(txt_files)} captions)"
|
|
)
|
|
|
|
filename = txt_files[index]
|
|
stem = os.path.splitext(filename)[0]
|
|
txt_path = os.path.join(folder_path, filename)
|
|
|
|
with open(txt_path, "r", encoding="utf-8") as f:
|
|
caption = f.read().strip()
|
|
|
|
save_prefix = f"{folder}/{stem}"
|
|
logging.info(f"LoadCaptionAtIndex: loaded {filename} (index {index})")
|
|
return (caption, save_prefix)
|
|
|
|
|
|
class LoadDatasetItemAtIndex:
|
|
"""Load image, caption, and save_prefix from a dataset folder at a given index.
|
|
|
|
Combines LoadImageAtIndex and LoadCaptionAtIndex into a single atomic node
|
|
to prevent race conditions where separate nodes could execute with different
|
|
index values during sequential queue processing.
|
|
|
|
Images and captions are paired by filename stem (e.g. abc.png + abc.txt).
|
|
Items are sorted alphabetically by caption filename.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"folder": (sorted(folder_paths.get_input_subfolders()),),
|
|
"index": ("INT", {"default": 0, "min": 0, "max": 999999, "control_after_generate": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING", "STRING")
|
|
RETURN_NAMES = ("image", "caption", "save_prefix")
|
|
FUNCTION = "load"
|
|
CATEGORY = "dataset"
|
|
|
|
def load(self, folder, index):
|
|
input_dir = folder_paths.get_input_directory()
|
|
folder_path = os.path.join(input_dir, folder)
|
|
|
|
if not os.path.isdir(folder_path):
|
|
raise ValueError(f"Folder not found: {folder_path}")
|
|
|
|
# Use caption files as the canonical list (they define the dataset items)
|
|
txt_files = sorted(
|
|
f for f in os.listdir(folder_path) if f.lower().endswith(".txt")
|
|
)
|
|
|
|
if not txt_files:
|
|
raise ValueError(f"No .txt caption files found in {folder_path}")
|
|
|
|
if index >= len(txt_files):
|
|
raise ValueError(
|
|
f"Index {index} out of range (folder has {len(txt_files)} items)"
|
|
)
|
|
|
|
txt_file = txt_files[index]
|
|
stem = os.path.splitext(txt_file)[0]
|
|
txt_path = os.path.join(folder_path, txt_file)
|
|
|
|
# Load caption
|
|
with open(txt_path, "r", encoding="utf-8") as f:
|
|
caption = f.read().strip()
|
|
|
|
save_prefix = f"{folder}/{stem}"
|
|
|
|
# Load matching image
|
|
valid_image_exts = [".png", ".jpg", ".jpeg", ".webp"]
|
|
img_tensor = None
|
|
for ext in valid_image_exts:
|
|
img_path = os.path.join(folder_path, stem + ext)
|
|
if os.path.exists(img_path):
|
|
img = Image.open(img_path).convert("RGB")
|
|
img_array = np.array(img).astype(np.float32) / 255.0
|
|
img_tensor = torch.from_numpy(img_array)[None,]
|
|
break
|
|
|
|
if img_tensor is None:
|
|
raise ValueError(
|
|
f"No matching image found for caption {txt_file} in {folder_path}"
|
|
)
|
|
|
|
logging.info(
|
|
f"LoadDatasetItemAtIndex: loaded {stem} (index {index}/{len(txt_files)})"
|
|
)
|
|
return (img_tensor, caption, save_prefix)
|
|
|
|
|
|
class SaveTextExact:
|
|
"""Save a text string to the output directory using the exact save_prefix
|
|
as the filename (no counter suffix). save_prefix may contain a subdirectory
|
|
component, e.g. 'test_sample/0a4597f94d8fe8752302c274b5d8a166'.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.output_dir = folder_paths.get_output_directory()
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"text": ("STRING", {"forceInput": True}),
|
|
"save_prefix": ("STRING", {"default": "output/text"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "save"
|
|
CATEGORY = "dataset"
|
|
|
|
def save(self, text, save_prefix):
|
|
save_prefix = save_prefix.replace("..", "").strip("/")
|
|
|
|
filepath = os.path.join(self.output_dir, save_prefix + ".txt")
|
|
dirpath = os.path.dirname(filepath)
|
|
os.makedirs(dirpath, exist_ok=True)
|
|
|
|
with open(filepath, "w", encoding="utf-8") as f:
|
|
f.write(text)
|
|
|
|
logging.info(f"SaveTextExact: saved {filepath}")
|
|
return {}
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"LoadCaptionDataset": LoadCaptionDataset,
|
|
"SaveImageExact": SaveImageExact,
|
|
"CountFilesInFolder": CountFilesInFolder,
|
|
"SequentialIndex": SequentialIndex,
|
|
"LoadImageAtIndex": LoadImageAtIndex,
|
|
"LoadCaptionAtIndex": LoadCaptionAtIndex,
|
|
"LoadDatasetItemAtIndex": LoadDatasetItemAtIndex,
|
|
"SaveTextExact": SaveTextExact,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"LoadCaptionDataset": "Load Caption Dataset",
|
|
"SaveImageExact": "Save Image (Exact Filename)",
|
|
"CountFilesInFolder": "Count Files In Folder",
|
|
"SequentialIndex": "Sequential Index",
|
|
"LoadImageAtIndex": "Load Image At Index",
|
|
"LoadCaptionAtIndex": "Load Caption At Index",
|
|
"LoadDatasetItemAtIndex": "Load Dataset Item At Index",
|
|
"SaveTextExact": "Save Text (Exact Filename)",
|
|
}
|