From bd5a8405194422d2f06c5edff183a5befb051327 Mon Sep 17 00:00:00 2001
From: layerdiffusion <19834515+lllyasviel@users.noreply.github.com>
Date: Mon, 5 Aug 2024 04:25:17 -0700
Subject: [PATCH] fix vae

---
 backend/diffusion_engine/base.py | 2 ++
 backend/diffusion_engine/sd15.py | 1 +
 backend/diffusion_engine/sdxl.py | 1 +
 3 files changed, 4 insertions(+)

diff --git a/backend/diffusion_engine/base.py b/backend/diffusion_engine/base.py
index a4b2a523..9aa30b80 100644
--- a/backend/diffusion_engine/base.py
+++ b/backend/diffusion_engine/base.py
@@ -31,6 +31,8 @@ class ForgeDiffusionEngine:
         self.current_lora_hash = str([])
         self.tiling_enabled = False
 
+        self.first_stage_model = None  # set this so that you can change VAE in UI
+
         # WebUI Dirty Legacy
         self.cond_stage_key = 'txt'
         self.is_sd3 = False
diff --git a/backend/diffusion_engine/sd15.py b/backend/diffusion_engine/sd15.py
index 741c32a0..2117c79f 100644
--- a/backend/diffusion_engine/sd15.py
+++ b/backend/diffusion_engine/sd15.py
@@ -52,6 +52,7 @@ class StableDiffusion(ForgeDiffusionEngine):
 
         # WebUI Legacy
         self.is_sd1 = True
+        self.first_stage_model = vae.first_stage_model
 
     def set_clip_skip(self, clip_skip):
         self.text_processing_engine.clip_skip = clip_skip
diff --git a/backend/diffusion_engine/sdxl.py b/backend/diffusion_engine/sdxl.py
index f255c3dd..6bb6ecd1 100644
--- a/backend/diffusion_engine/sdxl.py
+++ b/backend/diffusion_engine/sdxl.py
@@ -71,6 +71,7 @@ class StableDiffusionXL(ForgeDiffusionEngine):
 
         # WebUI Legacy
         self.is_sdxl = True
+        self.first_stage_model = vae.first_stage_model
 
     def set_clip_skip(self, clip_skip):
         self.text_processing_engine_l.clip_skip = clip_skip