From 8ab4927452b04dcd30847eaf92ea7a9f3b9c74e1 Mon Sep 17 00:00:00 2001 From: Muhammad Rizqi Nur Date: Wed, 2 Nov 2022 22:54:09 +0700 Subject: [PATCH] Fix model wasn't restored even when choosing "None" --- modules/sd_vae.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/modules/sd_vae.py b/modules/sd_vae.py index 71e7a6e63..7a79239f1 100644 --- a/modules/sd_vae.py +++ b/modules/sd_vae.py @@ -50,8 +50,8 @@ def delete_base_vae(): def restore_base_vae(model): - global base_vae, checkpoint_info if base_vae is not None and checkpoint_info == model.sd_checkpoint_info: + print("Restoring base VAE") load_vae_dict(model, base_vae) delete_base_vae() @@ -143,6 +143,7 @@ def load_vae(model, vae_file=None): 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} load_vae_dict(model, vae_dict_1) + store_base_vae(model) # If vae used is not in dict, update it # It will be removed on refresh though @@ -150,6 +151,9 @@ def load_vae(model, vae_file=None): if vae_opt not in vae_dict: vae_dict[vae_opt] = vae_file vae_list.append(vae_opt) + # shared.opts.data['sd_vae'] = vae_opt + else: + restore_base_vae(model) loaded_vae_file = vae_file @@ -166,12 +170,8 @@ def load_vae(model, vae_file=None): # don't call this from outside -def load_vae_dict(model, vae_dict_1=None): - if vae_dict_1: - store_base_vae(model) - model.first_stage_model.load_state_dict(vae_dict_1) - else: - restore_base_vae() +def load_vae_dict(model, vae_dict_1): + model.first_stage_model.load_state_dict(vae_dict_1) model.first_stage_model.to(devices.dtype_vae)