From a5bbcd215304e0c83ab2b9fe7f172f88536d7629 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Sat, 14 Jan 2023 19:56:09 +0300
Subject: fix bug with "Ignore selected VAE for..." option completely disabling
 VAE election rework VAE resolving code to be more simple

---
 scripts/xy_grid.py | 27 +++++++++++++--------------
 1 file changed, 13 insertions(+), 14 deletions(-)

(limited to 'scripts/xy_grid.py')

diff --git a/scripts/xy_grid.py b/scripts/xy_grid.py
index f04d9b7e..bd3087d4 100644
--- a/scripts/xy_grid.py
+++ b/scripts/xy_grid.py
@@ -125,24 +125,21 @@ def apply_upscale_latent_space(p, x, xs):
 
 
 def find_vae(name: str):
-    if name.lower() in ['auto', 'none']:
-        return name
+    if name.lower() in ['auto', 'automatic']:
+        return modules.sd_vae.unspecified
+    if name.lower() == 'none':
+        return None
     else:
-        vae_path = os.path.abspath(os.path.join(paths.models_path, 'VAE'))
-        found = glob.glob(os.path.join(vae_path, f'**/{name}.*pt'), recursive=True)
-        if found:
-            return found[0]
+        choices = [x for x in sorted(modules.sd_vae.vae_dict, key=lambda x: len(x)) if name.lower().strip() in x.lower()]
+        if len(choices) == 0:
+            print(f"No VAE found for {name}; using automatic")
+            return modules.sd_vae.unspecified
         else:
-            return 'auto'
+            return modules.sd_vae.vae_dict[choices[0]]
 
 
 def apply_vae(p, x, xs):
-    if x.lower().strip() == 'none':
-        modules.sd_vae.reload_vae_weights(shared.sd_model, vae_file='None')
-    else:
-        found = find_vae(x)
-        if found:
-            v = modules.sd_vae.reload_vae_weights(shared.sd_model, vae_file=found)
+    modules.sd_vae.reload_vae_weights(shared.sd_model, vae_file=find_vae(x))
 
 
 def apply_styles(p: StableDiffusionProcessingTxt2Img, x: str, _):
@@ -271,7 +268,9 @@ class SharedSettingsStackHelper(object):
   
     def __exit__(self, exc_type, exc_value, tb):
         modules.sd_models.reload_model_weights(self.model)
-        modules.sd_vae.reload_vae_weights(self.model, vae_file=find_vae(self.vae))
+
+        opts.data["sd_vae"] = self.vae
+        modules.sd_vae.reload_vae_weights(self.model)
 
         hypernetwork.load_hypernetwork(self.hypernetwork)
         hypernetwork.apply_strength()
-- 
cgit v1.2.1