"""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}