diff --git a/CHANGELOG.md b/CHANGELOG.md
index 435f7559b..9f94ba4f6 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -26,24 +26,28 @@ it is time to give credit to original [author](https://github.com/auTOMATIC1111)
   - **PWA** [SD.Next](https://github.com/vladmandic/automatic) now also includes valid manifest making it installable as PWA
   - **Gallery
   - **Gallery**: list, preview, search through all your images and videos!  
-    implemented as infinite-scroll with client-side-caching and lazy-loading while being fully async and non-blocking  
-    search or sort by path, name, size, width, height, mtime or any image metadata item, also with extended syntax like *width > 1000*  
-    *settings*: optional additional user-defined folders, thumbnails in fixed or variable aspect-ratio  
+    Implemented as infinite-scroll with client-side-caching and lazy-loading while being fully async and non-blocking  
+    Search or sort by path, name, size, width, height, mtime or any image metadata item, also with extended syntax like *width > 1000*  
+    *Settings*: optional additional user-defined folders, thumbnails in fixed or variable aspect-ratio  
   - [HiDiffusion](https://github.com/megvii-research/HiDiffusion):  
-    generate high-resolution images using your standard models without duplicates/distorsions AND improved performance  
-    for example, *SD15* can now go up to *2024x2048* and *SDXL* up to *4k* natively
-    simply enable checkbox in advanced menu and set desired resolution  
-    additional settings are available in *settings -> inference settings -> hidiffusion*  
-    and can also be set and used via *xyz grid*  
+    Generate high-resolution images using your standard models without duplicates/distorsions AND improved performance  
+    For example, *SD15* can now go up to *2024x2048* and *SDXL* up to *4k* natively
+    Simply enable checkbox in advanced menu and set desired resolution  
+    Additional settings are available in *settings -> inference settings -> hidiffusion*  
+    And can also be set and used via *xyz grid*  
   - **IP Adapter Masking**:  
-    powerful method of using masking with ip-adapters  
-    when combined with multiple ip-adapters, it allows for different inputs guidance for each segment of the input image  
-    *hint*: to create masks, you can use manually created masks or control->mask module with auto-segment to create masks and later upload them  
+    Powerful method of using masking with ip-adapters  
+    When combined with multiple ip-adapters, it allows for different inputs guidance for each segment of the input image  
+    *Hint*: to create masks, you can use manually created masks or control->mask module with auto-segment to create masks and later upload them  
   - **IP Adapter advanced layer configuration**:  
-    allows for more control over how each layer of ip-adapter is applied, requires a valid dict to be passed as input  
-    see [InstantStyle](https://github.com/InstantStyle/InstantStyle) for details  
+    Allows for more control over how each layer of ip-adapter is applied, requires a valid dict to be passed as input  
+    See [InstantStyle](https://github.com/InstantStyle/InstantStyle) for details  
   - **OneDiff**: new optimization/compile engine, thanks @aifartist  
-    as with all other compile engines, enable via *settings -> compute settings -> compile*  
+    As with all other compile engines, enable via *settings -> compute settings -> compile*  
+  - [ToDo](https://arxiv.org/html/2402.13573v2) Token Downsampling for Efficient Generation of High-Resolution Images  
+    Newer alternative method to [ToMe](https://github.com/dbolya/tomesd) that can provide speed-up with minimal quality loss  
+    Enable in *settings -> inference settings -> token merging*  
+    Also available in XYZ grid  
 - **UI**:
   - Faster **UI** load times
   - Theme types:  
diff --git a/TODO.md b/TODO.md
index dea5988e1..8750330b9 100644
--- a/TODO.md
+++ b/TODO.md
@@ -12,8 +12,6 @@ Main ToDo list can be found at [GitHub projects](https://github.com/users/vladma
 
 - stable diffusion 3.0
 - powerpaint: <https://github.com/zhuang2002/PowerPaint>
-- hyper-sd: <https://huggingface.co/ByteDance/Hyper-SD>
-  
 
 ### Features
 
diff --git a/html/sdnext-robot-2k.jpg b/html/sdnext-robot-2k.jpg
new file mode 100644
index 000000000..0a2578ccb
Binary files /dev/null and b/html/sdnext-robot-2k.jpg differ
diff --git a/javascript/extraNetworks.js b/javascript/extraNetworks.js
index 45ab26c45..96ac263a4 100644
--- a/javascript/extraNetworks.js
+++ b/javascript/extraNetworks.js
@@ -308,7 +308,6 @@ function quickApplyStyle() {
 }
 
 function quickSaveStyle() {
-  console.log('HERE');
   const tabname = getENActiveTab();
   const btnSave = gradioApp().getElementById(`${tabname}_extra_quicksave`);
   if (btnSave) btnSave.click();
diff --git a/modules/face/faceid.py b/modules/face/faceid.py
index 7639152f9..283594f54 100644
--- a/modules/face/faceid.py
+++ b/modules/face/faceid.py
@@ -68,8 +68,8 @@ def face_id(
     try:
         shared.prompt_styles.apply_styles_to_extra(p)
 
-        if not shared.opts.cuda_compile:
-            sd_models.apply_token_merging(p.sd_model, p.get_token_merging_ratio())
+        if shared.opts.cuda_compile_backend == 'none':
+            sd_models.apply_token_merging(p.sd_model)
             sd_hijack_freeu.apply_freeu(p, shared.backend == shared.Backend.ORIGINAL)
 
         script_callbacks.before_process_callback(p)
@@ -236,8 +236,8 @@ def face_id(
     finally:
         if faceid_model is not None and original_load_ip_adapter is not None:
             faceid_model.__class__.load_ip_adapter = original_load_ip_adapter
-        if not shared.opts.cuda_compile:
-            sd_models.apply_token_merging(p.sd_model, 0)
+        if shared.opts.cuda_compile_backend == 'none':
+            sd_models.remove_token_merging(p.sd_model)
         script_callbacks.after_process_callback(p)
 
     return processed_images
diff --git a/modules/generation_parameters_copypaste.py b/modules/generation_parameters_copypaste.py
index ecc603d7b..5688be7fd 100644
--- a/modules/generation_parameters_copypaste.py
+++ b/modules/generation_parameters_copypaste.py
@@ -284,10 +284,8 @@ infotext_to_setting_name_mapping = [
     ('UniPC variant', 'uni_pc_variant'),
     # Token Merging
     ('Mask weight', 'inpainting_mask_weight'),
-    ('Token merging ratio', 'token_merging_ratio'),
-    ('ToMe', 'token_merging_ratio'),
-    ('ToMe hires', 'token_merging_ratio_hr'),
-    ('ToMe img2img', 'token_merging_ratio_img2img'),
+    ('ToMe', 'tome_ratio'),
+    ('ToDo', 'todo_ratio'),
 ]
 
 
diff --git a/modules/processing.py b/modules/processing.py
index 230e924c4..c35610fdb 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -75,8 +75,6 @@ class Processed:
         self.all_negative_prompts = all_negative_prompts or p.all_negative_prompts or [self.negative_prompt]
         self.all_seeds = all_seeds or p.all_seeds or [self.seed]
         self.all_subseeds = all_subseeds or p.all_subseeds or [self.subseed]
-        self.token_merging_ratio = p.token_merging_ratio
-        self.token_merging_ratio_hr = p.token_merging_ratio_hr
         self.infotexts = infotexts or [info]
 
     def js(self):
@@ -114,9 +112,6 @@ class Processed:
     def infotext(self, p: StableDiffusionProcessing, index):
         return create_infotext(p, self.all_prompts, self.all_seeds, self.all_subseeds, comments=[], position_in_batch=index % self.batch_size, iteration=index // self.batch_size)
 
-    def get_token_merging_ratio(self, for_hr=False):
-        return self.token_merging_ratio_hr if for_hr else self.token_merging_ratio
-
 
 def process_images(p: StableDiffusionProcessing) -> Processed:
     debug(f'Process images: {vars(p)}')
@@ -161,8 +156,8 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
 
         shared.prompt_styles.apply_styles_to_extra(p)
         shared.prompt_styles.extract_comments(p)
-        if not shared.opts.cuda_compile:
-            sd_models.apply_token_merging(p.sd_model, p.get_token_merging_ratio())
+        if shared.opts.cuda_compile_backend == 'none':
+            sd_models.apply_token_merging(p.sd_model)
             sd_hijack_freeu.apply_freeu(p, shared.backend == shared.Backend.ORIGINAL)
 
         if p.width is not None:
@@ -194,8 +189,8 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
                 processed = process_images_inner(p)
 
     finally:
-        if not shared.opts.cuda_compile:
-            sd_models.apply_token_merging(p.sd_model, 0)
+        if shared.opts.cuda_compile_backend == 'none':
+            sd_models.remove_token_merging(p.sd_model)
 
         script_callbacks.after_process_callback(p)
 
diff --git a/modules/processing_class.py b/modules/processing_class.py
index b88ec6f68..3033b8122 100644
--- a/modules/processing_class.py
+++ b/modules/processing_class.py
@@ -63,8 +63,6 @@ class StableDiffusionProcessing:
         self.override_settings_restore_afterwards = override_settings_restore_afterwards
         self.is_using_inpainting_conditioning = False # a111 compatibility
         self.disable_extra_networks = False
-        self.token_merging_ratio = 0
-        self.token_merging_ratio_hr = 0
         # self.scripts = scripts.ScriptRunner() # set via property
         # self.script_args = script_args or [] # set via property
         self.per_script_args = {}
@@ -184,11 +182,6 @@ class StableDiffusionProcessing:
     def close(self):
         self.sampler = None # pylint: disable=attribute-defined-outside-init
 
-    def get_token_merging_ratio(self, for_hr=False):
-        if for_hr:
-            return self.token_merging_ratio_hr or shared.opts.token_merging_ratio_hr or self.token_merging_ratio or shared.opts.token_merging_ratio
-        return self.token_merging_ratio or shared.opts.token_merging_ratio
-
 
 class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
 
@@ -449,8 +442,6 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
         from modules import processing_original
         return processing_original.sample_img2img(self, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, prompts)
 
-    def get_token_merging_ratio(self, for_hr=False):
-        return self.token_merging_ratio or ("token_merging_ratio" in self.override_settings and shared.opts.token_merging_ratio) or shared.opts.token_merging_ratio_img2img or shared.opts.token_merging_ratio
 
 class StableDiffusionProcessingControl(StableDiffusionProcessingImg2Img):
     def __init__(self, **kwargs):
diff --git a/modules/processing_info.py b/modules/processing_info.py
index 218e1accd..5f60fb1a1 100644
--- a/modules/processing_info.py
+++ b/modules/processing_info.py
@@ -131,10 +131,8 @@ def create_infotext(p: StableDiffusionProcessing, all_prompts=None, all_seeds=No
         args['Sampler sigma noise'] = shared.opts.s_noise if shared.opts.s_noise != shared.opts.data_labels.get('s_noise').default else None
         args['Sampler sigma tmin'] = shared.opts.s_tmin if shared.opts.s_tmin != shared.opts.data_labels.get('s_tmin').default else None
     # tome
-    token_merging_ratio = p.get_token_merging_ratio()
-    token_merging_ratio_hr = p.get_token_merging_ratio(for_hr=True) if p.enable_hr else None
-    args['ToMe'] = token_merging_ratio if token_merging_ratio != 0 else None
-    args['ToMe hires'] = token_merging_ratio_hr if token_merging_ratio_hr != 0 else None
+    args['ToMe'] = shared.opts.tome_ratio if shared.opts.tome_ratio != 0 else None
+    args['ToDo'] = shared.opts.todo_ratio if shared.opts.todo_ratio != 0 else None
 
     args.update(p.extra_generation_params)
     params_text = ", ".join([k if k == v else f'{k}: {generation_parameters_copypaste.quote(v)}' for k, v in args.items() if v is not None])
diff --git a/modules/processing_original.py b/modules/processing_original.py
index 6e6d2bafe..0f61cf41b 100644
--- a/modules/processing_original.py
+++ b/modules/processing_original.py
@@ -135,10 +135,10 @@ def sample_txt2img(p: processing.StableDiffusionProcessingTxt2Img, conditioning,
                     p.sampler.initialize(p)
                 samples = samples[:, :, p.truncate_y//2:samples.shape[2]-(p.truncate_y+1)//2, p.truncate_x//2:samples.shape[3]-(p.truncate_x+1)//2]
                 noise = create_random_tensors(samples.shape[1:], seeds=seeds, subseeds=subseeds, subseed_strength=subseed_strength, p=p)
-                sd_models.apply_token_merging(p.sd_model, p.get_token_merging_ratio(for_hr=True))
+                sd_models.apply_token_merging(p.sd_model)
                 hypertile_set(p, hr=True)
                 samples = p.sampler.sample_img2img(p, samples, noise, conditioning, unconditional_conditioning, steps=p.hr_second_pass_steps or p.steps, image_conditioning=image_conditioning)
-                sd_models.apply_token_merging(p.sd_model, p.get_token_merging_ratio())
+                sd_models.apply_token_merging(p.sd_model)
             else:
                 p.ops.append('upscale')
         x = None
@@ -149,7 +149,7 @@ def sample_txt2img(p: processing.StableDiffusionProcessingTxt2Img, conditioning,
     return samples
 
 
-def sample_img2img(p, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, prompts):
+def sample_img2img(p, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, prompts): # pylint: disable=unused-argument
     hypertile_set(p)
     x = create_random_tensors([4, p.height // 8, p.width // 8], seeds=seeds, subseeds=subseeds, subseed_strength=p.subseed_strength, seed_resize_from_h=p.seed_resize_from_h, seed_resize_from_w=p.seed_resize_from_w, p=p)
     x *= p.initial_noise_multiplier
diff --git a/modules/sd_models.py b/modules/sd_models.py
index 060fe00ad..1e454c9cc 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -17,7 +17,6 @@ import torch
 import safetensors.torch
 import diffusers
 from omegaconf import OmegaConf
-import tomesd
 from transformers import logging as transformers_logging
 from ldm.util import instantiate_from_config
 from modules import paths, shared, shared_items, shared_state, modelloader, devices, script_callbacks, sd_vae, sd_unet, errors, hashes, sd_models_config, sd_models_compile, sd_hijack_accelerate
@@ -1581,31 +1580,75 @@ def unload_model_weights(op='model'):
             shared.log.debug(f'Unload weights {op}: {memory_stats()}')
 
 
-def apply_token_merging(sd_model, token_merging_ratio=0):
-    current_token_merging_ratio = getattr(sd_model, 'applied_token_merged_ratio', 0)
-    if token_merging_ratio is None or current_token_merging_ratio is None or current_token_merging_ratio == token_merging_ratio:
-        return
-    try:
-        if current_token_merging_ratio > 0:
-            tomesd.remove_patch(sd_model)
-    except Exception:
-        pass
-    if token_merging_ratio > 0:
+def apply_token_merging(sd_model):
+    current_tome = getattr(sd_model, 'applied_tome', 0)
+    current_todo = getattr(sd_model, 'applied_todo', 0)
+
+    if shared.opts.token_merging_method == 'ToMe' and shared.opts.tome_ratio > 0:
+        if current_tome == shared.opts.tome_ratio:
+            return
         if shared.opts.hypertile_unet_enabled and not shared.cmd_opts.experimental:
             shared.log.warning('Token merging not supported with HyperTile for UNet')
             return
         try:
+            import tomesd
             tomesd.apply_patch(
                 sd_model,
-                ratio=token_merging_ratio,
-                use_rand=False,  # can cause issues with some samplers
+                ratio=shared.opts.tome_ratio,
+                use_rand=False, # can cause issues with some samplers
                 merge_attn=True,
                 merge_crossattn=False,
                 merge_mlp=False
             )
-            shared.log.info(f'Applying token merging: ratio={token_merging_ratio}')
-            sd_model.applied_token_merged_ratio = token_merging_ratio
+            shared.log.info(f'Applying ToMe: ratio={shared.opts.tome_ratio}')
+            sd_model.applied_tome = shared.opts.tome_ratio
         except Exception:
             shared.log.warning(f'Token merging not supported: pipeline={sd_model.__class__.__name__}')
     else:
-        sd_model.applied_token_merged_ratio = 0
+        sd_model.applied_tome = 0
+
+    if shared.opts.token_merging_method == 'ToDo' and shared.opts.todo_ratio > 0:
+        if current_todo == shared.opts.todo_ratio:
+            return
+        if shared.opts.hypertile_unet_enabled and not shared.cmd_opts.experimental:
+            shared.log.warning('Token merging not supported with HyperTile for UNet')
+            return
+        try:
+            from modules.todo.todo_utils import patch_attention_proc
+            token_merge_args = {
+                        "ratio": shared.opts.todo_ratio,
+                        "merge_tokens": "keys/values",
+                        "merge_method": "downsample",
+                        "downsample_method": "nearest",
+                        "downsample_factor": 2,
+                        "timestep_threshold_switch": 0.0,
+                        "timestep_threshold_stop": 0.0,
+                        "downsample_factor_level_2": 1,
+                        "ratio_level_2": 0.0,
+                        }
+            patch_attention_proc(sd_model.unet, token_merge_args=token_merge_args)
+            shared.log.info(f'Applying ToDo: ratio={shared.opts.todo_ratio}')
+            sd_model.applied_todo = shared.opts.todo_ratio
+        except Exception:
+            shared.log.warning(f'Token merging not supported: pipeline={sd_model.__class__.__name__}')
+    else:
+        sd_model.applied_todo = 0
+
+
+def remove_token_merging(sd_model):
+    current_tome = getattr(sd_model, 'applied_tome', 0)
+    current_todo = getattr(sd_model, 'applied_todo', 0)
+    try:
+        if current_tome > 0:
+            import tomesd
+            tomesd.remove_patch(sd_model)
+            sd_model.applied_tome = 0
+    except Exception:
+        pass
+    try:
+        if current_todo > 0:
+            from modules.todo.todo_utils import remove_patch
+            remove_patch(sd_model)
+            sd_model.applied_todo = 0
+    except Exception:
+        pass
diff --git a/modules/shared.py b/modules/shared.py
index 62dedbfe9..06f5c6e41 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -468,9 +468,9 @@ options_templates.update(options_section(('cuda', "Compute Settings"), {
 
 options_templates.update(options_section(('advanced', "Inference Settings"), {
     "token_merging_sep": OptionInfo("<h2>Token merging</h2>", "", gr.HTML),
-    "token_merging_ratio": OptionInfo(0.0, "Token merging ratio for txt2img", gr.Slider, {"minimum": 0.0, "maximum": 0.9, "step": 0.1}),
-    "token_merging_ratio_img2img": OptionInfo(0.0, "Token merging ratio for img2img", gr.Slider, {"minimum": 0.0, "maximum": 0.9, "step": 0.1}),
-    "token_merging_ratio_hr": OptionInfo(0.0, "Token merging ratio for hires", gr.Slider, {"minimum": 0.0, "maximum": 0.9, "step": 0.1}),
+    "token_merging_method": OptionInfo("None", "Token merging method", gr.Radio, {"choices": ['None', 'ToMe', 'ToDo']}),
+    "tome_ratio": OptionInfo(0.0, "ToMe token merging ratio", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.05}),
+    "todo_ratio": OptionInfo(0.0, "ToDo token merging ratio", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.05}),
 
     "freeu_sep": OptionInfo("<h2>FreeU</h2>", "", gr.HTML),
     "freeu_enabled": OptionInfo(False, "FreeU"),
diff --git a/modules/todo/__init__.py b/modules/todo/__init__.py
new file mode 100644
index 000000000..0caaaac78
--- /dev/null
+++ b/modules/todo/__init__.py
@@ -0,0 +1,42 @@
+from modules.todo.todo_utils import patch_attention_proc
+
+
+def apply_todo(model, p, method='todo'):
+    mp = p.height * p.width / 1024 / 1024
+
+    if mp < 1.0: # 512px
+        downsample_factor = 2
+        ratio = 0.38
+        downsample_factor_level_2 = 1
+        ratio_level_2 = 0.0
+    elif mp < 1.1: # 1024+
+        downsample_factor = 2
+        ratio = 0.75
+        downsample_factor_level_2 = 1
+        ratio_level_2 = 0.0
+    elif mp < 2.3:
+        downsample_factor = 3
+        ratio = 0.89
+        downsample_factor_level_2 = 1
+        ratio_level_2 = 0.0
+    elif mp < 8:
+        downsample_factor = 4
+        ratio = 0.9375
+        downsample_factor_level_2 = 1
+        ratio_level_2 = 0.0
+    else:
+        return
+    merge_method = "downsample" if method == "todo" else "similarity"
+    merge_tokens = "keys/values" if method == "todo" else "all"
+    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(model.unet, token_merge_args=token_merge_args)
diff --git a/modules/todo/todo_merge.py b/modules/todo/todo_merge.py
new file mode 100644
index 000000000..bfbff5621
--- /dev/null
+++ b/modules/todo/todo_merge.py
@@ -0,0 +1,382 @@
+from typing import Optional, Tuple, Callable
+import math
+import torch
+import torch.nn.functional as F
+from diffusers.models.attention_processor import Attention
+from diffusers.utils import USE_PEFT_BACKEND
+from diffusers.utils.import_utils import is_xformers_available
+
+
+if is_xformers_available():
+    import xformers
+    import xformers.ops
+    xformers_is_available = True
+else:
+    xformers_is_available = False
+
+
+if hasattr(F, "scaled_dot_product_attention"):
+    torch2_is_available = True
+else:
+    torch2_is_available = False
+
+
+def init_generator(device: torch.device, fallback: torch.Generator = None):
+    """
+    Forks the current default random generator given device.
+    """
+    print(f"init_generator device = {device}")
+    if device.type == "cpu":
+        return torch.Generator(device="cpu").set_state(torch.get_rng_state())
+    elif device.type == "cuda":
+        return torch.Generator(device=device).set_state(torch.cuda.get_rng_state())
+    elif device.type == "cuda":
+        return torch.Generator(device=device).set_state(torch.mps.get_rng_state())
+    else:
+        if fallback is None:
+            return init_generator(torch.device("cpu"))
+        else:
+            return fallback
+
+
+def do_nothing(x: torch.Tensor, mode: str = None): # pylint: disable=unused-argument
+    return x
+
+
+def mps_gather_workaround(input, dim, index): # pylint: disable=redefined-builtin
+    if input.shape[-1] == 1:
+        return torch.gather(
+            input.unsqueeze(-1),
+            dim - 1 if dim < 0 else dim,
+            index.unsqueeze(-1)
+        ).squeeze(-1)
+    else:
+        return torch.gather(input, dim, index)
+
+
+def up_or_downsample(item, cur_w, cur_h, new_w, new_h, method):
+    batch_size = item.shape[0]
+
+    item = item.reshape(batch_size, cur_h, cur_w, -1)
+    item = item.permute(0, 3, 1, 2)
+    df = cur_h // new_h
+    if method in "max_pool":
+        item = F.max_pool2d(item, kernel_size=df, stride=df, padding=0)
+    elif method in "avg_pool":
+        item = F.avg_pool2d(item, kernel_size=df, stride=df, padding=0)
+    else:
+        item = F.interpolate(item, size=(new_h, new_w), mode=method)
+    item = item.permute(0, 2, 3, 1)
+    item = item.reshape(batch_size, new_h * new_w, -1)
+
+    return item
+
+
+def compute_merge(x: torch.Tensor, tome_info):
+    original_h, original_w = tome_info["size"]
+    original_tokens = original_h * original_w
+    downsample = int(math.ceil(math.sqrt(original_tokens // x.shape[1])))
+    dim = x.shape[-1]
+    if dim == 320:
+        cur_level = "level_1"
+        downsample_factor = tome_info['args']['downsample_factor']
+        ratio = tome_info['args']['ratio']
+    elif dim == 640:
+        cur_level = "level_2"
+        downsample_factor = tome_info['args']['downsample_factor_level_2']
+        ratio = tome_info['args']['ratio_level_2']
+    else:
+        cur_level = "other"
+        downsample_factor = 1
+        ratio = 0.0
+
+    args = tome_info["args"]
+
+    cur_h, cur_w = original_h // downsample, original_w // downsample
+    new_h, new_w = cur_h // downsample_factor, cur_w // downsample_factor
+
+    if tome_info['timestep'] / 1000 > tome_info['args']['timestep_threshold_switch']:
+        merge_method = args["merge_method"]
+    else:
+        merge_method = args["secondary_merge_method"]
+
+    if cur_level != "other" and tome_info['timestep'] / 1000 > tome_info['args']['timestep_threshold_stop']:
+        if merge_method == "downsample" and downsample_factor > 1:
+            m = lambda x: up_or_downsample(x, cur_w, cur_h, new_w, new_h, args["downsample_method"]) # pylint: disable=unnecessary-lambda-assignment
+            u = lambda x: up_or_downsample(x, new_w, new_h, cur_w, cur_h, args["downsample_method"]) # pylint: disable=unnecessary-lambda-assignment
+        elif merge_method == "similarity" and ratio > 0.0:
+            w = int(math.ceil(original_w / downsample))
+            h = int(math.ceil(original_h / downsample))
+            r = int(x.shape[1] * ratio)
+
+            # Re-init the generator if it hasn't already been initialized or device has changed.
+            if args["generator"] is None:
+                args["generator"] = init_generator(x.device)
+            elif args["generator"].device != x.device:
+                args["generator"] = init_generator(x.device, fallback=args["generator"])
+
+            # If the batch size is odd, then it's not possible for prompted and unprompted images to be in the same
+            # batch, which causes artifacts with use_rand, so force it to be off.
+            use_rand = False if x.shape[0] % 2 == 1 else args["use_rand"]
+            m, u = bipartite_soft_matching_random2d(x, w, h, args["sx"], args["sy"], r,
+                                                    no_rand=not use_rand, generator=args["generator"])
+        else:
+            m, u = (do_nothing, do_nothing)
+    else:
+        m, u = (do_nothing, do_nothing)
+
+    merge_fn, unmerge_fn = (m, u)
+
+    return merge_fn, unmerge_fn
+
+
+def bipartite_soft_matching_random2d(metric: torch.Tensor,
+                                     w: int,
+                                     h: int,
+                                     sx: int,
+                                     sy: int,
+                                     r: int,
+                                     no_rand: bool = False,
+                                     generator: torch.Generator = None) -> Tuple[Callable, Callable]:
+    """
+    Partitions the tokens into src and dst and merges r tokens from src to dst.
+    Dst tokens are partitioned by choosing one randomy in each (sx, sy) region.
+
+    Args:
+     - metric [B, N, C]: metric to use for similarity
+     - w: image width in tokens
+     - h: image height in tokens
+     - sx: stride in the x dimension for dst, must divide w
+     - sy: stride in the y dimension for dst, must divide h
+     - r: number of tokens to remove (by merging)
+     - no_rand: if true, disable randomness (use top left corner only)
+     - rand_seed: if no_rand is false, and if not None, sets random seed.
+    """
+    B, N, _ = metric.shape
+
+    if r <= 0:
+        return do_nothing, do_nothing
+
+    with torch.no_grad():
+        hsy, wsx = h // sy, w // sx
+
+        # For each sy by sx kernel, randomly assign one token to be dst and the rest src
+        if no_rand:
+            rand_idx = torch.zeros(hsy, wsx, 1, device=metric.device, dtype=torch.int64)
+        else:
+            rand_idx = torch.randint(sy * sx, size=(hsy, wsx, 1), device=generator.device, generator=generator).to(
+                metric.device)
+
+        # The image might not divide sx and sy, so we need to work on a view of the top left if the idx buffer instead
+        idx_buffer_view = torch.zeros(hsy, wsx, sy * sx, device=metric.device, dtype=torch.int64)
+        idx_buffer_view.scatter_(dim=2, index=rand_idx, src=-torch.ones_like(rand_idx, dtype=rand_idx.dtype))
+        idx_buffer_view = idx_buffer_view.view(hsy, wsx, sy, sx).transpose(1, 2).reshape(hsy * sy, wsx * sx)
+
+        # Image is not divisible by sx or sy so we need to move it into a new buffer
+        if (hsy * sy) < h or (wsx * sx) < w:
+            idx_buffer = torch.zeros(h, w, device=metric.device, dtype=torch.int64)
+            idx_buffer[:(hsy * sy), :(wsx * sx)] = idx_buffer_view
+        else:
+            idx_buffer = idx_buffer_view
+
+        # We set dst tokens to be -1 and src to be 0, so an argsort gives us dst|src indices
+        rand_idx = idx_buffer.reshape(1, -1, 1).argsort(dim=1)
+
+        # We're finished with these
+        del idx_buffer, idx_buffer_view
+
+        # rand_idx is currently dst|src, so split them
+        num_dst = hsy * wsx
+        a_idx = rand_idx[:, num_dst:, :]  # src
+        b_idx = rand_idx[:, :num_dst, :]  # dst
+
+        def split(x):
+            C = x.shape[-1]
+            src = torch.gather(x, dim=1, index=a_idx.expand(B, N - num_dst, C))
+            dst = torch.gather(x, dim=1, index=b_idx.expand(B, num_dst, C))
+            return src, dst
+
+        # Cosine similarity between A and B
+        metric = metric / metric.norm(dim=-1, keepdim=True)
+        a, b = split(metric)
+        scores = a @ b.transpose(-1, -2)
+
+        # Can't reduce more than the # tokens in src
+        r = min(a.shape[1], r)
+
+        # Find the most similar greedily
+        node_max, node_idx = scores.max(dim=-1)
+        edge_idx = node_max.argsort(dim=-1, descending=True)[..., None]
+
+        unm_idx = edge_idx[..., r:, :]  # Unmerged Tokens
+        src_idx = edge_idx[..., :r, :]  # Merged Tokens
+        dst_idx = torch.gather(node_idx[..., None], dim=-2, index=src_idx)
+
+    def merge(x: torch.Tensor, mode="mean") -> torch.Tensor:
+        src, dst = split(x)
+        n, t1, c = src.shape
+
+        unm = torch.gather(src, dim=-2, index=unm_idx.expand(n, t1 - r, c))
+        src = torch.gather(src, dim=-2, index=src_idx.expand(n, r, c))
+        dst = dst.scatter_reduce(-2, dst_idx.expand(n, r, c), src, reduce=mode)
+
+        return torch.cat([unm, dst], dim=1)
+
+    def unmerge(x: torch.Tensor) -> torch.Tensor:
+        unm_len = unm_idx.shape[1]
+        unm, dst = x[..., :unm_len, :], x[..., unm_len:, :]
+        _, _, c = unm.shape
+
+        src = torch.gather(dst, dim=-2, index=dst_idx.expand(B, r, c))
+
+        # Combine back to the original shape
+        out = torch.zeros(B, N, c, device=x.device, dtype=x.dtype)
+        out.scatter_(dim=-2, index=b_idx.expand(B, num_dst, c), src=dst)
+        out.scatter_(dim=-2,
+                     index=torch.gather(a_idx.expand(B, a_idx.shape[1], 1), dim=1, index=unm_idx).expand(B, unm_len, c),
+                     src=unm)
+        out.scatter_(dim=-2,
+                     index=torch.gather(a_idx.expand(B, a_idx.shape[1], 1), dim=1, index=src_idx).expand(B, r, c),
+                     src=src)
+
+        return out
+
+    return merge, unmerge
+
+
+class TokenMergeAttentionProcessor:
+    def __init__(self):
+        # priortize torch2's flash attention, if not fall back to xformers then regular attention
+        if torch2_is_available:
+            self.attn_method = "torch2"
+        elif xformers_is_available:
+            self.attn_method = "xformers"
+        else:
+            self.attn_method = "regular"
+
+    def torch2_attention(self, attn, query, key, value, attention_mask, batch_size):
+        inner_dim=key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+
+        return hidden_states
+
+    def xformers_attention(self, attn, query, key, value, attention_mask, batch_size):
+        query = attn.head_to_batch_dim(query).contiguous()
+        key = attn.head_to_batch_dim(key).contiguous()
+        value = attn.head_to_batch_dim(value).contiguous()
+
+        if attention_mask is not None:
+            attention_mask = attention_mask.reshape(batch_size * attn.heads, -1, attention_mask.shape[-1])
+
+        hidden_states = xformers.ops.memory_efficient_attention(
+            query, key, value, attn_bias=attention_mask, scale=attn.scale
+        )
+
+        hidden_states = attn.batch_to_head_dim(hidden_states)
+
+        return hidden_states
+
+
+    def regular_attention(self, attn, query, key, value, attention_mask, batch_size):
+        query = attn.head_to_batch_dim(query)
+        key = attn.head_to_batch_dim(key)
+        value = attn.head_to_batch_dim(value)
+
+        if attention_mask is not None:
+            attention_mask = attention_mask.reshape(batch_size * attn.heads, -1, attention_mask.shape[-1])
+
+        attention_probs = attn.get_attention_scores(query, key, attention_mask)
+        hidden_states = torch.bmm(attention_probs, value)
+        hidden_states = attn.batch_to_head_dim(hidden_states)
+
+        return hidden_states
+
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        temb: Optional[torch.FloatTensor] = None,
+        scale: float = 1.0,
+    ) -> torch.FloatTensor:
+        residual = hidden_states
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        )
+
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+            # scaled_dot_product_attention expects attention_mask shape to be
+            # (batch, heads, source_length, target_length)
+            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+        args = () if USE_PEFT_BACKEND else (scale,)
+
+        if self._tome_info['args']['merge_tokens'] == "all": # pylint: disable=no-member
+            merge_fn, unmerge_fn = compute_merge(hidden_states, self._tome_info) # pylint: disable=no-member
+            hidden_states = merge_fn(hidden_states)
+
+        query = attn.to_q(hidden_states, *args)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+        if self._tome_info['args']['merge_tokens'] == "keys/values": # pylint: disable=no-member
+            merge_fn, _ = compute_merge(encoder_hidden_states, self._tome_info) # pylint: disable=no-member
+            encoder_hidden_states = merge_fn(encoder_hidden_states)
+
+        key = attn.to_k(encoder_hidden_states, *args)
+        value = attn.to_v(encoder_hidden_states, *args)
+
+        if self.attn_method == "torch2":
+            hidden_states = self.torch2_attention(attn, query, key, value, attention_mask, batch_size)
+        elif self.attn_method == "xformers":
+            hidden_states = self.xformers_attention(attn, query, key, value, attention_mask, batch_size)
+        else:
+            hidden_states = self.regular_attention(attn, query, key, value, attention_mask, batch_size)
+
+        hidden_states = hidden_states.to(query.dtype)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states, *args)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if self._tome_info['args']['merge_tokens'] == "all": # pylint: disable=no-member
+            hidden_states = unmerge_fn(hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        return hidden_states
diff --git a/modules/todo/todo_utils.py b/modules/todo/todo_utils.py
new file mode 100644
index 000000000..34a24bb82
--- /dev/null
+++ b/modules/todo/todo_utils.py
@@ -0,0 +1,77 @@
+import torch
+import torch.nn.functional as F
+from diffusers.utils.import_utils import is_xformers_available
+from diffusers.models.attention_processor import AttnProcessor2_0, AttnProcessor
+from modules.todo.todo_merge import TokenMergeAttentionProcessor
+
+
+xformers_is_available = is_xformers_available()
+torch2_is_available = hasattr(F, "scaled_dot_product_attention")
+
+
+def hook_tome_model(model: torch.nn.Module):
+    """ Adds a forward pre hook to get the image size. This hook can be removed with remove_patch. """
+
+    def hook(module, args):
+        module._tome_info["size"] = (args[0].shape[2], args[0].shape[3]) # pylint: disable=protected-access
+        module._tome_info["timestep"] = args[1].item() # pylint: disable=protected-access
+        return None
+
+    model._tome_info["hooks"].append(model.register_forward_pre_hook(hook)) # pylint: disable=protected-access
+
+def remove_tome_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 = { # pylint: disable=protected-access
+        "size": None,
+        "timestep": None,
+        "hooks": [],
+        "args": {
+            "ratio": token_merge_args.get("ratio", 0.5),  # ratio of tokens to merge
+            "sx": token_merge_args.get("sx", 2),  # stride x for sim calculation
+            "sy": token_merge_args.get("sy", 2),  # stride y for sim calculation
+            "use_rand": token_merge_args.get("use_rand", True),
+            "generator": None,
+            "merge_tokens": token_merge_args.get("merge_tokens", "keys/values"),  # ["all", "keys/values"]
+            "merge_method": token_merge_args.get("merge_method", "downsample"),  # ["none","similarity", "downsample"]
+            "downsample_method": token_merge_args.get("downsample_method", "nearest-exact"), # native torch interpolation methods ["nearest", "linear", "bilinear", "bicubic", "nearest-exact"]
+            "downsample_factor": token_merge_args.get("downsample_factor", 2),  # amount to downsample by
+            "timestep_threshold_switch": token_merge_args.get("timestep_threshold_switch", 0.2), # timestep to switch to secondary method, 0.2 means 20% steps remaining
+            "timestep_threshold_stop": token_merge_args.get("timestep_threshold_stop", 0.0), # timestep to stop merging, 0.0 means stop at 0 steps remaining
+            "secondary_merge_method": token_merge_args.get("secondary_merge_method", "similarity"), # ["none", "similarity", "downsample"]
+            "downsample_factor_level_2": token_merge_args.get("downsample_factor_level_2", 1), # amount to downsample by at the 2nd down block of unet
+            "ratio_level_2": token_merge_args.get("ratio_level_2", 0.5), # ratio of tokens to merge at the 2nd down block of unet
+        }
+    }
+    hook_tome_model(unet)
+    attn_modules = [module for name, module in unet.named_modules() if module.__class__.__name__ == 'BasicTransformerBlock']
+
+    for _i, module in enumerate(attn_modules):
+        module.attn1.processor = TokenMergeAttentionProcessor()
+        module.attn1.processor._tome_info = unet._tome_info # pylint: disable=protected-access
+
+
+def remove_patch(pipe: torch.nn.Module):
+    """ Removes a patch from a ToMe Diffusion module if it was already patched. """
+
+    # this will remove our custom class
+    if torch2_is_available:
+        for _n, m in pipe.unet.named_modules():
+            if hasattr(m, "processor"):
+                m.processor = AttnProcessor2_0()
+
+    elif xformers_is_available:
+        pipe.enable_xformers_memory_efficient_attention()
+
+    else:
+        for _n, m in pipe.unet.named_modules():
+            if hasattr(m, "processor"):
+                m.processor = AttnProcessor()
diff --git a/scripts/xyz_grid.py b/scripts/xyz_grid.py
index ad63d73d5..444a6406c 100644
--- a/scripts/xyz_grid.py
+++ b/scripts/xyz_grid.py
@@ -267,8 +267,8 @@ axis_options = [
     AxisOption("[HDR] Maximize boundary", float, apply_field("hdr_max_boundry")),
     AxisOption("[HDR] Tint Color Hex", str, apply_field("hdr_color_picker")),
     AxisOption("[HDR] Tint Ratio", float, apply_field("hdr_tint_ratio")),
-    AxisOption("[ToMe] Token merging ratio (txt2img)", float, apply_override('token_merging_ratio')),
-    AxisOption("[ToMe] Token merging ratio (hires)", float, apply_override('token_merging_ratio_hr')),
+    AxisOption("[Token Merging] ToMe ratio", float, apply_setting('tome_ratio')),
+    AxisOption("[Token Merging] ToDo ratio", float, apply_setting('todo_ratio')),
     AxisOption("[FreeU] 1st stage backbone factor", float, apply_setting('freeu_b1')),
     AxisOption("[FreeU] 2nd stage backbone factor", float, apply_setting('freeu_b2')),
     AxisOption("[FreeU] 1st stage skip factor", float, apply_setting('freeu_s1')),
@@ -388,8 +388,8 @@ def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend
 class SharedSettingsStackHelper(object):
     vae = None
     schedulers_solver_order = None
-    token_merging_ratio_hr = None
-    token_merging_ratio = None
+    tome_ratio = None
+    todo_ratio = None
     sd_model_checkpoint = None
     sd_model_dict = None
     sd_vae_checkpoint = None
@@ -398,8 +398,8 @@ class SharedSettingsStackHelper(object):
         #Save overridden settings so they can be restored later.
         self.vae = shared.opts.sd_vae
         self.schedulers_solver_order = shared.opts.schedulers_solver_order
-        self.token_merging_ratio_hr = shared.opts.token_merging_ratio_hr
-        self.token_merging_ratio = shared.opts.token_merging_ratio
+        self.tome_ratio = shared.opts.tome_ratio
+        self.todo_ratio = shared.opts.todo_ratio
         self.sd_model_checkpoint = shared.opts.sd_model_checkpoint
         self.sd_model_dict = shared.opts.sd_model_dict
         self.sd_vae_checkpoint = shared.opts.sd_vae
@@ -408,8 +408,8 @@ class SharedSettingsStackHelper(object):
         #Restore overriden settings after plot generation.
         shared.opts.data["sd_vae"] = self.vae
         shared.opts.data["schedulers_solver_order"] = self.schedulers_solver_order
-        shared.opts.data["token_merging_ratio_hr"] = self.token_merging_ratio_hr
-        shared.opts.data["token_merging_ratio"] = self.token_merging_ratio
+        shared.opts.data["tome_ratio"] = self.tome_ratio
+        shared.opts.data["todo_ratio"] = self.todo_ratio
         if self.sd_model_dict != shared.opts.sd_model_dict:
             shared.opts.data["sd_model_dict"] = self.sd_model_dict
         if self.sd_model_checkpoint != shared.opts.sd_model_checkpoint: