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 <noreply@anthropic.com>

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())

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

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: " + (

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