Update app.py

app.py

@@ -1,8 +1,50 @@
 import gradio as gr
+from diffusers import StableDiffusionXLPipeline, EDMEulerScheduler
+from custom_pipeline import CosStableDiffusionXLInstructPix2PixPipeline
+import spaces 
+
+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")
+
+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")  # to avoid too much CPU/GPU communication
+
+EDMEulerScheduler.set_timesteps = set_timesteps_patched
+
+pipe_edit = CosStableDiffusionXLInstructPix2PixPipeline.from_single_file(
+    edit_file, num_in_channels=8
+)
+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.scheduler = EDMEulerScheduler(sigma_min=0.002, sigma_max=120.0, sigma_data=1.0, prediction_type="v_prediction")
+pipe_normal.to("cuda")
+
+@spaces.GPU
+def run_normal(prompt):
+    return pipe_normal(prompt, num_inference_steps=20).images[0]
+
+@spaces.GPU
+def run_edit(image, prompt):
+    resolution = 1024
+    image.resize((resolution, resolution))
+    return pipe_edit(prompt=prompt,image=image,height=resolution,width=resolution,num_inference_steps=20).images[0]
 
 with gr.Blocks() as demo:
-    gr.Markdown('''# CosXL unofficial demo
-    CosXL is a SDXL model tuned to produce full color range images. CosXL Edit allows you to perform edits on images. Both models have a [non-commercial community license](https://huggingface.co/stabilityai/cosxl/blob/main/LICENSE)
+    gr.Markdown('''# CosXL demo
+    Unofficial demo for CosXL, a SDXL model tuned to produce full color range images. CosXL Edit allows you to perform edits on images. Both have a [non-commercial community license](https://huggingface.co/stabilityai/cosxl/blob/main/LICENSE)
     ''')
     with gr.Tab("CosXL"):
       with gr.Group():
@@ -14,13 +56,22 @@ with gr.Blocks() as demo:
             pass
     with gr.Tab("CosXL Edit"):
       with gr.Group():
-          image_edit = gr.Image(label="Image you would like to edit")
+          image_edit = gr.Image(label="Image you would like to edit", type="pil")
           with gr.Row():
             prompt_edit = gr.Textbox(show_label=False, scale=4, placeholder="Edit instructions, e.g.: Make the day cloudy")
             button_edit = gr.Button("Generate", min_width=120)
           output_edit = gr.Image(label="Your result image", interactive=False)
           with gr.Accordion("Advanced Settings", open=False):
             pass
-      
+    button_normal.click(
+        fn=run_normal,
+        inputs=[prompt_normal],
+        outputs=[output_normal]
+    )
+    button_edit.click(
+        fn=run_edit,
+        inputs=[image_edit, prompt_edit],
+        outputs=[output_edit]
+    )
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(share=True)