Merge branch 'master' of github.com:AUTOMATIC1111/stable-diffusion-webui

1 files changed, 50 insertions, 13 deletions

diff --git a/modules/safe.py b/modules/safe.py
index 10460ad0..82d44be3 100644
--- a/modules/safe.py
+++ b/modules/safe.py
@@ -37,16 +37,16 @@ class RestrictedUnpickler(pickle.Unpickler):
 

         if module == 'collections' and name == 'OrderedDict':

             return getattr(collections, name)

-        if module == 'torch._utils' and name in ['_rebuild_tensor_v2', '_rebuild_parameter']:

+        if module == 'torch._utils' and name in ['_rebuild_tensor_v2', '_rebuild_parameter', '_rebuild_device_tensor_from_numpy']:

             return getattr(torch._utils, name)

-        if module == 'torch' and name in ['FloatStorage', 'HalfStorage', 'IntStorage', 'LongStorage', 'DoubleStorage', 'ByteStorage']:

+        if module == 'torch' and name in ['FloatStorage', 'HalfStorage', 'IntStorage', 'LongStorage', 'DoubleStorage', 'ByteStorage', 'float32']:

             return getattr(torch, name)

         if module == 'torch.nn.modules.container' and name in ['ParameterDict']:

             return getattr(torch.nn.modules.container, name)

-        if module == 'numpy.core.multiarray' and name == 'scalar':

-            return numpy.core.multiarray.scalar

-        if module == 'numpy' and name == 'dtype':

-            return numpy.dtype

+        if module == 'numpy.core.multiarray' and name in ['scalar', '_reconstruct']:

+            return getattr(numpy.core.multiarray, name)

+        if module == 'numpy' and name in ['dtype', 'ndarray']:

+            return getattr(numpy, name)

         if module == '_codecs' and name == 'encode':

             return encode

         if module == "pytorch_lightning.callbacks" and name == 'model_checkpoint':

@@ -80,7 +80,7 @@ def check_pt(filename, extra_handler):
         # new pytorch format is a zip file

         with zipfile.ZipFile(filename) as z:

             check_zip_filenames(filename, z.namelist())

-            

+

             # find filename of data.pkl in zip file: '<directory name>/data.pkl'

             data_pkl_filenames = [f for f in z.namelist() if data_pkl_re.match(f)]

             if len(data_pkl_filenames) == 0:

@@ -103,12 +103,12 @@ def check_pt(filename, extra_handler):
 

 

 def load(filename, *args, **kwargs):

-    return load_with_extra(filename, *args, **kwargs)

+    return load_with_extra(filename, extra_handler=global_extra_handler, *args, **kwargs)

 

 

 def load_with_extra(filename, extra_handler=None, *args, **kwargs):

     """

-    this functon is intended to be used by extensions that want to load models with

+    this function is intended to be used by extensions that want to load models with

     some extra classes in them that the usual unpickler would find suspicious.

 

     Use the extra_handler argument to specify a function that takes module and field name as text,

@@ -137,19 +137,56 @@ def load_with_extra(filename, extra_handler=None, *args, **kwargs):
     except pickle.UnpicklingError:

         print(f"Error verifying pickled file from {filename}:", file=sys.stderr)

         print(traceback.format_exc(), file=sys.stderr)

-        print(f"-----> !!!! The file is most likely corrupted !!!! <-----", file=sys.stderr)

-        print(f"You can skip this check with --disable-safe-unpickle commandline argument, but that is not going to help you.\n\n", file=sys.stderr)

+        print("-----> !!!! The file is most likely corrupted !!!! <-----", file=sys.stderr)

+        print("You can skip this check with --disable-safe-unpickle commandline argument, but that is not going to help you.\n\n", file=sys.stderr)

         return None

 

     except Exception:

         print(f"Error verifying pickled file from {filename}:", file=sys.stderr)

         print(traceback.format_exc(), file=sys.stderr)

-        print(f"\nThe file may be malicious, so the program is not going to read it.", file=sys.stderr)

-        print(f"You can skip this check with --disable-safe-unpickle commandline argument.\n\n", file=sys.stderr)

+        print("\nThe file may be malicious, so the program is not going to read it.", file=sys.stderr)

+        print("You can skip this check with --disable-safe-unpickle commandline argument.\n\n", file=sys.stderr)

         return None

 

     return unsafe_torch_load(filename, *args, **kwargs)

 

 

+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')

+```

+    """

+

+    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

+

+    def __exit__(self, exc_type, exc_val, exc_tb):

+        global global_extra_handler

+

+        global_extra_handler = None

+

+

 unsafe_torch_load = torch.load

 torch.load = load

+global_extra_handler = None

+