Claude
3adebb731f
Old project/workflow names and browser settings migrate on load; /funpack redirects to /fun. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016acRss69eH1235eRdFTZHY
295 lines
11 KiB
Python
295 lines
11 KiB
Python
"""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 = "_funui_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 publish(patcher, name: str, value) -> None:
|
|
"""Leave `value` in this model's transformer_options under `name`, where every
|
|
hook and wrapper of the run can read it. How one module tells another something
|
|
it knows (where the prompt's phrases sit) without either importing the other."""
|
|
_transformer_options(patcher)[name] = value
|
|
|
|
|
|
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):
|
|
# `funui_below` tells a hook that REPLACES the block (rather than
|
|
# calling original_block) that doing so would skip someone's hook.
|
|
if inner is None:
|
|
return {**extra, "funui_below": False}
|
|
original = extra["original_block"]
|
|
return {**extra, "funui_below": True,
|
|
"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
|
|
|
|
|
|
PROBE = "funui_probe"
|
|
# Set on late-branch guidance's weakened copy (which also carries PROBE). Hooks that must not
|
|
# run on it, or must tell it from a probe, read this.
|
|
WEAK_BRANCH = "funui_weak_branch"
|
|
|
|
|
|
# Set by a wrapper that moves the latent's frame of reference (camera move) on the first call
|
|
# it moves: an edit filed in the old frame must not be added to this call's input.
|
|
FRAME_CHANGE = "funui_frame_change"
|
|
|
|
|
|
def weak_branch(transformer_options) -> bool:
|
|
return bool((transformer_options or {}).get(WEAK_BRANCH))
|
|
|
|
|
|
def probing(transformer_options) -> bool:
|
|
"""A throwaway model call (a candidate scored and discarded, e.g. the
|
|
first-step seed search). Nothing may be learned from one: a feature that
|
|
captures from it banks a step the clip never took."""
|
|
return bool((transformer_options or {}).get(PROBE))
|
|
|
|
|
|
def last_step(transformer_options) -> bool:
|
|
"""Whether this denoise call is the schedule's final step -- True when that
|
|
cannot be told, so a capture taken "on the last step" still happens. Never
|
|
during a probe."""
|
|
if probing(transformer_options):
|
|
return False
|
|
where = current_step(transformer_options)
|
|
return where is None or where[0] >= where[1] - 1
|
|
|
|
|
|
def late_half(transformer_options, ahead: int = 0) -> float:
|
|
"""0 over the first half of the steps, rising to 1 at the last: the gate
|
|
every learning feature that steers late shares. By step POSITION, so it
|
|
holds on any schedule -- v4's sigma-based gate was open on 0 of 12 steps
|
|
of H3's default schedule, silently. 0 when the step cannot be told.
|
|
|
|
`ahead` reads a LATER step's gate: an edit carried into step i+1's input
|
|
(core/input_steer.py) is gated by the step it lands on."""
|
|
where = current_step(transformer_options)
|
|
if where is None:
|
|
return 0.0
|
|
index, total = where
|
|
return max(0.0, 2.0 * (index + ahead) / max(1, total) - 1.0)
|
|
|
|
|
|
def row_span(mask):
|
|
"""A [S] bool mask of ONE contiguous run -> (start, stop), or None.
|
|
|
|
Slicing a span is a view; indexing with a mask copies every row, which on
|
|
a 37k-token video is hundreds of MB per call. One host sync, so callers
|
|
cache the answer per sequence length.
|
|
"""
|
|
if mask is None:
|
|
return None
|
|
hit = mask.nonzero()
|
|
if not len(hit):
|
|
return None
|
|
lo, hi = int(hit[0]), int(hit[-1]) + 1
|
|
return (lo, hi) if hi - lo == len(hit) else None
|
|
|
|
|
|
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 as exc: # noqa: BLE001
|
|
from . import log
|
|
log.once(f"provider:target_rows:{_spec.id}:{type(exc).__name__}", log.WARNING, f"target_rows provider {_spec.id}",
|
|
f"failed while answering -- {type(exc).__name__}: {exc}") # skipped, but the callers' "could not find the rows" now has a cause
|
|
continue
|
|
if mask is not None:
|
|
return mask
|
|
return None
|
|
|
|
|
|
def without_compiler(patcher, key: str) -> None:
|
|
"""Run this model's sampling with ComfyUI's model compiler off.
|
|
|
|
The compiler (aimdo malloc-graph) records what each block allocates and
|
|
assumes it never changes. A hook that REPLACES a block's forward with one
|
|
allocating differently -- a repeated block, a second attention stream --
|
|
crashes it with "aimdo memory compile error" (v4, on rentals: block repeat,
|
|
REINS injection, shadow negative). Comfy-Org's own fix is
|
|
--disable-comfy-compiler; it is read live on every scope, so switching it
|
|
for the one sampling call is enough and the next run keeps the compiler.
|
|
|
|
Attached to the MODEL as an OUTER_SAMPLE wrapper, so it holds for any
|
|
sampler, not only FunUI's.
|
|
"""
|
|
from comfy.patcher_extension import WrappersMP
|
|
|
|
def outer(executor, *args, **kwargs):
|
|
import comfy.cli_args
|
|
cli = comfy.cli_args.args
|
|
had = hasattr(cli, "disable_comfy_compiler")
|
|
prior = getattr(cli, "disable_comfy_compiler", None)
|
|
cli.disable_comfy_compiler = True
|
|
try:
|
|
return executor(*args, **kwargs)
|
|
finally:
|
|
if had:
|
|
cli.disable_comfy_compiler = prior
|
|
else:
|
|
try:
|
|
del cli.disable_comfy_compiler
|
|
except AttributeError:
|
|
pass
|
|
|
|
patcher.add_wrapper_with_key(WrappersMP.OUTER_SAMPLE, key, outer)
|