wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-04-24 01:12:37 +08:00
Code Issues Actions 19 Packages Projects Releases Wiki Activity

Fix torch.load missing weights_only in LoadTrainingDataset

Browse Source 
The rest of the codebase uses weights_only=True. Also adds a path
check on folder_name.
This commit is contained in:
jakelodwick 2026-04-05 19:13:08 -06:00
parent 4b1444fc7a
commit 8088b347d0
1 changed files with 6 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
comfy_extras/nodes_dataset.py
View File

@ -1450,7 +1450,11 @@ class LoadTrainingDataset(io.ComfyNode):
@classmethod @classmethod
def execute(cls, folder_name): def execute(cls, folder_name):
# Get dataset directory # 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): if not os.path.exists(dataset_dir):
raise ValueError(f"Dataset directory not found: {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) shard_path = os.path.join(dataset_dir, shard_file)
with open(shard_path, "rb") as f: 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_latents.extend(shard_data["latents"])
all_conditioning.extend(shard_data["conditioning"]) all_conditioning.extend(shard_data["conditioning"])