option to enable taesd sdxl

app.py

@@ -2,6 +2,7 @@ from diffusers import (
     StableDiffusionXLPipeline,
     EulerDiscreteScheduler,
     UNet2DConditionModel,
+    AutoencoderTiny,
 )
 import torch
 import os
@@ -19,6 +20,7 @@ BASE = "stabilityai/stable-diffusion-xl-base-1.0"
 REPO = "ByteDance/SDXL-Lightning"
 # 1-step
 CHECKPOINT = "sdxl_lightning_2step_unet.safetensors"
+taesd_model = "madebyollin/taesdxl"
 
 # {
 #     "1-Step": ["sdxl_lightning_1step_unet_x0.safetensors", 1],
@@ -30,6 +32,8 @@ CHECKPOINT = "sdxl_lightning_2step_unet.safetensors"
 
 SFAST_COMPILE = os.environ.get("SFAST_COMPILE", "0") == "1"
 SAFETY_CHECKER = os.environ.get("SAFETY_CHECKER", "0") == "1"
+USE_TAESD = os.environ.get("USE_TAESD", "0") == "1"
+
 # check if MPS is available OSX only M1/M2/M3 chips
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -38,6 +42,7 @@ torch_dtype = torch.float16
 
 print(f"SAFETY_CHECKER: {SAFETY_CHECKER}")
 print(f"SFAST_COMPILE: {SFAST_COMPILE}")
+print(f"USE_TAESD: {USE_TAESD}")
 print(f"device: {device}")
 
 
@@ -49,6 +54,12 @@ pipe = StableDiffusionXLPipeline.from_pretrained(
     BASE, unet=unet, torch_dtype=torch.float16, variant="fp16", safety_checker=False
 ).to("cuda")
 
+if USE_TAESD:
+    pipe.vae = AutoencoderTiny.from_pretrained(
+        taesd_model, torch_dtype=torch_dtype, use_safetensors=True
+    ).to(device)
+
+
 # Ensure sampler uses "trailing" timesteps.
 pipe.scheduler = EulerDiscreteScheduler.from_config(
     pipe.scheduler.config, timestep_spacing="trailing"