feat(v5): randomizer and block repeat modules for H3, on shared chainable block hooks
Block repeat's video-only mode now really keeps audio single-pass: H3 blocks write in place, so v4 restored audio from an already-repeated stream. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
digital-garbage
co-authored by
Claude Opus 5.5
Claude Opus 5.5
parent
ec6e9893b4
commit
c3ad0dfd79
+64
@@ -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()
|
||||
|
||||
@@ -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
|
||||
+19
-2
@@ -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)
|
||||
|
||||
@@ -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
|
||||
+3
-3
@@ -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"]
|
||||
|
||||
@@ -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}
|
||||
|
||||
@@ -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
|
||||
@@ -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}
|
||||
@@ -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])
|
||||
@@ -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}
|
||||
@@ -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])
|
||||
Reference in New Issue
Block a user