Update app.py

app.py

@@ -6,7 +6,7 @@ import gradio as gr
 import numpy as np
 import torch
 from PIL import Image
-from diffusers import EDMEulerScheduler, StableDiffusionXLInstructPix2PixPipeline
+from diffusers import EDMEulerScheduler, StableDiffusionXLInstructPix2PixPipeline, StableDiffusionXLImg2ImgPipeline, AutoencoderKL
 from huggingface_hub import hf_hub_download, InferenceClient
 from diffusers import DiffusionPipeline
 
@@ -39,12 +39,20 @@ def set_timesteps_patched(self, num_inference_steps: int, device = None):
     self.sigmas = self.sigmas.to("cpu") 
 
 # Image Editor
+
+vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
+
+refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", vae=vae, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
+refiner.to("cuda")
+
 edit_file = hf_hub_download(repo_id="stabilityai/cosxl", filename="cosxl_edit.safetensors")
 EDMEulerScheduler.set_timesteps = set_timesteps_patched
 pipe_edit = StableDiffusionXLInstructPix2PixPipeline.from_single_file( 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")
 
+
+
 client1 = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 system_instructions1 = "<|system|>\nAct as Image Prompt Generation expert, Your task is to modify prompt by USER to more better and detailed prompt for Image Generation. \n  Ensure the prompt is deatiled, yet descriptive to generate an exceptional image that meets the user's expectations. \n Your task is to reply with final optimized prompt only. Reply with optimized prompt only.\n<|user|>\n"