From 893933e05ad267778111b4fad6d1ecb80937afdf Mon Sep 17 00:00:00 2001
From: hitomi <huahuohe@foxmail.com>
Date: Sun, 25 Dec 2022 20:49:25 +0800
Subject: Add memory cache for VAE weights

---
 modules/sd_vae.py | 31 +++++++++++++++++++++++++------
 1 file changed, 25 insertions(+), 6 deletions(-)

(limited to 'modules/sd_vae.py')

diff --git a/modules/sd_vae.py b/modules/sd_vae.py
index 3856418e..ac71d62d 100644
--- a/modules/sd_vae.py
+++ b/modules/sd_vae.py
@@ -1,5 +1,6 @@
 import torch
 import os
+import collections
 from collections import namedtuple
 from modules import shared, devices, script_callbacks
 from modules.paths import models_path
@@ -30,6 +31,7 @@ base_vae = None
 loaded_vae_file = None
 checkpoint_info = None
 
+checkpoints_loaded = collections.OrderedDict()
 
 def get_base_vae(model):
     if base_vae is not None and checkpoint_info == model.sd_checkpoint_info and model:
@@ -149,13 +151,30 @@ def load_vae(model, vae_file=None):
     global first_load, vae_dict, vae_list, loaded_vae_file
     # save_settings = False
 
+    cache_enabled = shared.opts.sd_vae_checkpoint_cache > 0
+
     if vae_file:
-        assert os.path.isfile(vae_file), f"VAE file doesn't exist: {vae_file}"
-        print(f"Loading VAE weights from: {vae_file}")
-        store_base_vae(model)
-        vae_ckpt = torch.load(vae_file, map_location=shared.weight_load_location)
-        vae_dict_1 = {k: v for k, v in vae_ckpt["state_dict"].items() if k[0:4] != "loss" and k not in vae_ignore_keys}
-        _load_vae_dict(model, vae_dict_1)
+        if cache_enabled and vae_file in checkpoints_loaded:
+            # use vae checkpoint cache
+            print(f"Loading VAE weights [{get_filename(vae_file)}] from cache")
+            store_base_vae(model)
+            _load_vae_dict(model, checkpoints_loaded[vae_file])
+        else:
+            assert os.path.isfile(vae_file), f"VAE file doesn't exist: {vae_file}"
+            print(f"Loading VAE weights from: {vae_file}")
+            store_base_vae(model)
+            vae_ckpt = torch.load(vae_file, map_location=shared.weight_load_location)
+            vae_dict_1 = {k: v for k, v in vae_ckpt["state_dict"].items() if k[0:4] != "loss" and k not in vae_ignore_keys}
+            _load_vae_dict(model, vae_dict_1)
+
+            if cache_enabled:
+                # cache newly loaded vae
+                checkpoints_loaded[vae_file] = vae_dict_1.copy()
+
+        # clean up cache if limit is reached
+        if cache_enabled:
+            while len(checkpoints_loaded) > shared.opts.sd_vae_checkpoint_cache + 1: # we need to count the current model
+                checkpoints_loaded.popitem(last=False)  # LRU
 
         # If vae used is not in dict, update it
         # It will be removed on refresh though
-- 
cgit v1.2.3


From cb255faec6e5f6b47b7632e6b7d450b9e2f6678b Mon Sep 17 00:00:00 2001
From: Lee Bousfield <ljbousfield@gmail.com>
Date: Sun, 8 Jan 2023 10:17:50 -0700
Subject: Add support for loading VAEs from safetensor files

---
 modules/sd_vae.py | 21 ++++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

(limited to 'modules/sd_vae.py')

diff --git a/modules/sd_vae.py b/modules/sd_vae.py
index ac71d62d..9fcfd9db 100644
--- a/modules/sd_vae.py
+++ b/modules/sd_vae.py
@@ -1,4 +1,5 @@
 import torch
+import safetensors.torch
 import os
 import collections
 from collections import namedtuple
@@ -72,8 +73,10 @@ def refresh_vae_list(vae_path=vae_path, model_path=model_path):
     candidates = [
         *glob.iglob(os.path.join(model_path, '**/*.vae.ckpt'), recursive=True),
         *glob.iglob(os.path.join(model_path, '**/*.vae.pt'), recursive=True),
+        *glob.iglob(os.path.join(model_path, '**/*.vae.safetensors'), recursive=True),
         *glob.iglob(os.path.join(vae_path, '**/*.ckpt'), recursive=True),
-        *glob.iglob(os.path.join(vae_path, '**/*.pt'), recursive=True)
+        *glob.iglob(os.path.join(vae_path, '**/*.pt'), recursive=True),
+        *glob.iglob(os.path.join(vae_path, '**/*.safetensors'), recursive=True),
     ]
     if shared.cmd_opts.vae_path is not None and os.path.isfile(shared.cmd_opts.vae_path):
         candidates.append(shared.cmd_opts.vae_path)
@@ -137,6 +140,12 @@ def resolve_vae(checkpoint_file=None, vae_file="auto"):
         if os.path.isfile(vae_file_try):
             vae_file = vae_file_try
             print(f"Using VAE found similar to selected model: {vae_file}")
+    # if still not found, try look for ".vae.safetensors" beside model
+    if vae_file == "auto":
+        vae_file_try = model_path + ".vae.safetensors"
+        if os.path.isfile(vae_file_try):
+            vae_file = vae_file_try
+            print(f"Using VAE found similar to selected model: {vae_file}")
     # No more fallbacks for auto
     if vae_file == "auto":
         vae_file = None
@@ -163,8 +172,14 @@ def load_vae(model, vae_file=None):
             assert os.path.isfile(vae_file), f"VAE file doesn't exist: {vae_file}"
             print(f"Loading VAE weights from: {vae_file}")
             store_base_vae(model)
-            vae_ckpt = torch.load(vae_file, map_location=shared.weight_load_location)
-            vae_dict_1 = {k: v for k, v in vae_ckpt["state_dict"].items() if k[0:4] != "loss" and k not in vae_ignore_keys}
+            _, extension = os.path.splitext(vae_file)
+            if extension.lower() == ".safetensors":
+                vae_ckpt = safetensors.torch.load_file(vae_file, device=shared.weight_load_location)
+            else:
+                vae_ckpt = torch.load(vae_file, map_location=shared.weight_load_location)
+            if "state_dict" in vae_ckpt:
+                vae_ckpt = vae_ckpt["state_dict"]
+            vae_dict_1 = {k: v for k, v in vae_ckpt.items() if k[0:4] != "loss" and k not in vae_ignore_keys}
             _load_vae_dict(model, vae_dict_1)
 
             if cache_enabled:
-- 
cgit v1.2.3


From 49c4509ce2302350210ff650fd26373518c46a79 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Mon, 9 Jan 2023 19:58:35 +0300
Subject: use existing function for loading VAE weights from file

---
 modules/sd_vae.py | 13 +++++--------
 1 file changed, 5 insertions(+), 8 deletions(-)

(limited to 'modules/sd_vae.py')

diff --git a/modules/sd_vae.py b/modules/sd_vae.py
index 9fcfd9db..0a49daa1 100644
--- a/modules/sd_vae.py
+++ b/modules/sd_vae.py
@@ -3,7 +3,7 @@ import safetensors.torch
 import os
 import collections
 from collections import namedtuple
-from modules import shared, devices, script_callbacks
+from modules import shared, devices, script_callbacks, sd_models
 from modules.paths import models_path
 import glob
 from copy import deepcopy
@@ -172,13 +172,8 @@ def load_vae(model, vae_file=None):
             assert os.path.isfile(vae_file), f"VAE file doesn't exist: {vae_file}"
             print(f"Loading VAE weights from: {vae_file}")
             store_base_vae(model)
-            _, extension = os.path.splitext(vae_file)
-            if extension.lower() == ".safetensors":
-                vae_ckpt = safetensors.torch.load_file(vae_file, device=shared.weight_load_location)
-            else:
-                vae_ckpt = torch.load(vae_file, map_location=shared.weight_load_location)
-            if "state_dict" in vae_ckpt:
-                vae_ckpt = vae_ckpt["state_dict"]
+
+            vae_ckpt = sd_models.read_state_dict(vae_file, map_location=shared.weight_load_location)
             vae_dict_1 = {k: v for k, v in vae_ckpt.items() if k[0:4] != "loss" and k not in vae_ignore_keys}
             _load_vae_dict(model, vae_dict_1)
 
@@ -210,10 +205,12 @@ def _load_vae_dict(model, vae_dict_1):
     model.first_stage_model.load_state_dict(vae_dict_1)
     model.first_stage_model.to(devices.dtype_vae)
 
+
 def clear_loaded_vae():
     global loaded_vae_file
     loaded_vae_file = None
 
+
 def reload_vae_weights(sd_model=None, vae_file="auto"):
     from modules import lowvram, devices, sd_hijack
 
-- 
cgit v1.2.3


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

---
 modules/sd_models.py |   6 +-
 modules/sd_vae.py    | 194 ++++++++++++++++++++-------------------------------
 modules/shared.py    |   4 +-
 scripts/xy_grid.py   |  27 ++++---
 4 files changed, 95 insertions(+), 136 deletions(-)

(limited to 'modules/sd_vae.py')

diff --git a/modules/sd_models.py b/modules/sd_models.py
index e5a0bc63..6a681cef 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -224,7 +224,7 @@ def read_state_dict(checkpoint_file, print_global_state=False, map_location=None
     return sd
 
 
-def load_model_weights(model, checkpoint_info: CheckpointInfo, vae_file="auto"):
+def load_model_weights(model, checkpoint_info: CheckpointInfo):
     sd_model_hash = checkpoint_info.calculate_shorthash()
 
     cache_enabled = shared.opts.sd_checkpoint_cache > 0
@@ -277,8 +277,8 @@ def load_model_weights(model, checkpoint_info: CheckpointInfo, vae_file="auto"):
 
     sd_vae.delete_base_vae()
     sd_vae.clear_loaded_vae()
-    vae_file = sd_vae.resolve_vae(checkpoint_info.filename, vae_file=vae_file)
-    sd_vae.load_vae(model, vae_file)
+    vae_file, vae_source = sd_vae.resolve_vae(checkpoint_info.filename)
+    sd_vae.load_vae(model, vae_file, vae_source)
 
 
 def enable_midas_autodownload():
diff --git a/modules/sd_vae.py b/modules/sd_vae.py
index 0a49daa1..6ea92711 100644
--- a/modules/sd_vae.py
+++ b/modules/sd_vae.py
@@ -9,23 +9,9 @@ import glob
 from copy import deepcopy
 
 
-model_dir = "Stable-diffusion"
-model_path = os.path.abspath(os.path.join(models_path, model_dir))
-vae_dir = "VAE"
-vae_path = os.path.abspath(os.path.join(models_path, vae_dir))
-
-
+vae_path = os.path.abspath(os.path.join(models_path, "VAE"))
 vae_ignore_keys = {"model_ema.decay", "model_ema.num_updates"}
-
-
-default_vae_dict = {"auto": "auto", "None": None, None: None}
-default_vae_list = ["auto", "None"]
-
-
-default_vae_values = [default_vae_dict[x] for x in default_vae_list]
-vae_dict = dict(default_vae_dict)
-vae_list = list(default_vae_list)
-first_load = True
+vae_dict = {}
 
 
 base_vae = None
@@ -64,100 +50,69 @@ def restore_base_vae(model):
 
 
 def get_filename(filepath):
-    return os.path.splitext(os.path.basename(filepath))[0]
-
-
-def refresh_vae_list(vae_path=vae_path, model_path=model_path):
-    global vae_dict, vae_list
-    res = {}
-    candidates = [
-        *glob.iglob(os.path.join(model_path, '**/*.vae.ckpt'), recursive=True),
-        *glob.iglob(os.path.join(model_path, '**/*.vae.pt'), recursive=True),
-        *glob.iglob(os.path.join(model_path, '**/*.vae.safetensors'), recursive=True),
-        *glob.iglob(os.path.join(vae_path, '**/*.ckpt'), recursive=True),
-        *glob.iglob(os.path.join(vae_path, '**/*.pt'), recursive=True),
-        *glob.iglob(os.path.join(vae_path, '**/*.safetensors'), recursive=True),
+    return os.path.basename(filepath)
+
+
+def refresh_vae_list():
+    vae_dict.clear()
+
+    paths = [
+        os.path.join(sd_models.model_path, '**/*.vae.ckpt'),
+        os.path.join(sd_models.model_path, '**/*.vae.pt'),
+        os.path.join(sd_models.model_path, '**/*.vae.safetensors'),
+        os.path.join(vae_path, '**/*.ckpt'),
+        os.path.join(vae_path, '**/*.pt'),
+        os.path.join(vae_path, '**/*.safetensors'),
     ]
-    if shared.cmd_opts.vae_path is not None and os.path.isfile(shared.cmd_opts.vae_path):
-        candidates.append(shared.cmd_opts.vae_path)
+
+    if shared.cmd_opts.ckpt_dir is not None and os.path.isdir(shared.cmd_opts.ckpt_dir):
+        paths += [
+            os.path.join(shared.cmd_opts.ckpt_dir, '**/*.vae.ckpt'),
+            os.path.join(shared.cmd_opts.ckpt_dir, '**/*.vae.pt'),
+            os.path.join(shared.cmd_opts.ckpt_dir, '**/*.vae.safetensors'),
+        ]
+
+    candidates = []
+    for path in paths:
+        candidates += glob.iglob(path, recursive=True)
+
     for filepath in candidates:
         name = get_filename(filepath)
-        res[name] = filepath
-    vae_list.clear()
-    vae_list.extend(default_vae_list)
-    vae_list.extend(list(res.keys()))
-    vae_dict.clear()
-    vae_dict.update(res)
-    vae_dict.update(default_vae_dict)
-    return vae_list
-
-
-def get_vae_from_settings(vae_file="auto"):
-    # else, we load from settings, if not set to be default
-    if vae_file == "auto" and shared.opts.sd_vae is not None:
-        # if saved VAE settings isn't recognized, fallback to auto
-        vae_file = vae_dict.get(shared.opts.sd_vae, "auto")
-        # if VAE selected but not found, fallback to auto
-        if vae_file not in default_vae_values and not os.path.isfile(vae_file):
-            vae_file = "auto"
-            print(f"Selected VAE doesn't exist: {vae_file}")
-    return vae_file
-
-
-def resolve_vae(checkpoint_file=None, vae_file="auto"):
-    global first_load, vae_dict, vae_list
-
-    # if vae_file argument is provided, it takes priority, but not saved
-    if vae_file and vae_file not in default_vae_list:
-        if not os.path.isfile(vae_file):
-            print(f"VAE provided as function argument doesn't exist: {vae_file}")
-            vae_file = "auto"
-    # for the first load, if vae-path is provided, it takes priority, saved, and failure is reported
-    if first_load and shared.cmd_opts.vae_path is not None:
-        if os.path.isfile(shared.cmd_opts.vae_path):
-            vae_file = shared.cmd_opts.vae_path
-            shared.opts.data['sd_vae'] = get_filename(vae_file)
-        else:
-            print(f"VAE provided as command line argument doesn't exist: {vae_file}")
-    # fallback to selector in settings, if vae selector not set to act as default fallback
-    if not shared.opts.sd_vae_as_default:
-        vae_file = get_vae_from_settings(vae_file)
-    # vae-path cmd arg takes priority for auto
-    if vae_file == "auto" and shared.cmd_opts.vae_path is not None:
-        if os.path.isfile(shared.cmd_opts.vae_path):
-            vae_file = shared.cmd_opts.vae_path
-            print(f"Using VAE provided as command line argument: {vae_file}")
-    # if still not found, try look for ".vae.pt" beside model
-    model_path = os.path.splitext(checkpoint_file)[0]
-    if vae_file == "auto":
-        vae_file_try = model_path + ".vae.pt"
-        if os.path.isfile(vae_file_try):
-            vae_file = vae_file_try
-            print(f"Using VAE found similar to selected model: {vae_file}")
-    # if still not found, try look for ".vae.ckpt" beside model
-    if vae_file == "auto":
-        vae_file_try = model_path + ".vae.ckpt"
-        if os.path.isfile(vae_file_try):
-            vae_file = vae_file_try
-            print(f"Using VAE found similar to selected model: {vae_file}")
-    # if still not found, try look for ".vae.safetensors" beside model
-    if vae_file == "auto":
-        vae_file_try = model_path + ".vae.safetensors"
-        if os.path.isfile(vae_file_try):
-            vae_file = vae_file_try
-            print(f"Using VAE found similar to selected model: {vae_file}")
-    # No more fallbacks for auto
-    if vae_file == "auto":
-        vae_file = None
-    # Last check, just because
-    if vae_file and not os.path.exists(vae_file):
-        vae_file = None
-
-    return vae_file
-
-
-def load_vae(model, vae_file=None):
-    global first_load, vae_dict, vae_list, loaded_vae_file
+        vae_dict[name] = filepath
+
+
+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
+
+    return None
+
+
+def resolve_vae(checkpoint_file):
+    if shared.cmd_opts.vae_path is not None:
+        return shared.cmd_opts.vae_path, 'from commandline argument'
+
+    vae_near_checkpoint = find_vae_near_checkpoint(checkpoint_file)
+    if vae_near_checkpoint is not None and (shared.opts.sd_vae_as_default or shared.opts.sd_vae == "auto"):
+        return vae_near_checkpoint, 'found near the checkpoint'
+
+    if shared.opts.sd_vae == "None":
+        return None, None
+
+    vae_from_options = vae_dict.get(shared.opts.sd_vae, None)
+    if vae_from_options is not None:
+        return vae_from_options, 'specified in settings'
+
+    if shared.opts.sd_vae != "Automatic":
+        print(f"Couldn't find VAE named {shared.opts.sd_vae}; using None instead")
+
+    return None, None
+
+
+def load_vae(model, vae_file=None, vae_source="from unknown source"):
+    global vae_dict, loaded_vae_file
     # save_settings = False
 
     cache_enabled = shared.opts.sd_vae_checkpoint_cache > 0
@@ -165,12 +120,12 @@ def load_vae(model, vae_file=None):
     if vae_file:
         if cache_enabled and vae_file in checkpoints_loaded:
             # use vae checkpoint cache
-            print(f"Loading VAE weights [{get_filename(vae_file)}] from cache")
+            print(f"Loading VAE weights {vae_source}: cached {get_filename(vae_file)}")
             store_base_vae(model)
             _load_vae_dict(model, checkpoints_loaded[vae_file])
         else:
-            assert os.path.isfile(vae_file), f"VAE file doesn't exist: {vae_file}"
-            print(f"Loading VAE weights from: {vae_file}")
+            assert os.path.isfile(vae_file), f"VAE {vae_source} doesn't exist: {vae_file}"
+            print(f"Loading VAE weights {vae_source}: {vae_file}")
             store_base_vae(model)
 
             vae_ckpt = sd_models.read_state_dict(vae_file, map_location=shared.weight_load_location)
@@ -191,14 +146,12 @@ def load_vae(model, vae_file=None):
         vae_opt = get_filename(vae_file)
         if vae_opt not in vae_dict:
             vae_dict[vae_opt] = vae_file
-            vae_list.append(vae_opt)
+
     elif loaded_vae_file:
         restore_base_vae(model)
 
     loaded_vae_file = vae_file
 
-    first_load = False
-
 
 # don't call this from outside
 def _load_vae_dict(model, vae_dict_1):
@@ -211,7 +164,10 @@ def clear_loaded_vae():
     loaded_vae_file = None
 
 
-def reload_vae_weights(sd_model=None, vae_file="auto"):
+unspecified = object()
+
+
+def reload_vae_weights(sd_model=None, vae_file=unspecified):
     from modules import lowvram, devices, sd_hijack
 
     if not sd_model:
@@ -219,7 +175,11 @@ def reload_vae_weights(sd_model=None, vae_file="auto"):
 
     checkpoint_info = sd_model.sd_checkpoint_info
     checkpoint_file = checkpoint_info.filename
-    vae_file = resolve_vae(checkpoint_file, vae_file=vae_file)
+
+    if vae_file == unspecified:
+        vae_file, vae_source = resolve_vae(checkpoint_file)
+    else:
+        vae_source = "from function argument"
 
     if loaded_vae_file == vae_file:
         return
@@ -231,7 +191,7 @@ def reload_vae_weights(sd_model=None, vae_file="auto"):
 
     sd_hijack.model_hijack.undo_hijack(sd_model)
 
-    load_vae(sd_model, vae_file)
+    load_vae(sd_model, vae_file, vae_source)
 
     sd_hijack.model_hijack.hijack(sd_model)
     script_callbacks.model_loaded_callback(sd_model)
@@ -239,5 +199,5 @@ def reload_vae_weights(sd_model=None, vae_file="auto"):
     if not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram:
         sd_model.to(devices.device)
 
-    print("VAE Weights loaded.")
+    print("VAE weights loaded.")
     return sd_model
diff --git a/modules/shared.py b/modules/shared.py
index e0ec3136..9756adea 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -83,7 +83,7 @@ parser.add_argument("--theme", type=str, help="launches the UI with light or dar
 parser.add_argument("--use-textbox-seed", action='store_true', help="use textbox for seeds in UI (no up/down, but possible to input long seeds)", default=False)
 parser.add_argument("--disable-console-progressbars", action='store_true', help="do not output progressbars to console", default=False)
 parser.add_argument("--enable-console-prompts", action='store_true', help="print prompts to console when generating with txt2img and img2img", default=False)
-parser.add_argument('--vae-path', type=str, help='Path to Variational Autoencoders model', default=None)
+parser.add_argument('--vae-path', type=str, help='Checkpoint to use as VAE; setting this argument disables all settings related to VAE', default=None)
 parser.add_argument("--disable-safe-unpickle", action='store_true', help="disable checking pytorch models for malicious code", default=False)
 parser.add_argument("--api", action='store_true', help="use api=True to launch the API together with the webui (use --nowebui instead for only the API)")
 parser.add_argument("--api-auth", type=str, help='Set authentication for API like "username:password"; or comma-delimit multiple like "u1:p1,u2:p2,u3:p3"', default=None)
@@ -383,7 +383,7 @@ options_templates.update(options_section(('sd', "Stable Diffusion"), {
     "sd_model_checkpoint": OptionInfo(None, "Stable Diffusion checkpoint", gr.Dropdown, lambda: {"choices": list_checkpoint_tiles()}, refresh=refresh_checkpoints),
     "sd_checkpoint_cache": OptionInfo(0, "Checkpoints to cache in RAM", gr.Slider, {"minimum": 0, "maximum": 10, "step": 1}),
     "sd_vae_checkpoint_cache": OptionInfo(0, "VAE Checkpoints to cache in RAM", gr.Slider, {"minimum": 0, "maximum": 10, "step": 1}),
-    "sd_vae": OptionInfo("auto", "SD VAE", gr.Dropdown, lambda: {"choices": sd_vae.vae_list}, refresh=sd_vae.refresh_vae_list),
+    "sd_vae": OptionInfo("Automatic", "SD VAE", gr.Dropdown, lambda: {"choices": ["Automatic", "None"] + list(sd_vae.vae_dict)}, refresh=sd_vae.refresh_vae_list),
     "sd_vae_as_default": OptionInfo(False, "Ignore selected VAE for stable diffusion checkpoints that have their own .vae.pt next to them"),
     "sd_hypernetwork": OptionInfo("None", "Hypernetwork", gr.Dropdown, lambda: {"choices": ["None"] + [x for x in hypernetworks.keys()]}, refresh=reload_hypernetworks),
     "sd_hypernetwork_strength": OptionInfo(1.0, "Hypernetwork strength", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.001}),
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.3


From f8c512478568293155539f616dce26c5e4495055 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Sat, 14 Jan 2023 20:00:12 +0300
Subject: typo?

---
 modules/sd_vae.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

(limited to 'modules/sd_vae.py')

diff --git a/modules/sd_vae.py b/modules/sd_vae.py
index 6ea92711..add5cecf 100644
--- a/modules/sd_vae.py
+++ b/modules/sd_vae.py
@@ -95,7 +95,7 @@ def resolve_vae(checkpoint_file):
         return shared.cmd_opts.vae_path, 'from commandline argument'
 
     vae_near_checkpoint = find_vae_near_checkpoint(checkpoint_file)
-    if vae_near_checkpoint is not None and (shared.opts.sd_vae_as_default or shared.opts.sd_vae == "auto"):
+    if vae_near_checkpoint is not None and (shared.opts.sd_vae_as_default or shared.opts.sd_vae == "Automatic"):
         return vae_near_checkpoint, 'found near the checkpoint'
 
     if shared.opts.sd_vae == "None":
-- 
cgit v1.2.3


From 3f887f7f61d69fa699a272166b79fdb787e9ce1d Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Mon, 16 Jan 2023 00:44:46 +0300
Subject: support old configs that say "auto" for ssd_vae change
 sd_vae_as_default to True by default as it's a more sensible setting

---
 modules/sd_vae.py | 6 ++++--
 modules/shared.py | 2 +-
 2 files changed, 5 insertions(+), 3 deletions(-)

(limited to 'modules/sd_vae.py')

diff --git a/modules/sd_vae.py b/modules/sd_vae.py
index add5cecf..e9c6bb40 100644
--- a/modules/sd_vae.py
+++ b/modules/sd_vae.py
@@ -94,8 +94,10 @@ def resolve_vae(checkpoint_file):
     if shared.cmd_opts.vae_path is not None:
         return shared.cmd_opts.vae_path, 'from commandline argument'
 
+    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 shared.opts.sd_vae == "Automatic"):
+    if vae_near_checkpoint is not None and (shared.opts.sd_vae_as_default or is_automatic):
         return vae_near_checkpoint, 'found near the checkpoint'
 
     if shared.opts.sd_vae == "None":
@@ -105,7 +107,7 @@ def resolve_vae(checkpoint_file):
     if vae_from_options is not None:
         return vae_from_options, 'specified in settings'
 
-    if shared.opts.sd_vae != "Automatic":
+    if is_automatic:
         print(f"Couldn't find VAE named {shared.opts.sd_vae}; using None instead")
 
     return None, None
diff --git a/modules/shared.py b/modules/shared.py
index f06ae610..c5fc250e 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -394,7 +394,7 @@ options_templates.update(options_section(('sd', "Stable Diffusion"), {
     "sd_checkpoint_cache": OptionInfo(0, "Checkpoints to cache in RAM", gr.Slider, {"minimum": 0, "maximum": 10, "step": 1}),
     "sd_vae_checkpoint_cache": OptionInfo(0, "VAE Checkpoints to cache in RAM", gr.Slider, {"minimum": 0, "maximum": 10, "step": 1}),
     "sd_vae": OptionInfo("Automatic", "SD VAE", gr.Dropdown, lambda: {"choices": ["Automatic", "None"] + list(sd_vae.vae_dict)}, refresh=sd_vae.refresh_vae_list),
-    "sd_vae_as_default": OptionInfo(False, "Ignore selected VAE for stable diffusion checkpoints that have their own .vae.pt next to them"),
+    "sd_vae_as_default": OptionInfo(True, "Ignore selected VAE for stable diffusion checkpoints that have their own .vae.pt next to them"),
     "sd_hypernetwork": OptionInfo("None", "Hypernetwork", gr.Dropdown, lambda: {"choices": ["None"] + [x for x in hypernetworks.keys()]}, refresh=reload_hypernetworks),
     "sd_hypernetwork_strength": OptionInfo(1.0, "Hypernetwork strength", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.001}),
     "inpainting_mask_weight": OptionInfo(1.0, "Inpainting conditioning mask strength", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),
-- 
cgit v1.2.3


From ff6a5bcec1ce25aa8f08b157ea957d764be23d8d Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Mon, 16 Jan 2023 01:28:20 +0300
Subject: bugfix for previous commit

---
 modules/sd_vae.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

(limited to 'modules/sd_vae.py')

diff --git a/modules/sd_vae.py b/modules/sd_vae.py
index e9c6bb40..b2af2ce7 100644
--- a/modules/sd_vae.py
+++ b/modules/sd_vae.py
@@ -107,7 +107,7 @@ def resolve_vae(checkpoint_file):
     if vae_from_options is not None:
         return vae_from_options, 'specified in settings'
 
-    if is_automatic:
+    if not is_automatic:
         print(f"Couldn't find VAE named {shared.opts.sd_vae}; using None instead")
 
     return None, None
-- 
cgit v1.2.3


From 6e08da2c315c346225aa834017f4e32cfc0de200 Mon Sep 17 00:00:00 2001
From: ddPn08 <pyn.g.1125@gmail.com>
Date: Tue, 17 Jan 2023 23:50:41 +0900
Subject: Add `--vae-dir` argument

---
 modules/sd_vae.py | 7 +++++++
 modules/shared.py | 1 +
 2 files changed, 8 insertions(+)

(limited to 'modules/sd_vae.py')

diff --git a/modules/sd_vae.py b/modules/sd_vae.py
index b2af2ce7..da1bf15c 100644
--- a/modules/sd_vae.py
+++ b/modules/sd_vae.py
@@ -72,6 +72,13 @@ def refresh_vae_list():
             os.path.join(shared.cmd_opts.ckpt_dir, '**/*.vae.safetensors'),
         ]
 
+    if shared.cmd_opts.vae_dir is not None and os.path.isdir(shared.cmd_opts.vae_dir):
+        paths += [
+            os.path.join(shared.cmd_opts.vae_dir, '**/*.ckpt'),
+            os.path.join(shared.cmd_opts.vae_dir, '**/*.pt'),
+            os.path.join(shared.cmd_opts.vae_dir, '**/*.safetensors'),
+        ]
+
     candidates = []
     for path in paths:
         candidates += glob.iglob(path, recursive=True)
diff --git a/modules/shared.py b/modules/shared.py
index a708f23c..a1345ad3 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -26,6 +26,7 @@ parser = argparse.ArgumentParser()
 parser.add_argument("--config", type=str, default=os.path.join(script_path, "configs/v1-inference.yaml"), help="path to config which constructs model",)
 parser.add_argument("--ckpt", type=str, default=sd_model_file, help="path to checkpoint of stable diffusion model; if specified, this checkpoint will be added to the list of checkpoints and loaded",)
 parser.add_argument("--ckpt-dir", type=str, default=None, help="Path to directory with stable diffusion checkpoints")
+parser.add_argument("--vae-dir", type=str, default=None, help="Path to directory with stable VAE files")
 parser.add_argument("--gfpgan-dir", type=str, help="GFPGAN directory", default=('./src/gfpgan' if os.path.exists('./src/gfpgan') else './GFPGAN'))
 parser.add_argument("--gfpgan-model", type=str, help="GFPGAN model file name", default=None)
 parser.add_argument("--no-half", action='store_true', help="do not switch the model to 16-bit floats")
-- 
cgit v1.2.3


From 0f5dbfffd0b7202a48e404d8e74b5cc9a3e5b135 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Thu, 19 Jan 2023 10:39:51 +0300
Subject: allow baking in VAE in checkpoint merger tab do not save config if
 it's the default for checkpoint merger tab change file naming scheme for
 checkpoint merger tab allow just saving A without any merging for checkpoint
 merger tab some stylistic changes for UI in checkpoint merger tab

---
 javascript/hints.js |   1 +
 javascript/ui.js    |   2 -
 modules/extras.py   | 112 +++++++++++++++++++++++++++++++---------------------
 modules/sd_vae.py   |   9 ++++-
 modules/shared.py   |   3 +-
 modules/ui.py       |  17 ++++++--
 style.css           |  15 ++++---
 7 files changed, 101 insertions(+), 58 deletions(-)

(limited to 'modules/sd_vae.py')

diff --git a/javascript/hints.js b/javascript/hints.js
index fa5e5ae8..e746e20d 100644
--- a/javascript/hints.js
+++ b/javascript/hints.js
@@ -92,6 +92,7 @@ titles = {
 
     "Weighted sum": "Result = A * (1 - M) + B * M",
     "Add difference": "Result = A + (B - C) * M",
+    "No interpolation": "Result = A",
 
 	"Initialization text": "If the number of tokens is more than the number of vectors, some may be skipped.\nLeave the textbox empty to start with zeroed out vectors",
     "Learning rate": "How fast should training go. Low values will take longer to train, high values may fail to converge (not generate accurate results) and/or may break the embedding (This has happened if you see Loss: nan in the training info textbox. If this happens, you need to manually restore your embedding from an older not-broken backup).\n\nYou can set a single numeric value, or multiple learning rates using the syntax:\n\n   rate_1:max_steps_1, rate_2:max_steps_2, ...\n\nEG:   0.005:100, 1e-3:1000, 1e-5\n\nWill train with rate of 0.005 for first 100 steps, then 1e-3 until 1000 steps, then 1e-5 for all remaining steps.",
diff --git a/javascript/ui.js b/javascript/ui.js
index 428375d4..37788a3e 100644
--- a/javascript/ui.js
+++ b/javascript/ui.js
@@ -176,8 +176,6 @@ function modelmerger(){
     var id = randomId()
     requestProgress(id, gradioApp().getElementById('modelmerger_results_panel'), null, function(){})
 
-    gradioApp().getElementById('modelmerger_result').innerHTML = ''
-
     var res = create_submit_args(arguments)
     res[0] = id
     return res
diff --git a/modules/extras.py b/modules/extras.py
index 034f28e4..fe701a0e 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -15,7 +15,7 @@ from typing import Callable, List, OrderedDict, Tuple
 from functools import partial
 from dataclasses import dataclass
 
-from modules import processing, shared, images, devices, sd_models, sd_samplers
+from modules import processing, shared, images, devices, sd_models, sd_samplers, sd_vae
 from modules.shared import opts
 import modules.gfpgan_model
 from modules.ui import plaintext_to_html
@@ -251,7 +251,8 @@ def run_pnginfo(image):
 
 def create_config(ckpt_result, config_source, a, b, c):
     def config(x):
-        return sd_models.find_checkpoint_config(x) if x else None
+        res = sd_models.find_checkpoint_config(x) if x else None
+        return res if res != shared.sd_default_config else None
 
     if config_source == 0:
         cfg = config(a) or config(b) or config(c)
@@ -274,10 +275,12 @@ def create_config(ckpt_result, config_source, a, b, c):
     shutil.copyfile(cfg, checkpoint_filename)
 
 
-def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_model_name, interp_method, multiplier, save_as_half, custom_name, checkpoint_format, config_source):
+chckpoint_dict_skip_on_merge = ["cond_stage_model.transformer.text_model.embeddings.position_ids"]
+
+
+def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_model_name, interp_method, multiplier, save_as_half, custom_name, checkpoint_format, config_source, bake_in_vae):
     shared.state.begin()
     shared.state.job = 'model-merge'
-    shared.state.job_count = 1
 
     def fail(message):
         shared.state.textinfo = message
@@ -293,41 +296,68 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
     def add_difference(theta0, theta1_2_diff, alpha):
         return theta0 + (alpha * theta1_2_diff)
 
+    def filename_weighed_sum():
+        a = primary_model_info.model_name
+        b = secondary_model_info.model_name
+        Ma = round(1 - multiplier, 2)
+        Mb = round(multiplier, 2)
+
+        return f"{Ma}({a}) + {Mb}({b})"
+
+    def filename_add_differnece():
+        a = primary_model_info.model_name
+        b = secondary_model_info.model_name
+        c = tertiary_model_info.model_name
+        M = round(multiplier, 2)
+
+        return f"{a} + {M}({b} - {c})"
+
+    def filename_nothing():
+        return primary_model_info.model_name
+
+    theta_funcs = {
+        "Weighted sum": (filename_weighed_sum, None, weighted_sum),
+        "Add difference": (filename_add_differnece, get_difference, add_difference),
+        "No interpolation": (filename_nothing, None, None),
+    }
+    filename_generator, theta_func1, theta_func2 = theta_funcs[interp_method]
+    shared.state.job_count = (1 if theta_func1 else 0) + (1 if theta_func2 else 0)
+
     if not primary_model_name:
         return fail("Failed: Merging requires a primary model.")
 
     primary_model_info = sd_models.checkpoints_list[primary_model_name]
 
-    if not secondary_model_name:
+    if theta_func2 and not secondary_model_name:
         return fail("Failed: Merging requires a secondary model.")
-    
-    secondary_model_info = sd_models.checkpoints_list[secondary_model_name]
 
-    theta_funcs = {
-        "Weighted sum": (None, weighted_sum),
-        "Add difference": (get_difference, add_difference),
-    }
-    theta_func1, theta_func2 = theta_funcs[interp_method]
+    secondary_model_info = sd_models.checkpoints_list[secondary_model_name] if theta_func2 else None
 
     if theta_func1 and not tertiary_model_name:
         return fail(f"Failed: Interpolation method ({interp_method}) requires a tertiary model.")
-    
+
     tertiary_model_info = sd_models.checkpoints_list[tertiary_model_name] if theta_func1 else None
 
     result_is_inpainting_model = False
 
-    shared.state.textinfo = f"Loading {secondary_model_info.filename}..."
-    print(f"Loading {secondary_model_info.filename}...")
-    theta_1 = sd_models.read_state_dict(secondary_model_info.filename, map_location='cpu')
+    if theta_func2:
+        shared.state.textinfo = f"Loading B"
+        print(f"Loading {secondary_model_info.filename}...")
+        theta_1 = sd_models.read_state_dict(secondary_model_info.filename, map_location='cpu')
+    else:
+        theta_1 = None
 
     if theta_func1:
-        shared.state.job_count += 1
-
+        shared.state.textinfo = f"Loading C"
         print(f"Loading {tertiary_model_info.filename}...")
         theta_2 = sd_models.read_state_dict(tertiary_model_info.filename, map_location='cpu')
 
+        shared.state.textinfo = 'Merging B and C'
         shared.state.sampling_steps = len(theta_1.keys())
         for key in tqdm.tqdm(theta_1.keys()):
+            if key in chckpoint_dict_skip_on_merge:
+                continue
+
             if 'model' in key:
                 if key in theta_2:
                     t2 = theta_2.get(key, torch.zeros_like(theta_1[key]))
@@ -345,12 +375,10 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
     theta_0 = sd_models.read_state_dict(primary_model_info.filename, map_location='cpu')
 
     print("Merging...")
-
-    chckpoint_dict_skip_on_merge = ["cond_stage_model.transformer.text_model.embeddings.position_ids"]
-
+    shared.state.textinfo = 'Merging A and B'
     shared.state.sampling_steps = len(theta_0.keys())
     for key in tqdm.tqdm(theta_0.keys()):
-        if 'model' in key and key in theta_1:
+        if theta_1 and 'model' in key and key in theta_1:
 
             if key in chckpoint_dict_skip_on_merge:
                 continue
@@ -358,7 +386,6 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
             a = theta_0[key]
             b = theta_1[key]
 
-            shared.state.textinfo = f'Merging layer {key}'
             # this enables merging an inpainting model (A) with another one (B);
             # where normal model would have 4 channels, for latenst space, inpainting model would
             # have another 4 channels for unmasked picture's latent space, plus one channel for mask, for a total of 9
@@ -378,34 +405,31 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
 
         shared.state.sampling_step += 1
 
-    # I believe this part should be discarded, but I'll leave it for now until I am sure
-    for key in theta_1.keys():
-        if 'model' in key and key not in theta_0:
+    del theta_1
+
+    bake_in_vae_filename = sd_vae.vae_dict.get(bake_in_vae, None)
+    if bake_in_vae_filename is not None:
+        print(f"Baking in VAE from {bake_in_vae_filename}")
+        shared.state.textinfo = 'Baking in VAE'
+        vae_dict = sd_vae.load_vae_dict(bake_in_vae_filename, map_location='cpu')
 
-            if key in chckpoint_dict_skip_on_merge:
-                continue
+        for key in vae_dict.keys():
+            theta_0_key = 'first_stage_model.' + key
+            if theta_0_key in theta_0:
+                theta_0[theta_0_key] = vae_dict[key].half() if save_as_half else vae_dict[key]
 
-            theta_0[key] = theta_1[key]
-            if save_as_half:
-                theta_0[key] = theta_0[key].half()
-    del theta_1
+        del vae_dict
 
     ckpt_dir = shared.cmd_opts.ckpt_dir or sd_models.model_path
 
-    filename = \
-        primary_model_info.model_name + '_' + str(round(1-multiplier, 2)) + '-' + \
-        secondary_model_info.model_name + '_' + str(round(multiplier, 2)) + '-' + \
-        interp_method.replace(" ", "_") + \
-        '-merged.' +  \
-        ("inpainting." if result_is_inpainting_model else "") + \
-        checkpoint_format
-
-    filename = filename if custom_name == '' else (custom_name + '.' + checkpoint_format)
+    filename = filename_generator() if custom_name == '' else custom_name
+    filename += ".inpainting" if result_is_inpainting_model else ""
+    filename += "." + checkpoint_format
 
     output_modelname = os.path.join(ckpt_dir, filename)
 
     shared.state.nextjob()
-    shared.state.textinfo = f"Saving to {output_modelname}..."
+    shared.state.textinfo = "Saving"
     print(f"Saving to {output_modelname}...")
 
     _, extension = os.path.splitext(output_modelname)
@@ -418,8 +442,8 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
 
     create_config(output_modelname, config_source, primary_model_info, secondary_model_info, tertiary_model_info)
 
-    print("Checkpoint saved.")
-    shared.state.textinfo = "Checkpoint saved to " + output_modelname
+    print(f"Checkpoint saved to {output_modelname}.")
+    shared.state.textinfo = "Checkpoint saved"
     shared.state.end()
 
     return [*[gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)], "Checkpoint saved to " + output_modelname]
diff --git a/modules/sd_vae.py b/modules/sd_vae.py
index da1bf15c..4ce238b8 100644
--- a/modules/sd_vae.py
+++ b/modules/sd_vae.py
@@ -120,6 +120,12 @@ def resolve_vae(checkpoint_file):
     return None, None
 
 
+def load_vae_dict(filename, map_location):
+    vae_ckpt = sd_models.read_state_dict(filename, map_location=map_location)
+    vae_dict_1 = {k: v for k, v in vae_ckpt.items() if k[0:4] != "loss" and k not in vae_ignore_keys}
+    return vae_dict_1
+
+
 def load_vae(model, vae_file=None, vae_source="from unknown source"):
     global vae_dict, loaded_vae_file
     # save_settings = False
@@ -137,8 +143,7 @@ def load_vae(model, vae_file=None, vae_source="from unknown source"):
             print(f"Loading VAE weights {vae_source}: {vae_file}")
             store_base_vae(model)
 
-            vae_ckpt = sd_models.read_state_dict(vae_file, map_location=shared.weight_load_location)
-            vae_dict_1 = {k: v for k, v in vae_ckpt.items() if k[0:4] != "loss" and k not in vae_ignore_keys}
+            vae_dict_1 = load_vae_dict(vae_file, map_location=shared.weight_load_location)
             _load_vae_dict(model, vae_dict_1)
 
             if cache_enabled:
diff --git a/modules/shared.py b/modules/shared.py
index 77e5e91c..29b28bff 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -20,10 +20,11 @@ from modules.paths import models_path, script_path, sd_path
 
 demo = None
 
+sd_default_config = os.path.join(script_path, "configs/v1-inference.yaml")
 sd_model_file = os.path.join(script_path, 'model.ckpt')
 default_sd_model_file = sd_model_file
 parser = argparse.ArgumentParser()
-parser.add_argument("--config", type=str, default=os.path.join(script_path, "configs/v1-inference.yaml"), help="path to config which constructs model",)
+parser.add_argument("--config", type=str, default=sd_default_config, help="path to config which constructs model",)
 parser.add_argument("--ckpt", type=str, default=sd_model_file, help="path to checkpoint of stable diffusion model; if specified, this checkpoint will be added to the list of checkpoints and loaded",)
 parser.add_argument("--ckpt-dir", type=str, default=None, help="Path to directory with stable diffusion checkpoints")
 parser.add_argument("--vae-dir", type=str, default=None, help="Path to directory with VAE files")
diff --git a/modules/ui.py b/modules/ui.py
index aeee7853..4e381a49 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -20,7 +20,7 @@ import numpy as np
 from PIL import Image, PngImagePlugin
 from modules.call_queue import wrap_gradio_gpu_call, wrap_queued_call, wrap_gradio_call
 
-from modules import sd_hijack, sd_models, localization, script_callbacks, ui_extensions, deepbooru
+from modules import sd_hijack, sd_models, localization, script_callbacks, ui_extensions, deepbooru, sd_vae
 from modules.ui_components import FormRow, FormGroup, ToolButton, FormHTML
 from modules.paths import script_path
 
@@ -1185,7 +1185,7 @@ def create_ui():
             with gr.Column(variant='compact'):
                 gr.HTML(value="<p style='margin-bottom: 2.5em'>A merger of the two checkpoints will be generated in your <b>checkpoint</b> directory.</p>")
 
-                with FormRow():
+                with FormRow(elem_id="modelmerger_models"):
                     primary_model_name = gr.Dropdown(modules.sd_models.checkpoint_tiles(), elem_id="modelmerger_primary_model_name", label="Primary model (A)")
                     create_refresh_button(primary_model_name, modules.sd_models.list_models, lambda: {"choices": modules.sd_models.checkpoint_tiles()}, "refresh_checkpoint_A")
 
@@ -1197,13 +1197,20 @@ def create_ui():
 
                 custom_name = gr.Textbox(label="Custom Name (Optional)", elem_id="modelmerger_custom_name")
                 interp_amount = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Multiplier (M) - set to 0 to get model A', value=0.3, elem_id="modelmerger_interp_amount")
-                interp_method = gr.Radio(choices=["Weighted sum", "Add difference"], value="Weighted sum", label="Interpolation Method", elem_id="modelmerger_interp_method")
+                interp_method = gr.Radio(choices=["No interpolation", "Weighted sum", "Add difference"], value="Weighted sum", label="Interpolation Method", elem_id="modelmerger_interp_method")
 
                 with FormRow():
                     checkpoint_format = gr.Radio(choices=["ckpt", "safetensors"], value="ckpt", label="Checkpoint format", elem_id="modelmerger_checkpoint_format")
                     save_as_half = gr.Checkbox(value=False, label="Save as float16", elem_id="modelmerger_save_as_half")
 
-                config_source = gr.Radio(choices=["A, B or C", "B", "C", "Don't"], value="A, B or C", label="Copy config from", type="index", elem_id="modelmerger_config_method")
+                with FormRow():
+                    with gr.Column():
+                        config_source = gr.Radio(choices=["A, B or C", "B", "C", "Don't"], value="A, B or C", label="Copy config from", type="index", elem_id="modelmerger_config_method")
+
+                    with gr.Column():
+                        with FormRow():
+                            bake_in_vae = gr.Dropdown(choices=["None"] + list(sd_vae.vae_dict), value="None", label="Bake in VAE", elem_id="modelmerger_bake_in_vae")
+                            create_refresh_button(bake_in_vae, sd_vae.refresh_vae_list, lambda: {"choices": ["None"] + list(sd_vae.vae_dict)}, "modelmerger_refresh_bake_in_vae")
 
                 with gr.Row():
                     modelmerger_merge = gr.Button(elem_id="modelmerger_merge", value="Merge", variant='primary')
@@ -1757,6 +1764,7 @@ def create_ui():
                 return [*[gr.Dropdown.update(choices=modules.sd_models.checkpoint_tiles()) for _ in range(4)], f"Error merging checkpoints: {e}"]
             return results
 
+        modelmerger_merge.click(fn=lambda: '', inputs=[], outputs=[modelmerger_result])
         modelmerger_merge.click(
             fn=wrap_gradio_gpu_call(modelmerger, extra_outputs=lambda: [gr.update() for _ in range(4)]),
             _js='modelmerger',
@@ -1771,6 +1779,7 @@ def create_ui():
                 custom_name,
                 checkpoint_format,
                 config_source,
+                bake_in_vae,
             ],
             outputs=[
                 primary_model_name,
diff --git a/style.css b/style.css
index 32ba4753..c10e32a1 100644
--- a/style.css
+++ b/style.css
@@ -641,6 +641,16 @@ canvas[key="mask"] {
     margin: 0.6em 0em 0.55em 0;
 }
 
+#modelmerger_results_container{
+    margin-top: 1em;
+    overflow: visible;
+}
+
+#modelmerger_models{
+    gap: 0;
+}
+
+
 #quicksettings .gr-button-tool{
     margin: 0;
 }
@@ -737,11 +747,6 @@ footer {
     line-height: 2.4em;
 }
 
-#modelmerger_results_container{
-    margin-top: 1em;
-    overflow: visible;
-}
-
 /* The following handles localization for right-to-left (RTL) languages like Arabic.
 The rtl media type will only be activated by the logic in javascript/localization.js.
 If you change anything above, you need to make sure it is RTL compliant by just running
-- 
cgit v1.2.3


From 5eee2ac39863f9e44591b50d0710dd2615416a13 Mon Sep 17 00:00:00 2001
From: Max Audron <audron@cocaine.farm>
Date: Wed, 25 Jan 2023 17:15:42 +0100
Subject: add data-dir flag and set all user data directories based on it

---
 modules/extensions.py                      |  2 +-
 modules/generation_parameters_copypaste.py |  4 ++--
 modules/gfpgan_model.py                    |  5 ++---
 modules/hashes.py                          |  4 +++-
 modules/interrogate.py                     |  2 +-
 modules/paths.py                           | 10 +++++++++-
 modules/processing.py                      |  3 ++-
 modules/sd_models.py                       |  6 +++---
 modules/sd_vae.py                          |  5 ++---
 modules/shared.py                          | 11 ++++++-----
 modules/textual_inversion/preprocess.py    |  5 ++---
 modules/ui.py                              |  6 +++---
 modules/ui_extensions.py                   |  2 +-
 modules/upscaler.py                        |  5 ++---
 14 files changed, 39 insertions(+), 31 deletions(-)

(limited to 'modules/sd_vae.py')

diff --git a/modules/extensions.py b/modules/extensions.py
index b522125c..92ee8144 100644
--- a/modules/extensions.py
+++ b/modules/extensions.py
@@ -7,7 +7,7 @@ import git
 from modules import paths, shared
 
 extensions = []
-extensions_dir = os.path.join(paths.script_path, "extensions")
+extensions_dir = os.path.join(paths.data_path, "extensions")
 extensions_builtin_dir = os.path.join(paths.script_path, "extensions-builtin")
 
 
diff --git a/modules/generation_parameters_copypaste.py b/modules/generation_parameters_copypaste.py
index 46e12dc6..35f72808 100644
--- a/modules/generation_parameters_copypaste.py
+++ b/modules/generation_parameters_copypaste.py
@@ -6,7 +6,7 @@ import re
 from pathlib import Path
 
 import gradio as gr
-from modules.shared import script_path
+from modules.paths import data_path, script_path
 from modules import shared, ui_tempdir, script_callbacks
 import tempfile
 from PIL import Image
@@ -289,7 +289,7 @@ Steps: 20, Sampler: Euler a, CFG scale: 7, Seed: 965400086, Size: 512x512, Model
 def connect_paste(button, paste_fields, input_comp, jsfunc=None):
     def paste_func(prompt):
         if not prompt and not shared.cmd_opts.hide_ui_dir_config:
-            filename = os.path.join(script_path, "params.txt")
+            filename = os.path.join(data_path, "params.txt")
             if os.path.exists(filename):
                 with open(filename, "r", encoding="utf8") as file:
                     prompt = file.read()
diff --git a/modules/gfpgan_model.py b/modules/gfpgan_model.py
index 1e2dbc32..fbe6215a 100644
--- a/modules/gfpgan_model.py
+++ b/modules/gfpgan_model.py
@@ -6,12 +6,11 @@ import facexlib
 import gfpgan
 
 import modules.face_restoration
-from modules import shared, devices, modelloader
-from modules.paths import models_path
+from modules import paths, shared, devices, modelloader
 
 model_dir = "GFPGAN"
 user_path = None
-model_path = os.path.join(models_path, model_dir)
+model_path = os.path.join(paths.models_path, model_dir)
 model_url = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth"
 have_gfpgan = False
 loaded_gfpgan_model = None
diff --git a/modules/hashes.py b/modules/hashes.py
index b85a7580..819362a3 100644
--- a/modules/hashes.py
+++ b/modules/hashes.py
@@ -4,8 +4,10 @@ import os.path
 
 import filelock
 
+from modules.paths import data_path
 
-cache_filename = "cache.json"
+
+cache_filename = os.path.join(data_path, "cache.json")
 cache_data = None
 
 
diff --git a/modules/interrogate.py b/modules/interrogate.py
index c72ff694..cbb80683 100644
--- a/modules/interrogate.py
+++ b/modules/interrogate.py
@@ -12,7 +12,7 @@ from torchvision import transforms
 from torchvision.transforms.functional import InterpolationMode
 
 import modules.shared as shared
-from modules import devices, paths, lowvram, modelloader, errors
+from modules import devices, paths, shared, lowvram, modelloader, errors
 
 blip_image_eval_size = 384
 clip_model_name = 'ViT-L/14'
diff --git a/modules/paths.py b/modules/paths.py
index 20b3e4d8..08e6f9b9 100644
--- a/modules/paths.py
+++ b/modules/paths.py
@@ -4,7 +4,15 @@ import sys
 import modules.safe
 
 script_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
-models_path = os.path.join(script_path, "models")
+
+# Parse the --data-dir flag first so we can use it as a base for our other argument default values
+parser = argparse.ArgumentParser()
+parser.add_argument("--data-dir", type=str, default=os.path.dirname(os.path.dirname(os.path.realpath(__file__))), help="base path where all user data is stored",)
+cmd_opts_pre = parser.parse_known_args()[0]
+data_path = cmd_opts_pre.data_dir
+models_path = os.path.join(data_path, "models")
+
+# data_path = cmd_opts_pre.data
 sys.path.insert(0, script_path)
 
 # search for directory of stable diffusion in following places
diff --git a/modules/processing.py b/modules/processing.py
index 262806a1..5072fc40 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -17,6 +17,7 @@ from modules import devices, prompt_parser, masking, sd_samplers, lowvram, gener
 from modules.sd_hijack import model_hijack
 from modules.shared import opts, cmd_opts, state
 import modules.shared as shared
+import modules.paths as paths
 import modules.face_restoration
 import modules.images as images
 import modules.styles
@@ -584,7 +585,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
             if not p.disable_extra_networks:
                 extra_networks.activate(p, extra_network_data)
 
-        with open(os.path.join(shared.script_path, "params.txt"), "w", encoding="utf8") as file:
+        with open(os.path.join(paths.data_path, "params.txt"), "w", encoding="utf8") as file:
             processed = Processed(p, [], p.seed, "")
             file.write(processed.infotext(p, 0))
 
diff --git a/modules/sd_models.py b/modules/sd_models.py
index 37dad18d..b2d48a51 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -12,13 +12,13 @@ import ldm.modules.midas as midas
 
 from ldm.util import instantiate_from_config
 
-from modules import shared, modelloader, devices, script_callbacks, sd_vae, sd_disable_initialization, errors, hashes, sd_models_config
+from modules import paths, shared, modelloader, devices, script_callbacks, sd_vae, sd_disable_initialization, errors, hashes, sd_models_config
 from modules.paths import models_path
 from modules.sd_hijack_inpainting import do_inpainting_hijack
 from modules.timer import Timer
 
 model_dir = "Stable-diffusion"
-model_path = os.path.abspath(os.path.join(models_path, model_dir))
+model_path = os.path.abspath(os.path.join(paths.models_path, model_dir))
 
 checkpoints_list = {}
 checkpoint_alisases = {}
@@ -307,7 +307,7 @@ def enable_midas_autodownload():
     location automatically.
     """
 
-    midas_path = os.path.join(models_path, 'midas')
+    midas_path = os.path.join(paths.models_path, 'midas')
 
     # stable-diffusion-stability-ai hard-codes the midas model path to
     # a location that differs from where other scripts using this model look.
diff --git a/modules/sd_vae.py b/modules/sd_vae.py
index 4ce238b8..9b00f76e 100644
--- a/modules/sd_vae.py
+++ b/modules/sd_vae.py
@@ -3,13 +3,12 @@ import safetensors.torch
 import os
 import collections
 from collections import namedtuple
-from modules import shared, devices, script_callbacks, sd_models
-from modules.paths import models_path
+from modules import paths, shared, devices, script_callbacks, sd_models
 import glob
 from copy import deepcopy
 
 
-vae_path = os.path.abspath(os.path.join(models_path, "VAE"))
+vae_path = os.path.abspath(os.path.join(paths.models_path, "VAE"))
 vae_ignore_keys = {"model_ema.decay", "model_ema.num_updates"}
 vae_dict = {}
 
diff --git a/modules/shared.py b/modules/shared.py
index 14be993d..474fcc42 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -14,7 +14,7 @@ import modules.memmon
 import modules.styles
 import modules.devices as devices
 from modules import localization, extensions, script_loading, errors, ui_components, shared_items
-from modules.paths import models_path, script_path
+from modules.paths import models_path, script_path, data_path
 
 
 demo = None
@@ -25,6 +25,7 @@ sd_model_file = os.path.join(script_path, 'model.ckpt')
 default_sd_model_file = sd_model_file
 
 parser = argparse.ArgumentParser()
+parser.add_argument("--data-dir", type=str, default=os.path.dirname(os.path.dirname(os.path.realpath(__file__))), help="base path where all user data is stored",)
 parser.add_argument("--config", type=str, default=sd_default_config, help="path to config which constructs model",)
 parser.add_argument("--ckpt", type=str, default=sd_model_file, help="path to checkpoint of stable diffusion model; if specified, this checkpoint will be added to the list of checkpoints and loaded",)
 parser.add_argument("--ckpt-dir", type=str, default=None, help="Path to directory with stable diffusion checkpoints")
@@ -35,7 +36,7 @@ parser.add_argument("--no-half", action='store_true', help="do not switch the mo
 parser.add_argument("--no-half-vae", action='store_true', help="do not switch the VAE model to 16-bit floats")
 parser.add_argument("--no-progressbar-hiding", action='store_true', help="do not hide progressbar in gradio UI (we hide it because it slows down ML if you have hardware acceleration in browser)")
 parser.add_argument("--max-batch-count", type=int, default=16, help="maximum batch count value for the UI")
-parser.add_argument("--embeddings-dir", type=str, default=os.path.join(script_path, 'embeddings'), help="embeddings directory for textual inversion (default: embeddings)")
+parser.add_argument("--embeddings-dir", type=str, default=os.path.join(data_path, 'embeddings'), help="embeddings directory for textual inversion (default: embeddings)")
 parser.add_argument("--textual-inversion-templates-dir", type=str, default=os.path.join(script_path, 'textual_inversion_templates'), help="directory with textual inversion templates")
 parser.add_argument("--hypernetwork-dir", type=str, default=os.path.join(models_path, 'hypernetworks'), help="hypernetwork directory")
 parser.add_argument("--localizations-dir", type=str, default=os.path.join(script_path, 'localizations'), help="localizations directory")
@@ -74,16 +75,16 @@ parser.add_argument("--use-cpu", nargs='+', help="use CPU as torch device for sp
 parser.add_argument("--listen", action='store_true', help="launch gradio with 0.0.0.0 as server name, allowing to respond to network requests")
 parser.add_argument("--port", type=int, help="launch gradio with given server port, you need root/admin rights for ports < 1024, defaults to 7860 if available", default=None)
 parser.add_argument("--show-negative-prompt", action='store_true', help="does not do anything", default=False)
-parser.add_argument("--ui-config-file", type=str, help="filename to use for ui configuration", default=os.path.join(script_path, 'ui-config.json'))
+parser.add_argument("--ui-config-file", type=str, help="filename to use for ui configuration", default=os.path.join(data_path, 'ui-config.json'))
 parser.add_argument("--hide-ui-dir-config", action='store_true', help="hide directory configuration from webui", default=False)
 parser.add_argument("--freeze-settings", action='store_true', help="disable editing settings", default=False)
-parser.add_argument("--ui-settings-file", type=str, help="filename to use for ui settings", default=os.path.join(script_path, 'config.json'))
+parser.add_argument("--ui-settings-file", type=str, help="filename to use for ui settings", default=os.path.join(data_path, 'config.json'))
 parser.add_argument("--gradio-debug",  action='store_true', help="launch gradio with --debug option")
 parser.add_argument("--gradio-auth", type=str, help='set gradio authentication like "username:password"; or comma-delimit multiple like "u1:p1,u2:p2,u3:p3"', default=None)
 parser.add_argument("--gradio-img2img-tool", type=str, help='does not do anything')
 parser.add_argument("--gradio-inpaint-tool", type=str, help="does not do anything")
 parser.add_argument("--opt-channelslast", action='store_true', help="change memory type for stable diffusion to channels last")
-parser.add_argument("--styles-file", type=str, help="filename to use for styles", default=os.path.join(script_path, 'styles.csv'))
+parser.add_argument("--styles-file", type=str, help="filename to use for styles", default=os.path.join(data_path, 'styles.csv'))
 parser.add_argument("--autolaunch", action='store_true', help="open the webui URL in the system's default browser upon launch", default=False)
 parser.add_argument("--theme", type=str, help="launches the UI with light or dark theme", default=None)
 parser.add_argument("--use-textbox-seed", action='store_true', help="use textbox for seeds in UI (no up/down, but possible to input long seeds)", default=False)
diff --git a/modules/textual_inversion/preprocess.py b/modules/textual_inversion/preprocess.py
index c0ac11d3..2239cb84 100644
--- a/modules/textual_inversion/preprocess.py
+++ b/modules/textual_inversion/preprocess.py
@@ -6,8 +6,7 @@ import sys
 import tqdm
 import time
 
-from modules import shared, images, deepbooru
-from modules.paths import models_path
+from modules import paths, shared, images, deepbooru
 from modules.shared import opts, cmd_opts
 from modules.textual_inversion import autocrop
 
@@ -199,7 +198,7 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pre
 
             dnn_model_path = None
             try:
-                dnn_model_path = autocrop.download_and_cache_models(os.path.join(models_path, "opencv"))
+                dnn_model_path = autocrop.download_and_cache_models(os.path.join(paths.models_path, "opencv"))
             except Exception as e:
                 print("Unable to load face detection model for auto crop selection. Falling back to lower quality haar method.", e)
 
diff --git a/modules/ui.py b/modules/ui.py
index 85ae62c7..0117df3e 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -21,7 +21,7 @@ from modules.call_queue import wrap_gradio_gpu_call, wrap_queued_call, wrap_grad
 
 from modules import sd_hijack, sd_models, localization, script_callbacks, ui_extensions, deepbooru, sd_vae, extra_networks, postprocessing, ui_components, ui_common, ui_postprocessing
 from modules.ui_components import FormRow, FormGroup, ToolButton, FormHTML
-from modules.paths import script_path
+from modules.paths import script_path, data_path
 
 from modules.shared import opts, cmd_opts, restricted_opts
 
@@ -1497,8 +1497,8 @@ def create_ui():
         with open(cssfile, "r", encoding="utf8") as file:
             css += file.read() + "\n"
 
-    if os.path.exists(os.path.join(script_path, "user.css")):
-        with open(os.path.join(script_path, "user.css"), "r", encoding="utf8") as file:
+    if os.path.exists(os.path.join(data_path, "user.css")):
+        with open(os.path.join(data_path, "user.css"), "r", encoding="utf8") as file:
             css += file.read() + "\n"
 
     if not cmd_opts.no_progressbar_hiding:
diff --git a/modules/ui_extensions.py b/modules/ui_extensions.py
index 742e745e..66a41865 100644
--- a/modules/ui_extensions.py
+++ b/modules/ui_extensions.py
@@ -132,7 +132,7 @@ def install_extension_from_url(dirname, url):
     normalized_url = normalize_git_url(url)
     assert len([x for x in extensions.extensions if normalize_git_url(x.remote) == normalized_url]) == 0, 'Extension with this URL is already installed'
 
-    tmpdir = os.path.join(paths.script_path, "tmp", dirname)
+    tmpdir = os.path.join(paths.data_path, "tmp", dirname)
 
     try:
         shutil.rmtree(tmpdir, True)
diff --git a/modules/upscaler.py b/modules/upscaler.py
index a5bf5acb..e2eaa730 100644
--- a/modules/upscaler.py
+++ b/modules/upscaler.py
@@ -11,7 +11,6 @@ from modules import modelloader, shared
 
 LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS)
 NEAREST = (Image.Resampling.NEAREST if hasattr(Image, 'Resampling') else Image.NEAREST)
-from modules.paths import models_path
 
 
 class Upscaler:
@@ -39,7 +38,7 @@ class Upscaler:
         self.mod_scale = None
 
         if self.model_path is None and self.name:
-            self.model_path = os.path.join(models_path, self.name)
+            self.model_path = os.path.join(shared.models_path, self.name)
         if self.model_path and create_dirs:
             os.makedirs(self.model_path, exist_ok=True)
 
@@ -143,4 +142,4 @@ class UpscalerNearest(Upscaler):
     def __init__(self, dirname=None):
         super().__init__(False)
         self.name = "Nearest"
-        self.scalers = [UpscalerData("Nearest", None, self)]
\ No newline at end of file
+        self.scalers = [UpscalerData("Nearest", None, self)]
-- 
cgit v1.2.3