Update app.py

app.py

@@ -6,14 +6,10 @@ import gradio as gr
 import torch
 import spaces
 
-vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16)
-#safety_checker=StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker"),
-#feature_extractor=CLIPFeatureExtractor.from_pretrained("openai/clip-vit-base-patch32")
-pretrain_model = "SG161222/Realistic_Vision_V5.1_noVAE"
+vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
 scheduler = DDIMScheduler.from_pretrained(pretrain_model, subfolder="scheduler")
-pipe = StableDiffusionPipeline.from_pretrained(pretrain_model, scheduler = scheduler, vae=vae, torch_dtype=torch.float16).to("cuda")
+pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", vae=vae, scheduler=scheduler, torch_dtype=torch.float16, use_safetensors=True, variant="fp16").to("cuda")
 
-# # Optional. enable_xformers_memory_efficient_attention can save memory usage and increase inference speed. enable_model_cpu_offload and enable_vae_tiling can save memory usage.
 #pipe.enable_model_cpu_offload()
 #pipe.enable_vae_tiling()