From b5c59388539955e4936490d000e8dca57211f8fb Mon Sep 17 00:00:00 2001 From: Won-Kyu Park Date: Tue, 22 Oct 2024 05:17:16 +0900 Subject: [PATCH] fix RestrictedUnpickler it does no harm to remove the pytorch_lightening dependency see also https://github.com/comfyanonymous/ComfyUI/commit/735ac4cf81862b21902b312930ebfc92eef63357 --- modules/safe.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/modules/safe.py b/modules/safe.py index af019ffd9..9014043c6 100644 --- a/modules/safe.py +++ b/modules/safe.py @@ -37,6 +37,9 @@ class RestrictedUnpickler(pickle.Unpickler): if res is not None: return res + class Empty: + pass + if module == 'collections' and name == 'OrderedDict': return getattr(collections, name) if module == 'torch._utils' and name in ['_rebuild_tensor_v2', '_rebuild_parameter', '_rebuild_device_tensor_from_numpy']: @@ -51,12 +54,8 @@ class RestrictedUnpickler(pickle.Unpickler): return getattr(numpy, name) if module == '_codecs' and name == 'encode': return encode - if module == "pytorch_lightning.callbacks" and name == 'model_checkpoint': - import pytorch_lightning.callbacks - return pytorch_lightning.callbacks.model_checkpoint - if module == "pytorch_lightning.callbacks.model_checkpoint" and name == 'ModelCheckpoint': - import pytorch_lightning.callbacks.model_checkpoint - return pytorch_lightning.callbacks.model_checkpoint.ModelCheckpoint + if module.startswith("pytorch_lightning"): + return Empty if module == "__builtin__" and name == 'set': return set