From 8f1ce2a1fe0ef471a82ba5b9ef47fc7d9a27edc2 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Sat, 3 Oct 2026 11:37:22 +0300 Subject: [PATCH] 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 --- modules/sampling/momentum/__init__.py | 98 +++++++++++++++++++ .../sampling/momentum/tests/test_momentum.py | 64 ++++++++++++ 2 files changed, 162 insertions(+) create mode 100644 modules/sampling/momentum/__init__.py create mode 100644 modules/sampling/momentum/tests/test_momentum.py diff --git a/modules/sampling/momentum/__init__.py b/modules/sampling/momentum/__init__.py new file mode 100644 index 00000000..873d9fce --- /dev/null +++ b/modules/sampling/momentum/__init__.py @@ -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} diff --git a/modules/sampling/momentum/tests/test_momentum.py b/modules/sampling/momentum/tests/test_momentum.py new file mode 100644 index 00000000..bea5d4a3 --- /dev/null +++ b/modules/sampling/momentum/tests/test_momentum.py @@ -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)