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

Merge cf43a3a63e into 3086026401

Browse Source
This commit is contained in:
Jake Lodwick 2026-04-18 15:26:07 -07:00 committed by GitHub
commit 219468b0d6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 8 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
comfy_extras/nodes_dataset.py
View File

@ -1450,7 +1450,13 @@ 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")
real_output_dir = os.path.realpath(output_dir)
real_dataset_dir = os.path.realpath(dataset_dir)
if os.path.commonpath((real_output_dir, real_dataset_dir)) != real_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 +1483,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"])