digital-garbageandClaude Sonnet 5.5
793990e840
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
367 lines
20 KiB
Python
367 lines
20 KiB
Python
"""Step 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 (and so is the sound, when the model has it). 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.
|
|
|
|
"Second run" keeps a whole second video going beside yours: its first input is the second seed's noise and every
|
|
later input is its own last answer re-noised, the way a few-step sampler chains its steps. At the named steps its
|
|
answer is blended into yours the same way. It costs one extra pass at every step up to the last one named.
|
|
|
|
"Simultaneous actions": with `[A|B]` in the prompt (or `A|B` between two shortcuts) and Alternating prompts >
|
|
"Read A and B separately" on, the named steps are answered
|
|
once for each option -- the same picture and noise, a prompt that holds only that option -- and the answers are
|
|
merged the same way, instead of the second-seed pass. The prompt outside those steps holds both ("A while B").
|
|
|
|
Costs one extra model pass per merged step. The second pass is marked as a probe, so no
|
|
learning feature records it.
|
|
"""
|
|
|
|
import random
|
|
import time
|
|
|
|
import math
|
|
|
|
import torch
|
|
from comfy.model_management import InterruptProcessingException
|
|
from comfy.patcher_extension import WrappersMP
|
|
|
|
from ..._core import dit_hooks, log, streams
|
|
|
|
ID = "noise_merge"
|
|
TITLE = "Step merging"
|
|
MOUNT = "generation.sampling"
|
|
STAGE = "sampling"
|
|
CATEGORY = "sampling"
|
|
STATUS = "experimental"
|
|
REQUIRES = ["temporal_latent"]
|
|
|
|
SETTINGS = {
|
|
"enabled": {
|
|
"type": "bool", "default": False,
|
|
"label": "Step merging",
|
|
"hint": "At the steps below, answers the step twice with different noise and blends the two. Costs one extra pass per merged step.",
|
|
},
|
|
"what": {
|
|
"type": "enum", "default": "answers", "options": [
|
|
{"value": "answers", "label": "Step answers"},
|
|
{"value": "parallel", "label": "Second run"},
|
|
{"value": "noise", "label": "Starting noise"},
|
|
],
|
|
"label": "Merge", "ui": "segmented",
|
|
"hint": "Blends in a second seed. Step answers: at chosen steps. Second run: a whole second video (slower). Starting noise: free.",
|
|
"when": {"enabled": True},
|
|
},
|
|
"steps": {
|
|
"type": "text", "default": "1",
|
|
"label": "Steps",
|
|
"hint": "Step numbers from 1, e.g. 1 or 2-3 or 1,4. Steps past the end of the schedule are ignored.",
|
|
"when": {"enabled": True, "what": ["answers", "parallel"]},
|
|
},
|
|
"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 your seed's picture, B the same step on a second seed. The first one named is blended in; the other is the base.",
|
|
"when": {"enabled": True, "what": ["answers", "parallel"]},
|
|
},
|
|
"strength": {
|
|
"type": "float", "default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05,
|
|
"label": "Merge strength", "ui": "slider",
|
|
"hint": "0 keeps your own result untouched, 1 replaces it with the second seed's. For starting noise, 0.5 is an even mix.",
|
|
"when": {"enabled": True},
|
|
},
|
|
"fresh": {
|
|
"type": "float", "default": 1.0, "min": 0.0, "max": 1.0, "step": 0.05,
|
|
"label": "Second pass new noise", "ui": "slider",
|
|
"hint": "How much of the step's noise the second pass swaps for the second seed's. 1: all of it. 0: no change.",
|
|
"when": {"enabled": True, "what": "answers"},
|
|
},
|
|
"sound": {
|
|
"type": "bool", "default": True,
|
|
"label": "Merge sound too",
|
|
"hint": "Blends the sound the same way as the picture. Off keeps only the base's own sound.",
|
|
"when": {"enabled": True, "what": ["answers", "parallel"]},
|
|
},
|
|
"seed": {
|
|
"type": "int", "default": -1, "min": -1, "max": 2147483647,
|
|
"label": "Second seed",
|
|
"hint": "The noise for picture B (or the second noise). -1 picks a new one every run; a number gives the same B each time.",
|
|
"when": {"enabled": True},
|
|
},
|
|
}
|
|
|
|
|
|
def _say(message, level=log.ALERT):
|
|
log.once(f"{ID}:{message}", level, "FunUI Step 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 second_input(x, a, sigma, seed, step, fresh):
|
|
"""The input for picture B. Rectified flow: x = (1 - sigma) picture + sigma noise. A's answer is the picture;
|
|
the noise is the step's own noise (what is left of x once A's picture is taken out) turned `fresh` of the
|
|
way toward the second seed's noise, at unit variance. fresh 1 is all second-seed noise, 0 hands back x."""
|
|
if sigma < 1e-4:
|
|
return x
|
|
own = (x.float() - (1.0 - sigma) * a.float()) / sigma
|
|
gen = torch.Generator().manual_seed(int(seed) + 7919 * int(step))
|
|
other = torch.randn(x.shape, generator=gen).to(device=x.device)
|
|
theta = fresh * torch.pi / 2
|
|
noise = float(torch.cos(torch.tensor(theta))) * own + float(torch.sin(torch.tensor(theta))) * other
|
|
return ((1.0 - sigma) * a.float() + sigma * noise).to(x.dtype)
|
|
|
|
|
|
SIMUL = "funui_simultaneous" # what Prompt markup leaves in transformer_options: {"options", "cond_len"}
|
|
|
|
|
|
def _mix(out_a, out_b, named, strength, b_first, with_sound):
|
|
"""Two full answers of one step -> (merged answer, how far apart they were), or None when either has no picture."""
|
|
split_a, split_b = streams.video_of(out_a, named), streams.video_of(out_b, named)
|
|
if split_a is None or split_b is None:
|
|
return None
|
|
(video_a, rebuild_a), (video_b, rebuild_b) = split_a, split_b
|
|
tail_a, tail_b = _sound(out_a, video_a), _sound(out_b, video_b)
|
|
diff = float((video_a - video_b).norm() / video_a.norm().clamp(min=1e-8))
|
|
if b_first:
|
|
base, other, base_tail, other_tail, rebuild = video_a, video_b, tail_a, tail_b, rebuild_a
|
|
else:
|
|
base, other, base_tail, other_tail, rebuild = video_b, video_a, tail_b, tail_a, rebuild_b
|
|
merged = rebuild(merge(base, other, strength).to(base.dtype))
|
|
if with_sound and base_tail is not None and other_tail is not None:
|
|
n = merged.shape[-1] - base_tail.shape[-1]
|
|
merged = torch.cat([merged[..., :n], merge(base_tail, other_tail, strength).to(merged.dtype)], dim=-1)
|
|
return merged, diff
|
|
|
|
|
|
def _second_x(x, out, named, sigma, seed, step, fresh, with_sound):
|
|
"""The input for a second pass over `x` after `out` answered it: the step's noise turned `fresh` of the way toward
|
|
the second seed's noise around out's picture (and sound), on x's own device and dtype. x itself when out has no picture."""
|
|
split_x, split_a = streams.video_of(x, named), streams.video_of(out, named)
|
|
if split_x is None or split_a is None:
|
|
return x
|
|
video_a, tail_a = split_a[0].detach(), _sound(out, split_a[0])
|
|
x_b = split_x[1](second_input(split_x[0], video_a, sigma, seed, step, fresh))
|
|
if tail_a is not None and with_sound:
|
|
n = x_b.shape[-1] - tail_a.shape[-1]
|
|
x_b = torch.cat([x_b[..., :n], second_input(x[..., n:], tail_a.detach(), sigma, seed + 1, step, fresh)], dim=-1)
|
|
return x_b
|
|
|
|
|
|
def _sound(packed, video):
|
|
"""The part of a packed [B, 1, N] result that follows the picture (H3's sound), or None."""
|
|
if getattr(packed, "dim", None) is None or packed.dim() != 3:
|
|
return None
|
|
n = video[0].numel()
|
|
return packed[..., n:] if packed.shape[-1] > n else None
|
|
|
|
|
|
def _resized(video, shape):
|
|
"""`video` [B, C, T, H, W] brought to `shape` (only the picture's height and width may differ), or None."""
|
|
if video.dim() != 5 or len(shape) != 5 or tuple(video.shape[:3]) != tuple(shape[:3]):
|
|
return None
|
|
return torch.nn.functional.interpolate(video.float(), size=tuple(shape[2:]), mode="trilinear",
|
|
align_corners=False).to(device=video.device, dtype=video.dtype)
|
|
|
|
|
|
def mix_noise(noise, other, strength):
|
|
"""cos(t) noise + sin(t) other, t = strength * 90 degrees: two independent unit-variance noises
|
|
stay unit variance (a plain average would shrink to 71% at an even mix, and the model reads that as a calmer, blurrier start)."""
|
|
theta = strength * torch.pi / 2
|
|
return (math.cos(theta) * noise + math.sin(theta) * other.to(device=noise.device, dtype=noise.dtype)).to(noise.dtype) # plain numbers: no tensor to put on a device
|
|
|
|
|
|
def install_noise(patcher, values, key):
|
|
strength = float(values.get("strength", 0.5))
|
|
if strength <= 0.0:
|
|
_say("off: strength is 0")
|
|
return None
|
|
chosen = int(values.get("seed", -1))
|
|
run_seed = {}
|
|
|
|
def sampler_sample(executor, model_wrap, sigmas, extra_args, callback, noise,
|
|
latent_image=None, denoise_mask=None, disable_pbar=False):
|
|
if not isinstance(noise, torch.Tensor):
|
|
_say("off this run: the starting noise is not a plain tensor")
|
|
return executor(model_wrap, sigmas, extra_args, callback, noise, latent_image, denoise_mask, disable_pbar)
|
|
second_stage = dit_hooks.stage_of(model_wrap) == 2 # Two stages: both stages mix with the SAME second seed
|
|
if not second_stage or "seed" not in run_seed:
|
|
run_seed["seed"] = chosen if chosen >= 0 else random.randrange(2 ** 31)
|
|
seed = run_seed["seed"]
|
|
other = torch.randn(noise.shape, generator=torch.Generator().manual_seed(seed)).to(noise)
|
|
if not second_stage:
|
|
log.info("FunUI Step merging", f"on: starting noise mixed with the second seed's ({seed}) at "
|
|
f"strength {strength:.2f}, no extra pass")
|
|
return executor(model_wrap, sigmas, extra_args, callback, mix_noise(noise, other, strength),
|
|
latent_image, denoise_mask, disable_pbar)
|
|
|
|
patcher.add_wrapper_with_key(WrappersMP.SAMPLER_SAMPLE, key, sampler_sample)
|
|
return f"starting noise, strength {strength:.2f}, second seed {'random' if chosen < 0 else chosen}"
|
|
|
|
|
|
def install(patcher, values, key):
|
|
if not values.get("enabled"):
|
|
return None
|
|
if values.get("what") == "noise":
|
|
return install_noise(patcher, values, key)
|
|
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:
|
|
_say("off: no valid steps named")
|
|
return None
|
|
b_first = values.get("order", "b_into_a") == "b_into_a"
|
|
with_sound = bool(values.get("sound", True))
|
|
fresh = float(values.get("fresh", 1.0))
|
|
parallel = values.get("what") == "parallel"
|
|
last = max(wanted)
|
|
chosen = int(values.get("seed", -1))
|
|
live = {"seed": max(chosen, 0), "merged": 0, "prev": None, "options": {}}
|
|
|
|
def start(executor, *args, **kwargs):
|
|
# Two stages is two sampler calls of one run: merges counted in the first still count, and the verdict comes after the last.
|
|
stage = dit_hooks.stage_of(getattr(executor, "class_obj", None))
|
|
live["seed"] = chosen if chosen >= 0 else random.randrange(2 ** 31)
|
|
live["prev"], live["options"] = None, {}
|
|
if stage != 2:
|
|
live["merged"] = 0
|
|
log.info("FunUI Step merging", f"on: {'second run, ' if parallel else ''}steps {values.get('steps')}, strength {strength:.2f}, "
|
|
f"second seed {live['seed']}")
|
|
try:
|
|
return executor(*args, **kwargs)
|
|
finally:
|
|
live["options"].clear()
|
|
if stage != 1 and not live["merged"]:
|
|
_say("Inactive | no step of this run matched the steps you named "
|
|
"(or the model was not called as a plain prompt step)")
|
|
|
|
patcher.add_wrapper_with_key(WrappersMP.OUTER_SAMPLE, key, start)
|
|
|
|
def apply_model(executor, x, t, *args, **kwargs):
|
|
named = streams.model_args(args, kwargs)
|
|
to = named.get("transformer_options") or {}
|
|
where = dit_hooks.current_step(to)
|
|
simul = to.get(SIMUL) if not parallel else None
|
|
context = named.get("c_crossattn")
|
|
if (simul and where is not None and where[0] in wanted and not dit_hooks.probing(to)
|
|
and all(int(i) == 0 for i in (to.get("cond_or_uncond") or [0]))):
|
|
if (isinstance(context, torch.Tensor) and context.dim() == 3
|
|
and int(context.shape[1]) == simul["cond_len"] and len(simul["options"]) > 1):
|
|
try:
|
|
# Strength 0 keeps the base's answer and strength 1 (with fully new noise, so the second run does
|
|
# not start from the first one's picture) is the other's alone: that one is the only pass made.
|
|
base_k, other_k = (0, 1) if b_first else (1, 0)
|
|
only = base_k if strength <= 0.0 else other_k if strength >= 1.0 and fresh >= 1.0 else None
|
|
outs, x_k = {}, x
|
|
for k in ((only,) if only is not None else (0, 1)):
|
|
args_k, kwargs_k = streams.with_context(args, kwargs, streams.option_context(
|
|
getattr(executor, "class_obj", None), simul["options"][k], context, live["options"], k))
|
|
args_k, kwargs_k = streams.with_options(args_k, kwargs_k, {**to, streams.OPTION_PINNED: True})
|
|
outs[k] = executor(x_k, t, *args_k, **kwargs_k)
|
|
if k == 0 and fresh > 0.0 and only is None:
|
|
# The second option starts from partly new noise: the step's own noise, turned `fresh`
|
|
# of the way toward the second seed's, around the first option's picture.
|
|
x_k = _second_x(x, outs[0], named, float(t.max()), live["seed"], where[0], fresh, with_sound)
|
|
if only is not None:
|
|
live["merged"] += 1
|
|
log.info("FunUI Step merging", f"step {where[0] + 1}: only {'AB'[only]} read "
|
|
f"(strength {strength:.2f}, new noise {fresh:.2f}), nothing to merge")
|
|
return outs[only]
|
|
mixed = _mix(outs[0], outs[1], named, strength, b_first, with_sound)
|
|
if mixed is not None:
|
|
live["merged"] += 1
|
|
log.info("FunUI Step merging", f"step {where[0] + 1}: both actions answered separately, the "
|
|
f"answers differ by {mixed[1]:.1%}, merged at {strength:.2f}")
|
|
return mixed[0]
|
|
except InterruptProcessingException:
|
|
raise
|
|
except Exception as exc: # noqa: BLE001
|
|
log.broke("FunUI Step merging", exc, "merging the simultaneous actions")
|
|
return executor(x, t, *args, **kwargs)
|
|
else:
|
|
_say("Simultaneous actions skipped: another feature changed the prompt's conditioning")
|
|
out = executor(x, t, *args, **kwargs)
|
|
if strength <= 0.0 or where is None or dit_hooks.probing(to) or where[0] > last or (where[0] not in wanted and not parallel):
|
|
return out # strength 0 keeps your own answer: no second pass to blend in
|
|
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
|
|
tail_a = _sound(out, video_a)
|
|
gen_step = where[0]
|
|
if parallel:
|
|
# The second run: each step's input is its own last answer (zero at the start) re-noised with the
|
|
# second seed's noise, the way a few-step sampler chains its steps.
|
|
prev = live["prev"] or (torch.zeros_like(video_a), torch.zeros_like(tail_a) if tail_a is not None else None)
|
|
if prev[0].shape != video_a.shape:
|
|
prev = (_resized(prev[0], video_a.shape), prev[1]) # Two stages: the second run moves to the full size too
|
|
if prev[0] is None:
|
|
_say("Inactive | the latent changed size during the run, so the second run was dropped")
|
|
return out
|
|
base_v, base_t, mix = prev[0], prev[1], 1.0
|
|
else:
|
|
base_v, base_t, mix = video_a.detach(), tail_a, fresh
|
|
x_b = split_x[1](second_input(split_x[0], base_v, sigma, live["seed"], gen_step, mix))
|
|
if tail_a is not None and with_sound:
|
|
n = x_b.shape[-1] - tail_a.shape[-1]
|
|
x_b = torch.cat([x_b[..., :n], second_input(x[..., n:], base_t.detach(), sigma, live["seed"] + 1,
|
|
gen_step, mix)], dim=-1)
|
|
args_b, kwargs_b = streams.with_options(args, kwargs, {**to, dit_hooks.PROBE: True})
|
|
began = time.perf_counter()
|
|
out_b = executor(x_b, t, *args_b, **kwargs_b)
|
|
took = time.perf_counter() - began
|
|
split_b = streams.video_of(out_b, named)
|
|
if split_b is None:
|
|
return out
|
|
video_b, rebuild_b = split_b
|
|
tail_b = _sound(out_b, video_b)
|
|
if parallel:
|
|
live["prev"] = (video_b.detach(), None if tail_b is None else tail_b.detach())
|
|
if gen_step not in wanted:
|
|
return out
|
|
live["merged"] += 1
|
|
diff = float((video_a - video_b).norm() / video_a.norm().clamp(min=1e-8))
|
|
log.info("FunUI Step merging",
|
|
f"step {where[0] + 1}: second answer differs from the first by {diff:.1%}, "
|
|
f"merged at {strength:.2f}{' (picture and sound)' if tail_a is not None and with_sound else ''}; "
|
|
f"the second pass took {took:.1f}s")
|
|
mixed = _mix(out, out_b, named, strength, b_first, with_sound)
|
|
return out if mixed is None else mixed[0]
|
|
except InterruptProcessingException:
|
|
raise # Cancel stops the run, it is not a failed merge
|
|
except Exception as exc: # noqa: BLE001
|
|
log.broke("FunUI Step 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 {'random' if chosen < 0 else chosen}")
|
|
|
|
|
|
PROVIDES = {"modifier": install}
|