diff --git a/conftest.py b/conftest.py index d0bf0bc5..4dd9e909 100644 --- a/conftest.py +++ b/conftest.py @@ -42,3 +42,67 @@ def comfyui(): if COMFYUI is None: pytest.skip("no ComfyUI source on this machine") return COMFYUI + + +class TinyH3: + """A REAL MiniMax H3 -- ComfyUI's own class, four blocks wide as a thumb. + + Random weights, CPU, milliseconds per forward. Every H3 modifier is tested + against this rather than a mock, because what breaks these features is the + real forward's shapes (joint sequence without a batch axis, tensor + containers around q/k/v, mod_segments) and a mock only has the shapes its + author remembered. + """ + + def __init__(self, blocks=4, seed=0): + import torch + import comfy.ops + from comfy.latent_formats import MiniMaxH3AV + from comfy.ldm.minimax.model import MiniMaxH3Model + from comfy.model_patcher import ModelPatcher + + torch.manual_seed(seed) + dm = MiniMaxH3Model( + hidden_size=64, num_layers=blocks, token_refiner_num_layers=1, + num_attention_heads=2, attention_head_dim=32, ffn_hidden_size=64, + text_dim=48, timestep_input_dim=32, time_embed_hidden_size=64, + time_embed_dim=32, rope_inv_freq_len=4, dtype=torch.float32, + operations=comfy.ops.manual_cast) + for p in dm.parameters(): + torch.nn.init.normal_(p, std=0.05) + p.requires_grad_(False) # the in-place rope kernel refuses autograd + dm.rope.inv_freq.copy_(torch.rand(4)) + + class Config: + latent_format = MiniMaxH3AV + + class Base(torch.nn.Module): + pass + + base = Base() + base.diffusion_model = dm + base.model_config = Config() + self.patcher = ModelPatcher(base, load_device=torch.device("cpu"), + offload_device=torch.device("cpu")) + self.video = torch.randn(1, 24, 2, 4, 4) + self.audio = torch.randn(1, 32, 2, 3) + self.context = torch.randn(1, 5, 48) + + def run(self, patcher=None, sigma=0.5, sigmas=None, negative=None): + """One denoise call through `patcher`'s transformer_options, as the + sampler would make it. Returns (video, audio).""" + import torch + patcher = patcher or self.patcher + to = dict(patcher.model_options.get("transformer_options", {})) + sched = sigmas if sigmas is not None else torch.tensor([1.0, sigma, 0.0]) + to.setdefault("sample_sigmas", sched) + to["sigmas"] = torch.tensor([sigma]) + dm = patcher.model.diffusion_model + with torch.inference_mode(): + return dm([self.video, self.audio], torch.tensor([sigma * 1000.0]), + self.context, transformer_options=to) + + +@pytest.fixture +def tiny_h3(comfyui): + return TinyH3() diff --git a/core/dit_hooks.py b/core/dit_hooks.py new file mode 100644 index 00000000..398ad863 --- /dev/null +++ b/core/dit_hooks.py @@ -0,0 +1,182 @@ +"""Block hooks and attention overrides that stack, come off again, and cannot end a run. + +ComfyUI gives a transformer two extension points that are NOT lists: + +* `transformer_options["patches_replace"]["dit"][("double_block", i)]` -- one + callable per block, which replaces that block's call. +* `transformer_options["optimized_attention_override"]` -- one callable for every + attention call in the model. + +One slot each. A second modifier that simply assigns the slot silently deletes the +first one's effect, and v4 learned that the slow way (temperature vanishing under +REINS, with no error). So everything here CHAINS: a hook is written as if it were +alone -- it calls `extra["original_block"](args)`, or `func(q, k, v, ...)` -- and +what it calls is really the rest of the chain. No hook ever needs to know another +exists. + +Model-agnostic on purpose. The two slot names are ComfyUI conventions shared by +several DiTs; which models actually honour them is a trait the MODEL module +announces (`dit_block_hooks`), never something this file decides. + +Removal: every installed callable carries `patching.TAG` and a pointer to what it +wrapped, so `patching.strip` unwinds ours and puts back whatever was underneath -- +including someone else's hook we chained through. +""" + +from . import patching + +PREV = "_funpack_prev" +BLOCK = "double_block" + + +def _transformer_options(patcher) -> dict: + options = patcher.model_options + to = options.get("transformer_options") + if not isinstance(to, dict): + to = {} + options["transformer_options"] = to + return to + + +def _guard(patcher, fn, neutral): + """Guarded when the patcher can guard (it is always a GuardedPatcher inside a + run); a bare patcher -- tests, a direct call -- gets the hook unchanged.""" + wrap = getattr(patcher, "guarded", None) + return wrap(fn, neutral) if callable(wrap) else fn + + +def _mark(fn, key, prev): + patching.tag(fn, key) + setattr(fn, PREV, prev) + return fn + + +def add_block_hook(patcher, key: str, index: int, hook) -> None: + """Put `hook(args, extra)` around block `index`, on top of whatever is there. + + `extra["original_block"]` inside the hook runs the rest of the chain, ending in + the real block. A failing hook is dropped for the rest of the run and the block + runs as though it had never been installed. + """ + to = _transformer_options(patcher) + replace = to.setdefault("patches_replace", {}) + dit = replace.setdefault("dit", {}) + slot = (BLOCK, int(index)) + inner = dit.get(slot) + + def below(extra): + if inner is None: + return extra + original = extra["original_block"] + return {**extra, "original_block": lambda a: inner(a, {**extra, "original_block": original})} + + def neutral(args, extra): + return below(extra)["original_block"](args) + + guarded = _guard(patcher, lambda args, extra: hook(args, below(extra)), neutral) + + def chained(args, extra): + return guarded(args, extra) + + dit[slot] = _mark(chained, key, inner) + + +def add_attention_override(patcher, key: str, override) -> None: + """Put `override(func, q, k, v, heads, **kw)` around every attention call. + + Calling `func(...)` inside it runs the rest of the chain, ending in the real + backend. The override sees plain tensors: ComfyUI unwraps its tensor + containers before an override is called. + """ + to = _transformer_options(patcher) + inner = to.get("optimized_attention_override") + + def rest(func): + if inner is None: + return func + return lambda *a, **kw: inner(func, *a, **kw) + + def neutral(func, *args, **kwargs): + return rest(func)(*args, **kwargs) + + guarded = _guard(patcher, lambda func, *a, **kw: override(rest(func), *a, **kw), neutral) + + def chained(func, *args, **kwargs): + return guarded(func, *args, **kwargs) + + to["optimized_attention_override"] = _mark(chained, key, inner) + + +def block_count(model) -> int: + """How many hookable blocks the model has, or 0 when it cannot be read.""" + root = getattr(getattr(model, "model", None), "diffusion_model", None) + blocks = getattr(root, "blocks", None) + try: + return len(blocks) + except TypeError: + return 0 + + +def parse_blocks(spec, n_blocks: int): + """"40" | "38-42" | "10,40,44" | "" -> (sorted indices, [problems]). + + Out-of-range and unparseable entries are reported, not silently dropped: a + typo that turns a setting into a no-op is exactly the inert knob this project + refuses to ship, so the caller says what was ignored. + """ + out, problems = set(), [] + for part in str(spec or "").replace(" ", "").split(","): + if not part: + continue + try: + if "-" in part[1:]: + a, b = part.split("-", 1) + lo, hi = sorted((int(a), int(b))) + else: + lo = hi = int(part) + except ValueError: + problems.append(f"{part!r} is not a block number or range") + continue + wanted = set(range(lo, hi + 1)) + kept = {i for i in wanted if 0 <= i < n_blocks} + if kept != wanted: + problems.append(f"{part!r} reaches past this model's blocks 0-{n_blocks - 1}") + out |= kept + return sorted(out), problems + + +def current_step(transformer_options): + """(step index, step count) of the denoise call in progress, or None. + + Core puts the whole schedule in `sample_sigmas` and this call's sigma in + `sigmas`; the index is exact, so a step window never has to guess a sigma + threshold per schedule. + """ + import torch + to = transformer_options or {} + sched, cur = to.get("sample_sigmas"), to.get("sigmas") + if sched is None or cur is None: + return None + cur = cur.reshape(-1)[0].to(sched.dtype) + hit = torch.nonzero(torch.isclose(sched, cur, rtol=1e-4)) + if not len(hit): + return None + return int(hit[0]), int(sched.shape[0]) - 1 + + +def target_rows(mod_segments, seq_len, device, stream): + """[seq_len] bool mask of the model's TARGET `stream` ("video" | "audio") + rows, or None when no model module can prove it for this call. + + Asked of whichever model module announces `target_rows`, so core never + learns a model's sequence layout; the first provider with an answer wins. + """ + from . import registry + for _spec, provider in registry.current().providers("target_rows"): + try: + mask = provider(mod_segments, seq_len, device, stream) + except Exception: # noqa: BLE001 + continue + if mask is not None: + return mask + return None diff --git a/core/patching.py b/core/patching.py index 052c1b36..b9639e90 100644 --- a/core/patching.py +++ b/core/patching.py @@ -74,8 +74,16 @@ def _strip_options(options: dict, prefix: str) -> int: removed += len(value) - len(keep) options[name] = keep elif callable(value) and _ours(value, prefix): - options.pop(name, None) - removed += 1 + # A chained hook remembers what it wrapped (dit_hooks.PREV): put that + # back rather than dropping it, or stripping ours deletes someone + # else's hook along with it. + while callable(value) and _ours(value, prefix): + value = getattr(value, "_funpack_prev", None) + removed += 1 + if value is None: + options.pop(name, None) + else: + options[name] = value elif isinstance(value, dict): removed += _strip_options(value, prefix) return removed @@ -221,5 +229,14 @@ class GuardedPatcher: return install + def guarded(self, fn, neutral): + """`fn` guarded like any hook installed through this patcher, for hooks + that go in through a slot rather than a `set_model_*` method (block + replacements, the attention override). `neutral` has the hook's own + signature and returns what the call would have without it.""" + if not self._guarding: + return fn + return guard(fn, self._key, neutral, self._dropped) + def __setattr__(self, name, value): setattr(self._patcher, name, value) diff --git a/core/tests/test_dit_hooks.py b/core/tests/test_dit_hooks.py new file mode 100644 index 00000000..b0d29c29 --- /dev/null +++ b/core/tests/test_dit_hooks.py @@ -0,0 +1,102 @@ +"""Chained block hooks / attention overrides, on a REAL H3 forward.""" + +import pytest +import torch + + +@pytest.fixture(autouse=True) +def _needs_comfy(comfyui): + """Imports comfy.""" + + +def _clone(tiny): + return tiny.patcher.clone() + + +def test_two_block_hooks_on_one_block_both_run_and_the_output_is_the_block(tiny_h3): + from core import dit_hooks + base = tiny_h3.run() + p = _clone(tiny_h3) + seen = [] + + def hook(name): + def h(args, extra): + seen.append(name) + return extra["original_block"](args) + return h + + dit_hooks.add_block_hook(p, "funpack.a", 1, hook("a")) + dit_hooks.add_block_hook(p, "funpack.b", 1, hook("b")) + out = tiny_h3.run(p) + assert seen == ["b", "a"], "the second hook replaced the first instead of chaining" + assert torch.allclose(out[0], base[0]) and torch.allclose(out[1], base[1]) + + +def test_two_attention_overrides_both_see_every_call(tiny_h3): + from core import dit_hooks + p = _clone(tiny_h3) + calls = {"a": 0, "b": 0} + + def ov(name): + def o(func, q, k, v, *a, **kw): + calls[name] += 1 + return func(q, k, v, *a, **kw) + return o + + dit_hooks.add_attention_override(p, "funpack.a", ov("a")) + dit_hooks.add_attention_override(p, "funpack.b", ov("b")) + tiny_h3.run(p) + # 4 DiT blocks + 1 token-refiner block + assert calls == {"a": 5, "b": 5} + + +def test_strip_removes_ours_and_puts_back_a_foreign_hook_underneath(tiny_h3): + from core import dit_hooks, patching + p = _clone(tiny_h3) + foreign = [] + + def theirs(args, extra): + foreign.append(1) + return extra["original_block"](args) + + dit = p.model_options["transformer_options"].setdefault("patches_replace", {}).setdefault("dit", {}) + dit[("double_block", 0)] = theirs + dit_hooks.add_block_hook(p, "funpack.x", 0, lambda a, e: e["original_block"](a)) + dit_hooks.add_attention_override(p, "funpack.x", lambda f, *a, **k: f(*a, **k)) + + assert patching.strip(p, "funpack.") == 2 + to = p.model_options["transformer_options"] + assert to["patches_replace"]["dit"][("double_block", 0)] is theirs + assert "optimized_attention_override" not in to + tiny_h3.run(p) + assert foreign == [1] + + +def test_a_failing_block_hook_is_dropped_and_the_run_matches_an_unhooked_one(tiny_h3): + from core import dit_hooks, patching + base = tiny_h3.run() + p = _clone(tiny_h3) + dropped = patching.Dropped() + guarded = patching.GuardedPatcher(p, "funpack.bad", dropped) + + def bad(args, extra): + raise RuntimeError("boom") + + def bad_attn(func, *a, **k): + raise RuntimeError("boom") + + dit_hooks.add_block_hook(guarded, "funpack.bad", 2, bad) + dit_hooks.add_attention_override(guarded, "funpack.bad", bad_attn) + out = tiny_h3.run(p) + assert "funpack.bad" in dropped + assert torch.allclose(out[0], base[0]) + + +@pytest.mark.parametrize("spec,expect,problems", [ + ("1", [1], 0), ("0-2", [0, 1, 2], 0), ("3,1", [1, 3], 0), + ("2-9", [2, 3], 1), ("x", [], 1), ("", [], 0), ("-1", [], 1), +]) +def test_parse_blocks_reports_what_it_ignores(spec, expect, problems): + from core import dit_hooks + got, said = dit_hooks.parse_blocks(spec, 4) + assert got == expect and len(said) == problems diff --git a/modules/_core.py b/modules/_core.py index cd0f443c..b988f960 100644 --- a/modules/_core.py +++ b/modules/_core.py @@ -11,8 +11,8 @@ leading underscore keeps it out of the module scan. """ try: # root on sys.path: tests, dev server - from core import chain, contract, log, media, patching, registry, relations, run, schema, traits + from core import chain, contract, dit_hooks, log, media, patching, registry, relations, run, schema, traits except ImportError: # inside ComfyUI: a subpackage of the pack - from ..core import chain, contract, log, media, patching, registry, relations, run, schema, traits + from ..core import chain, contract, dit_hooks, log, media, patching, registry, relations, run, schema, traits -__all__ = ["chain", "contract", "log", "media", "patching", "registry", "relations", "run", "schema", "traits"] +__all__ = ["chain", "contract", "dit_hooks", "log", "media", "patching", "registry", "relations", "run", "schema", "traits"] diff --git a/modules/models/minimax_h3/__init__.py b/modules/models/minimax_h3/__init__.py index cce0390e..f0174b75 100644 --- a/modules/models/minimax_h3/__init__.py +++ b/modules/models/minimax_h3/__init__.py @@ -28,6 +28,7 @@ from comfy_extras.nodes_minimax_h3 import temporal_shape from ..._core import traits as _traits from .pipeline import presets as _pipeline_presets +from .rows import target_rows as _target_rows has_block = _traits.has_block @@ -105,7 +106,7 @@ def probe_traits(keys) -> list: return [] from comfy.latent_formats import MiniMaxH3AV - found = ["audio_stream"] + found = ["audio_stream", "dit_block_hooks"] rank = _traits.LATENT_RANK.get(MiniMaxH3AV.latent_dimensions) if rank: found.append(rank) @@ -118,7 +119,10 @@ def traits(model): """What core cannot read off the model for itself.""" if not is_h3(model): return () - found = ["audio_stream"] + # dit_block_hooks: its forward honours patches_replace["dit"][("double_block", + # i)] and routes every attention call through optimized_attention, which is + # what core/dit_hooks.py hooks. + found = ["audio_stream", "dit_block_hooks"] if has_block(model, ADALN_CLASS): found.append("adaln_modalities") return found @@ -178,4 +182,5 @@ def empty_latent(model, width, height, length, batch_size=1): TRAITS = traits PROVIDES = {"empty_latent": empty_latent, "decode": decode, "detect": detect, - "pipeline_presets": _pipeline_presets, "probe_traits": probe_traits} + "pipeline_presets": _pipeline_presets, "probe_traits": probe_traits, + "target_rows": _target_rows} diff --git a/modules/models/minimax_h3/rows.py b/modules/models/minimax_h3/rows.py new file mode 100644 index 00000000..bd3a1876 --- /dev/null +++ b/modules/models/minimax_h3/rows.py @@ -0,0 +1,35 @@ +"""Which rows of H3's joint sequence are the TARGET video and audio. + +Upstream packs `[text | cond/ref | audio | video]` and emits one mod_segments +entry per stream, so target video is the LAST entry and target audio the one +before it. Matching by modality tag alone is wrong: tag 0 also marks the i2v +anchor pin, guide keyframes, reference media and the vision tokens inside the +text span, none of which a "video rows" feature may touch (fa04d6c audit). + +None -- never an all-False mask -- when the layout cannot be proven (the token +refiner's text-only calls, an upstream layout change), so callers no-op +instead of guessing. +""" + +import torch + +_TAG = {"video": (0, -1), "audio": (2, -2)} + + +def target_rows(mod_segments, seq_len, device, stream): + if stream not in _TAG: + return None + tag, at = _TAG[stream] + segs = list(mod_segments or ()) + if len(segs) < 2: + return None + a, b, row = segs[at] + try: + ok = bool(((row % 3) == tag).all()) if torch.is_tensor(row) else int(row) % 3 == tag + except Exception: # noqa: BLE001 + return None + if not ok or not (0 <= a < b <= seq_len): + return None + mask = torch.zeros(seq_len, dtype=torch.bool, device=device) + mask[a:b] = True + return mask diff --git a/modules/sampling/attention_temperature/__init__.py b/modules/sampling/attention_temperature/__init__.py new file mode 100644 index 00000000..1dfb5e0e --- /dev/null +++ b/modules/sampling/attention_temperature/__init__.py @@ -0,0 +1,95 @@ +"""Randomizer: loosen or tighten how firmly chosen blocks weigh the prompt. + +Scales the query before attention at the named blocks, which is softmax +temperature without ever building the score matrix: Q/T . K == (Q . K)/T. Above +1 the model is less sure which tokens matter and wanders; below 1 it commits +harder to what the prompt already made confident. Either way it can only +re-weigh what is in context, so unlike added noise it stays inside the prompt. + +Validated on H3 by eye (2026-09-20), with two conditions the user found: keep it +moderate, and keep it off blocks a steering module is already pushing on. +Known-good: 0.75x (sharper) at blocks 2-10, or blocks 29-46. + +Not rating-driven, on purpose: a plain dial. Coupling one hard-to-read mechanism +to another's data would make neither debuggable. +""" + +from ..._core import dit_hooks, log + +ID = "attention_temperature" +TITLE = "Randomizer" +MOUNT = "generation.sampling" +STAGE = "sampling" +CATEGORY = "sampling" +STATUS = "proven" +REQUIRES = ["dit_block_hooks"] + +SETTINGS = { + "enabled": { + "type": "bool", "default": False, + "label": "Randomizer", + "hint": "Makes the model more adventurous or more literal with your prompt.", + }, + "amount": { + "type": "float", "default": -0.25, "min": -0.9, "max": 3.0, "step": 0.05, + "label": "Amount", "ui": "slider", + "hint": "Below zero sticks closer to the prompt, above zero wanders. Keep it small.", + "when": {"enabled": True}, + }, + "blocks": { + "type": "text", "default": "2-10", + "label": "Blocks", + "hint": "e.g. 2-10 or 29-46. Keep away from blocks a steering feature uses.", + "when": {"enabled": True}, + }, +} + +# Carried on the per-block copy of transformer_options, so an attention call knows +# it belongs to a named block without any shared flag that could leak. +_SCOPE = "funpack_attention_temperature" + + +def install(patcher, values, key): + if not values.get("enabled"): + return None + amount = float(values.get("amount", 0.0)) + if amount == 0.0: + return None + blocks, problems = dit_hooks.parse_blocks(values.get("blocks"), dit_hooks.block_count(patcher)) + for problem in problems: + log.once(f"{ID}:{problem}", log.ALERT, "FunPack Randomizer", f"ignored {problem}") + if not blocks: + raise ValueError(f"no usable blocks in {values.get('blocks')!r}, so nothing would change") + + # The widget floor keeps this above zero; the clamp is for API callers. + temperature = max(1.0 + amount, 0.05) + + def scoped(args, extra): + to = dict(args.get("transformer_options") or {}) + to[_SCOPE] = temperature + return extra["original_block"]({**args, "transformer_options": to}) + + for index in blocks: + dit_hooks.add_block_hook(patcher, key, index, scoped) + + def override(func, q, k, v, *args, **kwargs): + t = (kwargs.get("transformer_options") or {}).get(_SCOPE) + return func(q if t is None else q / t, k, v, *args, **kwargs) + + dit_hooks.add_attention_override(patcher, key, override) + direction = "sharper" if amount < 0 else "looser" + return f"{temperature:g}x ({direction}) at blocks {_ranges(blocks)}" + + +def _ranges(blocks): + out, start = [], None + for i, b in enumerate(blocks): + if start is None: + start = b + if i + 1 == len(blocks) or blocks[i + 1] != b + 1: + out.append(str(start) if start == b else f"{start}-{b}") + start = None + return ",".join(out) + + +PROVIDES = {"modifier": install} diff --git a/modules/sampling/attention_temperature/tests/test_attention_temperature.py b/modules/sampling/attention_temperature/tests/test_attention_temperature.py new file mode 100644 index 00000000..7b5476ad --- /dev/null +++ b/modules/sampling/attention_temperature/tests/test_attention_temperature.py @@ -0,0 +1,75 @@ +"""The randomizer, through the real modifier loader onto a real H3 forward.""" + +import pytest +import torch + + +@pytest.fixture(autouse=True) +def _needs_comfy(comfyui): + """Imports comfy.""" + + +def _load(tiny, **values): + from modules.sampling.modifiers.nodes import FunPackLoadModifiers + settings = {"attention_temperature": {"enabled": True, **values}} + return FunPackLoadModifiers.execute(tiny.patcher, settings).result + + +def test_it_changes_the_output(tiny_h3): + base = tiny_h3.run() + patched, status = _load(tiny_h3, amount=-0.5, blocks="1-2") + assert "attention_temperature: 0.5x (sharper) at blocks 1-2" in status + + out = tiny_h3.run(patched) + assert not torch.allclose(out[0], base[0]), "temperature ran and changed nothing" + + +def test_blocks_outside_the_list_see_the_original_query(tiny_h3): + from core import dit_hooks + patched, _ = _load(tiny_h3, amount=1.0, blocks="3") + seen = [] + + def spy(func, q, k, v, *a, **kw): + seen.append(float(q.abs().mean())) + return func(q, k, v, *a, **kw) + + # Installed on top, so it sees q BEFORE the temperature override divides it, + # and a second spy underneath sees it after. + under = [] + + def spy_under(func, q, k, v, *a, **kw): + under.append(float(q.abs().mean())) + return func(q, k, v, *a, **kw) + + probe = patched.clone() + dit_hooks.add_attention_override(probe, "test.spy", spy) + tiny_h3.run(probe) + # Rebuild with the spy UNDER the temperature: clone the base, spy first. + low = tiny_h3.patcher.clone() + dit_hooks.add_attention_override(low, "test.spy", spy_under) + from modules.sampling import attention_temperature as m + m.install(low, {"enabled": True, "amount": 1.0, "blocks": "3"}, key="funpack.attention_temperature") + tiny_h3.run(low) + # call order: refiner, blocks 0..3 -> only the last is halved + ratios = [u / s for u, s in zip(under, seen)] + assert [round(r, 4) for r in ratios] == [1.0, 1.0, 1.0, 1.0, 0.5] + + +def test_off_or_zero_installs_nothing(tiny_h3): + from modules.sampling.modifiers.nodes import FunPackLoadModifiers + _p, status = FunPackLoadModifiers.execute( + tiny_h3.patcher, {"attention_temperature": {"enabled": True, "amount": 0.0}}).result + assert "attention_temperature" not in status.splitlines()[0] + + +def test_a_block_list_with_nothing_usable_is_refused_not_silently_inert(tiny_h3): + _p, status = _load(tiny_h3, amount=0.5, blocks="40-49") + assert "attention_temperature: failed to install" in status + + +def test_loading_twice_does_not_double_the_effect(tiny_h3): + once, _ = _load(tiny_h3, amount=-0.5, blocks="1") + from modules.sampling.modifiers.nodes import FunPackLoadModifiers + twice, status = FunPackLoadModifiers.execute( + once, {"attention_temperature": {"enabled": True, "amount": -0.5, "blocks": "1"}}).result + assert torch.allclose(tiny_h3.run(once)[0], tiny_h3.run(twice)[0]) diff --git a/modules/sampling/block_repeat/__init__.py b/modules/sampling/block_repeat/__init__.py new file mode 100644 index 00000000..1788261a --- /dev/null +++ b/modules/sampling/block_repeat/__init__.py @@ -0,0 +1,181 @@ +"""Block repeat: run chosen blocks a second time, inside one sampler step. + +`x' = x + f(x)`, then `x'' = x' + f(x')`. No VAE round trip, no re-noise, no +extra step: one more block forward per repeat, ~2% of a step per block on a +50-block model. + +What it does, from the v4 sweep (research_block_repeat_sweep): a repeat is a +TRAJECTORY RE-ROLL -- a different take on the same prompt -- not a sharpening +pass. The band 20-41 behaved best. The blocks after a repeated one were trained +on its normal output, so pushing further (more passes, wider spans) drifts out +of distribution. + +Two shapes: + +* **each** -- 31,31,32,32..40,40. Every block gets its own doubled output. +* **span** -- 31..40,31..40. The whole span loops, so every block still gets + input from its normal predecessor and there is ONE seam (40 -> 31) instead of + one per block. Span mode runs the span's blocks directly, so a hook another + module put on a block INSIDE the span does not run on the looped passes. + +"Video only" keeps the extra pass for video rows and gives text and audio the +single pass, because repeating the audio rows is the lead suspect for the +unprompted dialogue band 31-40 invents. Attention has already mixed the doubled +audio into the video before the restore; it protects what you hear, not what +the video attended to. +""" + +import torch + +from ..._core import dit_hooks, log + +ID = "block_repeat" +TITLE = "Block repeat" +MOUNT = "generation.sampling" +STAGE = "sampling" +CATEGORY = "sampling" +STATUS = "experimental" +REQUIRES = ["dit_block_hooks"] + +SETTINGS = { + "enabled": { + "type": "bool", "default": False, + "label": "Block repeat", + "hint": "Re-rolls part of the model's thinking for a different take on the same prompt.", + }, + "blocks": { + "type": "text", "default": "31-40", + "label": "Blocks", + "hint": "e.g. 31-40, or 20,35. Somewhere in 20-41 works best.", + "when": {"enabled": True}, + }, + "shape": { + "type": "enum", "default": "each", "options": [ + {"value": "each", "label": "Each block"}, + {"value": "span", "label": "Whole span"}, + ], + "label": "Repeat", "ui": "segmented", + "hint": "each: every block twice in place. span: the whole range loops once.", + "when": {"enabled": True}, + }, + "times": { + "type": "int", "default": 1, "min": 1, "max": 4, + "label": "Extra passes", + "hint": "1 = each block runs twice. More drifts further.", + "when": {"enabled": True}, + }, + "video_only": { + "type": "bool", "default": False, + "label": "Video only", + "hint": "Keeps the sound out of it. Try this if people start talking unprompted.", + "when": {"enabled": True}, + }, + "steps": { + "type": "int", "default": 0, "min": -50, "max": 50, + "label": "Which steps", + "hint": "0 = every step. 3 = only the last 3. -3 = only the first 3.", + "when": {"enabled": True}, + }, +} + + +def _say(message): + log.once(f"{ID}:{message}", log.ALERT, "FunPack Block repeat", message) + + +def _in_window(args, steps): + if steps == 0: + return True + where = dit_hooks.current_step(args.get("transformer_options")) + if where is None: + # Fails OPEN -- repeat every step -- and says so, rather than quietly + # never repeating at all. + _say("could not tell which step this is, so it repeats on every step this run") + return True + index, total = where + return index >= total - steps if steps > 0 else index < -steps + + +def _video_only(args, doubled, single): + """Doubled video rows, single-pass everything else.""" + if single.shape != doubled.shape: + return doubled + mask = dit_hooks.target_rows(args.get("mod_segments"), int(doubled.shape[0]), + doubled.device, "video") + if mask is None: + _say("could not find the video rows, so 'video only' repeated everything this run") + return doubled + return torch.where(mask.view(-1, *([1] * (doubled.dim() - 1))), doubled, single) + + +def install(patcher, values, key): + if not values.get("enabled"): + return None + n = dit_hooks.block_count(patcher) + blocks, problems = dit_hooks.parse_blocks(values.get("blocks"), n) + for problem in problems: + _say(f"ignored {problem}") + if not blocks: + raise ValueError(f"no usable blocks in {values.get('blocks')!r}, so nothing would repeat") + times = int(values.get("times", 1)) + steps = int(values.get("steps", 0)) + video_only = bool(values.get("video_only")) + span = values.get("shape") == "span" + + if span: + lo, hi = blocks[0], blocks[-1] + if blocks != list(range(lo, hi + 1)): + raise ValueError(f"span needs a range like 31-40, got {values.get('blocks')!r}") + real = patcher.model.diffusion_model.blocks + + def run_span(args, h): + for i in range(lo, hi + 1): + h = real[i](h, args["t_emb"], args["mod_segments"], args["rope_freqs"], + transformer_options=args.get("transformer_options")) + return h + + def head(args, extra): + once = run_span(args, args["img"]) + if not _in_window(args, steps): + return {"img": once} + # H3's blocks add their residual INTO the input tensor, so without a + # copy the next pass overwrites `once` and "video only" restores + # audio from an already-repeated stream (v4 shipped that bug). + h = once.clone() if video_only else once + for _ in range(times): + h = run_span(args, h) + return {"img": _video_only(args, h, once) if video_only else h} + + def skip(args, extra): + return {"img": args["img"]} # the head already ran this block + + dit_hooks.add_block_hook(patcher, key, lo, head) + for i in range(lo + 1, hi + 1): + dit_hooks.add_block_hook(patcher, key, i, skip) + what = f"blocks {lo}-{hi} loop {times + 1}x as one span" + else: + def repeat(args, extra): + once = extra["original_block"](args)["img"] + if not _in_window(args, steps): + return {"img": once} + # H3's blocks add their residual INTO the input tensor, so without a + # copy the next pass overwrites `once` and "video only" restores + # audio from an already-repeated stream (v4 shipped that bug). + h = once.clone() if video_only else once + for _ in range(times): + # A fresh dict: the caller's args must not see the intermediate. + h = extra["original_block"]({**args, "img": h})["img"] + return {"img": _video_only(args, h, once) if video_only else h} + + for i in blocks: + dit_hooks.add_block_hook(patcher, key, i, repeat) + what = f"blocks {','.join(map(str, blocks))} run {times + 1}x each" + + if video_only: + what += ", video only" + if steps: + what += f", {'last' if steps > 0 else 'first'} {abs(steps)} step(s)" + return what + + +PROVIDES = {"modifier": install} diff --git a/modules/sampling/block_repeat/tests/test_block_repeat.py b/modules/sampling/block_repeat/tests/test_block_repeat.py new file mode 100644 index 00000000..e9455a07 --- /dev/null +++ b/modules/sampling/block_repeat/tests/test_block_repeat.py @@ -0,0 +1,94 @@ +"""Block repeat on a real H3 forward.""" + +import pytest +import torch + + +@pytest.fixture(autouse=True) +def _needs_comfy(comfyui): + """Imports comfy.""" + + +def _load(tiny, patcher=None, **values): + from modules.sampling.modifiers.nodes import FunPackLoadModifiers + return FunPackLoadModifiers.execute( + patcher or tiny.patcher, {"block_repeat": {"enabled": True, **values}}).result + + +def _count_block_calls(tiny, patched, **run): + """How many times each real block ran during one forward.""" + counts = {} + handles = [] + for i, block in enumerate(tiny.patcher.model.diffusion_model.blocks): + handles.append(block.register_forward_pre_hook( + lambda m, a, i=i: counts.__setitem__(i, counts.get(i, 0) + 1))) + try: + tiny.run(patched, **run) + finally: + for h in handles: + h.remove() + return counts + + +def test_each_runs_every_named_block_times_plus_one(tiny_h3): + patched, status = _load(tiny_h3, blocks="1-2", times=2) + assert "block_repeat: blocks 1,2 run 3x each" in status + assert _count_block_calls(tiny_h3, patched) == {0: 1, 1: 3, 2: 3, 3: 1} + + +def test_span_loops_the_range_and_runs_each_block_once_per_loop(tiny_h3): + patched, status = _load(tiny_h3, blocks="1-2", shape="span", times=1) + assert "blocks 1-2 loop 2x as one span" in status + assert _count_block_calls(tiny_h3, patched) == {0: 1, 1: 2, 2: 2, 3: 1} + + +def test_span_output_equals_running_the_span_twice_by_hand(tiny_h3): + """The loop really feeds block 2's output back into block 1.""" + patched, _ = _load(tiny_h3, blocks="1-2", shape="span") + looped = tiny_h3.run(patched) + # By hand: blocks 0,1,2,1,2,3 in order. + from core import dit_hooks + manual = tiny_h3.patcher.clone() + real = tiny_h3.patcher.model.diffusion_model.blocks + + def twice(args, extra): + h = args["img"] + for i in (1, 2, 1, 2): + h = real[i](h, args["t_emb"], args["mod_segments"], args["rope_freqs"], + transformer_options=args["transformer_options"]) + return {"img": h} + + dit_hooks.add_block_hook(manual, "t", 1, twice) + dit_hooks.add_block_hook(manual, "t", 2, lambda a, e: {"img": a["img"]}) + assert torch.allclose(looped[0], tiny_h3.run(manual)[0]) + + +def test_a_scattered_list_is_refused_for_span(tiny_h3): + _p, status = _load(tiny_h3, blocks="0,2", shape="span") + assert "block_repeat: failed to install" in status + + +def test_step_window_repeats_only_on_the_last_steps(tiny_h3): + patched, _ = _load(tiny_h3, blocks="1", steps=1) + sched = torch.tensor([1.0, 0.7, 0.3, 0.0]) + early = _count_block_calls(tiny_h3, patched, sigma=0.7, sigmas=sched) + last = _count_block_calls(tiny_h3, patched, sigma=0.3, sigmas=sched) + assert early[1] == 1 and last[1] == 2 + + +def test_video_only_leaves_audio_as_a_single_pass_would(tiny_h3): + """Repeating the LAST block with video only: the audio branch must match a + run with no repeat at all, the video branch must not.""" + base = tiny_h3.run() + patched, _ = _load(tiny_h3, blocks="3", video_only=True) + out = tiny_h3.run(patched) + assert torch.allclose(out[1], base[1], atol=1e-6), "audio got the repeat" + assert not torch.allclose(out[0], base[0]), "video did not" + + +def test_video_only_span_leaves_audio_as_a_single_pass_would(tiny_h3): + base = tiny_h3.run() + patched, _ = _load(tiny_h3, blocks="2-3", shape="span", video_only=True) + out = tiny_h3.run(patched) + assert torch.allclose(out[1], base[1], atol=1e-6), "audio got the loop" + assert not torch.allclose(out[0], base[0])