commit 8116f97faba2b868bccf63c84cf3c828e2e40130
Author: Solo Yolo <solo@yolo.com>
Date:   Sun Mar 22 21:07:15 2026 +0900

    Initial commit: dataset utility nodes for ComfyUI
    
    Eight custom nodes for loading, iterating, and saving image/caption
    datasets in ComfyUI workflows.
    
    Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..43ae0e2
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,2 @@
+__pycache__/
+*.py[cod]
diff --git a/README.md b/README.md
new file mode 100644
index 0000000..2bae5bf
--- /dev/null
+++ b/README.md
@@ -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`.
diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..39a8c6b
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,3 @@
+from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
diff --git a/nodes.py b/nodes.py
new file mode 100644
index 0000000..df20136
--- /dev/null
+++ b/nodes.py
@@ -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)",
+}