From 8a241cf9be87b30e231cb1549ce8d77168d8f48e Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 25 Oct 2026 16:35:34 +0000 Subject: [PATCH] Noise merging: answer chosen steps twice with a second seed and blend the two pictures Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01VmbxSsiiz8wmt226dtbFXe --- CHANGELOG.md | 1 + modules/sampling/noise_merge/__init__.py | 138 ++++++++++++++++++ .../noise_merge/tests/test_noise_merge.py | 92 ++++++++++++ 3 files changed, 231 insertions(+) create mode 100644 modules/sampling/noise_merge/__init__.py create mode 100644 modules/sampling/noise_merge/tests/test_noise_merge.py diff --git a/CHANGELOG.md b/CHANGELOG.md index ad54120c..2f511e15 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ ### Added (experimental) +- **Noise merging** (Settings > Engine > Sampling): at the steps you name, the step is answered a second time on the same picture with different noise (a second seed), and the two pictures are blended at your strength. Merge order picks which one is the base. Costs one extra model pass per merged step. Untested on a real H3. - **SLA rising sparsity** (diffusion model loader, `sla_sparsity_start`): skip fewer attention blocks on early steps and more on late ones, in a straight line up to `sla_sparsity`. Off at 0. One real H3 render (4 steps, 0.70 to 0.90): quality held, perhaps slightly better. Idea from GSRS (arXiv 2610.21129); only the schedule is taken, not its other parts. ## [5.0.3] "Classy Chopper" - 2026-10-25 diff --git a/modules/sampling/noise_merge/__init__.py b/modules/sampling/noise_merge/__init__.py new file mode 100644 index 00000000..ba4d5a33 --- /dev/null +++ b/modules/sampling/noise_merge/__init__.py @@ -0,0 +1,138 @@ +"""Noise merging: at chosen steps, run the step a second time on the same picture with +different noise and blend the two answers. + +Step N is answered twice. Picture A is the run's own answer. Picture B is the answer to +the same step after A's picture is re-noised to this step's noise level with a second +seed (the sound part of the input is left alone). The two are then mixed: + + merged = base + strength * (other - base) + +`Merge order` picks the base. "B into A" keeps A's picture and lets B's variation in; +"A into B" starts from B and lets A in. Two independent answers averaged lose contrast +(the independent detail partly cancels), so the mix is rescaled to the base's own spread. + +Costs one extra model pass per merged step. The second pass is marked as a probe, so no +learning feature records it. +""" + +import torch +from comfy.patcher_extension import WrappersMP + +from ..._core import dit_hooks, log, streams + +ID = "noise_merge" +TITLE = "Noise merging" +MOUNT = "generation.sampling" +STAGE = "sampling" +CATEGORY = "sampling" +STATUS = "experimental" +REQUIRES = ["temporal_latent"] + +SETTINGS = { + "enabled": { + "type": "bool", "default": False, + "label": "Noise merging", + "hint": "At the steps below, answers the step twice with different noise and blends the two. Costs one extra pass per merged step.", + }, + "steps": { + "type": "text", "default": "1", + "label": "Steps", + "hint": "Step numbers from 1, e.g. 1 or 2-3 or 1,4. Steps past the schedule's end are ignored.", + "when": {"enabled": True}, + }, + "order": { + "type": "enum", "default": "b_into_a", "options": [ + {"value": "b_into_a", "label": "B into A"}, + {"value": "a_into_b", "label": "A into B"}, + ], + "label": "Merge order", "ui": "segmented", + "hint": "A is the picture your seed makes, B the same step on a second seed. The first named is blended in; the other is the base that keeps its overall look.", + "when": {"enabled": True}, + }, + "strength": { + "type": "float", "default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05, + "label": "Merge strength", "ui": "slider", + "hint": "0 keeps the base picture untouched, 1 replaces it with the other.", + "when": {"enabled": True}, + }, + "seed": { + "type": "int", "default": 1, "min": 0, "max": 2147483647, + "label": "Second seed", + "hint": "The noise for picture B. The same number gives the same B.", + "when": {"enabled": True}, + }, +} + + +def _say(message, level=log.ALERT): + log.once(f"{ID}:{message}", level, "FunUI Noise merging", message) + + +def merge(base, other, strength): + """base + strength * (other - base), rescaled so each clip keeps the base's spread + (std around its own mean) -- a plain mix of two independent answers is flatter.""" + mixed = base + strength * (other - base) + dims = tuple(range(1, base.dim())) + base_std = base.float().std(dim=dims, keepdim=True) + mixed_std = mixed.float().std(dim=dims, keepdim=True).clamp(min=1e-8) + mean = mixed.float().mean(dim=dims, keepdim=True) + return ((mixed.float() - mean) * (base_std / mixed_std) + mean).to(base.dtype) + + +def renoise(x0, sigma, seed, step): + """A rectified-flow input at `sigma` for picture x0 with fresh noise: (1 - sigma) x0 + sigma noise.""" + gen = torch.Generator().manual_seed(int(seed) + 7919 * int(step)) + noise = torch.randn(x0.shape, generator=gen).to(device=x0.device, dtype=x0.dtype) + return (1.0 - sigma) * x0 + sigma * noise + + +def install(patcher, values, key): + if not values.get("enabled"): + return None + wanted, problems = dit_hooks.parse_blocks(values.get("steps"), 100000) + wanted = {i - 1 for i in wanted if i >= 1} + for p in problems: + _say(f"ignored step entry: {p}") + strength = float(values.get("strength", 0.5)) + if not wanted or strength <= 0.0: + _say("off: " + ("no valid steps named" if not wanted else "strength is 0")) + return None + b_first = values.get("order", "b_into_a") == "b_into_a" + seed = int(values.get("seed", 1)) + + def apply_model(executor, x, t, *args, **kwargs): + named = streams.model_args(args, kwargs) + to = named.get("transformer_options") or {} + out = executor(x, t, *args, **kwargs) + where = dit_hooks.current_step(to) + if where is None or dit_hooks.probing(to) or where[0] not in wanted: + return out + if any(int(i) != 0 for i in (to.get("cond_or_uncond") or [0])): + return out # a negative-prompt call: guidance is for the prompt's picture + try: + split_x, split_a = streams.video_of(x, named), streams.video_of(out, named) + if split_x is None or split_a is None: + _say("off this run: could not find the picture in this model's latent") + return out + sigma = float(t.max()) + video_a, rebuild_a = split_a + x_b = split_x[1](renoise(video_a.detach(), sigma, seed, where[0]).to(split_x[0].dtype)) + args_b, kwargs_b = streams.with_options(args, kwargs, {**to, dit_hooks.PROBE: True}) + out_b = executor(x_b, t, *args_b, **kwargs_b) + split_b = streams.video_of(out_b, named) + if split_b is None: + return out + video_b, rebuild_b = split_b + if b_first: + return rebuild_a(merge(video_a, video_b, strength).to(video_a.dtype)) + return rebuild_b(merge(video_b, video_a, strength).to(video_b.dtype)) + except Exception as exc: # noqa: BLE001 + log.broke("FunUI Noise merging", exc, "merging the second answer") + return out + + patcher.add_wrapper_with_key(WrappersMP.APPLY_MODEL, key, apply_model) + return (f"steps {values.get('steps')}, {'B into A' if b_first else 'A into B'}, " + f"strength {strength:.2f}, second seed {seed}") + + +PROVIDES = {"modifier": install} diff --git a/modules/sampling/noise_merge/tests/test_noise_merge.py b/modules/sampling/noise_merge/tests/test_noise_merge.py new file mode 100644 index 00000000..63b2ba5a --- /dev/null +++ b/modules/sampling/noise_merge/tests/test_noise_merge.py @@ -0,0 +1,92 @@ +"""Noise merging: the maths and the wrapper on a packed H3 latent.""" + +import pytest +import torch + + +@pytest.fixture(autouse=True) +def _needs_comfy(comfyui): + """Imports comfy.""" + + +def _sched(steps=8): + return torch.linspace(1.0, 0.0, steps + 1) + + +def _to(i, steps=8): + s = _sched(steps) + return {"sample_sigmas": s, "sigmas": s[i:i + 1]} + + +def _load(tiny, **values): + from comfy.patcher_extension import WrappersMP + from modules.sampling.modifiers.nodes import FunUILoadModifiers + patched, status = FunUILoadModifiers.execute( + tiny.patcher, {"noise_merge": {"enabled": True, **values}}).result + wrap = [w for ws in patched.wrappers[WrappersMP.APPLY_MODEL].values() for w in ws][-1] + return wrap, status + + +def _packed(): + from conftest import packed_av + torch.manual_seed(2) + return packed_av(torch.randn(1, 4, 3, 8, 8), torch.randn(1, 8, 5)) + + +def test_merge_keeps_the_base_spread_and_strength_zero_is_the_base(): + base, other = torch.randn(2, 4, 3, 4, 4), torch.randn(2, 4, 3, 4, 4) * 3 + assert torch.allclose(merge_(base, other, 0.0), base, atol=1e-5) + mixed = merge_(base, other, 0.5) + assert torch.allclose(mixed.flatten(1).std(1), base.flatten(1).std(1), atol=1e-4) + assert not torch.allclose(mixed, base) + + +def merge_(*a): + from modules.sampling.noise_merge import merge + return merge(*a) + + +def test_renoise_is_seeded_and_lands_on_the_step_noise_level(): + from modules.sampling.noise_merge import renoise + x0 = torch.randn(1, 4, 3, 8, 8) + assert torch.equal(renoise(x0, 0.5, 1, 2), renoise(x0, 0.5, 1, 2)) + assert not torch.equal(renoise(x0, 0.5, 1, 2), renoise(x0, 0.5, 2, 2)) + assert torch.equal(renoise(x0, 0.0, 1, 2), x0) + + +def test_a_named_step_runs_the_model_twice_and_others_once(tiny_h3): + wrap, status = _load(tiny_h3, steps="2", strength=0.5) + assert "steps 2" in status + x0, shapes = _packed() + calls = [] + + def model(x, *a, **k): + calls.append(a[4].get("funui_probe")) + return x0 + 0.1 * x + + for i in range(4): + wrap(model, x0.clone(), _sched()[i:i + 1], None, None, None, _to(i), latent_shapes=shapes) + assert len(calls) == 5 and calls.count(True) == 1 + + +def test_the_answer_changes_only_in_the_picture_and_order_matters(tiny_h3): + from conftest import unpacked + x0, shapes = _packed() + model = lambda x, *a, **k: x * 0.5 + 0.3 * x0 + results = {} + for order in ("b_into_a", "a_into_b"): + wrap, _ = _load(tiny_h3, steps="1", order=order, strength=0.7) + results[order] = wrap(model, x0.clone(), _sched()[0:1], None, None, None, _to(0), latent_shapes=shapes) + plain = model(x0, None) + a, b = unpacked(results["b_into_a"], shapes), unpacked(results["a_into_b"], shapes) + p = unpacked(plain, shapes) + assert not torch.allclose(a[0], p[0]) and not torch.allclose(a[0], b[0]) + assert torch.allclose(a[1], p[1]) # sound is the model's own answer for A + + +def test_nothing_named_or_strength_zero_installs_nothing(tiny_h3): + from comfy.patcher_extension import WrappersMP + from modules.sampling.modifiers.nodes import FunUILoadModifiers + for bad in ({"steps": "x"}, {"strength": 0.0}): + patched, _ = FunUILoadModifiers.execute(tiny_h3.patcher, {"noise_merge": {"enabled": True, **bad}}).result + assert not patched.wrappers.get(WrappersMP.APPLY_MODEL)