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:

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.