Update app.py

app.py

@@ -10,7 +10,7 @@ import numpy as np
 import PIL.Image
 import spaces
 import torch
-from diffusers import AutoencoderKL, DiffusionPipeline
+from diffusers import DPMSolverMultistepScheduler, AutoencoderKL, DiffusionPipeline
 
 DESCRIPTION = "# SDXL"
 if not torch.cuda.is_available():
@@ -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,15 +49,25 @@ 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 = 'stabilityai/stable-diffusion-xl-base-1.0',
-    vaecall = 'madebyollin/sdxl-vae-fp16-fix',
+    model = 'SG161222/Realistic_Vision_V6.0_B1_noVAE',
+    vaecall = 'stabilityai/sd-vae-ft-mse',
     lora = 'amazonaws-la/juliette',
+    lora_scale: float = 0.7,
 ) -> PIL.Image.Image:
     if torch.cuda.is_available():
-        pipe = DiffusionPipeline.from_pretrained(model, 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)
@@ -134,6 +145,13 @@ with gr.Blocks(css="style.css") as demo:
         model = gr.Text(label='Modelo')
         vaecall = gr.Text(label='VAE')
         lora = gr.Text(label='LoRA')
+        lora_scale = gr.Slider(
+                label="Lora Scale",
+                minimum=0.01,
+                maximum=1,
+                step=0.01,
+                value=0.7,
+            )
         with gr.Row():
             prompt = gr.Text(
                 label="Prompt",
@@ -191,6 +209,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():
@@ -253,6 +272,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,
@@ -298,11 +324,13 @@ 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,
             vaecall,
             lora,
+            lora_scale,
         ],
         outputs=result,
         api_name="run",