Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from diffusers import StableDiffusionXLPipeline, EDMEulerScheduler, StableDiffusionXLInstructPix2PixPipeline
+from diffusers import StableDiffusionXLPipeline, EDMEulerScheduler, StableDiffusionXLInstructPix2PixPipeline, AutoencoderKL
 from custom_pipeline import CosStableDiffusionXLInstructPix2PixPipeline
 from huggingface_hub import hf_hub_download
 import numpy as np
@@ -40,13 +40,15 @@ def resize_image(image, resolution):
 
 EDMEulerScheduler.set_timesteps = set_timesteps_patched
 
+vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
+
 pipe_edit = StableDiffusionXLInstructPix2PixPipeline.from_single_file(
-    edit_file, num_in_channels=8, is_cosxl_edit=True, torch_dtype=torch.float16
+    edit_file, num_in_channels=8, is_cosxl_edit=True, vae=vae, torch_dtype=torch.float16,
 )
 pipe_edit.scheduler = EDMEulerScheduler(sigma_min=0.002, sigma_max=120.0, sigma_data=1.0, prediction_type="v_prediction")
 pipe_edit.to("cuda")
 
-pipe_normal = StableDiffusionXLPipeline.from_single_file(normal_file, torch_dtype=torch.float16)
+pipe_normal = StableDiffusionXLPipeline.from_single_file(normal_file, torch_dtype=torch.float16, vae=vae)
 pipe_normal.scheduler = EDMEulerScheduler(sigma_min=0.002, sigma_max=120.0, sigma_data=1.0, prediction_type="v_prediction")
 pipe_normal.to("cuda")