From 4fbb7a88dd77f3fdcedcb8a15e3ffb86645e637c Mon Sep 17 00:00:00 2001
From: Gazzoo-byte <73721238+Gazzoo-byte@users.noreply.github.com>
Date: Tue, 25 Apr 2023 20:29:27 +0100
Subject: [PATCH] Change 'Automatic' behaviour

Modifies the behaviour of 'Automatic' to check in vae_path for vae files with a filename that matches the checkpoint, failing this falls back to None
---
 modules/sd_vae.py | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/modules/sd_vae.py b/modules/sd_vae.py
index 7f3634a99..042e147d7 100644
--- a/modules/sd_vae.py
+++ b/modules/sd_vae.py
@@ -5,6 +5,7 @@ from copy import deepcopy
 from rich import print # pylint: disable=redefined-builtin
 import torch
 from modules import paths, shared, devices, script_callbacks, sd_models
+from pathlib import Path
 
 
 vae_ignore_keys = {"model_ema.decay", "model_ema.num_updates"}
@@ -98,8 +99,17 @@ def resolve_vae(checkpoint_file):
     is_automatic = shared.opts.sd_vae in {"Automatic", "auto"}  # "auto" for people with old config
 
     vae_near_checkpoint = find_vae_near_checkpoint(checkpoint_file)
-    if vae_near_checkpoint is not None and (shared.opts.sd_vae_as_default or is_automatic):
+    if vae_near_checkpoint is not None and (shared.opts.sd_vae_as_default):
         return vae_near_checkpoint, 'found near the checkpoint'
+    
+    if is_automatic:
+        for named_vae_location in [vae_path + "\\" + Path(checkpoint_file).stem + ".vae.pt", vae_path + "\\" + Path(checkpoint_file).stem + ".vae.ckpt", vae_path + "\\" + Path(checkpoint_file).stem + ".vae.safetensors"]:
+            if os.path.isfile(named_vae_location):
+                print(named_vae_location+' found in VAE dir')
+                return named_vae_location
+            else:
+                print(f"Couldn't find a VAE with a matching name in {vae_path}, using None instead")
+                return None, None
 
     if shared.opts.sd_vae == "None":
         return None, None