Add ToMe

app.py

@@ -4,7 +4,7 @@ import gradio as gr
 
 from generate import generate
 
-DEFAULT_NEGATIVE_PROMPT = "<bad_prompt>, ugly, unattractive, deformed, disfigured, mutated, blurry, distorted, noisy, grainy, glitch, worst quality"
+DEFAULT_NEGATIVE_PROMPT = "<fast_negative>"
 
 # base font stacks
 MONO_FONTS = ["monospace"]
@@ -60,6 +60,7 @@ with gr.Blocks(
         font=[gr.themes.GoogleFont("Inter"), *SANS_FONTS],
         font_mono=[gr.themes.GoogleFont("Ubuntu Mono"), *MONO_FONTS],
     ).set(
+        layout_gap="8px",
         block_shadow="0 0 #0000",
         block_shadow_dark="0 0 #0000",
         block_background_fill=gr.themes.colors.gray.c50,
@@ -67,28 +68,49 @@ with gr.Blocks(
     ),
 ) as demo:
     gr.HTML(read_file("./partials/intro.html"))
-    output_images = gr.Gallery(
-        elem_classes=["gallery"],
-        show_share_button=False,
-        interactive=False,
-        show_label=False,
-        label="Output",
-        format="png",
-        columns=2,
-    )
-    prompt = gr.Textbox(
-        placeholder="corgi, at the beach, cute, 8k",
-        show_label=False,
-        label="Prompt",
-        value=None,
-        lines=2,
-    )
-    generate_btn = gr.Button("Generate", variant="primary", elem_classes=[])
+
+    with gr.Group():
+        output_images = gr.Gallery(
+            elem_classes=["gallery"],
+            show_share_button=False,
+            interactive=False,
+            show_label=False,
+            label="Output",
+            format="png",
+            columns=2,
+        )
+        prompt = gr.Textbox(
+            placeholder="corgi, at the beach, cute, 8k",
+            show_label=False,
+            label="Prompt",
+            value=None,
+            lines=2,
+        )
+
+    with gr.Row():
+        generate_btn = gr.Button("Generate", variant="primary", scale=6, elem_classes=[])
+        random_btn = gr.Button(
+            elem_classes=["icon-button"],
+            variant="secondary",
+            elem_id="random",
+            min_width=0,
+            value="🎲",
+            scale=1,
+        )
+        clear_btn = gr.ClearButton(
+            elem_classes=["icon-button"],
+            components=[output_images],
+            variant="secondary",
+            elem_id="clear",
+            min_width=0,
+            value="🗑️",
+            scale=1,
+        )
 
     with gr.Accordion(
         elem_classes=["accordion"],
         elem_id="menu",
-        label="Menu",
+        label="Open menu",
         open=False,
     ):
         with gr.Tabs():
@@ -98,12 +120,12 @@ with gr.Blocks(
                         label="Negative Prompt",
                         value=DEFAULT_NEGATIVE_PROMPT,
                         placeholder="",
-                        lines=1,
+                        lines=2,
                     )
 
                     with gr.Row():
                         num_images = gr.Dropdown(
-                            choices=[1, 2, 3, 4],
+                            choices=list(range(1, 9)),
                             filterable=False,
                             label="Images",
                             value=1,
@@ -113,7 +135,7 @@ with gr.Blocks(
                             label="Width",
                             minimum=256,
                             maximum=1024,
-                            value=512,
+                            value=448,
                             step=32,
                             scale=2,
                         )
@@ -121,7 +143,7 @@ with gr.Blocks(
                             label="Height",
                             minimum=256,
                             maximum=1024,
-                            value=512,
+                            value=576,
                             step=32,
                             scale=2,
                         )
@@ -131,7 +153,7 @@ with gr.Blocks(
                             label="Guidance Scale",
                             minimum=1.0,
                             maximum=15.0,
-                            value=7.5,
+                            value=7,
                             step=0.1,
                         )
                         inference_steps = gr.Slider(
@@ -171,7 +193,7 @@ with gr.Blocks(
                                 "PNDM",
                             ],
                         )
-                        seed = gr.Number(label="Seed", value=42)
+                        seed = gr.Number(label="Seed", value=42, scale=1)
 
                     with gr.Row():
                         use_karras = gr.Checkbox(
@@ -184,39 +206,33 @@ with gr.Blocks(
                             elem_classes=["checkbox"],
                             label="Autoincrement",
                             value=True,
-                            scale=2,
-                        )
-                        random_seed_btn = gr.Button(
-                            "🎲 Random seed",
-                            variant="secondary",
-                            size="sm",
-                            scale=1,
+                            scale=4,
                         )
 
             with gr.TabItem("🛠️ Advanced"):
                 with gr.Group():
                     with gr.Row():
-                        deep_cache_interval = gr.Slider(
+                        deepcache_interval = gr.Slider(
                             label="DeepCache Interval",
                             minimum=1,
                             maximum=4,
-                            value=0,
-                            step=1,
-                        )
-                        deep_cache_branch = gr.Slider(
-                            label="DeepCache Branch",
-                            minimum=0,
-                            maximum=3,
-                            value=0,
+                            value=2,
                             step=1,
                         )
                         tgate_step = gr.Slider(
                             label="T-GATE Step",
                             minimum=0,
                             maximum=50,
-                            value=0,
+                            value=20,
                             step=1,
                         )
+                        tome_ratio = gr.Slider(
+                            label="ToMe Ratio",
+                            minimum=0.0,
+                            maximum=1.0,
+                            value=0.0,
+                            step=0.01,
+                        )
 
                     with gr.Row():
                         use_taesd = gr.Checkbox(
@@ -242,7 +258,7 @@ with gr.Blocks(
                 gr.Markdown(read_file("info.md"), elem_classes=["markdown"])
 
     # update the random seed using JavaScript
-    random_seed_btn.click(None, outputs=[seed], js="() => Math.floor(Math.random() * 2**32)")
+    random_btn.click(None, outputs=[seed], js="() => Math.floor(Math.random() * 2**32)")
 
     # ensure correct argument order
     generate_btn.click(
@@ -266,9 +282,9 @@ with gr.Blocks(
             use_clip_skip,
             truncate_prompts,
             increment_seed,
-            deep_cache_interval,
-            deep_cache_branch,
+            deepcache_interval,
             tgate_step,
+            tome_ratio,
         ],
     )
 

generate.py

@@ -7,6 +7,7 @@ from types import MethodType
 from warnings import filterwarnings
 
 import spaces
+import tomesd
 import torch
 from compel import Compel, DiffusersTextualInversionManager, ReturnedEmbeddingsType
 from DeepCache import DeepCacheSDHelper
@@ -54,27 +55,63 @@ class Loader:
             cls._instance.pipe = None
         return cls._instance
 
+    def _load_deepcache(self, interval=1):
+        has_deepcache = hasattr(self.pipe, "deepcache")
+
+        if has_deepcache and self.pipe.deepcache.params["cache_interval"] == interval:
+            return self.pipe.deepcache
+        if has_deepcache:
+            self.pipe.deepcache.disable()
+        else:
+            self.pipe.deepcache = DeepCacheSDHelper(pipe=self.pipe)
+
+        self.pipe.deepcache.set_params(cache_interval=interval)
+        self.pipe.deepcache.enable()
+        return self.pipe.deepcache
+
+    def _load_tgate(self):
+        has_tgate = hasattr(self.pipe, "tgate")
+        has_deepcache = hasattr(self.pipe, "deepcache")
+
+        if not has_tgate:
+            self.pipe.tgate = MethodType(
+                tgate_sd_deepcache if has_deepcache else tgate_sd,
+                self.pipe,
+            )
+
+        return self.pipe.tgate
+
     def _load_vae(self, model_name=None, taesd=False, dtype=None):
-        if taesd:
+        vae_type = type(self.pipe.vae)
+        is_kl = issubclass(vae_type, (AutoencoderKL, OptimizedModule))
+        is_tiny = issubclass(vae_type, AutoencoderTiny)
+
+        # by default all models use KL
+        if is_kl and taesd:
             # can't compile tiny VAE
-            return AutoencoderTiny.from_pretrained(
+            print("Switching to Tiny VAE...")
+            self.pipe.vae = AutoencoderTiny.from_pretrained(
                 pretrained_model_name_or_path="madebyollin/taesd",
                 use_safetensors=True,
                 torch_dtype=dtype,
             ).to(self.gpu)
+            return self.pipe.vae
+
+        if is_tiny and not taesd:
+            print("Switching to KL VAE...")
+            self.pipe.vae = torch.compile(
+                fullgraph=True,
+                mode="reduce-overhead",
+                model=AutoencoderKL.from_pretrained(
+                    pretrained_model_name_or_path=model_name,
+                    use_safetensors=True,
+                    torch_dtype=dtype,
+                    subfolder="vae",
+                ).to(self.gpu),
+            )
+        return self.pipe.vae
 
-        return torch.compile(
-            fullgraph=True,
-            mode="reduce-overhead",
-            model=AutoencoderKL.from_pretrained(
-                pretrained_model_name_or_path=model_name,
-                use_safetensors=True,
-                torch_dtype=dtype,
-                subfolder="vae",
-            ).to(self.gpu),
-        )
-
-    def load(self, model, scheduler, karras, taesd, dtype=None):
+    def load(self, model, scheduler, karras, taesd, deepcache_interval, dtype=None):
         model_lower = model.lower()
 
         schedulers = {
@@ -126,13 +163,9 @@ class Loader:
                 if not same_scheduler or not same_karras:
                     self.pipe.scheduler = schedulers[scheduler](**scheduler_kwargs)
 
-                # if compiled will be an OptimizedModule
-                vae_type = type(self.pipe.vae)
-                if (issubclass(vae_type, (AutoencoderKL, OptimizedModule)) and taesd) or (
-                    issubclass(vae_type, AutoencoderTiny) and not taesd
-                ):
-                    print(f"Switching to {'Tiny' if taesd else 'KL'} VAE...")
-                    self.pipe.vae = self._load_vae(model_lower, taesd, dtype)
+                self._load_vae(model_lower, taesd, dtype)
+                self._load_deepcache(interval=deepcache_interval)
+                self._load_tgate()
                 return self.pipe
             else:
                 print(f"Unloading {model_name.lower()}...")
@@ -149,7 +182,9 @@ class Loader:
 
         print(f"Loading {model_lower} with {'Tiny' if taesd else 'KL'} VAE...")
         self.pipe = StableDiffusionPipeline.from_pretrained(**pipe_kwargs).to(self.gpu)
-        self.pipe.vae = self._load_vae(model_lower, taesd, dtype)
+        self._load_vae(model_lower, taesd, dtype)
+        self._load_deepcache(interval=deepcache_interval)
+        self._load_tgate()
         self.pipe.load_textual_inversion(
             pretrained_model_name_or_path=list(EMBEDDINGS.keys()),
             tokens=list(EMBEDDINGS.values()),
@@ -157,26 +192,15 @@ class Loader:
         return self.pipe
 
 
+# applies tome to the pipeline
 @contextmanager
-def deep_cache(pipe, interval=1, branch=0, tgate_step=0):
-    if interval > 1:
-        helper = DeepCacheSDHelper(pipe=pipe)
-        helper.set_params(cache_interval=interval, cache_branch_id=branch)
-        helper.enable()
-
-        if tgate_step > 0:
-            pipe.deepcache = helper
-            pipe.tgate = MethodType(tgate_sd_deepcache, pipe)
-
-        try:
-            yield helper
-        finally:
-            helper.disable()
-    elif interval < 2 and tgate_step > 0:
-        pipe.tgate = MethodType(tgate_sd, pipe)
-        yield None
-    else:
-        yield None
+def token_merging(pipe, tome_ratio=0):
+    try:
+        if tome_ratio > 0:
+            tomesd.apply_patch(pipe, max_downsample=1, sx=2, sy=2, ratio=tome_ratio)
+        yield
+    finally:
+        tomesd.remove_patch(pipe)  # idempotent
 
 
 # parse prompts with arrays
@@ -194,7 +218,6 @@ def parse_prompt(prompt: str) -> list[str]:
         current_prompt = prompt
         for i, token in enumerate(combo):
             current_prompt = current_prompt.replace(f"[[{arrays[i]}]]", token.strip(), 1)
-
         prompts.append(current_prompt)
     return prompts
 
@@ -216,9 +239,9 @@ def generate(
     clip_skip=False,
     truncate_prompts=False,
     increment_seed=True,
-    deep_cache_interval=1,
-    deep_cache_branch=0,
+    deepcache_interval=1,
     tgate_step=0,
+    tome_ratio=0,
     Error=Exception,
 ):
     if not torch.cuda.is_available():
@@ -241,7 +264,7 @@ def generate(
 
     with torch.inference_mode():
         loader = Loader()
-        pipe = loader.load(model, scheduler, karras, taesd, dtype=TORCH_DTYPE)
+        pipe = loader.load(model, scheduler, karras, taesd, deepcache_interval, TORCH_DTYPE)
 
         # prompt embeds
         compel = Compel(
@@ -271,25 +294,21 @@ def generate(
                 [pos_embeds, neg_embeds]
             )
 
-            with deep_cache(
-                pipe,
-                interval=deep_cache_interval,
-                branch=deep_cache_branch,
-                tgate_step=tgate_step,
-            ):
-                pipe_kwargs = {
-                    "num_inference_steps": inference_steps,
-                    "negative_prompt_embeds": neg_embeds,
-                    "guidance_scale": guidance_scale,
-                    "prompt_embeds": pos_embeds,
-                    "generator": generator,
-                    "height": height,
-                    "width": width,
-                }
-                result = (
-                    pipe.tgate(**pipe_kwargs, gate_step=tgate_step)
-                    if tgate_step > 0
-                    else pipe(**pipe_kwargs)
+            with token_merging(pipe, tome_ratio=tome_ratio):
+                # cap the tgate step
+                gate_step = min(
+                    tgate_step if tgate_step > 0 else inference_steps,
+                    inference_steps,
+                )
+                result = pipe.tgate(
+                    num_inference_steps=inference_steps,
+                    negative_prompt_embeds=neg_embeds,
+                    guidance_scale=guidance_scale,
+                    prompt_embeds=pos_embeds,
+                    gate_step=gate_step,
+                    generator=generator,
+                    height=height,
+                    width=width,
                 )
                 images.append((result.images[0], str(current_seed)))
 

requirements.txt

@@ -1,13 +1,14 @@
 accelerate
 compel
-deepcache
+deepcache==0.1.1
 diffusers
 hf-transfer
-gradio
+gradio==4.39.0
 ruff
 scipy  # for LMS scheduler
 spaces
-tgate
+tgate==0.1.2
+tomesd==0.1.3
 torch
 torchvision
 transformers