Files

523 lines
25 KiB
Python

"""Attention explorer: each run tries a small attention change, and your ratings pick which ones stay.
The chosen blocks are split by depth into a few groups. Every run, each group draws
one of three softmax temperatures -- sharper, as is, looser -- from how that group's
earlier runs were rated (Thompson draws, like Shot memory). A like makes the change
that run used more likely for that group, a dislike less. Groups drift toward what
you liked and away from what you did not, and the prompt still leads: temperature
only re-weighs what is already in context (see Randomizer).
"As is" starts with a head start so a fresh key stays close to the plain model.
How far a change goes learns too: a liked run that changed something pulls the size
toward the size it used, a disliked one pushes it away (as Shot memory does for its
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
ID = "attention_bandit"
TITLE = "Attention explorer"
MOUNT = "generation.sampling"
STAGE = "sampling"
CATEGORY = "sampling"
STATUS = "experimental"
REQUIRES = ["block_scope"]
USES = ["taste_store"]
SETTINGS = {
"enabled": {
"type": "bool", "default": False,
"label": "Attention explorer",
"hint": "Each run tries a small attention change; your ratings decide which ones stay.",
},
"mode": {
"type": "enum", "default": "learned", "options": [
{"value": "learned", "label": "Learn from ratings"},
{"value": "manual", "label": "My value"},
],
"label": "How far", "ui": "segmented",
"hint": "Learn from ratings: each rating nudges the size of a change toward what you liked. My value: you set it below.",
"when": {"enabled": True},
},
"range": {
"type": "float", "default": 0.2, "min": 0.05, "max": 0.5, "step": 0.05,
"label": "Change size", "ui": "slider",
"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},
},
"sound": {
"type": "bool", "default": True,
"label": "Change the sound too",
"hint": "Off: only the picture's attention changes; the sound is left as the model makes it.",
"when": {"enabled": True},
},
"blocks": {
"type": "text", "model_blocks": True, "default": "2-10,29-46",
"label": "Blocks",
"hint": "e.g. 2-10,29-46. Split by depth into groups that learn separately.",
"when": {"enabled": True},
},
}
KIND = "attention_bandit"
GROUPS = 4
# 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
NEW_HEAD_START = 4.0 # the same for prompt focus and clip regions, which change less often until they have data
# A factor on the prompt's keys only is not the same size of change as a temperature: it moves the prompt's
# logits against everyone else's, and at the full size it wrecked a trained key. A fraction of the size.
PROMPT_SCALE = 0.3
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
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):
log.once(f"{ID}:{message}", level, "FunUI Attention explorer", message)
def split(blocks, groups=GROUPS):
"""Sorted block indices -> up to `groups` lists, by depth, sizes differing by at most one."""
n = min(groups, len(blocks))
size, extra = divmod(len(blocks), n)
out, at = [], 0
for i in range(n):
step = size + (1 if i < extra else 0)
out.append(list(blocks[at:at + step]))
at += step
return out
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 i, r in enumerate(rows):
p = r["rows"]
if field not in p or r["reward"] == 0 or not _same_signature(p, signature):
continue
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
def _sign(x):
return (x > 0) - (x < 0)
def learned_spread(rows, signature):
"""The size of a change, replayed from the rated runs in order: a liked run that changed
something pulls it toward the size it used, a disliked one pushes it away."""
m = START_SPREAD
for r in rows:
p = r["rows"]
if "spread" not in p or "arms" not in p or not r["reward"] or not _same_signature(p, signature):
continue
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
# size down from a bigger one, never up from a smaller one (that would ratchet to the cap).
m += LR * (used - m) if r["reward"] > 0 else LR * min(0.0, m - used)
m = min(max(m, SPREAD_MIN), LEARN_MAX)
return m
def tally(rows):
"""'N rated runs: L liked, D disliked' over everything this key holds for this kind."""
liked = sum(1 for r in rows if r["reward"] > 0)
disliked = sum(1 for r in rows if r["reward"] < 0)
return f"{len(rows)} rated runs ({liked} liked, {disliked} disliked)"
def ranges(blocks):
out, start = [], None
for i, b in enumerate(blocks):
if start is None:
start = b
if i + 1 == len(blocks) or blocks[i + 1] != b + 1:
out.append(str(start) if start == b else f"{start}-{b}")
start = None
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, head=STAY_HEAD_START):
"""One arm per group: the arm whose Beta draw (likes + 1, dislikes + 1) is highest."""
picks = []
for group in table:
best, best_theta = STAY, -1.0
for arm, (liked, disliked) in enumerate(group):
theta = rng.betavariate(liked + (head if arm == STAY else 0.0) + 1.0, disliked + 1.0)
if theta > best_theta:
best, best_theta = arm, theta
picks.append(best)
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, head=NEW_HEAD_START) if focus else None
rarms = draw(counts(rows, signature, len(REGIONS), "rarms", weights), rng, head=NEW_HEAD_START) 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 followed(factors, options, index, blocks):
"""(temperature, prompt gain) for one block, nudged by the block map when this model has one: each change
from 1 is spread by how much that block can take at this point of the run and softly limited."""
if (options or {}).get("funui_block_map") is None:
return factors
return tuple(max(0.05, 1.0 + dit_hooks.block_map_value(options, "attention", index, v - 1.0, blocks)) for v in factors) # never zero or negative: that would flip the query
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"]]
def spare_sound(q, temperature, options, paired, sound_width):
"""q cooled by `temperature` on every row but the sound's, or None when this call cannot say which
rows are the sound. A model with a stream of its own for the sound: that call is left whole. A packed
sequence: the sound's rows are put back as they were."""
if paired:
return q if int(q.shape[-1]) == sound_width else q / temperature
audio = _segments(options).get("audio")
if audio is None or q.ndim != 4 or audio[1] > q.shape[2]:
return None
out = q / temperature
out[:, :, audio[0]:audio[1]] = q[:, :, audio[0]:audio[1]]
return out
def paired_regions(q, factors, video_width, sound_width):
"""q with each clip region's rows asking differently, for a model whose picture and sound attend in
calls of their own (rows run frame by frame): the picture's self-attention is cut into BANDS equal
time bands, the sound's self-attention is the last region. Any other call comes back as it was."""
width = int(q.shape[-1])
if width == video_width:
edges = [int(q.shape[1]) * i // BANDS for i in range(BANDS + 1)]
spans = list(zip(edges[:-1], edges[1:]))
mine = factors[:BANDS]
elif width == sound_width:
spans, mine = [(0, int(q.shape[1]))], factors[BANDS:BANDS + 1]
else:
return q
if all(f == 1.0 for f in mine):
return q
q = q.clone()
for (a, b), factor in zip(spans, mine):
if factor != 1.0:
q[:, a:b] /= factor
return q
# Carried on the per-block copy of transformer_options, as Randomizer does, so an
# attention call knows its block's factors without a shared flag.
_SCOPE = "funui_attention_bandit"
def install(patcher, values, key):
if not values.get("enabled"):
return None
taste = registry.current().ask("taste_store", patcher)
if taste is None:
_say("off: set a Taste key first -- it is what your ratings teach")
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))
sound = bool(values.get("sound", True))
blocks, problems = dit_hooks.parse_blocks(values.get("blocks"), dit_hooks.block_count(patcher))
for problem in problems:
_say(f"ignored {problem}")
if not blocks:
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 = {"factors": {}, "regions": [1.0] * len(REGIONS), "scaled": 0, "calls": 0,
"text": 0, "text_wanted": False, "region_hits": 0, "region_wanted": False,
"spared": 0, "unsure": 0}
captured = taste.collect(patcher, key, KIND, keep=MAX_ROWS)
dm = getattr(patcher.model, "diffusion_model", None)
paired = hasattr(dm, "transformer_blocks") # picture and sound attend in separate calls, one stream each
video_width, sound_width = getattr(dm, "inner_dim", None), getattr(dm, "audio_inner_dim", None)
if paired:
signature.append(1) # ratings made on a packed-sequence model (H3) must not teach a paired one
silent = not paired and getattr(dm, "audio_latents_dim", 1) is None # a picture-only model (H3-Image): no sound to spare, no clip time to cut into regions
if silent:
regions, sound = False, True
signature.append(2) # and what was rated on a model with a sound stream must not teach one without
if focus and paired:
focus = False # the prompt's rows are not one run of the sequence here: that part is H3's
if not sound and not paired:
focus = False # the prompt's keys are read by the sound's rows too, and one call cannot tell them apart
def scoped_for(index):
def scoped(args, extra):
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] = followed(f, to, index, blocks)
return extra["original_block"]({**args, "transformer_options": to})
return scoped
for index in blocks:
dit_hooks.add_block_hook(patcher, key, index, scoped_for(index))
def override(func, q, k, v, *args, **kwargs):
options = kwargs.get("transformer_options") or {}
scope = options.get(_SCOPE)
live["calls"] += 1
if scope is None:
return func(q, k, v, *args, **kwargs)
temperature, prompt = scope
live["scaled"] += 1
if temperature != 1.0:
if sound:
q = q / temperature
else:
spared = spare_sound(q, temperature, options, paired, sound_width)
if spared is None:
live["unsure"] += 1
spared = q / temperature
else:
live["spared"] += 1
q = spared
if live["region_wanted"] and paired and q.ndim == 3 and q.shape == k.shape:
q = paired_regions(q, live["regions"], video_width, sound_width)
live["region_hits"] += 1
elif 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 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 + PROMPT_SCALE * spread * ARMS[a] for a in tarms] if focus else [1.0] * len(groups)
if regions and not sound:
rarms = list(rarms[:-1]) + [STAY]
rfac = [1.0 + spread * ARMS[a] for a in rarms] if regions else [1.0] * len(REGIONS)
live.update(spared=0, unsure=0, 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, PROMPT_SCALE * 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)
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 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 not sound and live["unsure"]:
_say(f"the sound could not be told apart from the picture in {live['unsure']} attention calls, "
"so those changed both.", 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", prompt keys scaled in {live['text']}" if live["text_wanted"] else "")
+ (f", clip regions in {live['region_hits']}" if live["region_wanted"] else "")
+ (f", sound left alone in {live['spared']}" if not sound else ""))
return out
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}"
+ (", prompt focus on" if focus else "") + (", clip regions on" if regions else "") + ("" if sound else ", sound kept as is"))
PROVIDES = {"modifier": install}