Upload folder using huggingface_hub

README.md

@@ -8,8 +8,9 @@ sdk_version: 4.19.2
 # ToDo: Token Downsampling for Efficient Generation of High-Resolution Images
 ---
 
-This is a demo for our recently proposed method, ["ToDo: Token Downsampling for Efficient Generation of High-Resolution Images"](https://arxiv.org/abs/2402.13573), compared against a popular token merging method, ToMe. 
+We provide a [HuggingFace Spaces demo](https://huggingface.co/spaces/aningineer/ToDo) for our recently proposed method, ["ToDo: Token Downsampling for Efficient Generation of High-Resolution Images"](https://arxiv.org/abs/2402.13573), and compare it against a popular token merging method called ToMe. 
 
+If you consider our research to be helpful, please consider citing us:
 ```
 @misc{smith2024todo,
       title={ToDo: Token Downsampling for Efficient Generation of High-Resolution Images}, 

app.py

@@ -3,10 +3,11 @@ import spaces
 import gradio as gr
 import torch
 import diffusers
-from utils import patch_attention_proc
+from utils import patch_attention_proc, remove_patch
 import math
 import numpy as np
 from PIL import Image
+from threading import Semaphore
 
 # Globals
 css = """
@@ -21,8 +22,12 @@ pipe = diffusers.StableDiffusionPipeline.from_pretrained("Lykon/DreamShaper").to
 pipe.scheduler = diffusers.EulerDiscreteScheduler.from_config(pipe.scheduler.config)
 pipe.safety_checker = None
 
+semaphore = Semaphore() # for preventing collisions of two simultaneous button presses
+
 @spaces.GPU
-def generate(prompt, seed, steps, height_width, negative_prompt, guidance_scale, method):
+def generate_baseline(prompt, seed, steps, height_width, negative_prompt, guidance_scale, method):
+
+    semaphore.acquire()
 
     downsample_factor = 2
     ratio = 0.38
@@ -42,8 +47,8 @@ def generate(prompt, seed, steps, height_width, negative_prompt, guidance_scale,
     elif height_width == 2048:
         downsample_factor = 4
         ratio = 0.9375
-        downsample_factor_level_2 = 2
-        ratio_level_2 = 0.75
+        downsample_factor_level_2 = 1
+        ratio_level_2 = 0.0
 
     token_merge_args = {"ratio": ratio,
                 "merge_tokens": merge_tokens,
@@ -56,17 +61,60 @@ def generate(prompt, seed, steps, height_width, negative_prompt, guidance_scale,
                 "ratio_level_2": ratio_level_2
                 }
 
-    l_r = torch.rand(1).item()
     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()
 
-    patch_attention_proc(pipe.unet, token_merge_args=token_merge_args)
+    result = f"Baseline image: {end_time_base-start_time_base:.2f} sec"
+
+    semaphore.release()
+
+    return base_img, result
+
+
+@spaces.GPU
+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,
@@ -75,9 +123,11 @@ def generate(prompt, seed, steps, height_width, negative_prompt, guidance_scale,
                         guidance_scale=guidance_scale).images[0]
     end_time_merge = time.time()
 
-    result = f"Baseline image: {end_time_base-start_time_base:.2f} sec  |  {'ToDo' if method == 'todo' else 'ToMe'} image: {end_time_merge-start_time_merge:.2f} sec"
+    result = f"{'ToDo' if method == 'todo' else 'ToMe'} image: {end_time_merge-start_time_merge:.2f} sec"
 
-    return base_img, merged_img, result
+    semaphore.release()
+
+    return merged_img, result
 
 
 
@@ -95,13 +145,18 @@ with gr.Blocks(css=css) as demo:
         steps = gr.Number(label="steps", value=20, precision=0)
         seed = gr.Number(label="seed", value=1, precision=0)
 
-    result = gr.Textbox(label="Result")
     with gr.Row():
-        base_image = gr.Image(label=f"baseline_image", type="pil", interactive=False)
-        output_image = gr.Image(label=f"output_image", type="pil", interactive=False)
-
-    gen = gr.Button("generate")
-    gen.click(generate, inputs=[prompt, seed, steps, height_width, negative_prompt,
-                                guidance_scale, method], outputs=[base_image, output_image, result])
+        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)

utils.py

@@ -25,6 +25,15 @@ def hook_tome_model(model: torch.nn.Module):
 
     model._tome_info["hooks"].append(model.register_forward_pre_hook(hook))
 
+def remove_patch(pipe: torch.nn.Module):
+    """ Removes a patch from a ToMe Diffusion module if it was already patched. """
+
+    if hasattr(pipe.unet, "_tome_info"):
+        del pipe.unet._tome_info
+
+    for n,m in pipe.unet.named_modules():
+        if hasattr(m, "processor"):
+            m.processor = AttnProcessor2_0()
 
 def patch_attention_proc(unet, token_merge_args={}):
     unet._tome_info = {