Update app.py

app.py

@@ -6,11 +6,13 @@ import gradio as gr
 import numpy as np
 import torch
 from PIL import Image
-from diffusers import DiffusionPipeline, StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler, StableDiffusionInstructPix2PixPipeline, AutoencoderKL
+from diffusers import DiffusionPipeline, StableDiffusionXLPipeline, EDMEulerScheduler, StableDiffusionXLInstructPix2PixPipeline, AutoencoderKL
 from custom_pipeline import CosStableDiffusionXLInstructPix2PixPipeline
 from huggingface_hub import hf_hub_download
 from huggingface_hub import InferenceClient
+from diffusers import StableDiffusion3Pipeline, SD3Transformer2DModel, FlowMatchEulerDiscreteScheduler
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.float16
 vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
 
@@ -21,6 +23,8 @@ pipe.set_adapters("lora")
 pipe.to("cuda")
 
 refiner = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", vae=vae, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
+refiner.load_lora_weights("KingNish/Better-Image-XL-Lora", weight_name="example-03.safetensors", adapter_name="lora")
+refiner.set_adapters("lora")
 refiner.to("cuda")
 
 help_text = """
@@ -30,11 +34,30 @@ To optimize image results:
 - Experiment with different **random seeds** and **CFG values** for varied outcomes.
 - **Rephrase your instructions** for potentially better results.
 - **Increase the number of steps** for enhanced edits.
-""" 
+"""
+
+def set_timesteps_patched(self, num_inference_steps: int, device = None):
+    self.num_inference_steps = num_inference_steps
+    
+    ramp = np.linspace(0, 1, self.num_inference_steps)
+    sigmas = torch.linspace(math.log(self.config.sigma_min), math.log(self.config.sigma_max), len(ramp)).exp().flip(0)
+    
+    sigmas = (sigmas).to(dtype=torch.float32, device=device)
+    self.timesteps = self.precondition_noise(sigmas)
+    
+    self.sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)])
+    self._step_index = None
+    self._begin_index = None
+    self.sigmas = self.sigmas.to("cpu") 
 
 # Image Editor
-pipe_edit = StableDiffusionInstructPix2PixPipeline.from_pretrained( "vinesmsuic/magicbrush-jul7", torch_dtype=torch.float16 ).to("cuda")
-pipe_edit.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe_edit.scheduler.config)
+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")
 
 # Generator
 @spaces.GPU(duration=30, queue=False)
@@ -45,7 +68,7 @@ def king(type ,
         randomize_seed: bool = False,
         seed: int = 25,
         text_cfg_scale: float = 7.3,
-        image_cfg_scale: float = 1.5,
+        image_cfg_scale: float = 1.7,
         width: int = 1024,
         height: int = 1024,
         guidance_scale: float = 6,
@@ -172,7 +195,7 @@ with gr.Blocks(css=css) as demo:
 
     with gr.Row():
         text_cfg_scale = gr.Number(value=7.3, step=0.1, label="Text CFG", interactive=True)
-        image_cfg_scale = gr.Number(value=1.5, step=0.1,label="Image CFG", interactive=True)
+        image_cfg_scale = gr.Number(value=1.7, step=0.1,label="Image CFG", interactive=True)
         guidance_scale = gr.Number(value=6.0, step=0.1, label="Image Generation Guidance Scale", interactive=True)
         steps = gr.Number(value=25, step=1, label="Steps", interactive=True)
         randomize_seed = gr.Radio(