first

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)

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