Mirrors
/
stable-diffusion-webui
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Store base_vae and loaded_vae_file in sd_model

This commit is contained in:
Uminosachi 2023-08-20 19:44:37 +09:00
parent 042e1d5d0b
commit 5159edbf0e
1 changed files with 8 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
modules/sd_models.py
View File

@ -462,7 +462,6 @@ class SdModelData:
def __init__(self):
self.sd_model = None
self.loaded_sd_models = []
self.loaded_vae_states = {}
self.was_loaded_at_least_once = False
self.lock = threading.Lock()
@ -489,24 +488,19 @@ class SdModelData:
def set_sd_model(self, v, already_loaded=False):
self.sd_model = v
if already_loaded:
sd_vae_state = self.loaded_vae_states.get(v.sd_model_hash, {})
sd_vae.base_vae = sd_vae_state.get("base_vae", None)
sd_vae.loaded_vae_file = sd_vae_state.get("loaded_vae_file", None)
sd_vae.checkpoint_info = sd_vae_state.get("checkpoint_info", None)
sd_vae.base_vae = getattr(v, "base_vae", None)
sd_vae.loaded_vae_file = getattr(v, "loaded_vae_file", None)
sd_vae.checkpoint_info = v.sd_checkpoint_info
try:
self.loaded_sd_models.remove(v)
self.loaded_vae_states.pop(v.sd_model_hash, {}).clear()
except ValueError:
pass
if v is not None:
setattr(v, "base_vae", sd_vae.base_vae)
setattr(v, "loaded_vae_file", sd_vae.loaded_vae_file)
self.loaded_sd_models.insert(0, v)
self.loaded_vae_states[v.sd_model_hash] = dict(
base_vae=sd_vae.base_vae,
loaded_vae_file=sd_vae.loaded_vae_file,
checkpoint_info=sd_vae.checkpoint_info,
)
model_data = SdModelData()
@ -661,7 +655,6 @@ def reuse_model_from_already_loaded(sd_model, checkpoint_info, timer):
if len(model_data.loaded_sd_models) > shared.opts.sd_checkpoints_limit > 0:
print(f"Unloading model {len(model_data.loaded_sd_models)} over the limit of {shared.opts.sd_checkpoints_limit}: {loaded_model.sd_checkpoint_info.title}")
model_data.loaded_sd_models.pop()
model_data.loaded_vae_states.pop(loaded_model.sd_model_hash, {}).clear()
send_model_to_trash(loaded_model)
timer.record("send model to trash")
@ -691,10 +684,9 @@ def reuse_model_from_already_loaded(sd_model, checkpoint_info, timer):
sd_model = model_data.loaded_sd_models.pop()
model_data.sd_model = sd_model
sd_vae_state = model_data.loaded_vae_states.pop(sd_model.sd_model_hash, {})
sd_vae.base_vae = sd_vae_state.get("base_vae", None)
sd_vae.loaded_vae_file = sd_vae_state.get("loaded_vae_file", None)
sd_vae.checkpoint_info = sd_vae_state.get("checkpoint_info", None)
sd_vae.base_vae = getattr(sd_model, "base_vae", None)
sd_vae.loaded_vae_file = getattr(sd_model, "loaded_vae_file", None)
sd_vae.checkpoint_info = sd_model.sd_checkpoint_info
print(f"Reusing loaded model {sd_model.sd_checkpoint_info.title} to load {checkpoint_info.title}")
return sd_model