Spaces:
Sleeping
Sleeping
| import time | |
| import spaces | |
| import gradio as gr | |
| import torch | |
| import diffusers | |
| from utils import patch_attention_proc, remove_patch | |
| import math | |
| import numpy as np | |
| from PIL import Image | |
| from threading import Semaphore | |
| # Globals | |
| css = """ | |
| h1 { | |
| text-align: center; | |
| display: block; | |
| } | |
| """ | |
| # Pipeline | |
| pipe = diffusers.StableDiffusionPipeline.from_pretrained("Lykon/DreamShaper").to("cuda", torch.float16) | |
| pipe.scheduler = diffusers.EulerDiscreteScheduler.from_config(pipe.scheduler.config) | |
| pipe.safety_checker = None | |
| semaphore = Semaphore() # for preventing collisions of two simultaneous button presses | |
| def generate_baseline(prompt, seed, steps, height_width, negative_prompt, guidance_scale, method): | |
| semaphore.acquire() | |
| downsample_factor = 2 | |
| ratio = 0.38 | |
| merge_method = "downsample" if method == "todo" else "similarity" | |
| merge_tokens = "keys/values" if method == "todo" else "all" | |
| if height_width == 1024: | |
| downsample_factor = 2 | |
| ratio = 0.75 | |
| downsample_factor_level_2 = 1 | |
| ratio_level_2 = 0.0 | |
| elif height_width == 1536: | |
| downsample_factor = 3 | |
| ratio = 0.89 | |
| downsample_factor_level_2 = 1 | |
| ratio_level_2 = 0.0 | |
| elif height_width == 2048: | |
| downsample_factor = 4 | |
| ratio = 0.9375 | |
| downsample_factor_level_2 = 1 | |
| ratio_level_2 = 0.0 | |
| token_merge_args = {"ratio": ratio, | |
| "merge_tokens": merge_tokens, | |
| "merge_method": merge_method, | |
| "downsample_method": "nearest", | |
| "downsample_factor": downsample_factor, | |
| "timestep_threshold_switch": 0.0, | |
| "timestep_threshold_stop": 0.0, | |
| "downsample_factor_level_2": downsample_factor_level_2, | |
| "ratio_level_2": ratio_level_2 | |
| } | |
| torch.manual_seed(seed) | |
| start_time_base = time.time() | |
| remove_patch(pipe) | |
| base_img = pipe(prompt, | |
| num_inference_steps=steps, height=height_width, width=height_width, | |
| negative_prompt=negative_prompt, | |
| guidance_scale=guidance_scale).images[0] | |
| end_time_base = time.time() | |
| result = f"Baseline Runtime: {end_time_base-start_time_base:.2f} sec" | |
| semaphore.release() | |
| return base_img, result | |
| def generate_merged(prompt, seed, steps, height_width, negative_prompt, guidance_scale, method): | |
| semaphore.acquire() | |
| downsample_factor = 2 | |
| ratio = 0.38 | |
| merge_method = "downsample" if method == "todo" else "similarity" | |
| merge_tokens = "keys/values" if method == "todo" else "all" | |
| if height_width == 1024: | |
| downsample_factor = 2 | |
| ratio = 0.75 | |
| downsample_factor_level_2 = 1 | |
| ratio_level_2 = 0.0 | |
| elif height_width == 1536: | |
| downsample_factor = 3 | |
| ratio = 0.89 | |
| downsample_factor_level_2 = 1 | |
| ratio_level_2 = 0.0 | |
| elif height_width == 2048: | |
| downsample_factor = 4 | |
| ratio = 0.9375 | |
| downsample_factor_level_2 = 1 | |
| ratio_level_2 = 0.0 | |
| token_merge_args = {"ratio": ratio, | |
| "merge_tokens": merge_tokens, | |
| "merge_method": merge_method, | |
| "downsample_method": "nearest", | |
| "downsample_factor": downsample_factor, | |
| "timestep_threshold_switch": 0.0, | |
| "timestep_threshold_stop": 0.0, | |
| "downsample_factor_level_2": downsample_factor_level_2, | |
| "ratio_level_2": ratio_level_2 | |
| } | |
| patch_attention_proc(pipe.unet, token_merge_args=token_merge_args) | |
| torch.manual_seed(seed) | |
| start_time_merge = time.time() | |
| merged_img = pipe(prompt, | |
| num_inference_steps=steps, height=height_width, width=height_width, | |
| negative_prompt=negative_prompt, | |
| guidance_scale=guidance_scale).images[0] | |
| end_time_merge = time.time() | |
| result = f"{'ToDo' if method == 'todo' else 'ToMe'} Runtime: {end_time_merge-start_time_merge:.2f} sec" | |
| semaphore.release() | |
| return merged_img, result | |
| with gr.Blocks(css=css) as demo: | |
| gr.Markdown("# ToDo: Token Downsampling for Efficient Generation of High-Resolution Images") | |
| prompt = gr.Textbox(interactive=True, label="prompt") | |
| negative_prompt = gr.Textbox(interactive=True, label="negative_prompt") | |
| with gr.Row(): | |
| method = gr.Dropdown(["todo", "tome"], value="todo", label="method", info="Choose Your Desired Method (Default: todo)") | |
| height_width = gr.Dropdown([1024, 1536, 2048], value=1024, label="height/width", info="Choose Your Desired Height/Width (Default: 1024)") | |
| with gr.Row(): | |
| guidance_scale = gr.Number(label="guidance_scale", value=7.5, precision=1) | |
| steps = gr.Number(label="steps", value=20, precision=0) | |
| seed = gr.Number(label="seed", value=1, precision=0) | |
| with gr.Row(): | |
| with gr.Column(): | |
| base_result = gr.Textbox(label="Baseline Runtime") | |
| base_image = gr.Image(label=f"baseline_image", type="pil", interactive=False) | |
| gen = gr.Button("Generate Baseline") | |
| gen.click(generate_baseline, inputs=[prompt, seed, steps, height_width, negative_prompt, | |
| guidance_scale, method], outputs=[base_image, base_result]) | |
| with gr.Column(): | |
| output_result = gr.Textbox(label="Runtime") | |
| output_image = gr.Image(label=f"image", type="pil", interactive=False) | |
| gen = gr.Button("Generate") | |
| gen.click(generate_merged, inputs=[prompt, seed, steps, height_width, negative_prompt, | |
| guidance_scale, method], outputs=[output_image, output_result]) | |
| demo.launch(share=True) |