Update app.py

app.py

@@ -9,7 +9,7 @@ import gradio as gr
 import numpy as np
 import PIL.Image
 import torch
-from diffusers import DiffusionPipeline
+from diffusers import AutoencoderKL, DiffusionPipeline
 
 DESCRIPTION = 'This space is an API service meant to be used by VideoChain and VideoQuest.\nWant to use this space for yourself? Please use the original code: [https://huggingface.co/spaces/hysts/SD-XL](https://huggingface.co/spaces/hysts/SD-XL)'
 if not torch.cuda.is_available():
@@ -23,13 +23,16 @@ SECRET_TOKEN = os.getenv('SECRET_TOKEN', 'default_secret')
 
 device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
 if torch.cuda.is_available():
+    vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
     pipe = DiffusionPipeline.from_pretrained(
         'stabilityai/stable-diffusion-xl-base-1.0',
+        vae=vae,
         torch_dtype=torch.float16,
         use_safetensors=True,
         variant='fp16')
     refiner = DiffusionPipeline.from_pretrained(
         'stabilityai/stable-diffusion-xl-refiner-1.0',
+        vae=vae,
         torch_dtype=torch.float16,
         use_safetensors=True,
         variant='fp16')