Remove no longer necessary parts and add vae_file safeguard

2 files changed, 3 insertions, 8 deletions

diff --git a/modules/sd_models.py b/modules/sd_models.py
index 80addf03..c59151e0 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -165,16 +165,9 @@ def load_model_weights(model, checkpoint_info, vae_file="auto"):
 

     cache_enabled = shared.opts.sd_checkpoint_cache > 0

 

-    if cache_enabled:

-        sd_vae.restore_base_vae(model)

-

-    vae_file = sd_vae.resolve_vae(checkpoint_file, vae_file=vae_file)

-

     if cache_enabled and checkpoint_info in checkpoints_loaded:

         # use checkpoint cache

-        vae_name = sd_vae.get_filename(vae_file) if vae_file else None

-        vae_message = f" with {vae_name} VAE" if vae_name else ""

-        print(f"Loading weights [{sd_model_hash}]{vae_message} from cache")

+        print(f"Loading weights [{sd_model_hash}] from cache")

         model.load_state_dict(checkpoints_loaded[checkpoint_info])

     else:

         # load from file

@@ -220,6 +213,7 @@ def load_model_weights(model, checkpoint_info, vae_file="auto"):
     model.sd_model_checkpoint = checkpoint_file

     model.sd_checkpoint_info = checkpoint_info

 

+    vae_file = sd_vae.resolve_vae(checkpoint_file, vae_file=vae_file)

     sd_vae.load_vae(model, vae_file)

 

 

diff --git a/modules/sd_vae.py b/modules/sd_vae.py
index 71e7a6e6..8bdb2c17 100644
--- a/modules/sd_vae.py
+++ b/modules/sd_vae.py
@@ -139,6 +139,7 @@ def load_vae(model, vae_file=None):
     # save_settings = False
 
     if vae_file:
+        assert os.path.isfile(vae_file), f"VAE file doesn't exist: {vae_file}"
         print(f"Loading VAE weights from: {vae_file}")
         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}