wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-06-24 00:39:30 +08:00
Code Issues Actions 18 Packages Projects Releases Wiki Activity

Allowlist dataset conditioning hook classes for safe loading

Browse Source 
Datasets built with CLIP hooks store comfy.hooks objects in conditioning;
register them as safe globals so they still load with weights_only=True.
This commit is contained in:
CFionaBF 2026-06-14 01:37:33 -04:00
parent 1b877e583f
commit 0878ec8e5f
1 changed files with 24 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
comfy_extras/nodes_dataset.py
View File

@ -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):