feat: momentum guidance module (running-average direction blend below a noise level, carried into the next step's input)

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
2026-10-03 11:37:22 +03:00
parent 397f36014a
commit 8f1ce2a1fe
2 changed files with 162 additions and 0 deletions
@@ -0,0 +1,98 @@
+"""Momentum Guidance: smooth the fine-motion steps with a running average of past directions.
+
+Each step's direction (where the model says the picture should move) is averaged into a running
+memory. Once the noise has fallen below a level, the current direction is blended toward that memory,
+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.
+"""
+
+import torch
+from comfy.patcher_extension import WrappersMP
+
+from ..._core import input_steer, log, streams
+
+ID = "momentum"
+TITLE = "Momentum guidance"
+MOUNT = "generation.sampling"
+STAGE = "guidance"
+CATEGORY = "guidance"
+STATUS = "experimental"
+REQUIRES = ["temporal_latent"]
+
+SETTINGS = {
+    "enabled": {
+        "type": "bool", "default": False,
+        "label": "Smooth the fine motion",
+        "hint": "Calms jitter in the last steps. May also soften motion.",
+    },
+    "strength": {
+        "type": "float", "default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05,
+        "label": "How much", "ui": "slider",
+        "hint": "0 = nothing, 1 = follow the running average completely.",
+        "when": {"enabled": True},
+    },
+    "decay": {
+        "type": "float", "default": 0.5, "min": 0.0, "max": 0.95, "step": 0.05,
+        "label": "How long it remembers", "ui": "slider",
+        "hint": "High values hold on to the noisy first steps and make garbage on short schedules.",
+        "when": {"enabled": True},
+    },
+    "below_sigma": {
+        "type": "float", "default": 0.975, "min": 0.0, "max": 1.0, "step": 0.025,
+        "label": "Starts at noise level", "ui": "slider",
+        "hint": "Acts once the noise has dropped below this.",
+        "when": {"enabled": True},
+    },
+}
+
+
+def blend(x, out, sigma, ema, decay, strength):
+    """(new answer, new memory) for one step on a video tensor. The direction is
+    d = (x - answer) / sigma; the answer is rebuilt from d blended toward the memory."""
+    d = ((x - out) / max(sigma, 1e-6)).detach()
+    ema = d if ema is None else decay * ema + (1.0 - decay) * d
+    mixed = d + strength * (ema - d)
+    return x - sigma * mixed, ema
+
+
+def install(patcher, values, key):
+    if not values.get("enabled"):
+        return None
+    strength = max(0.0, min(1.0, float(values.get("strength", 0.5))))
+    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}
+    steer = input_steer.Steer("Momentum")
+    steer.attach(patcher, key)
+
+    def apply_model(executor, x, t, *args, **kwargs):
+        named = streams.model_args(args, kwargs)
+        step = steer.begin(x, t, named.get("transformer_options"))
+        out = executor(step.x, t, *args, **kwargs)
+        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
+        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
+        # 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:
+            return out
+        return step.keep(out, rebuild(new.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}"
+
+
+PROVIDES = {"modifier": install}
@@ -0,0 +1,64 @@
+import pytest
+import torch
+
+
+@pytest.fixture(autouse=True)
+def _needs_comfy(comfyui):
+    """Imports comfy."""
+
+
+def test_strength_zero_changes_nothing_and_one_follows_the_memory():
+    from modules.sampling.momentum import blend
+    x, out = torch.randn(1, 4, 2, 4, 4), torch.randn(1, 4, 2, 4, 4)
+    same, _ = blend(x, out, 0.5, None, 0.5, 0.0)
+    assert torch.allclose(same, out, atol=1e-5)
+    ema = torch.randn_like(x)
+    full, _ = blend(x, out, 0.5, ema, 0.0, 1.0)            # decay 0: memory is this step's direction
+    assert torch.allclose(full, out, atol=1e-5)
+    mixed, new = blend(x, out, 0.5, ema, 0.5, 1.0)
+    d = (x - out) / 0.5
+    assert torch.allclose(new, 0.5 * ema + 0.5 * d, atol=1e-5)
+    assert torch.allclose(mixed, x - 0.5 * new, atol=1e-5)
+
+
+def _load(tiny, **v):
+    from comfy.patcher_extension import WrappersMP
+    from modules.sampling.modifiers.nodes import FunPackLoadModifiers
+    patched, status = FunPackLoadModifiers.execute(
+        tiny.patcher, {"momentum": {"enabled": True, **v}}).result
+    wrap = [w for ws in patched.wrappers[WrappersMP.APPLY_MODEL].values() for w in ws][-1]
+    return patched, wrap, status
+
+
+def _to(i, steps=4):
+    s = torch.linspace(1.0, 0.0, steps + 1)
+    return {"sample_sigmas": s, "sigmas": s[i:i + 1]}
+
+
+def test_it_acts_below_the_threshold_and_carries_the_edit_into_the_next_step(tiny_h3):
+    from conftest import packed_av, unpacked
+    _p, wrap, status = _load(tiny_h3, strength=1.0, decay=0.5, below_sigma=0.8)
+    torch.manual_seed(0)
+    x0, shapes = packed_av(torch.randn(1, 4, 3, 8, 8), torch.randn(1, 8, 5))
+    answers = [x0 * f for f in (0.9, 0.5, 0.2, 0.1)]       # directions that differ step to step
+    seen = []
+
+    def model(x, *a, **k):
+        seen.append(x)
+        return answers[len(seen) - 1]
+
+    ts = torch.linspace(1.0, 0.0, 5)
+    for i in range(4):
+        out = wrap(model, x0, ts[i:i + 1], None, None, None, _to(i), latent_shapes=shapes)
+        assert torch.equal(out, answers[i])                  # the answer itself is never edited
+    assert torch.equal(seen[1], x0)                          # sigma 0.75 is the first step below 0.8: nothing was filed before it
+    push = unpacked(seen[2] - x0, shapes)                    # step 1 filed an edit; step 2 got it
+    assert not torch.allclose(push[0], torch.zeros_like(push[0]))
+    assert torch.equal(push[1], torch.zeros_like(push[1]))   # sound untouched
+
+
+def test_off_installs_nothing(tiny_h3):
+    from comfy.patcher_extension import WrappersMP
+    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)