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