diff options
Diffstat (limited to 'modules/sd_vae.py')
-rw-r--r-- | modules/sd_vae.py | 16 |
1 files changed, 14 insertions, 2 deletions
diff --git a/modules/sd_vae.py b/modules/sd_vae.py index b2af2ce7..4ce238b8 100644 --- a/modules/sd_vae.py +++ b/modules/sd_vae.py @@ -72,6 +72,13 @@ def refresh_vae_list(): os.path.join(shared.cmd_opts.ckpt_dir, '**/*.vae.safetensors'), ] + if shared.cmd_opts.vae_dir is not None and os.path.isdir(shared.cmd_opts.vae_dir): + paths += [ + os.path.join(shared.cmd_opts.vae_dir, '**/*.ckpt'), + os.path.join(shared.cmd_opts.vae_dir, '**/*.pt'), + os.path.join(shared.cmd_opts.vae_dir, '**/*.safetensors'), + ] + candidates = [] for path in paths: candidates += glob.iglob(path, recursive=True) @@ -113,6 +120,12 @@ def resolve_vae(checkpoint_file): return None, None +def load_vae_dict(filename, map_location): + vae_ckpt = sd_models.read_state_dict(filename, map_location=map_location) + vae_dict_1 = {k: v for k, v in vae_ckpt.items() if k[0:4] != "loss" and k not in vae_ignore_keys} + return vae_dict_1 + + def load_vae(model, vae_file=None, vae_source="from unknown source"): global vae_dict, loaded_vae_file # save_settings = False @@ -130,8 +143,7 @@ def load_vae(model, vae_file=None, vae_source="from unknown source"): print(f"Loading VAE weights {vae_source}: {vae_file}") store_base_vae(model) - vae_ckpt = sd_models.read_state_dict(vae_file, map_location=shared.weight_load_location) - vae_dict_1 = {k: v for k, v in vae_ckpt.items() if k[0:4] != "loss" and k not in vae_ignore_keys} + vae_dict_1 = load_vae_dict(vae_file, map_location=shared.weight_load_location) _load_vae_dict(model, vae_dict_1) if cache_enabled: |