Merge branch 'dev' into tooltip

1 files changed, 5 insertions, 4 deletions

diff --git a/modules/safe.py b/modules/safe.py
index 82d44be3..e6c2f2c0 100644
--- a/modules/safe.py
+++ b/modules/safe.py
@@ -1,6 +1,5 @@
 # this code is adapted from the script contributed by anon from /h/

 

-import io

 import pickle

 import collections

 import sys

@@ -12,11 +11,9 @@ import _codecs
 import zipfile

 import re

 

-

 # PyTorch 1.13 and later have _TypedStorage renamed to TypedStorage

 TypedStorage = torch.storage.TypedStorage if hasattr(torch.storage, 'TypedStorage') else torch.storage._TypedStorage

 

-

 def encode(*args):

     out = _codecs.encode(*args)

     return out

@@ -27,7 +24,11 @@ class RestrictedUnpickler(pickle.Unpickler):
 

     def persistent_load(self, saved_id):

         assert saved_id[0] == 'storage'

-        return TypedStorage()

+

+        try:

+            return TypedStorage(_internal=True)

+        except TypeError:

+            return TypedStorage()  # PyTorch before 2.0 does not have the _internal argument

 

     def find_class(self, module, name):

         if self.extra_handler is not None: