ComfyUI-DatasetGenerator

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.