jumbo merge part two

modules/sd_models.py

@@ -637,7 +637,6 @@ def reload_model_weights(sd_model=None, info=None, reuse_dict=False):
     timer.record("config")
     if sd_model is None or checkpoint_config != sd_model.used_config:
         del sd_model
-        checkpoints_loaded.clear()
         if shared.backend == shared.Backend.ORIGINAL:
             load_model(checkpoint_info, already_loaded_state_dict=state_dict, timer=timer)
         else:
@@ -676,26 +675,23 @@ def unload_model_weights(sd_model=None, _info=None):
     return sd_model
 
 
-def apply_token_merging(sd_model, hr: bool):
+def apply_token_merging(sd_model, token_merging_ratio):
     """
     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
-
-    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
-    )
+    current_token_merging_ratio = getattr(sd_model, 'applied_token_merged_ratio', 0)
+    shared.log.debug(f'Appplying token merging: current={current_token_merging_ratio} target={token_merging_ratio}')
+    if current_token_merging_ratio == token_merging_ratio:
+        return
+    if current_token_merging_ratio > 0:
+        tomesd.remove_patch(sd_model)
+    if token_merging_ratio > 0:
+        tomesd.apply_patch(
+            sd_model,
+            ratio=token_merging_ratio,
+            use_rand=False,  # can cause issues with some samplers
+            merge_attn=True,
+            merge_crossattn=False,
+            merge_mlp=False
+        )
+    sd_model.applied_token_merged_ratio = token_merging_ratio