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
8eef9d8e
Commit
8eef9d8e
authored
Dec 25, 2022
by
AUTOMATIC
Browse files
Options
Browse Files
Download
Email Patches
Plain Diff
a way to add an exception to unpickler without explicitly calling load_with_extra
parent
c5bdba20
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
38 additions
and
1 deletion
+38
-1
safe.py
modules/safe.py
+38
-1
No files found.
modules/safe.py
View file @
8eef9d8e
...
@@ -103,7 +103,7 @@ def check_pt(filename, extra_handler):
...
@@ -103,7 +103,7 @@ def check_pt(filename, extra_handler):
def
load
(
filename
,
*
args
,
**
kwargs
):
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
):
def
load_with_extra
(
filename
,
extra_handler
=
None
,
*
args
,
**
kwargs
):
...
@@ -151,5 +151,42 @@ def load_with_extra(filename, extra_handler=None, *args, **kwargs):
...
@@ -151,5 +151,42 @@ def load_with_extra(filename, extra_handler=None, *args, **kwargs):
return
unsafe_torch_load
(
filename
,
*
args
,
**
kwargs
)
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
unsafe_torch_load
=
torch
.
load
torch
.
load
=
load
torch
.
load
=
load
global_extra_handler
=
None
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