From 1b877e583f4c5bc37d94be77c3faa09f2eb078a2 Mon Sep 17 00:00:00 2001 From: CFionaBF Date: Sun, 7 Jun 2026 11:20:20 -0400 Subject: [PATCH 1/2] Load training dataset tensors with weights_only --- comfy_extras/nodes_dataset.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy_extras/nodes_dataset.py b/comfy_extras/nodes_dataset.py index 0253b4b4f..73fe75b7f 100644 --- a/comfy_extras/nodes_dataset.py +++ b/comfy_extras/nodes_dataset.py @@ -1583,7 +1583,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"]) From 0878ec8e5fc86a8020e14f532fdeb4aa7bc297fc Mon Sep 17 00:00:00 2001 From: CFionaBF Date: Sun, 14 Jun 2026 01:37:33 -0400 Subject: [PATCH 2/2] Allowlist dataset conditioning hook classes for safe loading Datasets built with CLIP hooks store comfy.hooks objects in conditioning; register them as safe globals so they still load with weights_only=True. --- comfy_extras/nodes_dataset.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/comfy_extras/nodes_dataset.py b/comfy_extras/nodes_dataset.py index 73fe75b7f..83d68b9d2 100644 --- a/comfy_extras/nodes_dataset.py +++ b/comfy_extras/nodes_dataset.py @@ -10,6 +10,30 @@ from typing_extensions import override import folder_paths import node_helpers from comfy_api.latest import ComfyExtension, io +import comfy.hooks + +# Datasets produced by MakeTrainingDataset can embed comfy.hooks objects in their +# conditioning (clip.encode_from_tokens_scheduled stores a HookGroup under "hooks" +# when the CLIP has hooks applied). LoadTrainingDataset loads shards with +# weights_only=True, so register these known classes as safe globals; otherwise such +# datasets fail to load. Only data-bearing hooks (e.g. WeightHook) round-trip this way. +torch.serialization.add_safe_globals([ + comfy.hooks.HookGroup, + comfy.hooks.Hook, + comfy.hooks.WeightHook, + comfy.hooks.ObjectPatchHook, + comfy.hooks.AdditionalModelsHook, + comfy.hooks.TransformerOptionsHook, + comfy.hooks.InjectionsHook, + comfy.hooks._HookRef, + comfy.hooks.HookKeyframe, + comfy.hooks.HookKeyframeGroup, + comfy.hooks.EnumHookMode, + comfy.hooks.EnumHookType, + comfy.hooks.EnumWeightTarget, + comfy.hooks.EnumHookScope, + comfy.hooks.default_should_register, +]) def load_and_process_images(image_files, input_dir):