From c4e5ceac6412e21c0ff77e48635d2ccc2bfb8f0b Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Sat, 3 Oct 2026 12:16:26 +0300 Subject: [PATCH] fix: momentum leaves second-order samplers alone and says so; nudge at strength 0 only learns; more tests Co-Authored-By: Claude Sonnet 5.5 --- modules/sampling/momentum/__init__.py | 17 ++++++++++-- .../sampling/momentum/tests/test_momentum.py | 14 ++++++++++ modules/sampling/prompt_nudge/__init__.py | 2 +- .../prompt_nudge/tests/test_prompt_nudge.py | 27 +++++++++++++++++++ 4 files changed, 57 insertions(+), 3 deletions(-) diff --git a/modules/sampling/momentum/__init__.py b/modules/sampling/momentum/__init__.py index 8eb43183..452838e9 100644 --- a/modules/sampling/momentum/__init__.py +++ b/modules/sampling/momentum/__init__.py @@ -6,7 +6,8 @@ which calms jitter in the refinement steps. It is the late-step partner of ALG, 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, so the -last step is never edited; it is sized so the next input moves exactly as v4's blended direction would. +last step is never edited; it is sized so the NEXT model input moves as v4's blended direction would. The sampler's own latent does +not take the edit, so later steps feel about three quarters of v4's effect (v4 edits the sampler's step). Only the picture is touched, never the sound. """ @@ -73,13 +74,15 @@ def install(patcher, values, key): below = float(values.get("below_sigma", 0.975)) if strength <= 0.0: return "strength 0: nothing to do" - state = {"ema": {}, "last": {}} # by cond/uncond call: they are different directions + state = {"ema": {}, "last": {}, "index": {}, "multi": False} # 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() + state["index"].clear() + state["multi"] = False try: return executor(*args, **kwargs) finally: @@ -99,6 +102,16 @@ def install(patcher, values, key): out = executor(step.x, t, *args, **kwargs) packed = x.dim() == 3 where = dit_hooks.current_step(to) + if not packed and where is not None and not dit_hooks.probing(to): + now, was = where[0], state["index"].get(tuple(int(v) for v in ((to or {}).get("cond_or_uncond") or ()))) + state["index"][tuple(int(v) for v in ((to or {}).get("cond_or_uncond") or ()))] = now + if now == was: # the same step twice: a second-order corrector call + state["multi"] = True + log.once(f"{ID}:multi", log.ALERT, "FunPack Momentum", + "off for the rest of this run: the model is called more than once per step " + "(a second-order sampler); use euler-style sampling") + if state["multi"]: + return out if not step.steering or dit_hooks.probing(to) or (not packed and where is not None and where[0] >= where[1] - 1): return out # a probe, an off-schedule call, the last step: not a step to edit sigma = float(t.max()) diff --git a/modules/sampling/momentum/tests/test_momentum.py b/modules/sampling/momentum/tests/test_momentum.py index 199c8c3a..72027856 100644 --- a/modules/sampling/momentum/tests/test_momentum.py +++ b/modules/sampling/momentum/tests/test_momentum.py @@ -161,3 +161,17 @@ def test_a_second_pass_that_restarts_lower_still_starts_with_an_empty_memory(tin wrap = fresh run(4, first + [0.1]) assert all(torch.allclose(a, b) for a, b in zip(twice, seen)) + + +def test_a_second_order_sampler_on_a_plain_latent_is_said_and_left_alone(tiny_h3): + from core import log + log.new_run() + _p, wrap, _ = _load(tiny_h3, strength=1.0, decay=0.5, below_sigma=1.0) + x = torch.randn(1, 4, 3, 8, 8) + sig = torch.linspace(1.0, 0.0, 5) + outs = [] + for i in (0, 1, 1, 2, 2, 3): # heun: each step's corrector lands on the next sigma + o = x * (0.9 - 0.1 * len(outs)) + outs.append(wrap(lambda *a, o=o, **kw: o, x, sig[i:i + 1], None, None, None, _to(i))) + assert any("more than once per step" in e["message"] for e in log.history()) + assert torch.equal(outs[-1], x * (0.9 - 0.5)) # nothing edited once it is known diff --git a/modules/sampling/prompt_nudge/__init__.py b/modules/sampling/prompt_nudge/__init__.py index a816416e..0802f2a9 100644 --- a/modules/sampling/prompt_nudge/__init__.py +++ b/modules/sampling/prompt_nudge/__init__.py @@ -105,7 +105,7 @@ def install(patcher, values, key): return executor(x, t, *args, **kwargs) amount = strength * dit_hooks.late_half(to) d = live["dir"] - if dit_hooks.last_step(to) and not acted.get("yes") and not dit_hooks.probing(to): + if strength > 0.0 and dit_hooks.last_step(to) and not acted.get("yes") and not dit_hooks.probing(to): acted["yes"] = d is not None and amount > 0.0 if not acted.get("yes"): _say("Inactive | nothing was nudged this run: " + ( diff --git a/modules/sampling/prompt_nudge/tests/test_prompt_nudge.py b/modules/sampling/prompt_nudge/tests/test_prompt_nudge.py index 8bf0dfee..ac4b7b48 100644 --- a/modules/sampling/prompt_nudge/tests/test_prompt_nudge.py +++ b/modules/sampling/prompt_nudge/tests/test_prompt_nudge.py @@ -149,3 +149,30 @@ def test_a_run_that_starts_with_the_negative_prompt_does_not_learn_from_it(tiny_ to = {**_to(3), "cond_or_uncond": [1]} wrap(lambda x, t, *a, **k: seen.append(a[1]) or x, torch.zeros(1, 1, 4), s[3:4], None, c, None, to) assert torch.equal(seen[0], c) # no positive prompt yet: nothing learned, nothing nudged + + +def test_strength_zero_only_learns_and_does_not_cry_wolf(tiny_h3): + from core import log + _teach(8) + wrap, status = _load(tiny_h3, strength=0.0) + log.new_run() + before = len(log.history()) + c = torch.ones(1, 4, 8) + ts = torch.linspace(1.0, 0.0, 5) + seen = [] + for i in range(4): + wrap(lambda x, t, *a, **k: seen.append(a[1]) or x, torch.zeros(1, 1, 4), ts[i:i + 1], None, c, None, _to(i)) + assert all(torch.equal(s, c) for s in seen) and "learning only" in str(status) + assert not [e["message"] for e in log.history()[before:] if "nothing was nudged" in e["message"]] + + +def test_a_probe_does_not_teach_the_taste_key(tiny_h3): + from core import dit_hooks + from modules.system.taste import store + _teach(8) + wrap, _ = _load(tiny_h3, strength=0.1) + c = torch.ones(1, 4, 8) + ts = torch.linspace(1.0, 0.0, 5) + wrap(lambda x, t, *a, **k: x, torch.zeros(1, 1, 4), ts[3:4], None, c, None, {**_to(3), dit_hooks.PROBE: True}) + pending = store._dir("fox") / "prompt_taste.pending.pt" + assert not pending.exists() # nothing was captured from the probe