diff --git a/modules/sampling/shot_memory/__init__.py b/modules/sampling/shot_memory/__init__.py new file mode 100644 index 00000000..1c9f7df6 --- /dev/null +++ b/modules/sampling/shot_memory/__init__.py @@ -0,0 +1,271 @@ +"""Shot memory: the starting noise remembers the shots you liked. + +With the prompt fixed, what a like or dislike reacts to is mostly the SHOT -- where +things sit, how the frame is split. The model decides that from the coarse part of +the starting noise, so a rating lands on the thing that caused it instead of on an +averaged direction in a space thousands of dimensions wide. + +Each run's coarse noise (one value per 4x4-latent cell per channel, averaged over +the whole clip) waits on the taste key for a rating. The next run then either +reuses a liked shot -- its coarse noise blended into the fresh noise, the fine +noise always fresh so the details are new -- or starts fresh. Both choices learn +from ratings (Thompson draws): whether reusing beats fresh for this key, which +liked shot, and how strongly (one number, nudged toward what was liked). A liked +child is a shot in its own right, so good shots breed. + +The noise stays unit Gaussian: swapping one block mean for another of the same +variance changes nothing the model can tell apart from ordinary noise. Zero model +calls. Sound noise is never touched. A latent that is not empty (a second pass, an +anchor) is left alone, and the run says so. + +Unvalidated on a GPU. v4 lessons kept: reusing a shot's own seed stacks its layout +1.41x (the blend is capped per channel, only ever turned down), and a one-cell grid +has no spread (left alone). +""" + +import math +import random + +import torch +import torch.nn.functional as F +from comfy.patcher_extension import WrappersMP + +from ..._core import log, registry, streams + +ID = "shot_memory" +TITLE = "Shot memory" +MOUNT = "generation.sampling" +STAGE = "latent" +CATEGORY = "sampling" +STATUS = "experimental" +REQUIRES = ["temporal_latent"] + +SETTINGS = { + "enabled": { + "type": "bool", "default": False, + "label": "Shot memory", + "hint": "Starts from the layout of shots you liked, with new details.", + }, + "mode": { + "type": "enum", "default": "learned", "options": [ + {"value": "learned", "label": "Learn from ratings"}, + {"value": "manual", "label": "My value"}, + ], + "label": "How much", "ui": "segmented", + "hint": "Learn: how strongly to reuse is tuned by your ratings.", + "when": {"enabled": True}, + }, + "amount": { + "type": "float", "default": 0.7, "min": 0.0, "max": 0.95, "step": 0.05, + "label": "Reuse amount", "ui": "slider", + "hint": "How much of the liked layout carries over. 0 = none.", + "when": {"enabled": True, "mode": "manual"}, + }, +} + +KIND = "shot_memory" +# Latent pixels per coarse cell: 64 image px on H3's 16x VAE. +# ponytail: fixed; learn it per key if one cell size stops fitting every resolution. +CELL = 4 +MAX_ROWS = 64 +START_AMOUNT = 0.7 +EXPLORE = 0.15 +LR = 0.5 +AMOUNT_MIN, AMOUNT_MAX = 0.2, 0.95 + + +def _say(message, level=log.ALERT): + log.once(f"{ID}:{message}", level, "FunPack Shot memory", message) + + +def _sign(x): + return (x > 0) - (x < 0) + + +def coarse_of(video): + """[1, C, T, H, W] noise -> its coarse part [C, h, w], unit variance, or None.""" + if video.dim() != 5 or video.shape[0] != 1: + return None + _b, c, t, H, W = video.shape + h, w = H // CELL, W // CELL + if h < 2 or w < 2: # a layout needs cells to compare; one cell has no spread + return None + blocks = video[0, :, :, :h * CELL, :w * CELL].float().reshape(c, t, h, CELL, w, CELL) + return blocks.mean(dim=(1, 3, 5)) * math.sqrt(t * CELL * CELL) + + +def with_coarse(video, coarse): + """`video` with its coarse part replaced by `coarse` [C, h, w]; the fine part kept.""" + _b, c, t, _H, _W = video.shape + h, w = coarse.shape[-2:] + core = video[0, :, :, :h * CELL, :w * CELL].float().reshape(c, t, h, CELL, w, CELL) + delta = coarse.to(core) / math.sqrt(t * CELL * CELL) - core.mean(dim=(1, 3, 5)) + core = core + delta[:, None, :, None, :, None] + out = video.clone() + out[0, :, :, :h * CELL, :w * CELL] = core.reshape(c, t, h * CELL, w * CELL).to(video.dtype) + return out + + +def fit(coarse, h, w): + """A stored coarse map on this run's grid. Another resolution is resampled and + re-standardised: the layout survives, exact Gaussianity does not.""" + coarse = coarse.float() + if tuple(coarse.shape[-2:]) == (h, w): + return coarse + x = F.interpolate(coarse[None], size=(h, w), mode="bilinear", align_corners=False)[0] + return (x - x.mean()) / x.std().clamp(min=1e-6) + + +def pooled(cond): + if not torch.is_tensor(cond): + return None + p = cond.detach().float() + while p.dim() > 1: + p = p.mean(dim=0) + return p.cpu() + + +def shots(rows): + """The key's rated shots, oldest first, as plain dicts.""" + out = [] + for r in rows: + p = r["rows"] + if "coarse" in p and "id" in p: + out.append({"id": int(p["id"]), "parent": int(p["parent"]), "amount": float(p["amount"]), + "reward": float(r["reward"]), "coarse": p["coarse"], "cond": p.get("cond")}) + return out + + +def learned_amount(rows): + """How strongly to reuse, replayed from the rated reuses in order: a liked one pulls + the value toward the amount it used, a disliked one pushes it away.""" + m = START_AMOUNT + for r in rows: + if r["parent"] < 0: + continue + m += LR * _sign(r["reward"]) * (r["amount"] - m) + m = min(max(m, AMOUNT_MIN), AMOUNT_MAX) + return m + + +def choose_parent(rows, cond, channels, rng=random): + """-> (liked shot to reuse | None, why). Two Thompson draws: reuse-vs-fresh from how + reuses and fresh shots have been rated, then which liked shot, weighted by how alike + its prompt was.""" + by_id = {r["id"] for r in rows} + reuse_ab, fresh_ab, kids = [1.0, 1.0], [1.0, 1.0], {} + for r in rows: + s = _sign(r["reward"]) + if not s: + continue + (fresh_ab if r["parent"] < 0 else reuse_ab)[0 if s > 0 else 1] += 1.0 + if r["parent"] in by_id: + kids.setdefault(r["parent"], [1.0, 1.0])[0 if s > 0 else 1] += 1.0 + liked = [r for r in rows if r["reward"] > 0 and int(r["coarse"].shape[0]) == int(channels)] + if not liked: + return None, "fresh: no liked shot yet" + reuse, fresh_ = rng.betavariate(*reuse_ab), rng.betavariate(*fresh_ab) + if fresh_ >= reuse: + return None, f"fresh (reuse {reuse:.2f} < fresh {fresh_:.2f})" + best, best_theta = None, -1.0 + for r in liked: + a, b = kids.get(r["id"], [1.0, 1.0]) + theta = rng.betavariate(a + 1.0, b) # +1: its own like + stored = r["cond"] + if cond is not None and torch.is_tensor(stored) and stored.numel() == cond.numel(): + # ponytail: raw pooled-prompt cosine; sits high for most prompts under one key. + theta *= float(F.cosine_similarity(cond, stored.float(), dim=0).clamp(0.0, 1.0)) + if theta > best_theta: + best, best_theta = r, theta + return best, f"reusing a liked shot of {len(liked)} (reuse {reuse:.2f} > fresh {fresh_:.2f})" + + +def blend(own, parent_coarse, amount): + """The layout to start from: the liked shot's, `amount` of the way.""" + coarse = amount * fit(parent_coarse, *own.shape[-2:]) + math.sqrt(1.0 - amount ** 2) * own + # Unit variance only holds when the two layouts are unrelated. Rerunning the liked + # shot's own seed makes them the same pattern and the blend stacks it (1.41x at 0.7): + # back to a normal-strength layout, per channel. Only ever down: a few-cell grid's + # measured spread can read low by chance, and dividing by it would strengthen it. + return coarse / coarse.std(dim=(1, 2), keepdim=True).clamp(min=1.0) + + +def _prompt_of(guider): + try: + return pooled(guider.conds["positive"][0]["cross_attn"]) + except (AttributeError, KeyError, IndexError, TypeError): + return None + + +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_amount = min(max(float(values.get("amount", START_AMOUNT)), 0.0), AMOUNT_MAX) + live = {"rows": [], "plan": None} + + def fresh(): + live["rows"], live["plan"] = shots(taste.rows(KIND)), None + liked = sum(1 for r in live["rows"] if r["reward"] > 0) + log.once(f"{ID}:state", log.INFO, "FunPack Shot memory", + f"key {taste.key!r}: {liked} liked shot(s) of {len(live['rows'])} rated") + + captured = taste.collect(patcher, key, KIND, keep=MAX_ROWS, fresh=fresh) + + def plan(cond, channels): + if live["plan"] is None: + parent, why = choose_parent(live["rows"], cond, channels) + amount = 0.0 + if parent is not None: + amount = manual_amount if manual else min(max( + learned_amount(live["rows"]) + random.gauss(0.0, EXPLORE), AMOUNT_MIN), AMOUNT_MAX) + why += f", amount {amount:.2f}" + log.once(f"{ID}:plan", log.INFO, "FunPack Shot memory", why) + live["plan"] = (parent, amount) + return live["plan"] + + def sampler_sample(executor, model_wrap, sigmas, extra_args, callback, noise, + latent_image=None, denoise_mask=None, disable_pbar=False): + def run(n): + return executor(model_wrap, sigmas, extra_args, callback, n, + latent_image, denoise_mask, disable_pbar) + + named = {"latent_shapes": getattr(getattr(model_wrap, "inner_model", None), "latent_shapes", None)} + split = streams.video_of(noise, named) + if split is None: + _say("off this run: could not find the picture in this model's noise") + return run(noise) + video, rebuild = split + base = streams.video_of(latent_image, named) if latent_image is not None else None + if base is None or bool(torch.count_nonzero(base[0])): + _say("Inactive | the starting latent is not empty (a second pass or an anchor), " + "so shot memory leaves this run's noise alone") + return run(noise) + own = coarse_of(video) + if own is None: + _say(f"Inactive | the picture is under {2 * CELL}x{2 * CELL} latent px (two layout " + "cells a side) or batched, so shot memory leaves it alone") + return run(noise) + cond = _prompt_of(model_wrap) + parent, amount = plan(cond, int(own.shape[0])) + coarse = own + if parent is not None: + coarse = blend(own, parent["coarse"], amount) + video = with_coarse(video, coarse) + captured.update({ + "id": torch.tensor(random.getrandbits(62)), + "parent": torch.tensor(parent["id"] if parent else -1), + "amount": torch.tensor(float(amount)), "coarse": coarse.half().cpu()}) + if cond is not None: + captured["cond"] = cond.half() + return run(rebuild(video)) + + patcher.add_wrapper_with_key(WrappersMP.SAMPLER_SAMPLE, key, sampler_sample) + return ("learned" if not manual else f"manual {manual_amount:g}") + ": starts from a liked layout when one is drawn" + + +PROVIDES = {"modifier": install} diff --git a/modules/sampling/shot_memory/tests/test_shot_memory.py b/modules/sampling/shot_memory/tests/test_shot_memory.py new file mode 100644 index 00000000..6dbb734d --- /dev/null +++ b/modules/sampling/shot_memory/tests/test_shot_memory.py @@ -0,0 +1,173 @@ +"""Shot memory: the noise maths, the Thompson choice, and the wrapper on packed H3 noise.""" + +import math +import random +import types + +import pytest +import torch + + +@pytest.fixture(autouse=True) +def _needs_comfy(comfyui): + """Imports comfy.""" + + +@pytest.fixture(autouse=True) +def root(tmp_path, monkeypatch): + from modules.system.taste import store + monkeypatch.setattr(store, "ROOT", tmp_path / "taste") + monkeypatch.setattr(store, "current_prompt_id", lambda: "run-1") + + +SHAPES = [torch.Size([1, 4, 2, 8, 8]), torch.Size([1, 3, 5])] +N_VIDEO = 4 * 2 * 8 * 8 + + +def test_the_coarse_part_is_unit_variance_and_swapping_it_keeps_the_fine_part(): + from modules.sampling.shot_memory import coarse_of, with_coarse + torch.manual_seed(0) + v = torch.randn(1, 4, 6, 32, 32) + c = coarse_of(v) + assert c.shape == (4, 8, 8) and c.std() == pytest.approx(1.0, abs=0.15) + other = torch.randn_like(c) + w = with_coarse(v, other) + assert torch.allclose(coarse_of(w), other, atol=1e-4) + fine = lambda x: x[0, :, :, :32, :32].reshape(4, 6, 8, 4, 8, 4) - x[0].reshape(4, 6, 8, 4, 8, 4).mean(dim=(1, 3, 5), keepdim=True) + assert torch.allclose(fine(w), fine(v), atol=1e-5) + + +def test_one_cell_grids_are_left_alone(): + from modules.sampling.shot_memory import coarse_of + assert coarse_of(torch.randn(1, 4, 2, 4, 4)) is None and coarse_of(torch.randn(2, 4, 2, 16, 16)) is None + + +def test_the_blend_never_raises_the_layout_strength_even_for_the_same_seed(): + from modules.sampling.shot_memory import blend + c = torch.randn(4, 8, 8) + out = blend(c, c.clone(), 0.7) # rerunning the liked shot's own seed + assert out.std(dim=(1, 2)).max() <= 1.0 + 1e-5 + + +def _row(i, parent, reward, channels=4, amount=0.7): + return {"id": i, "parent": parent, "amount": amount, "reward": reward, + "coarse": torch.randn(channels, 8, 8), "cond": None} + + +def test_no_liked_shot_means_fresh_and_a_liked_one_can_be_reused(): + from modules.sampling.shot_memory import choose_parent + assert choose_parent([_row(1, -1, -1.0)], None, 4)[0] is None + rows = [_row(1, -1, 1.0)] + [_row(i, 1, 1.0) for i in range(2, 12)] # reuse keeps winning + hits = sum(choose_parent(rows, None, 4, random.Random(s))[0] is not None for s in range(50)) + assert hits > 40 + assert choose_parent([_row(1, -1, 1.0, channels=8)], None, 4)[0] is None # other model's channels + + +def test_the_learned_amount_moves_toward_liked_reuses_and_away_from_disliked(): + from modules.sampling.shot_memory import START_AMOUNT, learned_amount + assert learned_amount([_row(1, -1, 1.0)]) == START_AMOUNT # a fresh shot says nothing about amount + assert learned_amount([_row(2, 1, 1.0, amount=0.9)]) > START_AMOUNT + assert learned_amount([_row(2, 1, -1.0, amount=0.9)]) < START_AMOUNT + + +class Guider: + def __init__(self): + self.inner_model = types.SimpleNamespace(latent_shapes=SHAPES) + self.conds = {"positive": [{"cross_attn": torch.randn(1, 5, 8)}]} + + +class Executor: + def __init__(self): + self.noise = None + + def __call__(self, guider, sigmas, extra_args, callback, noise, *rest): + self.noise = noise + return noise + + +def _load(tiny, **values): + from comfy.patcher_extension import WrappersMP + from modules.sampling.modifiers.nodes import FunPackLoadModifiers + patched, _ = FunPackLoadModifiers.execute(tiny.patcher, { + "taste": {"key": "fox"}, "shot_memory": {"enabled": True, **values}}).result + sample = [w for ws in patched.wrappers[WrappersMP.SAMPLER_SAMPLE].values() for w in ws][-1] + outer = [w for ws in patched.wrappers[WrappersMP.OUTER_SAMPLE].values() for w in ws] + return sample, outer + + +def _go(sample, outer, noise, latent=None, guider=None): + ex = Executor() + latent = torch.zeros_like(noise) if latent is None else latent + + def call(): + return sample(ex, guider or Guider(), torch.tensor([1.0, 0.0]), {"model_options": {}}, None, + noise, latent, None, True) + + run = call + for o in outer: + run = (lambda w, inner: (lambda: w(lambda: inner())))(o, run) + run() + return ex.noise + + +def _noise(): + from conftest import packed_av + torch.manual_seed(4) + return packed_av(torch.randn(1, 4, 2, 8, 8), torch.randn(1, 3, 5))[0] + + +def test_a_first_run_is_fresh_noise_is_untouched_and_the_shot_waits_for_a_rating(tiny_h3): + from modules.system.taste import store + sample, outer = _load(tiny_h3) + noise = _noise() + out = _go(sample, outer, noise) + assert torch.equal(out, noise) + assert store.rate("run-1", "liked")["recorded"] == ["shot_memory"] + row = store.load("fox", "shot_memory")["rows"][0]["rows"] + assert row["coarse"].shape == (4, 2, 2) and int(row["parent"]) == -1 + + +def test_once_a_shot_is_liked_the_next_run_starts_from_its_layout_and_keeps_sound_and_fine_noise(tiny_h3): + from modules.sampling.shot_memory import coarse_of + from modules.system.taste import store + from conftest import unpacked + sample, outer = _load(tiny_h3, mode="manual", amount=0.95) + first = _noise() + _go(sample, outer, first) + store.rate("run-1", "liked") + liked = store.load("fox", "shot_memory")["rows"][0]["rows"]["coarse"].float() + for _ in range(40): # Thompson may draw "fresh"; keep asking + from conftest import packed_av + second = packed_av(torch.randn(1, 4, 2, 8, 8), torch.randn(1, 3, 5))[0] + out = _go(sample, outer, second) + if not torch.equal(out, second): + break + else: + pytest.fail("a liked shot was never reused") + v_out, a_out = unpacked(out, SHAPES) + v_in, a_in = unpacked(second, SHAPES) + assert torch.equal(a_out, a_in) # sound noise untouched + cos = lambda a, b: torch.nn.functional.cosine_similarity(a.flatten(), b.flatten(), dim=0) + assert cos(coarse_of(v_out), liked) > cos(coarse_of(v_in), liked) + + +def test_a_non_empty_latent_is_left_alone_and_says_so_and_records_nothing(tiny_h3): + from core import log + from modules.system.taste import store + sample, outer = _load(tiny_h3) + noise = _noise() + log.new_run() + out = _go(sample, outer, noise, latent=torch.ones_like(noise)) + assert torch.equal(out, noise) + assert any("not empty" in e["message"] for e in log.history()) + assert store.rate("run-1", "liked")["why"] # nothing was captured to rate + + +def test_no_taste_key_is_off_and_says_so(tiny_h3): + from comfy.patcher_extension import WrappersMP + from core import log + from modules.sampling.modifiers.nodes import FunPackLoadModifiers + log.new_run() + patched, _ = FunPackLoadModifiers.execute(tiny_h3.patcher, {"shot_memory": {"enabled": True}}).result + assert any("Taste key" in e["message"] for e in log.history()) + assert not patched.wrappers.get(WrappersMP.SAMPLER_SAMPLE)