diff options
| author | JC-Array <44535867+JC-Array@users.noreply.github.com> | 2022-10-11 22:33:15 +0000 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2022-10-11 22:33:15 +0000 |
| commit | 963d98639673098fa8df975dd380f1ef56fff3b5 (patch) | |
| tree | 21d41f53af03ce2b21de7947fc216784fb2f2364 /modules/safe.py | |
| parent | ff4ef13dd591ec52f196f344f47537695df95364 (diff) | |
| parent | 6be32b31d181e42c639dad3451229aa7b9cfd1cf (diff) | |
| download | stable-diffusion-webui-gfx803-963d98639673098fa8df975dd380f1ef56fff3b5.tar.gz stable-diffusion-webui-gfx803-963d98639673098fa8df975dd380f1ef56fff3b5.tar.bz2 stable-diffusion-webui-gfx803-963d98639673098fa8df975dd380f1ef56fff3b5.zip | |
Merge branch 'AUTOMATIC1111:master' into deepdanbooru_pre_process
Diffstat (limited to 'modules/safe.py')
| -rw-r--r-- | modules/safe.py | 17 |
1 files changed, 17 insertions, 0 deletions
diff --git a/modules/safe.py b/modules/safe.py index 05917463..20be16a5 100644 --- a/modules/safe.py +++ b/modules/safe.py @@ -10,6 +10,7 @@ import torch import numpy
import _codecs
import zipfile
+import re
# PyTorch 1.13 and later have _TypedStorage renamed to TypedStorage
@@ -54,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()
|