"""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 and the sound is left as the model makes it. Prompt focus then stays off on models that mix picture and sound in one sequence.", "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 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}