From ac5569f2a04be716aac2d06d9e3d164f19baa090 Mon Sep 17 00:00:00 2001 From: Kohaku-Blueleaf <59680068+KohakuBlueleaf@users.noreply.github.com> Date: Tue, 7 Jul 2026 19:46:25 +0800 Subject: [PATCH 1/3] Add dataset folder --- extra_model_paths.yaml.example | 1 + folder_paths.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/extra_model_paths.yaml.example b/extra_model_paths.yaml.example index 6a31d8a63..755b8d124 100644 --- a/extra_model_paths.yaml.example +++ b/extra_model_paths.yaml.example @@ -29,6 +29,7 @@ # upscale_models: models/upscale_models/ # latent_upscale_models: models/latent_upscale_models/ # custom_nodes: custom_nodes/ +# datasets: datasets/ # hypernetworks: models/hypernetworks/ # photomaker: models/photomaker/ # classifiers: models/classifiers/ diff --git a/folder_paths.py b/folder_paths.py index ee048b0f2..c91065ccc 100644 --- a/folder_paths.py +++ b/folder_paths.py @@ -40,6 +40,8 @@ folder_names_and_paths["latent_upscale_models"] = ([os.path.join(models_dir, "la folder_names_and_paths["custom_nodes"] = ([os.path.join(base_path, "custom_nodes")], set()) +folder_names_and_paths["datasets"] = ([os.path.join(base_path, "datasets")], set()) + folder_names_and_paths["hypernetworks"] = ([os.path.join(models_dir, "hypernetworks")], supported_pt_extensions) folder_names_and_paths["photomaker"] = ([os.path.join(models_dir, "photomaker")], supported_pt_extensions) From 990dded2c1e6b658fc7e53978fc13bc85b9b2e68 Mon Sep 17 00:00:00 2001 From: Kohaku-Blueleaf <59680068+KohakuBlueleaf@users.noreply.github.com> Date: Tue, 7 Jul 2026 19:46:33 +0800 Subject: [PATCH 2/3] Use correct dataset folder in dataset nodes --- comfy_extras/nodes_dataset.py | 79 +++++++++++++++++++++++++++++------ 1 file changed, 66 insertions(+), 13 deletions(-) diff --git a/comfy_extras/nodes_dataset.py b/comfy_extras/nodes_dataset.py index 73fe75b7f..92f263620 100644 --- a/comfy_extras/nodes_dataset.py +++ b/comfy_extras/nodes_dataset.py @@ -42,6 +42,62 @@ def load_and_process_images(image_files, input_dir): return output_images +def secure_subfolder_path(base_dir, folder_name): + """Resolve folder_name inside base_dir, rejecting anything that escapes it. + + Blocks '..', absolute paths, drive letters and symlink escapes using the + same realpath containment check as the core file endpoints. + """ + target = os.path.abspath(os.path.join(base_dir, folder_name)) + if not folder_paths.is_within_directory(base_dir, target): + raise ValueError(f"Invalid folder name {folder_name!r}: resolves outside of {base_dir}") + return target + + +def list_dataset_folders(): + """Relative paths of dataset folders found under all dataset roots. + + Any subfolder containing a metadata.json or *.safetensors shard counts as + a dataset; the walk doesn't descend into matched folders. + """ + found = set() + for root in folder_paths.get_folder_paths("datasets"): + if not os.path.isdir(root): + continue + for dirpath, subdirs, filenames in os.walk(root, followlinks=True): + if dirpath != root and ( + "metadata.json" in filenames + or any(f.endswith(".safetensors") for f in filenames) + ): + found.add(os.path.relpath(dirpath, root).replace(os.sep, "/")) + subdirs[:] = [] + return sorted(found) + + +def get_dataset_save_dir(folder_name): + """Resolve the folder to save a new dataset into, inside the default root. + + The folder is not created here; callers makedirs after validation. + """ + root = folder_paths.get_folder_paths("datasets")[0] + target = secure_subfolder_path(root, folder_name) + if os.path.realpath(target) == os.path.realpath(root): + raise ValueError("folder_name must name a subfolder of the datasets directory, e.g. 'my_dataset'.") + return target + + +def get_dataset_dir(folder_name): + """Find an existing dataset folder by relative name across all dataset roots.""" + roots = folder_paths.get_folder_paths("datasets") + for root in roots: + target = secure_subfolder_path(root, folder_name) + if os.path.realpath(target) == os.path.realpath(root): + raise ValueError("folder_name must name a subfolder of the datasets directory, e.g. 'my_dataset'.") + if os.path.isdir(target): + return target + raise ValueError(f"Dataset folder {folder_name!r} not found in: {', '.join(roots)}") + + class LoadImageDataSetFromFolderNode(io.ComfyNode): @classmethod def define_schema(cls): @@ -252,7 +308,7 @@ class SaveImageDataSetToFolderNode(io.ComfyNode): filename_prefix = filename_prefix[0] mode = mode[0] - output_dir = os.path.join(folder_paths.get_output_directory(), folder_name) + output_dir = secure_subfolder_path(folder_paths.get_output_directory(), folder_name) saved_files = save_images_to_folder(images, output_dir, filename_prefix, mode=='overwrite') logging.info(f"Saved {len(saved_files)} images to {output_dir}.") @@ -306,7 +362,7 @@ class SaveImageTextDataSetToFolderNode(io.ComfyNode): filename_prefix = filename_prefix[0] mode = mode[0] - output_dir = os.path.join(folder_paths.get_output_directory(), folder_name) + output_dir = secure_subfolder_path(folder_paths.get_output_directory(), folder_name) saved_files = save_images_to_folder(images, output_dir, filename_prefix, mode=='overwrite') # Save captions @@ -1443,7 +1499,7 @@ class SaveTrainingDataset(io.ComfyNode): io.String.Input( "folder_name", default="training_dataset", - tooltip="Name of folder to save dataset (inside output directory).", + tooltip="Name of folder to save the dataset into, inside the datasets directory. Subfolders like 'project/run1' are allowed.", ), io.Int.Input( "shard_size", @@ -1473,8 +1529,8 @@ class SaveTrainingDataset(io.ComfyNode): f"Something went wrong in dataset preparation." ) - # Create output directory - output_dir = os.path.join(folder_paths.get_output_directory(), folder_name) + # Create output directory (inside the datasets root, traversal-safe) + output_dir = get_dataset_save_dir(folder_name) os.makedirs(output_dir, exist_ok=True) # Prepare data pairs @@ -1533,10 +1589,10 @@ class LoadTrainingDataset(io.ComfyNode): description="Load encoded training dataset (latents + conditioning) from disk for use in training.", is_experimental=True, inputs=[ - io.String.Input( + io.Combo.Input( "folder_name", - default="training_dataset", - tooltip="Name of folder containing the saved dataset (inside output directory).", + options=list_dataset_folders(), + tooltip="Saved dataset to load, from the datasets directory.", ), ], outputs=[ @@ -1555,11 +1611,8 @@ class LoadTrainingDataset(io.ComfyNode): @classmethod def execute(cls, folder_name): - # Get dataset directory - dataset_dir = os.path.join(folder_paths.get_output_directory(), folder_name) - - if not os.path.exists(dataset_dir): - raise ValueError(f"Dataset directory not found: {dataset_dir}") + # Get dataset directory (searched across all dataset roots, traversal-safe) + dataset_dir = get_dataset_dir(folder_name) # Find all shard files shard_files = sorted( From ce792cc96c021e323159903e3ef90cf1ffacd565 Mon Sep 17 00:00:00 2001 From: Kohaku-Blueleaf <59680068+KohakuBlueleaf@users.noreply.github.com> Date: Tue, 7 Jul 2026 21:00:19 +0800 Subject: [PATCH 3/3] Enhance dataset folder discovery logic Refactor dataset folder discovery to avoid symlink loops and improve directory handling. --- comfy_extras/nodes_dataset.py | 36 +++++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/comfy_extras/nodes_dataset.py b/comfy_extras/nodes_dataset.py index 92f263620..b5adfc699 100644 --- a/comfy_extras/nodes_dataset.py +++ b/comfy_extras/nodes_dataset.py @@ -59,18 +59,54 @@ def list_dataset_folders(): Any subfolder containing a metadata.json or *.safetensors shard counts as a dataset; the walk doesn't descend into matched folders. + + Symlinked directories are followed, but symlink loops are avoided. """ found = set() + for root in folder_paths.get_folder_paths("datasets"): if not os.path.isdir(root): continue + + root = os.path.abspath(root) + seen_dirs = set() + for dirpath, subdirs, filenames in os.walk(root, followlinks=True): + try: + st = os.stat(dirpath) # follows symlinks + except OSError: + subdirs[:] = [] + continue + + dir_key = (st.st_dev, st.st_ino) + if dir_key in seen_dirs: + subdirs[:] = [] + continue + + seen_dirs.add(dir_key) + if dirpath != root and ( "metadata.json" in filenames or any(f.endswith(".safetensors") for f in filenames) ): found.add(os.path.relpath(dirpath, root).replace(os.sep, "/")) subdirs[:] = [] + continue + + kept_subdirs = [] + for name in subdirs: + child = os.path.join(dirpath, name) + try: + child_st = os.stat(child) # follows symlinks + except OSError: + continue + + child_key = (child_st.st_dev, child_st.st_ino) + if child_key not in seen_dirs: + kept_subdirs.append(name) + + subdirs[:] = kept_subdirs + return sorted(found)