From 8088b347d0aeff7ac0319996d351eec62450a7ba Mon Sep 17 00:00:00 2001 From: jakelodwick Date: Sun, 5 Apr 2026 19:13:08 -0600 Subject: [PATCH] Fix torch.load missing weights_only in LoadTrainingDataset The rest of the codebase uses weights_only=True. Also adds a path check on folder_name. --- comfy_extras/nodes_dataset.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/comfy_extras/nodes_dataset.py b/comfy_extras/nodes_dataset.py index 98ed25d7e..f7d9afbff 100644 --- a/comfy_extras/nodes_dataset.py +++ b/comfy_extras/nodes_dataset.py @@ -1450,7 +1450,11 @@ 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) + output_dir = folder_paths.get_output_directory() + dataset_dir = os.path.join(output_dir, folder_name) + # Prevent path traversal (e.g. folder_name="../../etc") + if not os.path.realpath(dataset_dir).startswith(os.path.realpath(output_dir)): + raise ValueError(f"Invalid folder_name: path traversal detected") if not os.path.exists(dataset_dir): raise ValueError(f"Dataset directory not found: {dataset_dir}") @@ -1477,7 +1481,7 @@ class LoadTrainingDataset(io.ComfyNode): shard_path = os.path.join(dataset_dir, shard_file) with open(shard_path, "rb") as f: - shard_data = torch.load(f) + shard_data = torch.load(f, weights_only=True) all_latents.extend(shard_data["latents"]) all_conditioning.extend(shard_data["conditioning"])