Attention explorer: the log reports how many attention calls were actually changed, and warns if none

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RnrTBt9b8CdJwub5roffF3

modules/sampling/attention_bandit/__init__.py

@@ -20,6 +20,8 @@ a real model.
 
 import random
 
+from comfy.patcher_extension import WrappersMP
+
 from ..._core import dit_hooks, log, registry
 
 ID = "attention_bandit"
@@ -188,7 +190,7 @@ def install(patcher, values, key):
         raise ValueError(f"no usable blocks in {values.get('blocks')!r}, so nothing would change")
     groups = split(blocks)
     signature = [len(groups), blocks[0], blocks[-1]]
-    live = {"temps": {}}
+    live = {"temps": {}, "scaled": 0, "calls": 0}
 
     def fresh():
         import torch
@@ -196,6 +198,7 @@ def install(patcher, values, key):
         picks = draw(counts(rows, signature, len(groups)))
         spread = manual_spread if manual else min(max(
             learned_spread(rows, signature) + random.gauss(0.0, EXPLORE), SPREAD_MIN), SPREAD_MAX)
+        live["scaled"] = live["calls"] = 0
         live["temps"] = {b: 1.0 + spread * ARMS[a] for g, a in zip(groups, picks) for b in g}
         captured.update({"arms": torch.tensor(picks), "sig": torch.tensor(signature),
                          "spread": torch.tensor(float(spread))})
@@ -219,9 +222,24 @@ def install(patcher, values, key):
 
     def override(func, q, k, v, *args, **kwargs):
         t = (kwargs.get("transformer_options") or {}).get(_SCOPE)
+        live["calls"] += 1
+        live["scaled"] += t is not None
         return func(q if t is None else q / t, k, v, *args, **kwargs)
 
     dit_hooks.add_attention_override(patcher, key, override)
+
+    def report(executor, *args, **kwargs):
+        out = executor(*args, **kwargs)
+        changed = any(t != 1.0 for t in live["temps"].values())
+        if changed and not live["scaled"]:
+            _say("changed nothing this run: no attention call reached the changed blocks "
+                 f"({live['calls']} calls seen). Another feature may have replaced attention.", log.WARNING)
+        elif changed:
+            log.once(f"{ID}:applied", log.INFO, "FunUI Attention explorer",
+                     f"applied to {live['scaled']} of {live['calls']} attention calls")
+        return out
+
+    patcher.add_wrapper_with_key(WrappersMP.OUTER_SAMPLE, f"{key}.report", report)
     size = f"size {manual_spread:g}" if manual else "size learned from ratings"
     return f"{len(groups)} depth groups over {len(blocks)} blocks, {size}"
 

modules/sampling/attention_bandit/tests/test_attention_bandit.py

@@ -160,3 +160,13 @@ def test_the_log_counts_rated_runs():
     rows = [_row(1.0, [1], [1, 0, 0]), _row(-1.0, [1], [1, 0, 0]), _row(-1.0, [1], [1, 0, 0])]
     assert tally(rows) == "3 rated runs (1 liked, 2 disliked)"
     assert tally([]) == "0 rated runs (0 liked, 0 disliked)"
+
+
+def test_the_log_says_how_many_attention_calls_were_changed(tiny_h3, monkeypatch):
+    from core import log
+    from modules.sampling import attention_bandit as m
+    patched, _ = _load(tiny_h3, mode="manual", blocks="0-3")
+    monkeypatch.setattr(m, "draw", lambda table, rng=random: [0] * len(table))
+    log.new_run()
+    _run(_outer(patched), lambda: tiny_h3.run(patched))
+    assert any("applied to" in e["message"] and "attention calls" in e["message"] for e in log.history())