From 549b0fc5267e9539f321f0891aa757619b7248cb Mon Sep 17 00:00:00 2001
From: Uminosachi <49424133+Uminosachi@users.noreply.github.com>
Date: Sun, 20 Aug 2023 23:06:51 +0900
Subject: Change where VAE state are stored in model

---
 modules/sd_vae.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

(limited to 'modules/sd_vae.py')

diff --git a/modules/sd_vae.py b/modules/sd_vae.py
index dbade067..ee118656 100644
--- a/modules/sd_vae.py
+++ b/modules/sd_vae.py
@@ -192,7 +192,7 @@ def load_vae_dict(filename, map_location):
 
 
 def load_vae(model, vae_file=None, vae_source="from unknown source"):
-    global vae_dict, loaded_vae_file
+    global vae_dict, base_vae, loaded_vae_file
     # save_settings = False
 
     cache_enabled = shared.opts.sd_vae_checkpoint_cache > 0
@@ -230,6 +230,8 @@ def load_vae(model, vae_file=None, vae_source="from unknown source"):
         restore_base_vae(model)
 
     loaded_vae_file = vae_file
+    model.base_vae = base_vae
+    model.loaded_vae_file = loaded_vae_file
 
 
 # don't call this from outside
-- 
cgit v1.2.3