nodes/sintspiden-ComfyUI-DatasetGenerator
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Initial commit: dataset utility nodes for ComfyUI

Browse Source 
Eight custom nodes for loading, iterating, and saving image/caption
datasets in ComfyUI workflows.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Solo Yolo
2026-03-22 21:07:15 +09:00
co-authored by Claude Opus 4.6
commit 8116f97fab
4 changed files with 557 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+2
View File
@@ -0,0 +1,2 @@
__pycache__/
*.py[cod]
README.md
+90
View File
@@ -0,0 +1,90 @@
# ComfyUI-DatasetGenerator
Custom nodes for ComfyUI that help you build, iterate over, and save image/caption datasets. Useful for preparing training data for fine-tuning diffusion models or for batch-processing existing datasets through ComfyUI workflows.
## Nodes
All nodes appear under the **dataset** category in the node menu.
### Load Caption Dataset
Loads all `.txt` caption files from a subfolder of `ComfyUI/input/`. Outputs parallel lists of captions, save prefixes, and corresponding source images (matched by filename stem).
| Output | Type | Description |
|--------|------|-------------|
| `captions` | STRING (list) | Text content of each `.txt` file |
| `save_prefixes` | STRING (list) | `folder/stem` paths for saving outputs |
| `images` | IMAGE (list) | Matching images (64x64 black placeholder if no image found) |
### Load Dataset Item At Index
Loads a single image + caption pair from a dataset folder at a given index. Pairs are matched by filename stem (e.g., `photo.png` + `photo.txt`). Items are sorted alphabetically by caption filename.
This is the recommended node for sequential processing — it loads both image and caption atomically, preventing race conditions that can occur when using separate index-based nodes.
| Output | Type | Description |
|--------|------|-------------|
| `image` | IMAGE | The image at the given index |
| `caption` | STRING | The caption text at the given index |
| `save_prefix` | STRING | `folder/stem` path for saving |
### Load Image At Index
Loads a single image from a folder at a given index. Images are sorted alphabetically.
### Load Caption At Index
Loads a single `.txt` caption from a folder at a given index. Captions are sorted alphabetically.
### Sequential Index
An index counter with `control_after_generate` support. Set the control to **increment** and queue N jobs to iterate through items sequentially. Connect `max_count` from **Count Files In Folder** to stop at the end of the dataset.
### Count Files In Folder
Counts image files (`.png`, `.jpg`, `.jpeg`, `.webp`) or text files (`.txt`) in a subfolder of `ComfyUI/input/`.
### Save Image (Exact Filename)
Saves an image to the output directory using the exact `save_prefix` as the filename — no counter suffix is appended. Supports subdirectories (e.g., `my_dataset/image_001`). Embeds workflow metadata in the PNG.
### Save Text (Exact Filename)
Saves a text string to the output directory using the exact `save_prefix` as the filename (`.txt` extension). Supports subdirectories.
## Installation
Clone this repo into your ComfyUI custom nodes directory:
```bash
cd ComfyUI/custom_nodes
git clone https://github.com/sintspiden/ComfyUI-DatasetGenerator.git
```
No additional dependencies are required beyond what ComfyUI already provides.
## Usage Example
A typical workflow for batch-captioning a folder of images:
1. Add **Count Files In Folder** — select your input subfolder, set file type to `images`
2. Add **Sequential Index** — connect `count` to `max_count`, set control to `increment`
3. Add **Load Image At Index** — connect the index output
4. Connect the image to your captioning workflow (e.g., Florence2, BLIP)
5. Add **Save Text (Exact Filename)** — connect the caption output and `filename_stem` as save prefix
6. Queue the number of jobs matching your image count
## Dataset Folder Structure
Place your dataset in a subfolder of `ComfyUI/input/`:
```
ComfyUI/input/my_dataset/
├── image_001.png
├── image_001.txt # caption for image_001
├── image_002.jpg
├── image_002.txt # caption for image_002
└── ...
```
Images and captions are paired by filename stem — `photo.png` matches `photo.txt`.
__init__.py
+3
View File
@@ -0,0 +1,3 @@
from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
nodes.py
+462
View File
@@ -0,0 +1,462 @@
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)",
}