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

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

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

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)

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

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)