Skip to content
Projects
Groups
Snippets
Help
Loading...
Help
Contribute to GitLab
Sign in / Register
Toggle navigation
S
stable-diffusion-webui
Project
Project
Details
Activity
Cycle Analytics
Repository
Repository
Files
Commits
Branches
Tags
Contributors
Graph
Compare
Charts
Issues
0
Issues
0
List
Board
Labels
Milestones
Merge Requests
0
Merge Requests
0
CI / CD
CI / CD
Pipelines
Jobs
Schedules
Charts
Wiki
Wiki
Snippets
Snippets
Members
Members
Collapse sidebar
Close sidebar
Activity
Graph
Charts
Create a new issue
Jobs
Commits
Issue Boards
Open sidebar
Administrator
stable-diffusion-webui
Commits
c3777777
Unverified
Commit
c3777777
authored
Dec 03, 2022
by
AUTOMATIC1111
Committed by
GitHub
Dec 03, 2022
Browse files
Options
Browse Files
Download
Plain Diff
Merge pull request #5327 from smirkingface/master
Fixed safety checker for ckpt files written with pytorch >=1.13
parents
4b3c5bc2
e4614778
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
11 additions
and
7 deletions
+11
-7
safe.py
modules/safe.py
+11
-7
No files found.
modules/safe.py
View file @
c3777777
...
...
@@ -62,14 +62,12 @@ class RestrictedUnpickler(pickle.Unpickler):
raise
Exception
(
f
"global '{module}/{name}' is forbidden"
)
allowed_zip_names
=
[
"archive/data.pkl"
,
"archive/version"
]
allowed_zip_names_re
=
re
.
compile
(
r"^
archive/data/\d+
$"
)
# Regular expression that accepts 'dirname/version', 'dirname/data.pkl', and 'dirname/data/<number>'
allowed_zip_names_re
=
re
.
compile
(
r"^
([^/]+)/((data/\d+)|version|(data\.pkl))
$"
)
data_pkl_re
=
re
.
compile
(
r"^([^/]+)/data\.pkl$"
)
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
...
...
@@ -82,8 +80,14 @@ 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
())
with
z
.
open
(
'archive/data.pkl'
)
as
file
:
# 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
:
raise
Exception
(
f
"data.pkl not found in {filename}"
)
if
len
(
data_pkl_filenames
)
>
1
:
raise
Exception
(
f
"Multiple data.pkl found in {filename}"
)
with
z
.
open
(
data_pkl_filenames
[
0
])
as
file
:
unpickler
=
RestrictedUnpickler
(
file
)
unpickler
.
extra_handler
=
extra_handler
unpickler
.
load
()
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment