Update app.py

app.py

@@ -23,6 +23,7 @@ USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE") == "1"
 ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD") == "1"
 ENABLE_REFINER = os.getenv("ENABLE_REFINER", "1") == "1"
 ENABLE_USE_LORA = os.getenv("ENABLE_USE_LORA", "1") == "1"
+ENABLE_USE_VAE = os.getenv("ENABLE_USE_VAE", "1") == "1"
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
@@ -48,6 +49,7 @@ def generate(
     guidance_scale_refiner: float = 5.0,
     num_inference_steps_base: int = 25,
     num_inference_steps_refiner: int = 25,
+    use_vae: bool = False,
     use_lora: bool = False,
     apply_refiner: bool = False,
     model = 'SG161222/Realistic_Vision_V6.0_B1_noVAE',
@@ -55,10 +57,16 @@ def generate(
     lora = 'amazonaws-la/juliette',
 ) -> PIL.Image.Image:
     if torch.cuda.is_available():
-        vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
-        pipe = DiffusionPipeline.from_pretrained(model, vae=vae, torch_dtype=torch.float16)
+
+        if not use_vae:  
+        pipe = DiffusionPipeline.from_pretrained(model, torch_dtype=torch.float16)
         pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
 
+        if use_vae:
+            vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
+            pipe = DiffusionPipeline.from_pretrained(model, vae=vae, torch_dtype=torch.float16)
+            pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
+        
         if use_lora:
             pipe.load_lora_weights(lora)
             pipe.fuse_lora(lora_scale=0.7)
@@ -193,6 +201,7 @@ with gr.Blocks(css="style.css") as demo:
                 step=32,
                 value=1024,
             )
+        use_vae = gr.Checkbox(label='Use VAE', value=False, visible=ENABLE_USE_VAE)
         use_lora = gr.Checkbox(label='Use Lora', value=False, visible=ENABLE_USE_LORA)
         apply_refiner = gr.Checkbox(label="Apply refiner", value=False, visible=ENABLE_REFINER)
         with gr.Row():
@@ -255,6 +264,13 @@ with gr.Blocks(css="style.css") as demo:
         queue=False,
         api_name=False,
     )
+    use_vae.change(
+        fn=lambda x: gr.update(visible=x),
+        inputs=use_vae,
+        outputs=vaecall,
+        queue=False,
+        api_name=False,
+    )
     use_lora.change(
         fn=lambda x: gr.update(visible=x),
         inputs=use_lora,
@@ -300,6 +316,7 @@ with gr.Blocks(css="style.css") as demo:
             guidance_scale_refiner,
             num_inference_steps_base,
             num_inference_steps_refiner,
+            use_vae,
             use_lora,
             apply_refiner,
             model,