fix vae

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

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

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