safe

backend/misc/checkpoint_pickle.py

@@ -1,22 +0,0 @@
-import pickle
-
-load = pickle.load
-
-
-class Empty:
-    pass
-
-
-class RestrictedUnpickler(pickle.Unpickler):
-
-    def find_class(self, module: str, name: str):
-        if module.startswith("pytorch_lightning"):
-            return Empty
-
-        if module.startswith(("collections", "torch", "numpy", "__builtin__")):
-            return super().find_class(module, name)
-
-        raise NotImplementedError(f'"{module}.{name}" is forbidden')
-
-
-Unpickler = RestrictedUnpickler

backend/utils.py

@@ -6,9 +6,9 @@ import safetensors
 import torch
 from einops import rearrange, repeat
 
-import backend.misc.checkpoint_pickle
 from backend.args import args
 from backend.operations_gguf import ParameterGGUF
+from modules import safe
 
 MMAP_TORCH_FILES = args.mmap_torch_files
 DISABLE_MMAP = args.disable_mmap
@@ -46,7 +46,7 @@ def load_torch_file(ckpt: str, safe_load=False, device=None, *, return_metadata=
         except Exception as e:
             if len(e.args) > 0:
                 if "HeaderTooLarge" in e.args[0] or "MetadataIncompleteBuffer" in e.args[0]:
-                    raise ValueError('\nModel: "{}" is corrupt or invalid...'.format(ckpt))
+                    raise ValueError(f'\nModel: "{ckpt}" is corrupt or invalid...\nPlease download the model again')
             raise e
 
     elif ckpt.lower().endswith(".gguf"):
@@ -59,7 +59,7 @@ def load_torch_file(ckpt: str, safe_load=False, device=None, *, return_metadata=
         torch_args = {}
 
         if not safe_load:
-            torch_args["pickle_module"] = backend.misc.checkpoint_pickle
+            torch_args["pickle_module"] = safe
         else:
             torch_args["weights_only"] = True
             if MMAP_TORCH_FILES:

modules/safe.py

@@ -1,6 +1,10 @@
 import pickle
+
 import torch
 
+load = pickle.load
+unsafe_torch_load = torch.load
+
 
 class Empty:
     pass
@@ -19,41 +23,17 @@ class RestrictedUnpickler(pickle.Unpickler):
 
 
 class Extra:
-    """
-    A class for temporarily setting the global handler for when you can't explicitly call load_with_extra
-    (because it's not your code making the torch.load call). The intended use is like this:
-
-    ```
-    import torch
-    from modules import safe
-
-    def handler(module, name):
-        if module == "torch" and name in ["float64", "float16"]:
-            return getattr(torch, name)
-
-        return None
-
-    with safe.Extra(handler):
-        x = torch.load("model.pt")
-    ```
-    """
+    global_extra_handler = None
 
     def __init__(self, handler):
         self.handler = handler
 
     def __enter__(self):
-        global global_extra_handler
-
-        assert global_extra_handler is None, "already inside an Extra() block"
-        global_extra_handler = self.handler
+        assert Extra.global_extra_handler is None, "already inside an Extra() block"
+        Extra.global_extra_handler = self.handler
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        global global_extra_handler
+        Extra.global_extra_handler = None
 
-        global_extra_handler = None
-
-
-unsafe_torch_load = torch.load
-global_extra_handler = None
 
 Unpickler = RestrictedUnpickler