Noise merging: answer chosen steps twice with a second seed and blend the two pictures
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VmbxSsiiz8wmt226dtbFXe
This commit is contained in:
Claude
@@ -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}
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user