Remove ToMe

README.md

@@ -53,7 +53,7 @@ Gradio app for Stable Diffusion 1.5 including:
 * 100+ styles from sdxl_prompt_styler
 * Compel prompt weighting
 * Multiple samplers with Karras scheduling
-* DeepCache, ToMe, FreeU, and Clip Skip available
+* DeepCache, FreeU, and Clip Skip available
 * Real-ESRGAN upscaling
 * Optional tiny autoencoder
 

app.py

@@ -177,11 +177,12 @@ with gr.Blocks(
                             maximum=50,
                             step=1,
                         )
-                        seed = gr.Number(
-                            value=Config.SEED,
-                            label="Seed",
-                            minimum=-1,
-                            maximum=(2**64) - 1,
+                        deepcache_interval = gr.Slider(
+                            value=Config.DEEPCACHE_INTERVAL,
+                            label="DeepCache",
+                            minimum=1,
+                            maximum=4,
+                            step=1,
                         )
 
                     with gr.Row():
@@ -212,39 +213,31 @@ with gr.Blocks(
                             filterable=False,
                             label="Aspect Ratio",
                         )
-                        scale = gr.Dropdown(
-                            choices=[(f"{s}x", s) for s in Config.SCALES],
-                            filterable=False,
-                            value=Config.SCALE,
-                            label="Scale",
-                        )
 
                     with gr.Row():
+                        file_format = gr.Dropdown(
+                            choices=["png", "jpeg", "webp"],
+                            label="File Format",
+                            filterable=False,
+                            value="png",
+                        )
                         num_images = gr.Dropdown(
                             choices=list(range(1, 5)),
                             value=Config.NUM_IMAGES,
                             filterable=False,
                             label="Images",
                         )
-                        file_format = gr.Dropdown(
-                            choices=["png", "jpeg", "webp"],
-                            label="File Format",
+                        scale = gr.Dropdown(
+                            choices=[(f"{s}x", s) for s in Config.SCALES],
                             filterable=False,
-                            value="png",
-                        )
-                        deepcache_interval = gr.Slider(
-                            value=Config.DEEPCACHE_INTERVAL,
-                            label="DeepCache",
-                            minimum=1,
-                            maximum=4,
-                            step=1,
+                            value=Config.SCALE,
+                            label="Scale",
                         )
-                        tome_ratio = gr.Slider(
-                            value=Config.TOME_RATIO,
-                            label="ToMe Ratio",
-                            minimum=0.0,
-                            maximum=0.5,
-                            step=0.01,
+                        seed = gr.Number(
+                            value=Config.SEED,
+                            label="Seed",
+                            minimum=-1,
+                            maximum=(2**64) - 1,
                         )
 
                     with gr.Row():
@@ -440,7 +433,6 @@ with gr.Blocks(
             truncate_prompts,
             increment_seed,
             deepcache_interval,
-            tome_ratio,
             scale,
         ],
     )

cli.py

@@ -29,7 +29,6 @@ def main():
     parser.add_argument("--scheduler", type=str, metavar="STR", default=Config.SCHEDULER)
     parser.add_argument("--guidance", type=float, metavar="FLOAT", default=Config.GUIDANCE_SCALE)
     parser.add_argument("--steps", type=int, metavar="INT", default=Config.INFERENCE_STEPS)
-    parser.add_argument("--tome", type=float, metavar="FLOAT", default=Config.TOME_RATIO)
     parser.add_argument("--strength", type=float, metavar="FLOAT", default=Config.DENOISING_STRENGTH)
     parser.add_argument("--image", type=str, metavar="STR")
     parser.add_argument("--taesd", action="store_true")
@@ -63,7 +62,6 @@ def main():
         args.truncate,
         args.no_increment,
         args.deepcache,
-        args.tome,
         args.scale,
     )
     save_images(images, args.filename)

lib/config.py

@@ -43,7 +43,6 @@ Config = SimpleNamespace(
     INFERENCE_STEPS=30,
     DENOISING_STRENGTH=0.6,
     DEEPCACHE_INTERVAL=2,
-    TOME_RATIO=0.0,
     SCALE=1,
     SCALES=[1, 2, 4],
 )

lib/inference.py

@@ -2,14 +2,12 @@ import json
 import os
 import re
 import time
-from contextlib import contextmanager
 from datetime import datetime
 from itertools import product
 from typing import Callable
 
 import numpy as np
 import spaces
-import tomesd
 import torch
 from compel import Compel, DiffusersTextualInversionManager, ReturnedEmbeddingsType
 from compel.prompt_parser import PromptParser
@@ -25,17 +23,6 @@ with open("./data/styles.json") as f:
     styles = json.load(f)
 
 
-# applies tome to the pipeline
-@contextmanager
-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
 def parse_prompt(prompt: str) -> list[str]:
     arrays = re.findall(r"\[\[(.*?)\]\]", prompt)
@@ -106,7 +93,6 @@ def generate(
     truncate_prompts=False,
     increment_seed=True,
     deepcache=1,
-    tome_ratio=0,
     scale=1,
     Info: Callable[[str], None] = None,
     Error=Exception,
@@ -216,15 +202,14 @@ def generate(
                 kwargs["strength"] = denoising_strength
                 kwargs["image"] = prepare_image(image_prompt, (width, height))
 
-            with token_merging(pipe, tome_ratio=tome_ratio):
-                try:
-                    image = pipe(**kwargs).images[0]
-                    if scale > 1:
-                        image = upscaler.predict(image)
-                    images.append((image, str(current_seed)))
-                finally:
-                    pipe.unload_textual_inversion()
-                    torch.cuda.empty_cache()
+            try:
+                image = pipe(**kwargs).images[0]
+                if scale > 1:
+                    image = upscaler.predict(image)
+                images.append((image, str(current_seed)))
+            finally:
+                pipe.unload_textual_inversion()
+                torch.cuda.empty_cache()
 
             if increment_seed:
                 current_seed += 1

requirements.txt

@@ -8,7 +8,6 @@ gradio==4.41.0
 ruff==0.5.7
 scipy  # for LMS scheduler
 spaces
-tomesd==0.1.3
 torch==2.3.1
 torchvision==0.18.1
 transformers==4.43.4

usage.md

@@ -71,10 +71,6 @@ Denoising strength is essentially how much the generation will differ from the i
 * `3`: balanced
 * `4`: more speed
 
-#### ToMe
-
-[Token merging](https://github.com/dbolya/tomesd) (Bolya & Hoffman 2023) reduces the number of tokens processed by the model. Set `Ratio` to the desired reduction factor. Only necessary to speed up generation on older GPUs.
-
 #### FreeU
 
 [FreeU](https://github.com/ChenyangSi/FreeU) (Si et al. 2023) re-weights the contributions sourced from the U-Net’s skip connections and backbone feature maps to potentially improve image quality.