Merge branch 'master' into feature/scale_to

1 files changed, 22 insertions, 1 deletions

diff --git a/modules/safe.py b/modules/safe.py
index 4d06f2a5..20be16a5 100644
--- a/modules/safe.py
+++ b/modules/safe.py
@@ -10,6 +10,11 @@ import torch
 import numpy

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

@@ -20,7 +25,7 @@ def encode(*args):
 class RestrictedUnpickler(pickle.Unpickler):

     def persistent_load(self, saved_id):

         assert saved_id[0] == 'storage'

-        return torch.storage._TypedStorage()

+        return TypedStorage()

 

     def find_class(self, module, name):

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

@@ -50,11 +55,27 @@ class RestrictedUnpickler(pickle.Unpickler):
         raise pickle.UnpicklingError(f"global '{module}/{name}' is forbidden")

 

 

+allowed_zip_names = ["archive/data.pkl", "archive/version"]

+allowed_zip_names_re = re.compile(r"^archive/data/\d+$")

+

+

+def check_zip_filenames(filename, names):

+    for name in names:

+        if name in allowed_zip_names:

+            continue

+        if allowed_zip_names_re.match(name):

+            continue

+

+        raise Exception(f"bad file inside {filename}: {name}")

+

+

 def check_pt(filename):

     try:

 

         # new pytorch format is a zip file

         with zipfile.ZipFile(filename) as z:

+            check_zip_filenames(filename, z.namelist())

+

             with z.open('archive/data.pkl') as file:

                 unpickler = RestrictedUnpickler(file)

                 unpickler.load()