From e8ba35e7e9719017d8c16df0cd96131f0e5a2b44 Mon Sep 17 00:00:00 2001
From: vladmandic <mandic00@live.com>
Date: Fri, 13 Mar 2026 13:00:55 +0100
Subject: [PATCH] merge: modules/token_merge.py

---
 modules/token_merge.py | 46 +++++++++++++++++++++++++-----------------
 1 file changed, 28 insertions(+), 18 deletions(-)

diff --git a/modules/token_merge.py b/modules/token_merge.py
index 0cfaffe0d..800898f80 100644
--- a/modules/token_merge.py
+++ b/modules/token_merge.py
@@ -1,16 +1,26 @@
-from modules.logger import log
 from modules import shared
 
 
-def apply_token_merging(sd_model):
+def apply_token_merging(sd_model, p=None):
+    def _opt(key):
+        if p is not None:
+            val = getattr(p, key, None)
+            if val is not None:
+                return val
+        return getattr(shared.opts, key, None)
+
     current_tome = getattr(sd_model, 'applied_tome', 0)
     current_todo = getattr(sd_model, 'applied_todo', 0)
+    method = _opt('token_merging_method')
+    tome = _opt('tome_ratio')
+    todo = _opt('todo_ratio')
+    hypertile = _opt('hypertile_unet_enabled')
 
-    if shared.opts.token_merging_method == 'ToMe' and shared.opts.tome_ratio > 0:
-        if current_tome == shared.opts.tome_ratio:
+    if method == 'ToMe' and tome > 0:
+        if current_tome == tome:
             return
-        if shared.opts.hypertile_unet_enabled and not shared.cmd_opts.experimental:
-            log.warning('Token merging not supported with HyperTile for UNet')
+        if hypertile and not shared.cmd_opts.experimental:
+            shared.log.warning('Token merging not supported with HyperTile for UNet')
             return
         try:
             import installer
@@ -18,29 +28,29 @@ def apply_token_merging(sd_model):
             import tomesd
             tomesd.apply_patch(
                 sd_model,
-                ratio=shared.opts.tome_ratio,
+                ratio=tome,
                 use_rand=False, # can cause issues with some samplers
                 merge_attn=True,
                 merge_crossattn=False,
                 merge_mlp=False
             )
-            log.info(f'Applying ToMe: ratio={shared.opts.tome_ratio}')
-            sd_model.applied_tome = shared.opts.tome_ratio
+            shared.log.info(f'Applying ToMe: ratio={tome}')
+            sd_model.applied_tome = tome
         except Exception:
-            log.warning(f'Token merging not supported: pipeline={sd_model.__class__.__name__}')
+            shared.log.warning(f'Token merging not supported: pipeline={sd_model.__class__.__name__}')
     else:
         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:
+    if method == 'ToDo' and todo > 0:
+        if current_todo == todo:
             return
-        if shared.opts.hypertile_unet_enabled and not shared.cmd_opts.experimental:
-            log.warning('Token merging not supported with HyperTile for UNet')
+        if hypertile 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,
+                        "ratio": todo,
                         "merge_tokens": "keys/values",
                         "merge_method": "downsample",
                         "downsample_method": "nearest",
@@ -51,10 +61,10 @@ def apply_token_merging(sd_model):
                         "ratio_level_2": 0.0,
                         }
             patch_attention_proc(sd_model.unet, token_merge_args=token_merge_args)
-            log.info(f'Applying ToDo: ratio={shared.opts.todo_ratio}')
-            sd_model.applied_todo = shared.opts.todo_ratio
+            shared.log.info(f'Applying ToDo: ratio={todo}')
+            sd_model.applied_todo = todo
         except Exception:
-            log.warning(f'Token merging not supported: pipeline={sd_model.__class__.__name__}')
+            shared.log.warning(f'Token merging not supported: pipeline={sd_model.__class__.__name__}')
     else:
         sd_model.applied_todo = 0