From 7fd3a4e6d7b1c70461eed0c8a7dc4f2412cdaf1c Mon Sep 17 00:00:00 2001 From: Micky Brunetti Date: Tue, 9 May 2023 15:35:57 +0200 Subject: [PATCH] files in vae folder with same name as a checkpoint can be found too --- modules/sd_vae.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/modules/sd_vae.py b/modules/sd_vae.py index 9b00f76e9..4d2026e1c 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