diff options
author | Micky Brunetti <micky2be@users.noreply.github.com> | 2023-05-09 13:35:57 +0000 |
---|---|---|
committer | GitHub <noreply@github.com> | 2023-05-09 13:35:57 +0000 |
commit | 7fd3a4e6d7b1c70461eed0c8a7dc4f2412cdaf1c (patch) | |
tree | 911225dea445e2c9e0af8575a510f7dfa9481b07 | |
parent | 5ab7f213bec2f816f9c5644becb32eb72c8ffb89 (diff) | |
download | stable-diffusion-webui-gfx803-7fd3a4e6d7b1c70461eed0c8a7dc4f2412cdaf1c.tar.gz stable-diffusion-webui-gfx803-7fd3a4e6d7b1c70461eed0c8a7dc4f2412cdaf1c.tar.bz2 stable-diffusion-webui-gfx803-7fd3a4e6d7b1c70461eed0c8a7dc4f2412cdaf1c.zip |
files in vae folder with same name as a checkpoint can be found too
-rw-r--r-- | modules/sd_vae.py | 11 |
1 files changed, 7 insertions, 4 deletions
diff --git a/modules/sd_vae.py b/modules/sd_vae.py index 9b00f76e..4d2026e1 100644 --- a/modules/sd_vae.py +++ b/modules/sd_vae.py @@ -88,10 +88,13 @@ def refresh_vae_list(): def find_vae_near_checkpoint(checkpoint_file): - checkpoint_path = os.path.splitext(checkpoint_file)[0] - for vae_location in [checkpoint_path + ".vae.pt", checkpoint_path + ".vae.ckpt", checkpoint_path + ".vae.safetensors"]: - if os.path.isfile(vae_location): - return vae_location + checkpoint_path = os.path.basename(checkpoint_file).split('.', 1)[0] + print(f"checkpoint: {checkpoint_path}") + for vae_file in vae_dict.values(): + vae_path = os.path.basename(vae_file).split('.', 1)[0] + print(f"vae: {vae_path}") + if vae_path == checkpoint_path: + return vae_file return None |