ComfyUI-DatasetGenerator

Custom nodes for ComfyUI that help you 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 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.

Input	Type	Description
folder	dropdown	Subfolder of ComfyUI/input/
index	INT	Item index (supports control_after_generate)

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

Save Dataset Item

Saves an image and/or caption 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). Both image and caption inputs are optional — connect whichever outputs you need to save. The image is saved as .png with embedded workflow metadata, the caption as .txt.

Input	Type	Description
save_prefix	STRING	Output path relative to ComfyUI/output/
image	IMAGE (optional)	Image to save as .png
caption	STRING (optional)	Text to save as .txt

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-processing a dataset:

1. Add Load Dataset Item At Index — select your input subfolder, set control to increment
2. Connect image to your processing workflow (e.g., Florence2 captioning, img2img)
3. Add Save Dataset Item — connect the processed image/caption and save_prefix
4. Queue the number of jobs matching your dataset size

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.