nodes/olivv-cs-ComfyUI-FunPack
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Attention explorer: prompt weighting, prompt focus, per-prompt recipes, clip regions

Browse Source 
Regions rescale rows of q on a copy; checked against SLA's override contract.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RnrTBt9b8CdJwub5roffF3
This commit is contained in:
Claude
2026-10-25 22:10:19 +00:00
parent 5a98b2ed35
commit 01b8a37dff
3 changed files with 408 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/sampling/attention_bandit/__init__.py
+240 -55
View File
@@ -13,13 +13,27 @@ toward the size it used, a disliked one pushes it away (as Shot memory does for
amount), with a little exploration. Or type your own size.
Zero extra model calls; the speed cost is a division on the query.
Three more levers, all learned from the same ratings:
- Prompt focus: each group also draws a gain on the prompt tokens' keys, so attention to
the prompt can be sharpened or loosened separately from everything else.
- Prompt weighting: a rating counts more for runs of a similar prompt (pooled prompt
embedding, steep cosine) and less for other prompts, so what you tune on one prompt
shows up on that prompt first.
- Clip regions: the picture's four time bands and the sound each draw their own sharper or
looser setting, so a part you liked can stay while another changes.
- Recipes: when this prompt has a combination you liked, it is reused, with some groups
re-drawn so it keeps exploring.
Not proven: with a handful of ratings the per-group counts are noisy, and whether
depth groups are the right unit is a guess. Experimental until it has been rated on
a real model.
"""
import math
import random
import torch
import torch.nn.functional as F
from comfy.patcher_extension import WrappersMP
from ..._core import dit_hooks, log, registry
@@ -54,6 +68,18 @@ SETTINGS = {
"hint": "How much sharper or looser a change is. Keep it small.",
"when": {"enabled": True, "mode": "manual"},
},
"prompt_focus": {
"type": "bool", "default": True,
"label": "Prompt focus",
"hint": "Also learns how strongly each group attends to your prompt, apart from everything else.",
"when": {"enabled": True},
},
"regions": {
"type": "bool", "default": True,
"label": "Clip regions",
"hint": "Also learns the clip's four time bands and its sound separately, so a good part can stay while another changes.",
"when": {"enabled": True},
},
"blocks": {
"type": "text", "default": "2-10,29-46",
"label": "Blocks",
@@ -64,16 +90,26 @@ SETTINGS = {
KIND = "attention_bandit"
GROUPS = 4
# Each arm is a temperature as 1 + RANGE * ARMS[i]: sharper, as is, looser.
# Each arm is a factor as 1 + size * ARMS[i]: lower, as is, higher.
ARMS = (-1.0, 0.0, 1.0)
STAY = 1 # index of "as is"
STAY_HEAD_START = 2.0 # extra likes "as is" starts with, so a new key is stable
MAX_ROWS = 10 ** 9 # no cap: a rated run is under 1 KB
MAX_ROWS = 10 ** 9 # no cap: a rated run is a few KB at most
START_SPREAD = 0.2
SPREAD_MIN, SPREAD_MAX = 0.05, 0.5
LEARN_MAX = 0.35 # what ratings can grow the size to; 0.5x/1.5x on many blocks wrecked a picture on H3
LEARN_MAX = 0.35 # what ratings can grow the size to; 0.5x/1.5x on many blocks wrecked a picture on H3
LR = 0.5
EXPLORE = 0.05
# ponytail: pooled-prompt cosine sits high for most prompts, so the curve is steep and the
# constants are guesses. The log says how many rated runs count as this prompt.
SHARP = 100.0 # similarity = exp(SHARP * (cos - 1)); cos 0.99 -> 0.37
FLOOR = 0.1 # what a different prompt's rating still counts for
UNKNOWN = 0.5 # a rating stored without its prompt
SAME = 0.5 # similarity from which a run counts as this prompt
RECIPE_BAR = 0.65 # a liked combination is reused when its draw beats this
REGIONS = ("start of the clip", "second quarter", "third quarter", "end of the clip", "sound")
BANDS = 4 # equal time bands of the picture; the sound is the last region
RECIPE_EXPLORE = 0.25 # per group: chance a reused recipe re-draws that group
def _say(message, level=log.ALERT):
@@ -92,17 +128,35 @@ def split(blocks, groups=GROUPS):
return out
def counts(rows, signature, n_groups):
"""[group][arm] -> [likes, dislikes] over the rated runs made with this same block layout."""
def similarity(cond, stored):
"""1 = the same prompt, toward 0 = a different one; None when either is unknown."""
if cond is None or not torch.is_tensor(stored) or stored.numel() != cond.numel():
return None
cos = float(F.cosine_similarity(cond.float().flatten(), stored.float().flatten(), dim=0))
return math.exp(SHARP * (min(cos, 1.0) - 1.0))
def weights_of(rows, cond):
"""([how much each rated run counts now], [its similarity or None])."""
sims = [similarity(cond, r["rows"].get("cond")) for r in rows]
return [UNKNOWN if s is None else FLOOR + (1.0 - FLOOR) * s for s in sims], sims
def _same_signature(p, signature):
return "sig" in p and [int(x) for x in p["sig"].tolist()] == signature
def counts(rows, signature, n_groups, field="arms", weights=None):
"""[group][arm] -> [likes, dislikes] over the rated runs made with this block layout,
each counting its weight (1 when none given)."""
table = [[[0.0, 0.0] for _ in ARMS] for _ in range(n_groups)]
for r in rows:
for i, r in enumerate(rows):
p = r["rows"]
if "arms" not in p or "sig" not in p or r["reward"] == 0:
if field not in p or r["reward"] == 0 or not _same_signature(p, signature):
continue
if [int(x) for x in p["sig"].tolist()] != signature:
continue
for g, arm in enumerate(int(a) for a in p["arms"].tolist()):
table[g][arm][0 if r["reward"] > 0 else 1] += 1.0
w = 1.0 if weights is None else weights[i]
for g, arm in enumerate(int(a) for a in p[field].tolist()):
table[g][arm][0 if r["reward"] > 0 else 1] += w
return table
@@ -116,11 +170,9 @@ def learned_spread(rows, signature):
m = START_SPREAD
for r in rows:
p = r["rows"]
if "spread" not in p or "sig" not in p or "arms" not in p or not r["reward"]:
if "spread" not in p or "arms" not in p or not r["reward"] or not _same_signature(p, signature):
continue
if [int(x) for x in p["sig"].tolist()] != signature:
continue
if all(int(a) == STAY for a in p["arms"].tolist()):
if all(int(a) == STAY for f in ("arms", "tarms", "rarms") if f in p for a in p[f].tolist()):
continue # nothing changed, so nothing to blame on the size
used = float(p["spread"])
# A like pulls toward the size it used. A dislike can only say "too much": it pushes the
@@ -137,17 +189,6 @@ def tally(rows):
return f"{len(rows)} rated runs ({liked} liked, {disliked} disliked)"
def describe(groups, picks, spread):
"""Which blocks went sharper or looser, exactly; groups left as is are not mentioned."""
moved = {}
for g, a in zip(groups, picks):
if a != STAY:
moved.setdefault(a, []).extend(g)
parts = [f"{'sharper' if ARMS[a] < 0 else 'looser'} {1.0 + spread * ARMS[a]:g}x at blocks {ranges(sorted(bs))}"
for a, bs in sorted(moved.items())]
return "; ".join(parts) or "no change this run"
def ranges(blocks):
out, start = [], None
for i, b in enumerate(blocks):
@@ -159,6 +200,18 @@ def ranges(blocks):
return ",".join(out)
def describe(groups, picks, spread, words=("sharper", "looser"), what="", blocks=True):
"""Which blocks (or regions) went down or up, exactly; groups left as is are not mentioned."""
moved = {}
for g, a in zip(groups, picks):
if a != STAY:
moved.setdefault(a, []).extend(g)
parts = [f"{what}{words[0] if ARMS[a] < 0 else words[1]} {1.0 + spread * ARMS[a]:g}x "
+ (f"at blocks {ranges(sorted(bs))}" if blocks else f"in {', '.join(bs)}")
for a, bs in sorted(moved.items())]
return "; ".join(parts)
def draw(table, rng=random):
"""One arm per group: the arm whose Beta draw (likes + 1, dislikes + 1) is highest."""
picks = []
@@ -173,8 +226,95 @@ def draw(table, rng=random):
return picks
def pick_recipe(rows, sims, signature, focus, regions=False, rng=random):
"""The arms of a combination liked on THIS prompt, or None. Each combination seen on this
prompt gets a Beta draw from its likes and dislikes; the best one is reused when it beats the bar."""
combos = {}
for r, s in zip(rows, sims):
p = r["rows"]
if s is None or s < SAME or not r["reward"] or "arms" not in p or "spread" not in p \
or not _same_signature(p, signature) or (focus and "tarms" not in p) \
or (regions and "rarms" not in p):
continue
key = (tuple(int(a) for a in p["arms"].tolist()),
tuple(int(a) for a in p["tarms"].tolist()) if focus else None,
tuple(int(a) for a in p["rarms"].tolist()) if regions else None)
c = combos.setdefault(key, [0, 0, START_SPREAD])
c[0 if r["reward"] > 0 else 1] += 1
if r["reward"] > 0:
c[2] = float(p["spread"])
best, best_theta = None, -1.0
for key, (liked, disliked, spread) in combos.items():
if not liked:
continue
theta = rng.betavariate(1.0 + liked, 1.0 + disliked)
if theta > best_theta:
best, best_theta = (key, spread, liked, disliked), theta
return best if best_theta > RECIPE_BAR else None
def make_plan(rows, cond, signature, n_groups, focus, manual_spread=None, regions=False, rng=random):
"""What this run does: {arms, tarms | None, rarms | None, spread, recipe | None, same, unknown}."""
weights, sims = weights_of(rows, cond)
arms = draw(counts(rows, signature, n_groups, "arms", weights), rng)
tarms = draw(counts(rows, signature, n_groups, "tarms", weights), rng) if focus else None
rarms = draw(counts(rows, signature, len(REGIONS), "rarms", weights), rng) if regions else None
spread = learned_spread(rows, signature)
recipe = pick_recipe(rows, sims, signature, focus, regions, rng)
if recipe is not None:
(r_arms, r_tarms, r_rarms), spread, liked, disliked = recipe
for mine, theirs in ((arms, r_arms), (tarms, r_tarms), (rarms, r_rarms)):
for i in range(len(mine or ())):
if rng.random() >= RECIPE_EXPLORE:
mine[i] = theirs[i]
recipe = (liked, disliked)
spread = manual_spread if manual_spread is not None else min(max(
spread + rng.gauss(0.0, EXPLORE), SPREAD_MIN), LEARN_MAX)
return {"arms": arms, "tarms": tarms, "rarms": rarms, "spread": spread, "recipe": recipe,
"same": sum(1 for s in sims if s is not None and s >= SAME),
"unknown": sum(1 for s in sims if s is None)}
def _prompt_of(guider):
try:
p = guider.conds["positive"][0]["cross_attn"]
except (AttributeError, KeyError, IndexError, TypeError):
return None
if not torch.is_tensor(p):
return None
p = p.detach().float()
while p.dim() > 1:
p = p.mean(dim=0)
return p.cpu()
def _segments(options):
"""The packed sequence's {kind: (start, stop)}, or {} when this model does not say."""
try:
return {kind: (int(a), int(b)) for a, b, kind in options["minimax_h3_layout"].segments}
except (KeyError, AttributeError, TypeError, ValueError):
return {}
def _text_rows(options):
"""How many leading rows of the packed sequence are the prompt, or 0."""
a, b = _segments(options).get("text", (0, 0))
return b - a if a == 0 else 0
def region_spans(options):
"""[(start, stop)] per REGIONS entry: the target picture in equal time bands (its rows run
frame by frame), then the target sound. [] when the layout is unknown."""
seg = _segments(options)
if "video" not in seg or "audio" not in seg:
return []
a, b = seg["video"]
edges = [a + (b - a) * i // BANDS for i in range(BANDS + 1)]
return [(edges[i], edges[i + 1]) for i in range(BANDS)] + [seg["audio"]]
# Carried on the per-block copy of transformer_options, as Randomizer does, so an
# attention call knows its block's temperature without a shared flag.
# attention call knows its block's factors without a shared flag.
_SCOPE = "funui_attention_bandit"
@@ -187,6 +327,8 @@ def install(patcher, values, key):
return None
manual = values.get("mode") == "manual"
manual_spread = min(max(float(values.get("range", START_SPREAD)), SPREAD_MIN), SPREAD_MAX)
focus = bool(values.get("prompt_focus", True))
regions = bool(values.get("regions", True))
blocks, problems = dit_hooks.parse_blocks(values.get("blocks"), dit_hooks.block_count(patcher))
for problem in problems:
_say(f"ignored {problem}")
@@ -194,30 +336,18 @@ 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": {}, "scaled": 0, "calls": 0}
live = {"factors": {}, "regions": [1.0] * len(REGIONS), "scaled": 0, "calls": 0,
"text": 0, "text_wanted": False, "region_hits": 0, "region_wanted": False}
def fresh():
import torch
rows = taste.rows(KIND)
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), LEARN_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))})
log.once(f"{ID}:run", log.INFO, "FunUI Attention explorer",
f"key {taste.key!r}, {tally(rows)}: {describe(groups, picks, spread)}")
captured = taste.collect(patcher, key, KIND, keep=MAX_ROWS, fresh=fresh)
captured = taste.collect(patcher, key, KIND, keep=MAX_ROWS)
def scoped_for(index):
def scoped(args, extra):
t = live["temps"].get(index)
if t is None or t == 1.0:
f = live["factors"].get(index)
if f is None or (f == (1.0, 1.0) and not live["region_wanted"]):
return extra["original_block"](args)
to = dict(args.get("transformer_options") or {})
to[_SCOPE] = t
to[_SCOPE] = f
return extra["original_block"]({**args, "transformer_options": to})
return scoped
@@ -225,27 +355,82 @@ def install(patcher, values, key):
dit_hooks.add_block_hook(patcher, key, index, scoped_for(index))
def override(func, q, k, v, *args, **kwargs):
t = (kwargs.get("transformer_options") or {}).get(_SCOPE)
options = kwargs.get("transformer_options") or {}
scope = options.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)
if scope is None:
return func(q, k, v, *args, **kwargs)
temperature, prompt = scope
live["scaled"] += 1
if temperature != 1.0:
q = q / temperature
if live["region_wanted"] and q.ndim == 4:
spans = region_spans(options)
if spans and spans[-1][1] <= q.shape[2]:
q = q.clone()
for (a, b), factor in zip(spans, live["regions"]):
if factor != 1.0:
q[:, :, a:b] /= factor # only those rows ask differently: the rest of the clip is untouched
live["region_hits"] += 1
if prompt != 1.0 and kwargs.get("skip_reshape") and k.ndim == 4:
n = _text_rows(options)
if 0 < n < k.shape[2]:
k = k.clone()
k[:, :, :n] *= prompt # the prompt's keys only: its share of attention sharpens or loosens
live["text"] += 1
return func(q, k, v, *args, **kwargs)
dit_hooks.add_attention_override(patcher, key, override)
def report(executor, *args, **kwargs):
def run(executor, *args, **kwargs):
rows = taste.rows(KIND)
cond = _prompt_of(getattr(executor, "class_obj", None))
plan = make_plan(rows, cond, signature, len(groups), focus, manual_spread if manual else None, regions)
arms, tarms, rarms, spread = plan["arms"], plan["tarms"], plan["rarms"], plan["spread"]
tfac = [1.0 + spread * ARMS[a] for a in tarms] if focus else [1.0] * len(groups)
rfac = [1.0 + spread * ARMS[a] for a in rarms] if regions else [1.0] * len(REGIONS)
live.update(scaled=0, calls=0, text=0, region_hits=0, regions=rfac,
region_wanted=any(f != 1.0 for f in rfac), text_wanted=any(t != 1.0 for t in tfac),
factors={b: (1.0 + spread * ARMS[a], tf)
for g, a, tf in zip(groups, arms, tfac) for b in g})
captured.update({"arms": torch.tensor(arms), "sig": torch.tensor(signature),
"spread": torch.tensor(float(spread))})
if focus:
captured["tarms"] = torch.tensor(tarms)
if regions:
captured["rarms"] = torch.tensor(rarms)
if cond is not None:
captured["cond"] = cond.half()
known = "prompt unknown" if cond is None else f"{plan['same']} of them this prompt"
how = f"recipe liked {plan['recipe'][0]}x, disliked {plan['recipe'][1]}x" if plan["recipe"] else "exploring"
changes = "; ".join(x for x in (
describe(groups, arms, spread),
describe(groups, tarms, spread, ("weaker", "stronger"), "prompt ") if focus else "",
describe([[n] for n in REGIONS], rarms, spread, what="clip regions ", blocks=False) if regions else "") if x)
log.once(f"{ID}:run", log.INFO, "FunUI Attention explorer",
f"key {taste.key!r}, {tally(rows)}, {known}, {how}: {changes or 'no change this run'}")
out = executor(*args, **kwargs)
changed = any(t != 1.0 for t in live["temps"].values())
if changed and not live["scaled"]:
wanted = bool(changes)
if wanted 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:
elif live["text_wanted"] and not live["text"]:
_say("prompt focus changed nothing this run: this model does not say where the prompt sits in "
"its sequence, so only the plain attention changes were applied.", log.WARNING)
elif live["region_wanted"] and not live["region_hits"]:
_say("clip regions changed nothing this run: this model does not say where the picture and "
"sound sit in its sequence.", log.WARNING)
elif wanted:
log.once(f"{ID}:applied", log.INFO, "FunUI Attention explorer",
f"applied to {live['scaled']} of {live['calls']} attention calls")
f"applied to {live['scaled']} of {live['calls']} attention calls"
+ (f", prompt keys scaled in {live['text']}" if live["text_wanted"] else "")
+ (f", clip regions in {live['region_hits']}" if live["region_wanted"] else ""))
return out
patcher.add_wrapper_with_key(WrappersMP.OUTER_SAMPLE, f"{key}.report", report)
patcher.add_wrapper_with_key(WrappersMP.OUTER_SAMPLE, f"{key}.run", run)
size = f"size {manual_spread:g}" if manual else "size learned from ratings"
return f"{len(groups)} depth groups over {len(blocks)} blocks, {size}"
return (f"{len(groups)} depth groups over {len(blocks)} blocks, {size}"
+ (", prompt focus on" if focus else "") + (", clip regions on" if regions else ""))
PROVIDES = {"modifier": install}
modules/sampling/attention_bandit/tests/test_attention_bandit.py
+167 -1
View File
@@ -74,7 +74,7 @@ def test_dislikes_of_small_changes_never_grow_the_size():
def test_the_log_names_exact_blocks_and_skips_the_unchanged():
from modules.sampling.attention_bandit import describe
groups = [[2, 3, 4], [9, 10, 29, 30], [34, 35], [41, 42]]
assert describe(groups, [1, 1, 1, 1], 0.2) == "no change this run"
assert describe(groups, [1, 1, 1, 1], 0.2) == ""
assert describe(groups, [0, 1, 2, 0], 0.2) == "sharper 0.8x at blocks 2-4,41-42; looser 1.2x at blocks 34-35"
@@ -176,3 +176,169 @@ def test_the_log_says_how_many_attention_calls_were_changed(tiny_h3, monkeypatch
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())
# ── prompt weighting, recipes, prompt focus, clip regions ─────────────────────
def _cond(seed, n=16):
g = torch.Generator().manual_seed(seed)
return torch.randn(n, generator=g)
def test_the_same_prompt_counts_fully_another_prompt_counts_for_the_floor_and_unknown_in_between():
from modules.sampling.attention_bandit import FLOOR, UNKNOWN, similarity, weights_of
a, b = _cond(1), _cond(2)
assert similarity(a, a) == pytest.approx(1.0)
assert similarity(a, b) < 0.01 and similarity(a, None) is None and similarity(None, a) is None
rows = [{"reward": 1.0, "rows": {"cond": a.half()}}, {"reward": 1.0, "rows": {"cond": b.half()}},
{"reward": 1.0, "rows": {}}]
w, sims = weights_of(rows, a)
assert w[0] == pytest.approx(1.0, abs=1e-3) and w[1] == pytest.approx(FLOOR, abs=0.02) and w[2] == UNKNOWN
def test_counts_use_the_weights():
from modules.sampling.attention_bandit import counts
rows = [_row(1.0, [0], [1, 0, 0]), _row(1.0, [0], [1, 0, 0])]
assert counts(rows, [1, 0, 0], 1, weights=[1.0, 0.1])[0][0][0] == pytest.approx(1.1)
rows[0]["rows"]["tarms"] = torch.tensor([2])
assert counts(rows, [1, 0, 0], 1, "tarms")[0][2] == [1.0, 0.0] # only the row that has the field
def _combo_row(reward, cond, arms=(0, 1), tarms=(2, 1), rarms=(1, 1, 1, 1, 1), spread=0.3, sig=(2, 0, 9)):
return {"reward": reward, "rows": {"arms": torch.tensor(list(arms)), "tarms": torch.tensor(list(tarms)),
"rarms": torch.tensor(list(rarms)), "sig": torch.tensor(list(sig)),
"spread": torch.tensor(spread), "cond": cond.half()}}
def test_a_combination_liked_on_this_prompt_is_reused_and_one_liked_on_another_is_not():
from modules.sampling.attention_bandit import make_plan
here, other = _cond(1), _cond(2)
rows = [_combo_row(1.0, here) for _ in range(8)]
rng = random.Random(0)
plans = [make_plan(rows, here, [2, 0, 9], 2, True, regions=True, rng=rng) for _ in range(40)]
reused = [p for p in plans if p["recipe"]]
assert len(reused) > 30 and all(p["same"] == 8 for p in plans)
assert sum(p["arms"][0] == 0 for p in reused) > 0.6 * len(reused) # mostly the recipe, some groups re-drawn
away = [make_plan(rows, other, [2, 0, 9], 2, True, regions=True, rng=rng) for _ in range(40)]
assert not any(p["recipe"] for p in away) and all(p["same"] == 0 for p in away)
def test_disliked_combinations_are_not_a_recipe():
from modules.sampling.attention_bandit import make_plan
here = _cond(1)
rows = [_combo_row(-1.0, here) for _ in range(8)]
assert not any(make_plan(rows, here, [2, 0, 9], 2, True, regions=True, rng=random.Random(i))["recipe"] for i in range(30))
def _force(monkeypatch, arms=1, tarms=1, rarms=1):
from modules.sampling import attention_bandit as m
def fixed(table, rng=random):
n = len(table)
return [{4: arms, 5: rarms}.get(n, arms)] * n if n == 5 else [arms] * n
seq = iter([])
calls = []
def draw(table, rng=random):
calls.append(len(table))
# the first draw is the plain groups, the second the prompt groups, the third the regions
pick = (arms, tarms, rarms)[min(len(calls) - 1, 2)]
return [pick] * len(table)
monkeypatch.setattr(m, "draw", draw)
def test_prompt_focus_alone_changes_the_output_and_says_where(tiny_h3, monkeypatch):
from core import log
base = tiny_h3.run()
patched, status = _load(tiny_h3, mode="manual", range=0.5, blocks="0-3", regions=False)
assert "prompt focus on" in status
_force(monkeypatch, arms=1, tarms=2)
log.new_run()
outs = _run(_outer(patched), lambda: tiny_h3.run(patched))
assert not torch.equal(outs[0], base[0])
messages = [e["message"] for e in log.history()]
assert any("prompt stronger" in m and "blocks 0-3" in m for m in messages)
assert any("prompt keys scaled in" in m for m in messages)
def test_with_prompt_focus_off_nothing_about_the_prompt_is_drawn_or_stored(tiny_h3, monkeypatch):
from modules.system.taste import store
patched, _ = _load(tiny_h3, blocks="0-3", prompt_focus=False, regions=False)
_force(monkeypatch, arms=1)
outs = _run(_outer(patched), lambda: tiny_h3.run(patched))
assert torch.equal(outs[0], tiny_h3.run()[0])
store.rate("run-1", "liked")
assert "tarms" not in store.load("fox", "attention_bandit")["rows"][0]["rows"]
def test_clip_regions_change_only_the_rows_they_name(tiny_h3, monkeypatch):
from modules.sampling.attention_bandit import REGIONS, region_spans
captured = {}
def spy(func, q, k, v, *a, **kw):
options = kw.get("transformer_options") or {}
# the prompt refiner also runs with the layout in its options, on the prompt rows alone
if "minimax_h3_layout" in options and q.shape[2] > captured.get("rows", 0):
captured["spans"] = region_spans(options)
captured["rows"] = q.shape[2]
return func(q, k, v, *a, **kw)
from core import dit_hooks
probe = tiny_h3.patcher.clone()
dit_hooks.add_attention_override(probe, "test.spy", spy)
tiny_h3.run(probe)
spans = captured["spans"]
assert len(spans) == len(REGIONS)
flat = sorted(spans)
assert all(a[1] <= b[0] for a, b in zip(flat, flat[1:])) # no row in two regions
assert all(0 <= a < b <= captured["rows"] for a, b in spans)
def test_clip_regions_change_the_output_and_stay_off_when_all_as_is(tiny_h3, monkeypatch):
from core import log
base = tiny_h3.run()
patched, _ = _load(tiny_h3, mode="manual", range=0.5, blocks="0-3", prompt_focus=False)
_force(monkeypatch, arms=1, rarms=0)
# prompt focus is off, so the draws are groups then regions
from modules.sampling import attention_bandit as m
picks = iter([1, 0])
monkeypatch.setattr(m, "draw", lambda table, rng=random: [next(picks)] * len(table))
log.new_run()
outs = _run(_outer(patched), lambda: tiny_h3.run(patched))
assert not torch.equal(outs[0], base[0])
messages = [e["message"] for e in log.history()]
assert any("clip regions sharper" in m and "sound" in m for m in messages)
assert any("clip regions in" in m for m in messages)
def test_the_real_sla_override_underneath_still_gets_a_normal_call(tiny_h3, monkeypatch):
"""SLA hands anything off its fast path to dense and records a kernel failure if a call is odd.
Regions rescale rows of q in place on a copy, so shape, dtype and layout must be what SLA saw before."""
from core import dit_hooks
from modules.loaders import sla_attention as sla
seen = []
def spy(func, q, k, v, *a, **kw):
# strides are left out on purpose: the copy of q is compact, which SLA reads as well as the fused-buffer view
seen.append((tuple(q.shape), q.dtype, tuple(k.shape), kw.get("skip_reshape"), kw.get("mask")))
return func(q, k, v, *a, **kw)
def run_with(on):
probe = tiny_h3.patcher.clone()
state = sla.new_state()
to = dit_hooks._transformer_options(probe)
cfg = {**sla.SLA_DEFAULTS, "engine": "triton", "dense_steps": frozenset(), "min_seq_len": 0}
to["optimized_attention_override"] = sla.make_override(state, cfg)
dit_hooks.add_attention_override(probe, "test.spy", spy) # sits on top: sees exactly what SLA would get
patched, _ = _load(type("T", (), {"patcher": probe})(), mode="manual", range=0.5, blocks="0-3")
picks = iter([1, 1, 0 if on else 1])
from modules.sampling import attention_bandit as m
monkeypatch.setattr(m, "draw", lambda table, rng=random: [next(picks)] * len(table))
seen.clear()
_run(_outer(patched), lambda: tiny_h3.run(patched))
return list(seen), state
plain, state_plain = run_with(False)
scaled, state_scaled = run_with(True)
assert plain and scaled == plain # the same calls, same shapes, strides and flags
assert state_scaled["failed"] is None and state_scaled["calls"] == state_plain["calls"]