From 56680cd84ab68a283772cf697f8a72408a3f4167 Mon Sep 17 00:00:00 2001
From: papuSpartan <macabeg@icloud.com>
Date: Sat, 1 Apr 2023 02:07:08 -0500
Subject: [PATCH 01/19] first

---
 launch.py            | 6 ++++++
 modules/cmd_args.py  | 4 ++++
 modules/sd_models.py | 8 ++++++++
 3 files changed, 18 insertions(+)

diff --git a/launch.py b/launch.py
index 68e08114d..94bba5ca0 100644
--- a/launch.py
+++ b/launch.py
@@ -238,12 +238,14 @@ def prepare_environment():
     k_diffusion_repo = os.environ.get('K_DIFFUSION_REPO', 'https://github.com/crowsonkb/k-diffusion.git')
     codeformer_repo = os.environ.get('CODEFORMER_REPO', 'https://github.com/sczhou/CodeFormer.git')
     blip_repo = os.environ.get('BLIP_REPO', 'https://github.com/salesforce/BLIP.git')
+    tomesd_repo = os.environ.get('TOMESD_REPO', 'https://github.com/dbolya/tomesd.git')
 
     stable_diffusion_commit_hash = os.environ.get('STABLE_DIFFUSION_COMMIT_HASH', "cf1d67a6fd5ea1aa600c4df58e5b47da45f6bdbf")
     taming_transformers_commit_hash = os.environ.get('TAMING_TRANSFORMERS_COMMIT_HASH', "24268930bf1dce879235a7fddd0b2355b84d7ea6")
     k_diffusion_commit_hash = os.environ.get('K_DIFFUSION_COMMIT_HASH', "5b3af030dd83e0297272d861c19477735d0317ec")
     codeformer_commit_hash = os.environ.get('CODEFORMER_COMMIT_HASH', "c5b4593074ba6214284d6acd5f1719b6c5d739af")
     blip_commit_hash = os.environ.get('BLIP_COMMIT_HASH', "48211a1594f1321b00f14c9f7a5b4813144b2fb9")
+    tomesd_commit_hash = os.environ.get('TOMESD_COMMIT_HASH', "4f936c257e10848e0399fc6d0484a1761812092a")
 
     if not args.skip_python_version_check:
         check_python_version()
@@ -280,6 +282,10 @@ def prepare_environment():
         elif platform.system() == "Linux":
             run_pip(f"install {xformers_package}", "xformers")
 
+    if (not is_installed("tomesd") or args.reinstall_tomesd) and args.token_merging:
+        git_clone(tomesd_repo, repo_dir('tomesd'), "tomesd", tomesd_commit_hash)
+        run_pip(f"install {repo_dir('tomesd')}")
+
     if not is_installed("pyngrok") and args.ngrok:
         run_pip("install pyngrok", "ngrok")
 
diff --git a/modules/cmd_args.py b/modules/cmd_args.py
index 81c0b82a3..4314f97b1 100644
--- a/modules/cmd_args.py
+++ b/modules/cmd_args.py
@@ -101,3 +101,7 @@ parser.add_argument("--no-gradio-queue", action='store_true', help="Disables gra
 parser.add_argument("--skip-version-check", action='store_true', help="Do not check versions of torch and xformers")
 parser.add_argument("--no-hashing", action='store_true', help="disable sha256 hashing of checkpoints to help loading performance", default=False)
 parser.add_argument("--no-download-sd-model", action='store_true', help="don't download SD1.5 model even if no model is found in --ckpt-dir", default=False)
+
+# token merging / tomesd
+parser.add_argument("--token-merging", action='store_true', help="Provides generation speedup by merging redundant tokens. (compatible with --xformers)", default=False)
+parser.add_argument("--token-merging-ratio", type=float, help="Adjusts ratio of merged to untouched tokens. Range: (0.0-1.0], Defaults to 0.5", default=0.5)
diff --git a/modules/sd_models.py b/modules/sd_models.py
index 6ea874dfc..0b74aa0f0 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -9,6 +9,7 @@ from omegaconf import OmegaConf
 from os import mkdir
 from urllib import request
 import ldm.modules.midas as midas
+import tomesd
 
 from ldm.util import instantiate_from_config
 
@@ -430,6 +431,13 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None, time_taken_
     try:
         with sd_disable_initialization.DisableInitialization(disable_clip=clip_is_included_into_sd):
             sd_model = instantiate_from_config(sd_config.model)
+
+            if shared.cmd_opts.token_merging:
+                ratio = shared.cmd_opts.token_merging_ratio
+
+                tomesd.apply_patch(sd_model, ratio=ratio)
+                print(f"Model accelerated using {(ratio * 100)}% token merging via tomesd.")
+                timer.record("token merging")
     except Exception as e:
         pass
 

From ef8c0440512f2fae9ceeb977b73f881c0afbb90a Mon Sep 17 00:00:00 2001
From: papuSpartan <macabeg@icloud.com>
Date: Sat, 1 Apr 2023 03:21:23 -0500
Subject: [PATCH 02/19] forgot to add reinstall arg back earlier since args
 moved out of shared

---
 modules/cmd_args.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/modules/cmd_args.py b/modules/cmd_args.py
index 4314f97b1..80df94658 100644
--- a/modules/cmd_args.py
+++ b/modules/cmd_args.py
@@ -105,3 +105,4 @@ parser.add_argument("--no-download-sd-model", action='store_true', help="don't d
 # token merging / tomesd
 parser.add_argument("--token-merging", action='store_true', help="Provides generation speedup by merging redundant tokens. (compatible with --xformers)", default=False)
 parser.add_argument("--token-merging-ratio", type=float, help="Adjusts ratio of merged to untouched tokens. Range: (0.0-1.0], Defaults to 0.5", default=0.5)
+parser.add_argument("--reinstall-tomesd", action='store_true', help="Reinstalls tomesd", default=False)

From 26ab018253cb078630fcdde47dbaee85f2466145 Mon Sep 17 00:00:00 2001
From: papuSpartan <macabeg@icloud.com>
Date: Sat, 1 Apr 2023 03:31:22 -0500
Subject: [PATCH 03/19] delay import

---
 modules/sd_models.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/modules/sd_models.py b/modules/sd_models.py
index 0b74aa0f0..2c05ec17b 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -9,7 +9,6 @@ from omegaconf import OmegaConf
 from os import mkdir
 from urllib import request
 import ldm.modules.midas as midas
-import tomesd
 
 from ldm.util import instantiate_from_config
 
@@ -433,6 +432,7 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None, time_taken_
             sd_model = instantiate_from_config(sd_config.model)
 
             if shared.cmd_opts.token_merging:
+                import tomesd
                 ratio = shared.cmd_opts.token_merging_ratio
 
                 tomesd.apply_patch(sd_model, ratio=ratio)

From 8c88bf40060c86ba508646c6d7ddc21e389be846 Mon Sep 17 00:00:00 2001
From: papuSpartan <macabeg@icloud.com>
Date: Sat, 1 Apr 2023 14:12:12 -0500
Subject: [PATCH 04/19] use pypi package for tomesd intead of manually cloning
 repo

---
 launch.py           | 7 ++-----
 modules/cmd_args.py | 1 -
 2 files changed, 2 insertions(+), 6 deletions(-)

diff --git a/launch.py b/launch.py
index 94bba5ca0..846c4c200 100644
--- a/launch.py
+++ b/launch.py
@@ -238,14 +238,12 @@ def prepare_environment():
     k_diffusion_repo = os.environ.get('K_DIFFUSION_REPO', 'https://github.com/crowsonkb/k-diffusion.git')
     codeformer_repo = os.environ.get('CODEFORMER_REPO', 'https://github.com/sczhou/CodeFormer.git')
     blip_repo = os.environ.get('BLIP_REPO', 'https://github.com/salesforce/BLIP.git')
-    tomesd_repo = os.environ.get('TOMESD_REPO', 'https://github.com/dbolya/tomesd.git')
 
     stable_diffusion_commit_hash = os.environ.get('STABLE_DIFFUSION_COMMIT_HASH', "cf1d67a6fd5ea1aa600c4df58e5b47da45f6bdbf")
     taming_transformers_commit_hash = os.environ.get('TAMING_TRANSFORMERS_COMMIT_HASH', "24268930bf1dce879235a7fddd0b2355b84d7ea6")
     k_diffusion_commit_hash = os.environ.get('K_DIFFUSION_COMMIT_HASH', "5b3af030dd83e0297272d861c19477735d0317ec")
     codeformer_commit_hash = os.environ.get('CODEFORMER_COMMIT_HASH', "c5b4593074ba6214284d6acd5f1719b6c5d739af")
     blip_commit_hash = os.environ.get('BLIP_COMMIT_HASH', "48211a1594f1321b00f14c9f7a5b4813144b2fb9")
-    tomesd_commit_hash = os.environ.get('TOMESD_COMMIT_HASH', "4f936c257e10848e0399fc6d0484a1761812092a")
 
     if not args.skip_python_version_check:
         check_python_version()
@@ -282,9 +280,8 @@ def prepare_environment():
         elif platform.system() == "Linux":
             run_pip(f"install {xformers_package}", "xformers")
 
-    if (not is_installed("tomesd") or args.reinstall_tomesd) and args.token_merging:
-        git_clone(tomesd_repo, repo_dir('tomesd'), "tomesd", tomesd_commit_hash)
-        run_pip(f"install {repo_dir('tomesd')}")
+    if not is_installed("tomesd") and args.token_merging:
+        run_pip(f"install tomesd")
 
     if not is_installed("pyngrok") and args.ngrok:
         run_pip("install pyngrok", "ngrok")
diff --git a/modules/cmd_args.py b/modules/cmd_args.py
index 80df94658..4314f97b1 100644
--- a/modules/cmd_args.py
+++ b/modules/cmd_args.py
@@ -105,4 +105,3 @@ parser.add_argument("--no-download-sd-model", action='store_true', help="don't d
 # token merging / tomesd
 parser.add_argument("--token-merging", action='store_true', help="Provides generation speedup by merging redundant tokens. (compatible with --xformers)", default=False)
 parser.add_argument("--token-merging-ratio", type=float, help="Adjusts ratio of merged to untouched tokens. Range: (0.0-1.0], Defaults to 0.5", default=0.5)
-parser.add_argument("--reinstall-tomesd", action='store_true', help="Reinstalls tomesd", default=False)

From a609bd56b4206460d1df3c3022025fc78b66718f Mon Sep 17 00:00:00 2001
From: papuSpartan <macabeg@icloud.com>
Date: Sat, 1 Apr 2023 22:18:35 -0500
Subject: [PATCH 05/19] Transition to using settings through UI instead of cmd
 line args. Added feature to only apply to hr-fix. Install package using
 requirements_versions.txt

---
 launch.py                 |  3 ---
 modules/processing.py     | 35 +++++++++++++++++++++++++++++++
 modules/sd_models.py      |  7 -------
 modules/shared.py         | 44 +++++++++++++++++++++++++++++++++++++++
 requirements_versions.txt |  1 +
 5 files changed, 80 insertions(+), 10 deletions(-)

diff --git a/launch.py b/launch.py
index 846c4c200..68e08114d 100644
--- a/launch.py
+++ b/launch.py
@@ -280,9 +280,6 @@ def prepare_environment():
         elif platform.system() == "Linux":
             run_pip(f"install {xformers_package}", "xformers")
 
-    if not is_installed("tomesd") and args.token_merging:
-        run_pip(f"install tomesd")
-
     if not is_installed("pyngrok") and args.ngrok:
         run_pip("install pyngrok", "ngrok")
 
diff --git a/modules/processing.py b/modules/processing.py
index 6d9c6a8de..e115aadd1 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -29,6 +29,7 @@ from ldm.models.diffusion.ddpm import LatentDepth2ImageDiffusion
 
 from einops import repeat, rearrange
 from blendmodes.blend import blendLayers, BlendType
+import tomesd
 
 # some of those options should not be changed at all because they would break the model, so I removed them from options.
 opt_C = 4
@@ -500,9 +501,28 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
             if k == 'sd_vae':
                 sd_vae.reload_vae_weights()
 
+        if opts.token_merging:
+
+            if p.hr_second_pass_steps < 1 and not opts.token_merging_hr_only:
+                tomesd.apply_patch(
+                    p.sd_model,
+                    ratio=opts.token_merging_ratio,
+                    max_downsample=opts.token_merging_maximum_down_sampling,
+                    sx=opts.token_merging_stride_x,
+                    sy=opts.token_merging_stride_y,
+                    use_rand=opts.token_merging_random,
+                    merge_attn=opts.token_merging_merge_attention,
+                    merge_crossattn=opts.token_merging_merge_cross_attention,
+                    merge_mlp=opts.token_merging_merge_mlp
+                )
+
         res = process_images_inner(p)
 
     finally:
+        # undo model optimizations made by tomesd
+        if opts.token_merging:
+            tomesd.remove_patch(p.sd_model)
+
         # restore opts to original state
         if p.override_settings_restore_afterwards:
             for k, v in stored_opts.items():
@@ -938,6 +958,21 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
         x = None
         devices.torch_gc()
 
+        # apply token merging optimizations from tomesd for high-res pass
+        # check if hr_only so we don't redundantly apply patch
+        if opts.token_merging and opts.token_merging_hr_only:
+            tomesd.apply_patch(
+                self.sd_model,
+                ratio=opts.token_merging_ratio,
+                max_downsample=opts.token_merging_maximum_down_sampling,
+                sx=opts.token_merging_stride_x,
+                sy=opts.token_merging_stride_y,
+                use_rand=opts.token_merging_random,
+                merge_attn=opts.token_merging_merge_attention,
+                merge_crossattn=opts.token_merging_merge_cross_attention,
+                merge_mlp=opts.token_merging_merge_mlp
+            )
+
         samples = self.sampler.sample_img2img(self, samples, noise, conditioning, unconditional_conditioning, steps=self.hr_second_pass_steps or self.steps, image_conditioning=image_conditioning)
 
         return samples
diff --git a/modules/sd_models.py b/modules/sd_models.py
index 2c05ec17b..87c49b83c 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -431,13 +431,6 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None, time_taken_
         with sd_disable_initialization.DisableInitialization(disable_clip=clip_is_included_into_sd):
             sd_model = instantiate_from_config(sd_config.model)
 
-            if shared.cmd_opts.token_merging:
-                import tomesd
-                ratio = shared.cmd_opts.token_merging_ratio
-
-                tomesd.apply_patch(sd_model, ratio=ratio)
-                print(f"Model accelerated using {(ratio * 100)}% token merging via tomesd.")
-                timer.record("token merging")
     except Exception as e:
         pass
 
diff --git a/modules/shared.py b/modules/shared.py
index 5fd0eecbd..d7379e24f 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -427,6 +427,50 @@ options_templates.update(options_section((None, "Hidden options"), {
     "sd_checkpoint_hash": OptionInfo("", "SHA256 hash of the current checkpoint"),
 }))
 
+options_templates.update(options_section(('token_merging', 'Token Merging'), {
+    "token_merging": OptionInfo(
+        False, "Enable redundant token merging via tomesd. (currently incompatible with controlnet extension)",
+        gr.Checkbox
+    ),
+    "token_merging_ratio": OptionInfo(
+        0.5, "Merging Ratio",
+        gr.Slider, {"minimum": 0, "maximum": 0.9, "step": 0.1}
+    ),
+    "token_merging_hr_only": OptionInfo(
+        True, "Apply only to high-res fix pass. Disabling can yield a ~20-35% speedup on contemporary resolutions.",
+        gr.Checkbox
+    ),
+    # More advanced/niche settings:
+    "token_merging_random": OptionInfo(
+        True, "Use random perturbations - Disabling might help with certain samplers",
+        gr.Checkbox
+    ),
+    "token_merging_merge_attention": OptionInfo(
+        True, "Merge attention",
+        gr.Checkbox
+    ),
+     "token_merging_merge_cross_attention": OptionInfo(
+        False, "Merge cross attention",
+        gr.Checkbox
+    ),
+    "token_merging_merge_mlp": OptionInfo(
+        False, "Merge mlp",
+        gr.Checkbox
+    ),
+    "token_merging_maximum_down_sampling": OptionInfo(
+        1, "Maximum down sampling",
+        gr.Dropdown, lambda: {"choices": ["1", "2", "4", "8"]}
+    ),
+    "token_merging_stride_x": OptionInfo(
+        2, "Stride - X",
+        gr.Slider, {"minimum": 2, "maximum": 8, "step": 2}
+    ),
+    "token_merging_stride_y": OptionInfo(
+        2, "Stride - Y",
+        gr.Slider, {"minimum": 2, "maximum": 8, "step": 2}
+    )
+}))
+
 options_templates.update()
 
 
diff --git a/requirements_versions.txt b/requirements_versions.txt
index df65431a3..045230ab9 100644
--- a/requirements_versions.txt
+++ b/requirements_versions.txt
@@ -28,3 +28,4 @@ torchsde==0.2.5
 safetensors==0.3.0
 httpcore<=0.15
 fastapi==0.94.0
+tomesd>=0.1
\ No newline at end of file

From c707b7df95a61b66a05be94e805e1be9a432e294 Mon Sep 17 00:00:00 2001
From: papuSpartan <macabeg@icloud.com>
Date: Sat, 1 Apr 2023 23:47:10 -0500
Subject: [PATCH 06/19] remove excess condition

---
 modules/processing.py | 29 +++++++++++++++--------------
 1 file changed, 15 insertions(+), 14 deletions(-)

diff --git a/modules/processing.py b/modules/processing.py
index e115aadd1..557355725 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -501,26 +501,26 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
             if k == 'sd_vae':
                 sd_vae.reload_vae_weights()
 
-        if opts.token_merging:
-
-            if p.hr_second_pass_steps < 1 and not opts.token_merging_hr_only:
-                tomesd.apply_patch(
-                    p.sd_model,
-                    ratio=opts.token_merging_ratio,
-                    max_downsample=opts.token_merging_maximum_down_sampling,
-                    sx=opts.token_merging_stride_x,
-                    sy=opts.token_merging_stride_y,
-                    use_rand=opts.token_merging_random,
-                    merge_attn=opts.token_merging_merge_attention,
-                    merge_crossattn=opts.token_merging_merge_cross_attention,
-                    merge_mlp=opts.token_merging_merge_mlp
-                )
+        if opts.token_merging and not opts.token_merging_hr_only:
+            print("applying token merging to all passes")
+            tomesd.apply_patch(
+                p.sd_model,
+                ratio=opts.token_merging_ratio,
+                max_downsample=opts.token_merging_maximum_down_sampling,
+                sx=opts.token_merging_stride_x,
+                sy=opts.token_merging_stride_y,
+                use_rand=opts.token_merging_random,
+                merge_attn=opts.token_merging_merge_attention,
+                merge_crossattn=opts.token_merging_merge_cross_attention,
+                merge_mlp=opts.token_merging_merge_mlp
+            )
 
         res = process_images_inner(p)
 
     finally:
         # undo model optimizations made by tomesd
         if opts.token_merging:
+            print('removing token merging model optimizations')
             tomesd.remove_patch(p.sd_model)
 
         # restore opts to original state
@@ -961,6 +961,7 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
         # apply token merging optimizations from tomesd for high-res pass
         # check if hr_only so we don't redundantly apply patch
         if opts.token_merging and opts.token_merging_hr_only:
+            print("applying token merging for high-res pass")
             tomesd.apply_patch(
                 self.sd_model,
                 ratio=opts.token_merging_ratio,

From 5c8e53d5e98da0eabf384318955c57842d612c07 Mon Sep 17 00:00:00 2001
From: papuSpartan <macabeg@icloud.com>
Date: Tue, 4 Apr 2023 02:26:44 -0500
Subject: [PATCH 07/19] Allow different merge ratios to be used for each pass.
 Make toggle cmd flag work again. Remove ratio flag. Remove warning about
 controlnet being incompatible

---
 modules/cmd_args.py   |  3 +--
 modules/processing.py | 44 +++++++++++++++----------------------------
 modules/sd_models.py  | 29 +++++++++++++++++++++++++++-
 modules/shared.py     |  6 +++++-
 4 files changed, 49 insertions(+), 33 deletions(-)

diff --git a/modules/cmd_args.py b/modules/cmd_args.py
index 4314f97b1..8e5a7fabf 100644
--- a/modules/cmd_args.py
+++ b/modules/cmd_args.py
@@ -103,5 +103,4 @@ parser.add_argument("--no-hashing", action='store_true', help="disable sha256 ha
 parser.add_argument("--no-download-sd-model", action='store_true', help="don't download SD1.5 model even if no model is found in --ckpt-dir", default=False)
 
 # token merging / tomesd
-parser.add_argument("--token-merging", action='store_true', help="Provides generation speedup by merging redundant tokens. (compatible with --xformers)", default=False)
-parser.add_argument("--token-merging-ratio", type=float, help="Adjusts ratio of merged to untouched tokens. Range: (0.0-1.0], Defaults to 0.5", default=0.5)
+parser.add_argument("--token-merging", action='store_true', help="Provides speed and memory improvements by merging redundant tokens. This has a more pronounced effect on higher resolutions.", default=False)
diff --git a/modules/processing.py b/modules/processing.py
index 557355725..670a7a28a 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -501,26 +501,16 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
             if k == 'sd_vae':
                 sd_vae.reload_vae_weights()
 
-        if opts.token_merging and not opts.token_merging_hr_only:
-            print("applying token merging to all passes")
-            tomesd.apply_patch(
-                p.sd_model,
-                ratio=opts.token_merging_ratio,
-                max_downsample=opts.token_merging_maximum_down_sampling,
-                sx=opts.token_merging_stride_x,
-                sy=opts.token_merging_stride_y,
-                use_rand=opts.token_merging_random,
-                merge_attn=opts.token_merging_merge_attention,
-                merge_crossattn=opts.token_merging_merge_cross_attention,
-                merge_mlp=opts.token_merging_merge_mlp
-            )
+        if (opts.token_merging or cmd_opts.token_merging) and not opts.token_merging_hr_only:
+            print("\nApplying token merging\n")
+            sd_models.apply_token_merging(sd_model=p.sd_model, hr=False)
 
         res = process_images_inner(p)
 
     finally:
         # undo model optimizations made by tomesd
-        if opts.token_merging:
-            print('removing token merging model optimizations')
+        if opts.token_merging or cmd_opts.token_merging:
+            print('\nRemoving token merging model optimizations\n')
             tomesd.remove_patch(p.sd_model)
 
         # restore opts to original state
@@ -959,20 +949,16 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
         devices.torch_gc()
 
         # apply token merging optimizations from tomesd for high-res pass
-        # check if hr_only so we don't redundantly apply patch
-        if opts.token_merging and opts.token_merging_hr_only:
-            print("applying token merging for high-res pass")
-            tomesd.apply_patch(
-                self.sd_model,
-                ratio=opts.token_merging_ratio,
-                max_downsample=opts.token_merging_maximum_down_sampling,
-                sx=opts.token_merging_stride_x,
-                sy=opts.token_merging_stride_y,
-                use_rand=opts.token_merging_random,
-                merge_attn=opts.token_merging_merge_attention,
-                merge_crossattn=opts.token_merging_merge_cross_attention,
-                merge_mlp=opts.token_merging_merge_mlp
-            )
+        # check if hr_only so we are not redundantly patching
+        if (cmd_opts.token_merging or opts.token_merging) and (opts.token_merging_hr_only or opts.token_merging_ratio_hr != opts.token_merging_ratio):
+            # case where user wants to use separate merge ratios
+            if not opts.token_merging_hr_only:
+                # clean patch done by first pass. (clobbering the first patch might be fine? this might be excessive)
+                print('Temporarily reverting token merging optimizations in preparation for next pass')
+                tomesd.remove_patch(self.sd_model)
+
+            print("\nApplying token merging for high-res pass\n")
+            sd_models.apply_token_merging(sd_model=self.sd_model, hr=True)
 
         samples = self.sampler.sample_img2img(self, samples, noise, conditioning, unconditional_conditioning, steps=self.hr_second_pass_steps or self.steps, image_conditioning=image_conditioning)
 
diff --git a/modules/sd_models.py b/modules/sd_models.py
index 87c49b83c..696a2333d 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -16,6 +16,7 @@ from modules import paths, shared, modelloader, devices, script_callbacks, sd_va
 from modules.paths import models_path
 from modules.sd_hijack_inpainting import do_inpainting_hijack
 from modules.timer import Timer
+import tomesd
 
 model_dir = "Stable-diffusion"
 model_path = os.path.abspath(os.path.join(paths.models_path, model_dir))
@@ -545,4 +546,30 @@ def unload_model_weights(sd_model=None, info=None):
 
     print(f"Unloaded weights {timer.summary()}.")
 
-    return sd_model
\ No newline at end of file
+    return sd_model
+
+
+def apply_token_merging(sd_model, hr: bool):
+    """
+    Applies speed and memory optimizations from tomesd.
+
+    Args:
+        hr (bool): True if called in the context of a high-res pass
+    """
+
+    ratio = shared.opts.token_merging_ratio
+    if hr:
+        ratio = shared.opts.token_merging_ratio_hr
+        print("effective hr pass merge ratio is "+str(ratio))
+
+    tomesd.apply_patch(
+        sd_model,
+        ratio=ratio,
+        max_downsample=shared.opts.token_merging_maximum_down_sampling,
+        sx=shared.opts.token_merging_stride_x,
+        sy=shared.opts.token_merging_stride_y,
+        use_rand=shared.opts.token_merging_random,
+        merge_attn=shared.opts.token_merging_merge_attention,
+        merge_crossattn=shared.opts.token_merging_merge_cross_attention,
+        merge_mlp=shared.opts.token_merging_merge_mlp
+    )
diff --git a/modules/shared.py b/modules/shared.py
index d7379e24f..c7572e980 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -429,7 +429,7 @@ options_templates.update(options_section((None, "Hidden options"), {
 
 options_templates.update(options_section(('token_merging', 'Token Merging'), {
     "token_merging": OptionInfo(
-        False, "Enable redundant token merging via tomesd. (currently incompatible with controlnet extension)",
+        0.5, "Enable redundant token merging via tomesd. This can provide significant speed and memory improvements.",
         gr.Checkbox
     ),
     "token_merging_ratio": OptionInfo(
@@ -440,6 +440,10 @@ options_templates.update(options_section(('token_merging', 'Token Merging'), {
         True, "Apply only to high-res fix pass. Disabling can yield a ~20-35% speedup on contemporary resolutions.",
         gr.Checkbox
     ),
+    "token_merging_ratio_hr": OptionInfo(
+        0.5, "Merging Ratio (high-res pass) - If 'Apply only to high-res' is enabled, this will always be the ratio used.",
+        gr.Slider, {"minimum": 0, "maximum": 0.9, "step": 0.1}
+    ),
     # More advanced/niche settings:
     "token_merging_random": OptionInfo(
         True, "Use random perturbations - Disabling might help with certain samplers",

From ab195ab0da78aa9aa6304948ae2c76e45dab04eb Mon Sep 17 00:00:00 2001
From: papuSpartan <macabeg@icloud.com>
Date: Tue, 4 Apr 2023 02:31:57 -0500
Subject: [PATCH 08/19] bump tomesd package version

---
 requirements_versions.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/requirements_versions.txt b/requirements_versions.txt
index 045230ab9..035227152 100644
--- a/requirements_versions.txt
+++ b/requirements_versions.txt
@@ -28,4 +28,4 @@ torchsde==0.2.5
 safetensors==0.3.0
 httpcore<=0.15
 fastapi==0.94.0
-tomesd>=0.1
\ No newline at end of file
+tomesd>=0.1.1
\ No newline at end of file

From cf5a5773bfd1ca6a3c35232881661ec1ff4b67d7 Mon Sep 17 00:00:00 2001
From: papuSpartan <macabeg@icloud.com>
Date: Tue, 4 Apr 2023 02:39:13 -0500
Subject: [PATCH 09/19] :p

---
 modules/shared.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/modules/shared.py b/modules/shared.py
index c7572e980..568acdc4c 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -429,7 +429,7 @@ options_templates.update(options_section((None, "Hidden options"), {
 
 options_templates.update(options_section(('token_merging', 'Token Merging'), {
     "token_merging": OptionInfo(
-        0.5, "Enable redundant token merging via tomesd. This can provide significant speed and memory improvements.",
+        False, "Enable redundant token merging via tomesd. This can provide significant speed and memory improvements.",
         gr.Checkbox
     ),
     "token_merging_ratio": OptionInfo(

From 1c1106260300ca3956d9619875e28278b148adab Mon Sep 17 00:00:00 2001
From: papuSpartan <macabeg@icloud.com>
Date: Mon, 10 Apr 2023 03:37:15 -0500
Subject: [PATCH 10/19] add token merging options to infotext when necessary.
 Bump tomesd version

---
 modules/generation_parameters_copypaste.py | 37 ++++++++++++++++++++++
 modules/processing.py                      | 22 ++++++++++---
 modules/shared.py                          |  2 +-
 requirements_versions.txt                  |  2 +-
 4 files changed, 57 insertions(+), 6 deletions(-)

diff --git a/modules/generation_parameters_copypaste.py b/modules/generation_parameters_copypaste.py
index 6df76858f..a7ede5344 100644
--- a/modules/generation_parameters_copypaste.py
+++ b/modules/generation_parameters_copypaste.py
@@ -282,6 +282,32 @@ Steps: 20, Sampler: Euler a, CFG scale: 7, Seed: 965400086, Size: 512x512, Model
         res["Hires resize-1"] = 0
         res["Hires resize-2"] = 0
 
+    # Infer additional override settings for token merging
+    print("inferring settings for tomesd")
+    token_merging_ratio = res.get("Token merging ratio", None)
+    token_merging_ratio_hr = res.get("Token merging ratio hr", None)
+
+    if token_merging_ratio is not None or token_merging_ratio_hr is not None:
+        res["Token merging"] = 'True'
+
+        if token_merging_ratio is None:
+            res["Token merging hr only"] = 'True'
+        else:
+            res["Token merging hr only"] = 'False'
+
+        if res.get("Token merging random", None) is None:
+            res["Token merging random"] = 'False'
+        if res.get("Token merging merge attention", None) is None:
+            res["Token merging merge attention"] = 'True'
+        if res.get("Token merging merge cross attention", None) is None:
+            res["Token merging merge cross attention"] = 'False'
+        if res.get("Token merging merge mlp", None) is None:
+            res["Token merging merge mlp"] = 'False'
+        if res.get("Token merging stride x", None) is None:
+            res["Token merging stride x"] = '2'
+        if res.get("Token merging stride y", None) is None:
+            res["Token merging stride y"] = '2'
+
     restore_old_hires_fix_params(res)
 
     return res
@@ -304,6 +330,17 @@ infotext_to_setting_name_mapping = [
     ('UniPC skip type', 'uni_pc_skip_type'),
     ('UniPC order', 'uni_pc_order'),
     ('UniPC lower order final', 'uni_pc_lower_order_final'),
+    ('Token merging', 'token_merging'),
+    ('Token merging ratio', 'token_merging_ratio'),
+    ('Token merging hr only', 'token_merging_hr_only'),
+    ('Token merging ratio hr', 'token_merging_ratio_hr'),
+    ('Token merging random', 'token_merging_random'),
+    ('Token merging merge attention', 'token_merging_merge_attention'),
+    ('Token merging merge cross attention', 'token_merging_merge_cross_attention'),
+    ('Token merging merge mlp', 'token_merging_merge_mlp'),
+    ('Token merging maximum downsampling', 'token_merging_maximum_downsampling'),
+    ('Token merging stride x', 'token_merging_stride_x'),
+    ('Token merging stride y', 'token_merging_stride_y')
 ]
 
 
diff --git a/modules/processing.py b/modules/processing.py
index 670a7a28a..95058d0b7 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -31,6 +31,12 @@ from einops import repeat, rearrange
 from blendmodes.blend import blendLayers, BlendType
 import tomesd
 
+# add a logger for the processing module
+logger = logging.getLogger(__name__)
+# manually set output level here since there is no option to do so yet through launch options
+# logging.basicConfig(level=logging.DEBUG, format='%(asctime)s %(levelname)s %(name)s %(message)s')
+
+
 # some of those options should not be changed at all because they would break the model, so I removed them from options.
 opt_C = 4
 opt_f = 8
@@ -477,6 +483,14 @@ def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments=None, iter
         "Conditional mask weight": getattr(p, "inpainting_mask_weight", shared.opts.inpainting_mask_weight) if p.is_using_inpainting_conditioning else None,
         "Clip skip": None if clip_skip <= 1 else clip_skip,
         "ENSD": None if opts.eta_noise_seed_delta == 0 else opts.eta_noise_seed_delta,
+        "Token merging ratio": None if not (opts.token_merging or cmd_opts.token_merging) or opts.token_merging_hr_only else opts.token_merging_ratio,
+        "Token merging ratio hr": None if not (opts.token_merging or cmd_opts.token_merging) else opts.token_merging_ratio_hr,
+        "Token merging random": None if opts.token_merging_random is False else opts.token_merging_random,
+        "Token merging merge attention": None if opts.token_merging_merge_attention is True else opts.token_merging_merge_attention,
+        "Token merging merge cross attention": None if opts.token_merging_merge_cross_attention is False else opts.token_merging_merge_cross_attention,
+        "Token merging merge mlp": None if opts.token_merging_merge_mlp is False else opts.token_merging_merge_mlp,
+        "Token merging stride x": None if opts.token_merging_stride_x == 2 else opts.token_merging_stride_x,
+        "Token merging stride y": None if opts.token_merging_stride_y == 2 else opts.token_merging_stride_y
     }
 
     generation_params.update(p.extra_generation_params)
@@ -502,16 +516,16 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
                 sd_vae.reload_vae_weights()
 
         if (opts.token_merging or cmd_opts.token_merging) and not opts.token_merging_hr_only:
-            print("\nApplying token merging\n")
             sd_models.apply_token_merging(sd_model=p.sd_model, hr=False)
+            logger.debug('Token merging applied')
 
         res = process_images_inner(p)
 
     finally:
         # undo model optimizations made by tomesd
         if opts.token_merging or cmd_opts.token_merging:
-            print('\nRemoving token merging model optimizations\n')
             tomesd.remove_patch(p.sd_model)
+            logger.debug('Token merging model optimizations removed')
 
         # restore opts to original state
         if p.override_settings_restore_afterwards:
@@ -954,11 +968,11 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
             # case where user wants to use separate merge ratios
             if not opts.token_merging_hr_only:
                 # clean patch done by first pass. (clobbering the first patch might be fine? this might be excessive)
-                print('Temporarily reverting token merging optimizations in preparation for next pass')
                 tomesd.remove_patch(self.sd_model)
+                logger.debug('Temporarily removed token merging optimizations in preparation for next pass')
 
-            print("\nApplying token merging for high-res pass\n")
             sd_models.apply_token_merging(sd_model=self.sd_model, hr=True)
+            logger.debug('Applied token merging for high-res pass')
 
         samples = self.sampler.sample_img2img(self, samples, noise, conditioning, unconditional_conditioning, steps=self.hr_second_pass_steps or self.steps, image_conditioning=image_conditioning)
 
diff --git a/modules/shared.py b/modules/shared.py
index 568acdc4c..d9db7317a 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -446,7 +446,7 @@ options_templates.update(options_section(('token_merging', 'Token Merging'), {
     ),
     # More advanced/niche settings:
     "token_merging_random": OptionInfo(
-        True, "Use random perturbations - Disabling might help with certain samplers",
+        False, "Use random perturbations - Can improve outputs for certain samplers. For others, it may cause visaul artifacting.",
         gr.Checkbox
     ),
     "token_merging_merge_attention": OptionInfo(
diff --git a/requirements_versions.txt b/requirements_versions.txt
index 035227152..f972f975f 100644
--- a/requirements_versions.txt
+++ b/requirements_versions.txt
@@ -28,4 +28,4 @@ torchsde==0.2.5
 safetensors==0.3.0
 httpcore<=0.15
 fastapi==0.94.0
-tomesd>=0.1.1
\ No newline at end of file
+tomesd>=0.1.2
\ No newline at end of file

From c510cfd24b995f8267df3391cc961ac398922ba4 Mon Sep 17 00:00:00 2001
From: papuSpartan <30642826+papuSpartan@users.noreply.github.com>
Date: Mon, 10 Apr 2023 03:43:56 -0500
Subject: [PATCH 11/19] Update shared.py

fix typo
---
 modules/shared.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/modules/shared.py b/modules/shared.py
index d9db7317a..edf11e435 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -446,7 +446,7 @@ options_templates.update(options_section(('token_merging', 'Token Merging'), {
     ),
     # More advanced/niche settings:
     "token_merging_random": OptionInfo(
-        False, "Use random perturbations - Can improve outputs for certain samplers. For others, it may cause visaul artifacting.",
+        False, "Use random perturbations - Can improve outputs for certain samplers. For others, it may cause visual artifacting.",
         gr.Checkbox
     ),
     "token_merging_merge_attention": OptionInfo(

From a9902ca33119d6fae0a3623424bfc7ab86f2095a Mon Sep 17 00:00:00 2001
From: papuSpartan <30642826+papuSpartan@users.noreply.github.com>
Date: Mon, 10 Apr 2023 04:03:01 -0500
Subject: [PATCH 12/19] Update generation_parameters_copypaste.py

---
 modules/generation_parameters_copypaste.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/modules/generation_parameters_copypaste.py b/modules/generation_parameters_copypaste.py
index a7ede5344..ba2ca5ed7 100644
--- a/modules/generation_parameters_copypaste.py
+++ b/modules/generation_parameters_copypaste.py
@@ -283,7 +283,6 @@ Steps: 20, Sampler: Euler a, CFG scale: 7, Seed: 965400086, Size: 512x512, Model
         res["Hires resize-2"] = 0
 
     # Infer additional override settings for token merging
-    print("inferring settings for tomesd")
     token_merging_ratio = res.get("Token merging ratio", None)
     token_merging_ratio_hr = res.get("Token merging ratio hr", None)
 

From dff60e2e74964a8b02b75ecd8cf8007ef67a9712 Mon Sep 17 00:00:00 2001
From: papuSpartan <30642826+papuSpartan@users.noreply.github.com>
Date: Mon, 10 Apr 2023 04:10:50 -0500
Subject: [PATCH 13/19] Update sd_models.py

---
 modules/sd_models.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/modules/sd_models.py b/modules/sd_models.py
index 696a2333d..efcf730d6 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -560,7 +560,6 @@ def apply_token_merging(sd_model, hr: bool):
     ratio = shared.opts.token_merging_ratio
     if hr:
         ratio = shared.opts.token_merging_ratio_hr
-        print("effective hr pass merge ratio is "+str(ratio))
 
     tomesd.apply_patch(
         sd_model,

From e960781511eb175943be09b314ac2be46b6fc684 Mon Sep 17 00:00:00 2001
From: papuSpartan <30642826+papuSpartan@users.noreply.github.com>
Date: Wed, 3 May 2023 13:12:43 -0500
Subject: [PATCH 14/19] fix maximum downsampling option

---
 modules/generation_parameters_copypaste.py | 4 +++-
 modules/processing.py                      | 1 +
 modules/shared.py                          | 5 +----
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/modules/generation_parameters_copypaste.py b/modules/generation_parameters_copypaste.py
index 34c1b8608..83382e932 100644
--- a/modules/generation_parameters_copypaste.py
+++ b/modules/generation_parameters_copypaste.py
@@ -306,6 +306,8 @@ Steps: 20, Sampler: Euler a, CFG scale: 7, Seed: 965400086, Size: 512x512, Model
             res["Token merging stride x"] = '2'
         if res.get("Token merging stride y", None) is None:
             res["Token merging stride y"] = '2'
+        if res.get("Token merging maximum down sampling", None) is None:
+            res["Token merging maximum down sampling"] = '1'
 
     restore_old_hires_fix_params(res)
 
@@ -341,7 +343,7 @@ infotext_to_setting_name_mapping = [
     ('Token merging merge attention', 'token_merging_merge_attention'),
     ('Token merging merge cross attention', 'token_merging_merge_cross_attention'),
     ('Token merging merge mlp', 'token_merging_merge_mlp'),
-    ('Token merging maximum downsampling', 'token_merging_maximum_downsampling'),
+    ('Token merging maximum down sampling', 'token_merging_maximum_down_sampling'),
     ('Token merging stride x', 'token_merging_stride_x'),
     ('Token merging stride y', 'token_merging_stride_y'),
     ('RNG', 'randn_source'),
diff --git a/modules/processing.py b/modules/processing.py
index d5d1da5af..6807a301e 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -495,6 +495,7 @@ def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments=None, iter
         "Token merging merge mlp": None if opts.token_merging_merge_mlp is False else opts.token_merging_merge_mlp,
         "Token merging stride x": None if opts.token_merging_stride_x == 2 else opts.token_merging_stride_x,
         "Token merging stride y": None if opts.token_merging_stride_y == 2 else opts.token_merging_stride_y,
+        "Token merging maximum down sampling": None if opts.token_merging_maximum_down_sampling == 1 else opts.token_merging_maximum_down_sampling,
         "Init image hash": getattr(p, 'init_img_hash', None),
         "RNG": opts.randn_source if opts.randn_source != "GPU" else None,
         "NGMS": None if p.s_min_uncond == 0 else p.s_min_uncond,
diff --git a/modules/shared.py b/modules/shared.py
index 7b81ffc96..a7a72dd57 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -488,10 +488,7 @@ options_templates.update(options_section(('token_merging', 'Token Merging'), {
         False, "Merge mlp",
         gr.Checkbox
     ),
-    "token_merging_maximum_down_sampling": OptionInfo(
-        1, "Maximum down sampling",
-        gr.Dropdown, lambda: {"choices": ["1", "2", "4", "8"]}
-    ),
+    "token_merging_maximum_down_sampling": OptionInfo(1, "Maximum down sampling", gr.Radio, lambda: {"choices": ['1', '2', '4', '8']}),
     "token_merging_stride_x": OptionInfo(
         2, "Stride - X",
         gr.Slider, {"minimum": 2, "maximum": 8, "step": 2}

From f0efc8c211fc2d2c2f8caf6e2f92501922d18c99 Mon Sep 17 00:00:00 2001
From: papuSpartan <30642826+papuSpartan@users.noreply.github.com>
Date: Wed, 3 May 2023 21:10:31 -0500
Subject: [PATCH 15/19] not being cast properly every time, swap to ints

---
 modules/shared.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/modules/shared.py b/modules/shared.py
index a7a72dd57..eb06909cc 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -488,7 +488,7 @@ options_templates.update(options_section(('token_merging', 'Token Merging'), {
         False, "Merge mlp",
         gr.Checkbox
     ),
-    "token_merging_maximum_down_sampling": OptionInfo(1, "Maximum down sampling", gr.Radio, lambda: {"choices": ['1', '2', '4', '8']}),
+    "token_merging_maximum_down_sampling": OptionInfo(1, "Maximum down sampling", gr.Radio, lambda: {"choices": [1, 2, 4, 8]}),
     "token_merging_stride_x": OptionInfo(
         2, "Stride - X",
         gr.Slider, {"minimum": 2, "maximum": 8, "step": 2}

From 55e52c878ab669d5b11b001a4152ee1a3b8d4880 Mon Sep 17 00:00:00 2001
From: papuSpartan <30642826+papuSpartan@users.noreply.github.com>
Date: Sat, 13 May 2023 09:24:56 -0500
Subject: [PATCH 16/19] remove command line option

---
 modules/cmd_args.py   |  1 -
 modules/processing.py | 10 +++++-----
 2 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/modules/cmd_args.py b/modules/cmd_args.py
index 46043e333..f4a4ab36f 100644
--- a/modules/cmd_args.py
+++ b/modules/cmd_args.py
@@ -102,5 +102,4 @@ parser.add_argument("--no-gradio-queue", action='store_true', help="Disables gra
 parser.add_argument("--skip-version-check", action='store_true', help="Do not check versions of torch and xformers")
 parser.add_argument("--no-hashing", action='store_true', help="disable sha256 hashing of checkpoints to help loading performance", default=False)
 parser.add_argument("--no-download-sd-model", action='store_true', help="don't download SD1.5 model even if no model is found in --ckpt-dir", default=False)
-parser.add_argument("--token-merging", action='store_true', help="Provides speed and memory improvements by merging redundant tokens. This has a more pronounced effect on higher resolutions.", default=False)
 parser.add_argument('--subpath', type=str, help='customize the subpath for gradio, use with reverse proxy')
diff --git a/modules/processing.py b/modules/processing.py
index 8ba3a96b5..6828e8981 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -496,8 +496,8 @@ def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments=None, iter
         "Conditional mask weight": getattr(p, "inpainting_mask_weight", shared.opts.inpainting_mask_weight) if p.is_using_inpainting_conditioning else None,
         "Clip skip": None if clip_skip <= 1 else clip_skip,
         "ENSD": None if opts.eta_noise_seed_delta == 0 else opts.eta_noise_seed_delta,
-        "Token merging ratio": None if not (opts.token_merging or cmd_opts.token_merging) or opts.token_merging_hr_only else opts.token_merging_ratio,
-        "Token merging ratio hr": None if not (opts.token_merging or cmd_opts.token_merging) else opts.token_merging_ratio_hr,
+        "Token merging ratio": None if not opts.token_merging or opts.token_merging_hr_only else opts.token_merging_ratio,
+        "Token merging ratio hr": None if not opts.token_merging else opts.token_merging_ratio_hr,
         "Token merging random": None if opts.token_merging_random is False else opts.token_merging_random,
         "Token merging merge attention": None if opts.token_merging_merge_attention is True else opts.token_merging_merge_attention,
         "Token merging merge cross attention": None if opts.token_merging_merge_cross_attention is False else opts.token_merging_merge_cross_attention,
@@ -538,7 +538,7 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
             if k == 'sd_vae':
                 sd_vae.reload_vae_weights()
 
-        if (opts.token_merging or cmd_opts.token_merging) and not opts.token_merging_hr_only:
+        if opts.token_merging and not opts.token_merging_hr_only:
             sd_models.apply_token_merging(sd_model=p.sd_model, hr=False)
             logger.debug('Token merging applied')
 
@@ -546,7 +546,7 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
 
     finally:
         # undo model optimizations made by tomesd
-        if opts.token_merging or cmd_opts.token_merging:
+        if opts.token_merging:
             tomesd.remove_patch(p.sd_model)
             logger.debug('Token merging model optimizations removed')
 
@@ -1004,7 +1004,7 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
 
         # apply token merging optimizations from tomesd for high-res pass
         # check if hr_only so we are not redundantly patching
-        if (cmd_opts.token_merging or opts.token_merging) and (opts.token_merging_hr_only or opts.token_merging_ratio_hr != opts.token_merging_ratio):
+        if opts.token_merging and (opts.token_merging_hr_only or opts.token_merging_ratio_hr != opts.token_merging_ratio):
             # case where user wants to use separate merge ratios
             if not opts.token_merging_hr_only:
                 # clean patch done by first pass. (clobbering the first patch might be fine? this might be excessive)

From ac83627a31daac06f4d48b0e7db223ef807fe8e5 Mon Sep 17 00:00:00 2001
From: papuSpartan <30642826+papuSpartan@users.noreply.github.com>
Date: Sat, 13 May 2023 10:23:42 -0500
Subject: [PATCH 17/19] heavily simplify

---
 modules/generation_parameters_copypaste.py | 36 -------------------
 modules/processing.py                      | 35 ++++++++----------
 modules/sd_models.py                       | 11 +++---
 modules/shared.py                          | 42 +++-------------------
 4 files changed, 23 insertions(+), 101 deletions(-)

diff --git a/modules/generation_parameters_copypaste.py b/modules/generation_parameters_copypaste.py
index fb56254fc..a0a98bbcb 100644
--- a/modules/generation_parameters_copypaste.py
+++ b/modules/generation_parameters_copypaste.py
@@ -282,33 +282,6 @@ Steps: 20, Sampler: Euler a, CFG scale: 7, Seed: 965400086, Size: 512x512, Model
         res["Hires resize-1"] = 0
         res["Hires resize-2"] = 0
 
-    # Infer additional override settings for token merging
-    token_merging_ratio = res.get("Token merging ratio", None)
-    token_merging_ratio_hr = res.get("Token merging ratio hr", None)
-
-    if token_merging_ratio is not None or token_merging_ratio_hr is not None:
-        res["Token merging"] = 'True'
-
-        if token_merging_ratio is None:
-            res["Token merging hr only"] = 'True'
-        else:
-            res["Token merging hr only"] = 'False'
-
-        if res.get("Token merging random", None) is None:
-            res["Token merging random"] = 'False'
-        if res.get("Token merging merge attention", None) is None:
-            res["Token merging merge attention"] = 'True'
-        if res.get("Token merging merge cross attention", None) is None:
-            res["Token merging merge cross attention"] = 'False'
-        if res.get("Token merging merge mlp", None) is None:
-            res["Token merging merge mlp"] = 'False'
-        if res.get("Token merging stride x", None) is None:
-            res["Token merging stride x"] = '2'
-        if res.get("Token merging stride y", None) is None:
-            res["Token merging stride y"] = '2'
-        if res.get("Token merging maximum down sampling", None) is None:
-            res["Token merging maximum down sampling"] = '1'
-
     restore_old_hires_fix_params(res)
 
     # Missing RNG means the default was set, which is GPU RNG
@@ -335,17 +308,8 @@ infotext_to_setting_name_mapping = [
     ('UniPC skip type', 'uni_pc_skip_type'),
     ('UniPC order', 'uni_pc_order'),
     ('UniPC lower order final', 'uni_pc_lower_order_final'),
-    ('Token merging', 'token_merging'),
     ('Token merging ratio', 'token_merging_ratio'),
-    ('Token merging hr only', 'token_merging_hr_only'),
     ('Token merging ratio hr', 'token_merging_ratio_hr'),
-    ('Token merging random', 'token_merging_random'),
-    ('Token merging merge attention', 'token_merging_merge_attention'),
-    ('Token merging merge cross attention', 'token_merging_merge_cross_attention'),
-    ('Token merging merge mlp', 'token_merging_merge_mlp'),
-    ('Token merging maximum down sampling', 'token_merging_maximum_down_sampling'),
-    ('Token merging stride x', 'token_merging_stride_x'),
-    ('Token merging stride y', 'token_merging_stride_y'),
     ('RNG', 'randn_source'),
     ('NGMS', 's_min_uncond')
 ]
diff --git a/modules/processing.py b/modules/processing.py
index 6828e8981..32ff61e92 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -34,7 +34,7 @@ import tomesd
 # add a logger for the processing module
 logger = logging.getLogger(__name__)
 # manually set output level here since there is no option to do so yet through launch options
-# logging.basicConfig(level=logging.DEBUG, format='%(asctime)s %(levelname)s %(name)s %(message)s')
+logging.basicConfig(level=logging.DEBUG, format='%(asctime)s %(levelname)s %(name)s %(message)s')
 
 
 # some of those options should not be changed at all because they would break the model, so I removed them from options.
@@ -496,15 +496,8 @@ def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments=None, iter
         "Conditional mask weight": getattr(p, "inpainting_mask_weight", shared.opts.inpainting_mask_weight) if p.is_using_inpainting_conditioning else None,
         "Clip skip": None if clip_skip <= 1 else clip_skip,
         "ENSD": None if opts.eta_noise_seed_delta == 0 else opts.eta_noise_seed_delta,
-        "Token merging ratio": None if not opts.token_merging or opts.token_merging_hr_only else opts.token_merging_ratio,
-        "Token merging ratio hr": None if not opts.token_merging else opts.token_merging_ratio_hr,
-        "Token merging random": None if opts.token_merging_random is False else opts.token_merging_random,
-        "Token merging merge attention": None if opts.token_merging_merge_attention is True else opts.token_merging_merge_attention,
-        "Token merging merge cross attention": None if opts.token_merging_merge_cross_attention is False else opts.token_merging_merge_cross_attention,
-        "Token merging merge mlp": None if opts.token_merging_merge_mlp is False else opts.token_merging_merge_mlp,
-        "Token merging stride x": None if opts.token_merging_stride_x == 2 else opts.token_merging_stride_x,
-        "Token merging stride y": None if opts.token_merging_stride_y == 2 else opts.token_merging_stride_y,
-        "Token merging maximum down sampling": None if opts.token_merging_maximum_down_sampling == 1 else opts.token_merging_maximum_down_sampling,
+        "Token merging ratio": None if opts.token_merging_ratio == 0 else opts.token_merging_ratio,
+        "Token merging ratio hr": None if not p.enable_hr or opts.token_merging_ratio_hr == 0 else opts.token_merging_ratio_hr,
         "Init image hash": getattr(p, 'init_img_hash', None),
         "RNG": opts.randn_source if opts.randn_source != "GPU" else None,
         "NGMS": None if p.s_min_uncond == 0 else p.s_min_uncond,
@@ -538,15 +531,15 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
             if k == 'sd_vae':
                 sd_vae.reload_vae_weights()
 
-        if opts.token_merging and not opts.token_merging_hr_only:
+        if opts.token_merging_ratio > 0:
             sd_models.apply_token_merging(sd_model=p.sd_model, hr=False)
-            logger.debug('Token merging applied')
+            logger.debug(f"Token merging applied to first pass. Ratio: '{opts.token_merging_ratio}'")
 
         res = process_images_inner(p)
 
     finally:
         # undo model optimizations made by tomesd
-        if opts.token_merging:
+        if opts.token_merging_ratio > 0:
             tomesd.remove_patch(p.sd_model)
             logger.debug('Token merging model optimizations removed')
 
@@ -1003,19 +996,21 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
         devices.torch_gc()
 
         # apply token merging optimizations from tomesd for high-res pass
-        # check if hr_only so we are not redundantly patching
-        if opts.token_merging and (opts.token_merging_hr_only or opts.token_merging_ratio_hr != opts.token_merging_ratio):
-            # case where user wants to use separate merge ratios
-            if not opts.token_merging_hr_only:
-                # clean patch done by first pass. (clobbering the first patch might be fine? this might be excessive)
+        if opts.token_merging_ratio_hr > 0:
+            # in case the user has used separate merge ratios
+            if opts.token_merging_ratio > 0:
                 tomesd.remove_patch(self.sd_model)
-                logger.debug('Temporarily removed token merging optimizations in preparation for next pass')
+                logger.debug('Adjusting token merging ratio for high-res pass')
 
             sd_models.apply_token_merging(sd_model=self.sd_model, hr=True)
-            logger.debug('Applied token merging for high-res pass')
+            logger.debug(f"Applied token merging for high-res pass. Ratio: '{opts.token_merging_ratio_hr}'")
 
         samples = self.sampler.sample_img2img(self, samples, noise, conditioning, unconditional_conditioning, steps=self.hr_second_pass_steps or self.steps, image_conditioning=image_conditioning)
 
+        if opts.token_merging_ratio_hr > 0 or opts.token_merging_ratio > 0:
+            tomesd.remove_patch(self.sd_model)
+            logger.debug('Removed token merging optimizations from model')
+
         self.is_hr_pass = False
 
         return samples
diff --git a/modules/sd_models.py b/modules/sd_models.py
index 4787193c0..4c9a0a1fb 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -596,11 +596,8 @@ def apply_token_merging(sd_model, hr: bool):
     tomesd.apply_patch(
         sd_model,
         ratio=ratio,
-        max_downsample=shared.opts.token_merging_maximum_down_sampling,
-        sx=shared.opts.token_merging_stride_x,
-        sy=shared.opts.token_merging_stride_y,
-        use_rand=shared.opts.token_merging_random,
-        merge_attn=shared.opts.token_merging_merge_attention,
-        merge_crossattn=shared.opts.token_merging_merge_cross_attention,
-        merge_mlp=shared.opts.token_merging_merge_mlp
+        use_rand=False,  # can cause issues with some samplers
+        merge_attn=True,
+        merge_crossattn=False,
+        merge_mlp=False
     )
diff --git a/modules/shared.py b/modules/shared.py
index 4b3465855..0d96c14cb 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -459,47 +459,13 @@ options_templates.update(options_section((None, "Hidden options"), {
 }))
 
 options_templates.update(options_section(('token_merging', 'Token Merging'), {
-    "token_merging": OptionInfo(
-        False, "Enable redundant token merging via tomesd. This can provide significant speed and memory improvements.",
-        gr.Checkbox
+    "token_merging_ratio_hr": OptionInfo(
+        0, "Merging Ratio (high-res pass)",
+        gr.Slider, {"minimum": 0, "maximum": 0.9, "step": 0.1}
     ),
     "token_merging_ratio": OptionInfo(
-        0.5, "Merging Ratio",
+        0, "Merging Ratio",
         gr.Slider, {"minimum": 0, "maximum": 0.9, "step": 0.1}
-    ),
-    "token_merging_hr_only": OptionInfo(
-        True, "Apply only to high-res fix pass. Disabling can yield a ~20-35% speedup on contemporary resolutions.",
-        gr.Checkbox
-    ),
-    "token_merging_ratio_hr": OptionInfo(
-        0.5, "Merging Ratio (high-res pass) - If 'Apply only to high-res' is enabled, this will always be the ratio used.",
-        gr.Slider, {"minimum": 0, "maximum": 0.9, "step": 0.1}
-    ),
-    # More advanced/niche settings:
-    "token_merging_random": OptionInfo(
-        False, "Use random perturbations - Can improve outputs for certain samplers. For others, it may cause visual artifacting.",
-        gr.Checkbox
-    ),
-    "token_merging_merge_attention": OptionInfo(
-        True, "Merge attention",
-        gr.Checkbox
-    ),
-     "token_merging_merge_cross_attention": OptionInfo(
-        False, "Merge cross attention",
-        gr.Checkbox
-    ),
-    "token_merging_merge_mlp": OptionInfo(
-        False, "Merge mlp",
-        gr.Checkbox
-    ),
-    "token_merging_maximum_down_sampling": OptionInfo(1, "Maximum down sampling", gr.Radio, lambda: {"choices": [1, 2, 4, 8]}),
-    "token_merging_stride_x": OptionInfo(
-        2, "Stride - X",
-        gr.Slider, {"minimum": 2, "maximum": 8, "step": 2}
-    ),
-    "token_merging_stride_y": OptionInfo(
-        2, "Stride - Y",
-        gr.Slider, {"minimum": 2, "maximum": 8, "step": 2}
     )
 }))
 

From 917faa5325371e51d68f7d6f7b15ea4466bd5adf Mon Sep 17 00:00:00 2001
From: papuSpartan <30642826+papuSpartan@users.noreply.github.com>
Date: Sat, 13 May 2023 10:26:09 -0500
Subject: [PATCH 18/19] move to stable-diffusion tab

---
 modules/shared.py | 12 ++----------
 1 file changed, 2 insertions(+), 10 deletions(-)

diff --git a/modules/shared.py b/modules/shared.py
index 0d96c14cb..e49e9b74c 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -350,6 +350,8 @@ options_templates.update(options_section(('sd', "Stable Diffusion"), {
     "CLIP_stop_at_last_layers": OptionInfo(1, "Clip skip", gr.Slider, {"minimum": 1, "maximum": 12, "step": 1}),
     "upcast_attn": OptionInfo(False, "Upcast cross attention layer to float32"),
     "randn_source": OptionInfo("GPU", "Random number generator source. Changes seeds drastically. Use CPU to produce the same picture across different vidocard vendors.", gr.Radio, {"choices": ["GPU", "CPU"]}),
+    "token_merging_ratio_hr": OptionInfo(0, "Merging Ratio (high-res pass)", gr.Slider, {"minimum": 0, "maximum": 0.9, "step": 0.1}),
+    "token_merging_ratio": OptionInfo(0, "Merging Ratio", gr.Slider, {"minimum": 0, "maximum": 0.9, "step": 0.1})
 }))
 
 options_templates.update(options_section(('compatibility', "Compatibility"), {
@@ -458,16 +460,6 @@ options_templates.update(options_section((None, "Hidden options"), {
     "sd_checkpoint_hash": OptionInfo("", "SHA256 hash of the current checkpoint"),
 }))
 
-options_templates.update(options_section(('token_merging', 'Token Merging'), {
-    "token_merging_ratio_hr": OptionInfo(
-        0, "Merging Ratio (high-res pass)",
-        gr.Slider, {"minimum": 0, "maximum": 0.9, "step": 0.1}
-    ),
-    "token_merging_ratio": OptionInfo(
-        0, "Merging Ratio",
-        gr.Slider, {"minimum": 0, "maximum": 0.9, "step": 0.1}
-    )
-}))
 
 options_templates.update()
 

From c2fdb44880e07f43aee2f7edc1dc36a9516501e8 Mon Sep 17 00:00:00 2001
From: papuSpartan <30642826+papuSpartan@users.noreply.github.com>
Date: Sat, 13 May 2023 11:11:02 -0500
Subject: [PATCH 19/19] fix for img2img

---
 modules/processing.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/modules/processing.py b/modules/processing.py
index 32ff61e92..94fe2625c 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -34,7 +34,7 @@ import tomesd
 # add a logger for the processing module
 logger = logging.getLogger(__name__)
 # manually set output level here since there is no option to do so yet through launch options
-logging.basicConfig(level=logging.DEBUG, format='%(asctime)s %(levelname)s %(name)s %(message)s')
+# logging.basicConfig(level=logging.DEBUG, format='%(asctime)s %(levelname)s %(name)s %(message)s')
 
 
 # some of those options should not be changed at all because they would break the model, so I removed them from options.
@@ -478,6 +478,7 @@ def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments=None, iter
     index = position_in_batch + iteration * p.batch_size
 
     clip_skip = getattr(p, 'clip_skip', opts.CLIP_stop_at_last_layers)
+    enable_hr = getattr(p, 'enable_hr', False)
 
     generation_params = {
         "Steps": p.steps,
@@ -497,7 +498,7 @@ def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments=None, iter
         "Clip skip": None if clip_skip <= 1 else clip_skip,
         "ENSD": None if opts.eta_noise_seed_delta == 0 else opts.eta_noise_seed_delta,
         "Token merging ratio": None if opts.token_merging_ratio == 0 else opts.token_merging_ratio,
-        "Token merging ratio hr": None if not p.enable_hr or opts.token_merging_ratio_hr == 0 else opts.token_merging_ratio_hr,
+        "Token merging ratio hr": None if not enable_hr or opts.token_merging_ratio_hr == 0 else opts.token_merging_ratio_hr,
         "Init image hash": getattr(p, 'init_img_hash', None),
         "RNG": opts.randn_source if opts.randn_source != "GPU" else None,
         "NGMS": None if p.s_min_uncond == 0 else p.s_min_uncond,