Spaces:
Running
on
Zero
Running
on
Zero
| import gradio as gr | |
| from diffusers import StableDiffusionXLPipeline, EDMEulerScheduler, StableDiffusionXLInstructPix2PixPipeline, AutoencoderKL | |
| from custom_pipeline import CosStableDiffusionXLInstructPix2PixPipeline | |
| from huggingface_hub import hf_hub_download | |
| import numpy as np | |
| import math | |
| import spaces | |
| import torch | |
| from PIL import Image | |
| 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 | |
| def resize_image(image, resolution): | |
| original_width, original_height = image.size | |
| if original_width > original_height: | |
| new_width = resolution | |
| new_height = int((resolution / original_width) * original_height) | |
| else: | |
| new_height = resolution | |
| new_width = int((resolution / original_height) * original_width) | |
| resized_img = image.resize((new_width, new_height), Image.Resampling.LANCZOS) | |
| return resized_img | |
| 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") | |
| 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") | |
| def run_normal(prompt, negative_prompt="", guidance_scale=7, steps=20, progress=gr.Progress(track_tqdm=True)): | |
| return pipe_normal(prompt, negative_prompt=negative_prompt, guidance_scale=guidance_scale, num_inference_steps=steps).images[0] | |
| def run_edit(image, prompt, negative_prompt="", guidance_scale=7, steps=20, progress=gr.Progress(track_tqdm=True)): | |
| image = resize_image(image, 1024) | |
| print("Image resized to ", image.size) | |
| width, height = image.size | |
| #image.resize((resolution, resolution)) | |
| return pipe_edit(prompt=prompt,image=image,height=height,width=width,negative_prompt=negative_prompt, guidance_scale=guidance_scale,num_inference_steps=steps).images[0] | |
| css = ''' | |
| .gradio-container{ | |
| max-width: 768px !important; | |
| margin: 0 auto; | |
| } | |
| ''' | |
| normal_examples = ["portrait photo of a girl, photograph, highly detailed face, depth of field, moody light, golden hour, style by Dan Winters, Russell James, Steve McCurry, centered, extremely detailed, Nikon D850, award winning photography", "backlit photography of a dog", "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k", "A photo of beautiful mountain with realistic sunset and blue lake, highly detailed, masterpiece"] | |
| edit_examples = [["mountain.png", "make it a cloudy day"], ["painting.png", "make the earring fancier"]] | |
| with gr.Blocks(css=css) as demo: | |
| 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 Edit"): | |
| with gr.Group(): | |
| 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): | |
| negative_prompt_edit = gr.Textbox(label="Negative Prompt") | |
| guidance_scale_edit = gr.Number(label="Guidance Scale", value=7) | |
| steps_edit = gr.Slider(label="Steps", minimum=10, maximum=50, value=20) | |
| gr.Examples(examples=edit_examples, fn=run_edit, inputs=[image_edit, prompt_edit], outputs=[output_edit], cache_examples=True) | |
| with gr.Tab("CosXL"): | |
| with gr.Group(): | |
| with gr.Row(): | |
| prompt_normal = gr.Textbox(show_label=False, scale=4, placeholder="Your prompt, e.g.: backlit photography of a dog") | |
| button_normal = gr.Button("Generate", min_width=120) | |
| output_normal = gr.Image(label="Your result image", interactive=False) | |
| with gr.Accordion("Advanced Settings", open=False): | |
| negative_prompt_normal = gr.Textbox(label="Negative Prompt") | |
| guidance_scale_normal = gr.Number(label="Guidance Scale", value=7) | |
| steps_normal = gr.Slider(label="Steps", minimum=10, maximum=50, value=20) | |
| gr.Examples(examples=normal_examples, fn=run_normal, inputs=[prompt_normal], outputs=[output_normal], cache_examples="lazy") | |
| gr.on( | |
| triggers=[ | |
| button_normal.click, | |
| prompt_normal.submit | |
| ], | |
| fn=run_normal, | |
| inputs=[prompt_normal, negative_prompt_normal, guidance_scale_normal, steps_normal], | |
| outputs=[output_normal], | |
| ) | |
| gr.on( | |
| triggers=[ | |
| button_edit.click, | |
| prompt_edit.submit | |
| ], | |
| fn=run_edit, | |
| inputs=[image_edit, prompt_edit, negative_prompt_edit, guidance_scale_edit, steps_edit], | |
| outputs=[output_edit] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(share=True) | |