Update app.py

app.py

@@ -14,8 +14,8 @@ from PIL import Image
 import torch
 from diffusers import AutoencoderKL, StableDiffusionXLPipeline, UNet2DConditionModel
 from diffusers import EulerAncestralDiscreteScheduler
-from diffusers import DPMSolverMultistepScheduler
-from diffusers import AsymmetricAutoencoderKL
+#from diffusers import DPMSolverMultistepScheduler
+#from diffusers import AsymmetricAutoencoderKL
 from typing import Tuple
 import paramiko
 import gc
@@ -92,7 +92,7 @@ def load_and_prepare_model(model_id):
     model_dtypes = {"ford442/RealVisXL_V5.0_BF16": torch.bfloat16,}
     dtype = model_dtypes.get(model_id, torch.bfloat16)  # Default to float32 if not found
     #vaeX = AutoencoderKL.from_pretrained("ford442/sdxl-vae-bf16", safety_checker=None)
-    #vaeX = AutoencoderKL.from_pretrained("stabilityai/sdxl-vae", safety_checker=None, use_safetensors=False)
+    vaeXL = AutoencoderKL.from_pretrained("stabilityai/sdxl-vae", safety_checker=None, use_safetensors=False).to(device=device, dtype=torch.bfloat16)
     #vae = AutoencoderKL.from_pretrained('cross-attention/asymmetric-autoencoder-kl-x-2',use_safetensors=False)
     #vae = AutoencoderKL.from_single_file('https://huggingface.co/ford442/sdxl-vae-bf16/mySLR/myslrVAE_v10.safetensors')
     #vaeX = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse",use_safetensors=True)
@@ -239,6 +239,7 @@ def generate_30(
     randomize_seed: bool = False,
     use_resolution_binning: bool = True, 
     num_images: int = 1,  
+    juggernaut: bool = True,
     progress=gr.Progress(track_tqdm=True)  # Add progress as a keyword argument
 ):
     torch.backends.cudnn.benchmark = False
@@ -246,6 +247,8 @@ def generate_30(
     gc.collect()
     global models
     pipe = models[model_choice]
+    if juggernaut == False:
+        pipe.vae=vaeXL
     seed = int(randomize_seed_fn(seed, randomize_seed))
     generator = torch.Generator(device='cuda').manual_seed(seed)
     #prompt, negative_prompt = apply_style(style_selection, prompt, negative_prompt)
@@ -295,7 +298,8 @@ def generate_60(
     num_inference_steps: int = 250,
     randomize_seed: bool = False,
     use_resolution_binning: bool = True, 
-    num_images: int = 1,
+    num_images: int = 1,  
+    juggernaut: bool = True,
     progress=gr.Progress(track_tqdm=True)  # Add progress as a keyword argument
 ):
     torch.backends.cudnn.benchmark = True
@@ -303,6 +307,8 @@ def generate_60(
     gc.collect()
     global models
     pipe = models[model_choice]
+    if juggernaut == False:
+        pipe.vae=vaeXL
     seed = int(randomize_seed_fn(seed, randomize_seed))
     generator = torch.Generator(device='cuda').manual_seed(seed)
     #prompt, negative_prompt = apply_style(style_selection, prompt, negative_prompt)
@@ -353,6 +359,7 @@ def generate_90(
     randomize_seed: bool = False,
     use_resolution_binning: bool = True, 
     num_images: int = 1,  
+    juggernaut: bool = True,
     progress=gr.Progress(track_tqdm=True)  # Add progress as a keyword argument
 ):
     torch.backends.cudnn.benchmark = True
@@ -360,6 +367,8 @@ def generate_90(
     gc.collect()
     global models
     pipe = models[model_choice]
+    if juggernaut == False:
+        pipe.vae=vaeXL
     seed = int(randomize_seed_fn(seed, randomize_seed))
     generator = torch.Generator(device='cuda').manual_seed(seed)
     #prompt, negative_prompt = apply_style(style_selection, prompt, negative_prompt)
@@ -479,6 +488,7 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
             value=0,
         )
         randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+        juggernaut = gr.Checkbox(label="Use Juggernaut VAE", value=True)
         with gr.Row():
             width = gr.Slider(
                 label="Width",
@@ -541,7 +551,8 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
             guidance_scale,
             num_inference_steps,
             randomize_seed, 
-            num_images, 
+            num_images,
+            juggernaut,
         ],
         outputs=[result, seed],
     )
@@ -564,7 +575,8 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
             guidance_scale,
             num_inference_steps,
             randomize_seed, 
-            num_images, 
+            num_images,
+            juggernaut,
         ],
         outputs=[result, seed],
     )
@@ -587,7 +599,8 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
             guidance_scale,
             num_inference_steps,
             randomize_seed, 
-            num_images, 
+            num_images,
+            juggernaut,
         ],
         outputs=[result, seed],
     )