Merge branch 'dev' into text-drag-fix

1 files changed, 4 insertions, 7 deletions

diff --git a/modules/sd_vae.py b/modules/sd_vae.py
index 521e485a..e4ff2994 100644
--- a/modules/sd_vae.py
+++ b/modules/sd_vae.py
@@ -1,8 +1,5 @@
-import torch
-import safetensors.torch
 import os
 import collections
-from collections import namedtuple
 from modules import paths, shared, devices, script_callbacks, sd_models
 import glob
 from copy import deepcopy
@@ -88,10 +85,10 @@ def refresh_vae_list():
 
 
 def find_vae_near_checkpoint(checkpoint_file):
-    checkpoint_path = os.path.splitext(checkpoint_file)[0]
-    for vae_location in [f"{checkpoint_path}.vae.pt", f"{checkpoint_path}.vae.ckpt", f"{checkpoint_path}.vae.safetensors"]:
-        if os.path.isfile(vae_location):
-            return vae_location
+    checkpoint_path = os.path.basename(checkpoint_file).rsplit('.', 1)[0]
+    for vae_file in vae_dict.values():
+        if os.path.basename(vae_file).startswith(checkpoint_path):
+            return vae_file
 
     return None