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
abc1e79a
Commit
abc1e79a
authored
Nov 03, 2022
by
Muhammad Rizqi Nur
Browse files
Options
Browse Files
Download
Email Patches
Plain Diff
Fix base VAE caching was done after loading VAE, also add safeguard
parent
8ab49274
Show whitespace changes
Inline
Side-by-side
Showing
2 changed files
with
9 additions
and
11 deletions
+9
-11
sd_models.py
modules/sd_models.py
+1
-0
sd_vae.py
modules/sd_vae.py
+8
-11
No files found.
modules/sd_models.py
View file @
abc1e79a
...
@@ -220,6 +220,7 @@ def load_model_weights(model, checkpoint_info, vae_file="auto"):
...
@@ -220,6 +220,7 @@ def load_model_weights(model, checkpoint_info, vae_file="auto"):
model
.
sd_model_checkpoint
=
checkpoint_file
model
.
sd_model_checkpoint
=
checkpoint_file
model
.
sd_checkpoint_info
=
checkpoint_info
model
.
sd_checkpoint_info
=
checkpoint_info
sd_vae
.
clear_loaded_vae
()
sd_vae
.
load_vae
(
model
,
vae_file
)
sd_vae
.
load_vae
(
model
,
vae_file
)
...
...
modules/sd_vae.py
View file @
abc1e79a
...
@@ -15,7 +15,7 @@ vae_path = os.path.abspath(os.path.join(models_path, vae_dir))
...
@@ -15,7 +15,7 @@ vae_path = os.path.abspath(os.path.join(models_path, vae_dir))
vae_ignore_keys
=
{
"model_ema.decay"
,
"model_ema.num_updates"
}
vae_ignore_keys
=
{
"model_ema.decay"
,
"model_ema.num_updates"
}
default_vae_dict
=
{
"auto"
:
"auto"
,
"None"
:
"None"
}
default_vae_dict
=
{
"auto"
:
"auto"
,
"None"
:
None
,
None
:
None
}
default_vae_list
=
[
"auto"
,
"None"
]
default_vae_list
=
[
"auto"
,
"None"
]
...
@@ -39,6 +39,7 @@ def get_base_vae(model):
...
@@ -39,6 +39,7 @@ def get_base_vae(model):
def
store_base_vae
(
model
):
def
store_base_vae
(
model
):
global
base_vae
,
checkpoint_info
global
base_vae
,
checkpoint_info
if
checkpoint_info
!=
model
.
sd_checkpoint_info
:
if
checkpoint_info
!=
model
.
sd_checkpoint_info
:
assert
not
loaded_vae_file
,
"Trying to store non-base VAE!"
base_vae
=
model
.
first_stage_model
.
state_dict
()
.
copy
()
base_vae
=
model
.
first_stage_model
.
state_dict
()
.
copy
()
checkpoint_info
=
model
.
sd_checkpoint_info
checkpoint_info
=
model
.
sd_checkpoint_info
...
@@ -50,9 +51,11 @@ def delete_base_vae():
...
@@ -50,9 +51,11 @@ def delete_base_vae():
def
restore_base_vae
(
model
):
def
restore_base_vae
(
model
):
global
loaded_vae_file
if
base_vae
is
not
None
and
checkpoint_info
==
model
.
sd_checkpoint_info
:
if
base_vae
is
not
None
and
checkpoint_info
==
model
.
sd_checkpoint_info
:
print
(
"Restoring base VAE"
)
print
(
"Restoring base VAE"
)
load_vae_dict
(
model
,
base_vae
)
load_vae_dict
(
model
,
base_vae
)
loaded_vae_file
=
None
delete_base_vae
()
delete_base_vae
()
...
@@ -140,10 +143,10 @@ def load_vae(model, vae_file=None):
...
@@ -140,10 +143,10 @@ def load_vae(model, vae_file=None):
if
vae_file
:
if
vae_file
:
print
(
f
"Loading VAE weights from: {vae_file}"
)
print
(
f
"Loading VAE weights from: {vae_file}"
)
store_base_vae
(
model
)
vae_ckpt
=
torch
.
load
(
vae_file
,
map_location
=
shared
.
weight_load_location
)
vae_ckpt
=
torch
.
load
(
vae_file
,
map_location
=
shared
.
weight_load_location
)
vae_dict_1
=
{
k
:
v
for
k
,
v
in
vae_ckpt
[
"state_dict"
]
.
items
()
if
k
[
0
:
4
]
!=
"loss"
and
k
not
in
vae_ignore_keys
}
vae_dict_1
=
{
k
:
v
for
k
,
v
in
vae_ckpt
[
"state_dict"
]
.
items
()
if
k
[
0
:
4
]
!=
"loss"
and
k
not
in
vae_ignore_keys
}
load_vae_dict
(
model
,
vae_dict_1
)
load_vae_dict
(
model
,
vae_dict_1
)
store_base_vae
(
model
)
# If vae used is not in dict, update it
# If vae used is not in dict, update it
# It will be removed on refresh though
# It will be removed on refresh though
...
@@ -157,15 +160,6 @@ def load_vae(model, vae_file=None):
...
@@ -157,15 +160,6 @@ def load_vae(model, vae_file=None):
loaded_vae_file
=
vae_file
loaded_vae_file
=
vae_file
"""
# Save current VAE to VAE settings, maybe? will it work?
if save_settings:
if vae_file is None:
vae_opt = "None"
# shared.opts.sd_vae = vae_opt
"""
first_load
=
False
first_load
=
False
...
@@ -174,6 +168,9 @@ def load_vae_dict(model, vae_dict_1):
...
@@ -174,6 +168,9 @@ def load_vae_dict(model, vae_dict_1):
model
.
first_stage_model
.
load_state_dict
(
vae_dict_1
)
model
.
first_stage_model
.
load_state_dict
(
vae_dict_1
)
model
.
first_stage_model
.
to
(
devices
.
dtype_vae
)
model
.
first_stage_model
.
to
(
devices
.
dtype_vae
)
def
clear_loaded_vae
():
global
loaded_vae_file
loaded_vae_file
=
None
def
reload_vae_weights
(
sd_model
=
None
,
vae_file
=
"auto"
):
def
reload_vae_weights
(
sd_model
=
None
,
vae_file
=
"auto"
):
from
modules
import
lowvram
,
devices
,
sd_hijack
from
modules
import
lowvram
,
devices
,
sd_hijack
...
...
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