Update app.py

app.py

@@ -1,15 +1,17 @@
 from __future__ import annotations
-import uuid
-import os
 import math
 import random
 import spaces
 import gradio as gr
+import numpy as np
 import torch
 from PIL import Image
-from diffusers import StableDiffusionInstructPix2PixPipeline
+from diffusers import StableDiffusionXLPipeline, EDMEulerScheduler, StableDiffusionXLInstructPix2PixPipeline, AutoencoderKL
+from custom_pipeline import CosStableDiffusionXLInstructPix2PixPipeline
+from huggingface_hub import hf_hub_download
 from huggingface_hub import InferenceClient
 
+
 help_text = """
 To optimize image editing results:
 - Adjust the **Image CFG weight** if the image isn't changing enough or is changing too much. Lower it to allow bigger changes, or raise it to preserve original details.
@@ -20,14 +22,38 @@ To optimize image editing results:
 - For better facial details, especially if they're small, **crop the image** to enlarge the face's presence.
 """
 
-model_id = "timbrooks/instruct-pix2pix"
+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") 
+
+
+edit_file = hf_hub_download(repo_id="stabilityai/cosxl", filename="cosxl_edit.safetensors")
+normal_file = hf_hub_download(repo_id="stabilityai/cosxl", filename="cosxl.safetensors")
+
+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, 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")
 
 from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
 
 if not torch.cuda.is_available():
     DESCRIPTION += "\n<p>Running on CPU 🥶 This demo may not work on CPU.</p>"
-    
-MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "4096"))
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
@@ -45,8 +71,6 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
         seed = random.randint(0, 999999)
     return seed
 
-pipe2 = StableDiffusionInstructPix2PixPipeline.from_pretrained(model_id, torch_dtype=torch.float16, safety_checker=None).to("cuda")
-
 @spaces.GPU(duration=30, queue=False)
 def king(type = "Image Editing",
         input_image = None,
@@ -86,14 +110,14 @@ def king(type = "Image Editing",
         image_cfg_scale = image_cfg_scale
         input_image = input_image
 
-        steps=steps*6
+        steps=steps*3
         generator = torch.manual_seed(seed)
-        output_image = pipe2(
+        output_image = pipe_edit(
             instruction, image=input_image,
             guidance_scale=text_cfg_scale, image_guidance_scale=image_cfg_scale,
             num_inference_steps=steps, generator=generator).images[0]
         return seed, output_image
-    
+        
 def response(instruction, input_image=None):
     client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
         
@@ -131,28 +155,28 @@ def get_example():
     case = [
         [
             "Image Generation",
-            None,
             "A Super Car",
+            None,
         ],
         [
             "Image Editing",
-            "./supercar.png",
             "make it red",
+            "./supercar.png",
         ],
         [
             "Image Editing",
-            "./red_car.png",
             "add some snow",
+            "./red_car.png",
         ],
         [
             "Image Generation",
-            None,
             "Ironman flying in front of Ststue of liberty",
+            None,
         ],
         [
             "Image Generation",
-            None,
             "Beautiful Eiffel Tower at Night",
+             None,
         ],
     ]
     return case
@@ -195,7 +219,10 @@ with gr.Blocks(css=css) as demo:
 
     instruction.change(fn=response, inputs=[instruction,input_image], outputs=type, queue=False)
         
-    generate_button.click(
+    gr.on(triggers=[
+            generate_button.click,
+            instruction.submit
+        ],
             fn=king,
             inputs=[type,
                 input_image,