Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
523 lines
25 KiB
Python
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}
|