From ca2d23b3b135aa105afa1c70a19b6f7cf7bf8962 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Sat, 3 Oct 2026 11:49:47 +0300 Subject: [PATCH] fix: momentum memory per window/call, probes and the last step excluded, exact v4 step size, freed after the run; prompt nudge keeps the sound's channels, v4's absolute strength, teaches only the positive prompt, says when it never acted Co-Authored-By: Claude Sonnet 5.5 --- modules/sampling/momentum/__init__.py | 54 ++++++++++++++----- .../sampling/momentum/tests/test_momentum.py | 52 ++++++++++++++++++ modules/sampling/prompt_nudge/__init__.py | 41 ++++++++++---- .../prompt_nudge/tests/test_prompt_nudge.py | 40 +++++++++++++- modules/sampling/score_slider/__init__.py | 14 ++++- 5 files changed, 174 insertions(+), 27 deletions(-) diff --git a/modules/sampling/momentum/__init__.py b/modules/sampling/momentum/__init__.py index 873d9fce..ee49b723 100644 --- a/modules/sampling/momentum/__init__.py +++ b/modules/sampling/momentum/__init__.py @@ -5,14 +5,15 @@ memory. Once the noise has fallen below a level, the current direction is blende which calms jitter in the refinement steps. It is the late-step partner of ALG, which acts on the early ones (arXiv:2602.20360). It can damp motion; that is the price of averaging. -The edit is carried into the next step's input (core/input_steer.py), never onto the last answer, and -only the picture is touched, never the sound. +The edit is carried into the next step's input (core/input_steer.py), never onto the last answer, so the +last step is never edited; it is sized so the next input moves exactly as v4's blended direction would. +Only the picture is touched, never the sound. """ import torch from comfy.patcher_extension import WrappersMP -from ..._core import input_steer, log, streams +from ..._core import dit_hooks, input_steer, log, streams ID = "momentum" TITLE = "Momentum guidance" @@ -58,6 +59,13 @@ def blend(x, out, sigma, ema, decay, strength): return x - sigma * mixed, ema +def carry_factor(sigma, sigma_next, carried=True): + """What to multiply an x0 edit by so the next input moves as the blended direction would have: + v4 moves x by (sigma - sigma_next) * dd; an x0 edit of -sigma * dd lands in x_next scaled by + (sigma - sigma_next) / sigma when applied to the answer, and by (1 - sigma_next) when carried in.""" + return (sigma - sigma_next) / max(sigma * (1.0 - sigma_next) if carried else sigma, 1e-6) + + def install(patcher, values, key): if not values.get("enabled"): return None @@ -65,31 +73,51 @@ def install(patcher, values, key): decay = max(0.0, min(0.999, float(values.get("decay", 0.5)))) below = float(values.get("below_sigma", 0.975)) if strength <= 0.0: - return None - state = {"ema": None, "last": None} + return "strength 0: nothing to do" + state = {"ema": {}, "last": {}} # by window and by cond/uncond call: they are different directions steer = input_steer.Steer("Momentum") steer.attach(patcher, key) + def forget(executor, *args, **kwargs): + state["ema"].clear() + state["last"].clear() + try: + return executor(*args, **kwargs) + finally: + state["ema"].clear() # a latent-sized tensor must not outlive the run + + patcher.add_wrapper_with_key(WrappersMP.OUTER_SAMPLE, f"{key}.forget", forget) + def apply_model(executor, x, t, *args, **kwargs): named = streams.model_args(args, kwargs) - step = steer.begin(x, t, named.get("transformer_options")) + to = named.get("transformer_options") + step = steer.begin(x, t, to) out = executor(step.x, t, *args, **kwargs) + if not step.steering: # a probe, an off-schedule call, the last step: not a step of the run + return out sigma = float(t.max()) - if state["last"] is not None and sigma > state["last"] + 1e-4: - state["ema"] = None # sigma went back up: a new run - state["last"] = sigma + which = steer._key(to or {}) + if state["last"].get(which) is not None and sigma > state["last"][which] + 1e-4: + state["ema"].pop(which, None) # sigma went back up: a new pass + state["last"][which] = sigma seen_out, seen_in = streams.video_of(out, named), streams.video_of(step.x, named) if seen_out is None or seen_in is None: log.once(f"{ID}:layout", log.ALERT, "FunPack Momentum", "off this run: could not find the picture in this model's latent") return out video, rebuild = seen_out + ema = state["ema"].get(which) + if ema is not None and ema.shape != video.shape: + ema = None # The memory always accumulates, so the window opens on real history, not a cold start. - new, state["ema"] = blend(seen_in[0].float(), video.float(), sigma, state["ema"], decay, - strength if sigma < below else 0.0) - if sigma >= below or not step.steering: + new, state["ema"][which] = blend(seen_in[0].float(), video.float(), sigma, ema, decay, + strength if sigma < below else 0.0) + if sigma >= below: return out - return step.keep(out, rebuild(new.to(video.dtype))) + where = dit_hooks.current_step(to) + sigma_next = float(to["sample_sigmas"][where[0] + 1]) if where else 0.0 + edit = video.float() + carry_factor(sigma, sigma_next, x.dim() == 3) * (new - video.float()) + return step.keep(out, rebuild(edit.to(video.dtype))) patcher.add_wrapper_with_key(WrappersMP.APPLY_MODEL, key, apply_model) return f"strength {strength:g}, memory {decay:g}, from noise level {below:g}" diff --git a/modules/sampling/momentum/tests/test_momentum.py b/modules/sampling/momentum/tests/test_momentum.py index bea5d4a3..d474e5e5 100644 --- a/modules/sampling/momentum/tests/test_momentum.py +++ b/modules/sampling/momentum/tests/test_momentum.py @@ -62,3 +62,55 @@ def test_off_installs_nothing(tiny_h3): from modules.sampling.modifiers.nodes import FunPackLoadModifiers patched, _ = FunPackLoadModifiers.execute(tiny_h3.patcher, {"momentum": {"enabled": False}}).result assert not patched.wrappers.get(WrappersMP.APPLY_MODEL) + + +def test_the_carried_edit_is_sized_to_move_the_next_input_like_v4s_blended_direction(): + from modules.sampling.momentum import carry_factor + sigma, nxt = 0.75, 0.5 + # an x0 edit of -sigma*dd, scaled by this and carried in (x (1 - nxt)), moves x by (sigma - nxt)*dd + assert abs(carry_factor(sigma, nxt) * sigma * (1 - nxt) - (sigma - nxt) * 1.0) < 1e-9 + assert abs(carry_factor(sigma, nxt, carried=False) * sigma - (sigma - nxt)) < 1e-9 + + +def _run_calls(wrap, shapes, x0, n, steps=4, ts_start=0, **extra): + seen = [] + ts = torch.linspace(1.0, 0.0, steps + 1) + for i in range(ts_start, ts_start + n): + def model(x, *a, **k): + seen.append(x) + return x0 * (0.9 - 0.2 * len(seen)) + wrap(model, x0, ts[i:i + 1], None, None, None, _to(i, steps), latent_shapes=shapes, **extra) + return seen + + +def test_a_probe_call_does_not_feed_the_memory(tiny_h3): + from conftest import packed_av + from core import dit_hooks + torch.manual_seed(0) + x0, shapes = packed_av(torch.randn(1, 4, 3, 8, 8), torch.randn(1, 8, 5)) + clean = _run_calls(_load(tiny_h3, strength=1.0, decay=0.5, below_sigma=1.0)[1], shapes, x0, 4) + wrap = _load(tiny_h3, strength=1.0, decay=0.5, below_sigma=1.0)[1] + ts = torch.linspace(1.0, 0.0, 5) + probe_to = {**_to(0), dit_hooks.PROBE: True} + for _ in range(3): + wrap(lambda x, *a, **k: x0 * 0.1, x0, ts[0:1], None, None, None, probe_to, latent_shapes=shapes) + after = _run_calls(wrap, shapes, x0, 4) + assert all(torch.allclose(a, b) for a, b in zip(clean, after)) + + +def test_each_window_and_each_call_kind_has_its_own_memory_and_a_new_run_starts_empty(tiny_h3): + from conftest import packed_av + torch.manual_seed(0) + patched, wrap, _ = _load(tiny_h3, strength=1.0, decay=0.5, below_sigma=1.0) + big, sb = packed_av(torch.randn(1, 4, 3, 8, 8), torch.randn(1, 8, 5)) + small, ss = packed_av(torch.randn(1, 4, 2, 8, 8), torch.randn(1, 8, 5)) + + class Win: + def __init__(self, idx): + self.index_list = idx + ts = torch.linspace(1.0, 0.0, 5) + for i in range(3): # two windows of different sizes interleaved: must not collide + for x0, shapes, idx in ((big, sb, [0, 1, 2]), (small, ss, [3, 4])): + to = {**_to(i), "context_window": Win(idx)} + wrap(lambda x, *a, **k: x0 * 0.5, x0, ts[i:i + 1], None, None, None, to, latent_shapes=shapes) + assert not patched.model_options["funpack_dropped"] # no size clash: the guard would have dropped it diff --git a/modules/sampling/prompt_nudge/__init__.py b/modules/sampling/prompt_nudge/__init__.py index 772f66e0..d2e54fc6 100644 --- a/modules/sampling/prompt_nudge/__init__.py +++ b/modules/sampling/prompt_nudge/__init__.py @@ -3,7 +3,8 @@ prompts have in common. The same learned direction as the Taste slider (liked prompts relative to everything rated), applied the cheap way: the words the model reads are nudged in place, once per step, with no extra model passes. It -grows from nothing at the halfway point to full strength at the last step. Words only: reference-image +grows from nothing at the halfway point as the steps go on (the last step reaches half to three +quarters of it, depending on the schedule length). Words only: reference-image rows are left alone. Needs 3 liked clips on the taste key. Shares the slider's learning (`prompt_taste`), so turning on either teaches both. @@ -13,7 +14,7 @@ import torch from comfy.patcher_extension import WrappersMP from ..._core import dit_hooks, log, registry, streams -from ..score_slider import KIND, MIN_LIKED, NAME, direction +from ..score_slider import KIND, MIN_LIKED, NAME, cond_row, direction ID = "prompt_nudge" TITLE = "Taste prompt nudge" @@ -48,13 +49,23 @@ def _say(message): log.once(f"{ID}:{message}", log.ALERT, "FunPack Taste prompt nudge", message) -def nudged(c, words, d, amount): - """`c` with each word row moved along unit direction `d` by `amount` x the words' own size.""" - size = torch.linalg.vector_norm(c[:, words], dim=-1, dtype=torch.float32).mean() - step = (d.to(c.device) * amount * size).to(c.dtype) +def nudged(c, words, d, amount, width=None): + """`c` with each word row moved by `amount` along unit direction `d` (v4's absolute size: the + strength is the row shift itself). `width` limits it to the first channels: an audio+video model's + words carry the sound's channels after the picture's, and those stay as they were.""" + step = (d.to(c.device) * amount).to(c.dtype) + if width is not None: + step = torch.cat([step[:width], torch.zeros_like(step[width:])]) return c + step * words.view(1, -1, 1).to(c.dtype) +def picture_width(patcher, c): + """How many leading channels of the words belong to the picture, or None when all do.""" + dm = getattr(getattr(patcher, "model", None), "diffusion_model", None) + video, audio = getattr(dm, "cross_attention_dim", None), getattr(dm, "audio_cross_attention_dim", None) + return int(video) if video and audio and int(c.shape[-1]) == int(video) + int(audio) else None + + def install(patcher, values, key): if not values.get("enabled"): return None @@ -68,8 +79,10 @@ def install(patcher, values, key): def fresh(): live.clear() + acted["yes"] = False captured = taste.collect(patcher, key, KIND, fresh=fresh) + acted = {"yes": False} def apply_model(executor, x, t, *args, **kwargs): named = streams.model_args(args, kwargs) @@ -82,20 +95,28 @@ def install(patcher, values, key): words = torch.ones(c.shape[1], dtype=torch.bool, device=c.device) words = words.to(c.device) to = named.get("transformer_options") - if not dit_hooks.probing(to): - captured[NAME] = c[0][words].float().mean(0).detach() + row = cond_row(c, to) + if not dit_hooks.probing(to) and row is not None: + captured[NAME] = c[row][words].float().mean(0).detach() if "dir" not in live: - live["dir"], how = direction(taste.rows(KIND), c[0][words].float().mean(0), similar) + live["dir"], how = direction(taste.rows(KIND), c[row or 0][words].float().mean(0), similar) log.once(f"{ID}:state", log.INFO, "FunPack Taste prompt nudge", f"key {taste.key!r}: {how}") amount = strength * dit_hooks.late_half(to) d = live["dir"] + if dit_hooks.last_step(to) and not acted["yes"] and not dit_hooks.probing(to): + acted["yes"] = d is not None and amount > 0.0 + if not acted["yes"]: + _say("Inactive | nothing was nudged this run: " + ( + "no direction learned yet (needs liked clips)" if d is None else + "the schedule is too short for the late-step gate to open")) if d is None or amount <= 0.0: return executor(x, t, *args, **kwargs) + acted["yes"] = True if d.numel() != c.shape[-1]: _say("off this run: the learned direction was taught on a different model's text width; " "rate a few clips on this model") return executor(x, t, *args, **kwargs) - named = {**named, "c_crossattn": nudged(c, words, d, amount)} + named = {**named, "c_crossattn": nudged(c, words, d, amount, picture_width(patcher, c))} return executor(x, t, **named) patcher.add_wrapper_with_key(WrappersMP.APPLY_MODEL, key, apply_model) diff --git a/modules/sampling/prompt_nudge/tests/test_prompt_nudge.py b/modules/sampling/prompt_nudge/tests/test_prompt_nudge.py index 23f1469b..7a7f6dda 100644 --- a/modules/sampling/prompt_nudge/tests/test_prompt_nudge.py +++ b/modules/sampling/prompt_nudge/tests/test_prompt_nudge.py @@ -26,8 +26,7 @@ def test_words_move_along_the_direction_by_a_share_of_their_own_size(): d = torch.tensor([1.0, 0, 0, 0]) words = torch.tensor([True, True, False]) out = nudged(c, words, d, 0.1) - size = float(c[0, 0].norm()) - assert torch.allclose(out[0, 0, 0], c[0, 0, 0] + 0.1 * size) + assert torch.allclose(out[0, 0, 0], c[0, 0, 0] + 0.1) # the strength is the row shift itself, as in v4 assert torch.equal(out[0, 2], c[0, 2]) # a reference image's row is left alone assert torch.equal(out[0, :, 1:], c[0, :, 1:]) @@ -91,3 +90,40 @@ def test_off_installs_nothing(tiny_h3): from modules.sampling.modifiers.nodes import FunPackLoadModifiers patched, _ = FunPackLoadModifiers.execute(tiny_h3.patcher, {"prompt_nudge": {"enabled": False}}).result assert not patched.wrappers.get(WrappersMP.APPLY_MODEL) + + +def test_the_sounds_text_channels_are_left_alone(): + from modules.sampling.prompt_nudge import nudged, picture_width + + class Dm: + cross_attention_dim, audio_cross_attention_dim = 6, 2 + + class P: + model = type("M", (), {"diffusion_model": Dm()})() + + c = torch.ones(1, 3, 8) + width = picture_width(P(), c) + assert width == 6 and picture_width(P(), torch.ones(1, 3, 5)) is None + out = nudged(c, torch.ones(3, dtype=torch.bool), torch.ones(8) / 8 ** 0.5, 1.0, width) + assert (out[..., :6] > 1).all() and torch.equal(out[..., 6:], c[..., 6:]) + + +def test_a_split_negative_call_does_not_teach_the_negative_prompt(tiny_h3): + from modules.sampling.score_slider import cond_row + c = torch.ones(2, 3, 4) + assert cond_row(c, {"cond_or_uncond": [1, 0]}) == 1 + assert cond_row(c[:1], {"cond_or_uncond": [1]}) is None + assert cond_row(c, {}) == 0 + + +def test_a_run_that_never_nudged_says_so(tiny_h3): + from core import log + _teach(8) + wrap, _ = _load(tiny_h3, strength=0.1) + log.new_run() + c = torch.ones(1, 4, 8) + for i in range(2): # a two-step schedule: the late gate never opens + s = torch.linspace(1.0, 0.0, 3) + wrap(lambda x, t, *a, **k: x, torch.zeros(1, 1, 4), s[i:i + 1], None, c, None, + {"sample_sigmas": s, "sigmas": s[i:i + 1]}) + assert any("Inactive" in e["message"] and "too short" in e["message"] for e in log.history()) diff --git a/modules/sampling/score_slider/__init__.py b/modules/sampling/score_slider/__init__.py index 775ed1a3..c3c0e01b 100644 --- a/modules/sampling/score_slider/__init__.py +++ b/modules/sampling/score_slider/__init__.py @@ -82,6 +82,15 @@ def direction(rows, current=None, similar=False): return F.normalize(d, dim=0), f"from {len(liked)} liked of {len(pooled)} rated" +def cond_row(c, to): + """The batch row holding the positive prompt, or None when this call carries only the negative. + A split CFG call (cond and uncond in separate calls) must not teach the negative prompt.""" + kinds = list((to or {}).get("cond_or_uncond") or []) + if not kinds: + return 0 + return next((i for i, k in enumerate(kinds) if k == 0 and i < c.shape[0]), None) + + def install(patcher, values, key): if not values.get("enabled"): return None @@ -112,8 +121,9 @@ def install(patcher, values, key): if words is None: words = torch.ones(c.shape[1], dtype=torch.bool, device=c.device) words = words.to(c.device) - pooled = c[0][words].float().mean(0).detach() - if not dit_hooks.probing(named.get("transformer_options")): + row = cond_row(c, named.get("transformer_options")) + pooled = c[row or 0][words].float().mean(0).detach() + if not dit_hooks.probing(named.get("transformer_options")) and row is not None: captured[NAME] = pooled if "dir" not in live: live["dir"], how = direction(taste.rows(KIND), pooled, similar)