digital-garbageandClaude Sonnet 5.5
bc8631f683
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
329 lines
13 KiB
Python
329 lines
13 KiB
Python
"""Tagging what we install, so it can be removed again.
|
|
|
|
ComfyUI offers two shapes of hook. Some are keyed -- `add_wrapper_with_key` takes
|
|
a name and removing it is exact. Others are plain lists on `model_options`, where
|
|
`set_model_sampler_pre_cfg_function` appends an anonymous callable and nothing
|
|
records who put it there.
|
|
|
|
For the second kind, install-then-forget means a model run through the loader
|
|
twice carries two copies and applies the effect at double strength, reporting
|
|
once. That is the accumulation this project has already paid for: v4 registered
|
|
hooks on shared blocks and only a restart cleared them.
|
|
|
|
So everything we install is tagged, and everything tagged is stripped before the
|
|
next install. The walk below is generic -- it looks for our tag rather than for a
|
|
list of known hook names, because a list of hook names is a thing to keep current
|
|
and the day it falls behind is the day a hook stops being removable.
|
|
"""
|
|
|
|
from typing import Any
|
|
|
|
TAG = "_funpack_key"
|
|
|
|
|
|
def tag(fn, key: str):
|
|
"""Mark a callable as ours, so `strip` can find it later."""
|
|
try:
|
|
setattr(fn, TAG, key)
|
|
except AttributeError:
|
|
# A builtin or a slotted object cannot carry the mark, and something
|
|
# unremovable must not be installed silently.
|
|
raise TypeError(
|
|
f"{fn!r} cannot be tagged, so it could never be removed again. "
|
|
f"Install a plain function or a closure.")
|
|
return fn
|
|
|
|
|
|
def _ours(value: Any, prefix: str) -> bool:
|
|
key = getattr(value, TAG, None)
|
|
return isinstance(key, str) and key.startswith(prefix)
|
|
|
|
|
|
def strip(patcher, prefix: str) -> int:
|
|
"""Remove every tagged hook under `prefix`, and nothing else.
|
|
|
|
Covers both shapes: keyed wrappers and callbacks, and the anonymous function
|
|
lists on model_options. Returns how many were removed, so a caller can say
|
|
that it replaced something rather than leaving it to be inferred.
|
|
"""
|
|
removed = 0
|
|
|
|
for holder in (getattr(patcher, "wrappers", None), getattr(patcher, "callbacks", None)):
|
|
if not isinstance(holder, dict):
|
|
continue
|
|
for by_key in holder.values():
|
|
if not isinstance(by_key, dict):
|
|
continue
|
|
for key in [k for k in by_key if isinstance(k, str) and k.startswith(prefix)]:
|
|
by_key.pop(key, None)
|
|
removed += 1
|
|
|
|
options = getattr(patcher, "model_options", None)
|
|
if isinstance(options, dict):
|
|
removed += _strip_options(options, prefix)
|
|
|
|
# add_object_patch entries (a module attribute swapped at load time) are
|
|
# keyed by attribute path, so a clone carries ours forward too; without this
|
|
# a feature switched OFF kept its patch on a model that went round again.
|
|
objects = getattr(patcher, "object_patches", None)
|
|
if isinstance(objects, dict):
|
|
for name in [n for n, v in objects.items() if _ours(v, prefix)]:
|
|
objects.pop(name, None)
|
|
removed += 1
|
|
return removed
|
|
|
|
|
|
def _strip_options(options: dict, prefix: str) -> int:
|
|
"""Walk model_options for tagged callables, in lists or on their own."""
|
|
removed = 0
|
|
for name, value in list(options.items()):
|
|
if isinstance(value, list):
|
|
keep = [item for item in value if not _ours(item, prefix)]
|
|
if len(keep) != len(value):
|
|
removed += len(value) - len(keep)
|
|
options[name] = keep
|
|
elif callable(value) and _ours(value, prefix):
|
|
# 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)
|
|
elif _ours(value, prefix):
|
|
# A tagged object that is not a callable (a handler core looks up by key): it has
|
|
# nothing underneath to put back.
|
|
options.pop(name, None)
|
|
removed += 1
|
|
return removed
|
|
|
|
|
|
# --- keeping a module from killing a run ----------------------------------
|
|
#
|
|
# Catching a module that fails to INSTALL is easy and was already done. The one
|
|
# that matters is a module that installs fine and raises on step 3 of 30: the
|
|
# generation dies after the GPU time is spent, and on a rental that is the whole
|
|
# cost of the run. A modifier is an opinion about the picture. It is never worth
|
|
# the picture.
|
|
#
|
|
# So a failing hook is dropped for the rest of the run and the run continues
|
|
# without it. Loudly -- a dropped modifier is the difference between "I asked for
|
|
# this" and "this happened", so it is said once, with the traceback, and carried
|
|
# in the run's status.
|
|
|
|
#: Set by core/control.py: told when a module's hook raised, so the fault outlives the run.
|
|
on_fault = None
|
|
|
|
|
|
class Dropped:
|
|
"""What was disabled during a run, and why. One per run, not per process."""
|
|
|
|
def __init__(self):
|
|
self.reasons: dict = {}
|
|
|
|
def clear(self) -> None:
|
|
"""Forget what was dropped. Called when a generation starts: this run's record starts
|
|
empty (a module that FAILED also stays off across runs: that is
|
|
core/control.py's quarantine, not this record) -- and because ComfyUI's cache hands the same
|
|
object to both."""
|
|
self.reasons.clear()
|
|
|
|
def record(self, key: str, exc: BaseException) -> bool:
|
|
"""True the first time this key fails, so callers report once."""
|
|
if key in self.reasons:
|
|
return False
|
|
self.reasons[key] = f"{type(exc).__name__}: {exc}"
|
|
if on_fault is not None:
|
|
try:
|
|
on_fault(key, exc)
|
|
except Exception: # noqa: BLE001 -- bookkeeping must not end a run
|
|
pass
|
|
return True
|
|
|
|
def __contains__(self, key) -> bool:
|
|
return key in self.reasons
|
|
|
|
def __bool__(self) -> bool:
|
|
return bool(self.reasons)
|
|
|
|
def items(self):
|
|
return self.reasons.items()
|
|
|
|
|
|
def guard(fn, key: str, neutral, dropped: Dropped):
|
|
"""`fn`, but unable to end a run.
|
|
|
|
`neutral(*args, **kwargs)` produces the result the caller would have got had
|
|
the hook never been installed -- it differs per hook shape, which is why it
|
|
is supplied rather than guessed. After the first failure the hook is skipped
|
|
outright: a modifier that raises once will raise every step, and thirty
|
|
tracebacks say nothing the first one did not.
|
|
"""
|
|
import traceback
|
|
|
|
def guarded(*args, **kwargs):
|
|
if key in dropped:
|
|
return neutral(*args, **kwargs)
|
|
try:
|
|
return fn(*args, **kwargs)
|
|
except Exception as exc: # noqa: BLE001 -- the whole point
|
|
if dropped.record(key, exc):
|
|
from . import log
|
|
log.warning(
|
|
key,
|
|
"failed during sampling and is now OFF for the rest of this run; "
|
|
"the run continues without it\n"
|
|
+ "".join(traceback.format_exception(exc)).rstrip())
|
|
return neutral(*args, **kwargs)
|
|
|
|
# Carried over so a guarded hook is still ours to strip.
|
|
marker = getattr(fn, TAG, None)
|
|
if marker is not None:
|
|
guarded.__dict__[TAG] = marker
|
|
return guarded
|
|
|
|
|
|
class _Executor:
|
|
"""ComfyUI's wrapper executor, seen through a glass: calls and attributes pass
|
|
through, and what the call raised or returned is remembered."""
|
|
|
|
def __init__(self, executor):
|
|
self._executor = executor
|
|
self.raised = None
|
|
self.done = False
|
|
self.result = None
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
try:
|
|
self.result = self._executor(*args, **kwargs)
|
|
except BaseException as exc:
|
|
self.raised = exc
|
|
raise
|
|
self.done = True
|
|
return self.result
|
|
|
|
def __getattr__(self, name):
|
|
return getattr(self._executor, name)
|
|
|
|
|
|
def guard_wrapper(fn, key: str, dropped: Dropped):
|
|
"""`guard` for `wrapper(executor, *args)` hooks, which the generic one gets
|
|
wrong twice: the executor IS the rest of the run, so
|
|
|
|
* what the executor raises (an interrupt, an OOM, the next module's bug) is
|
|
not this module's failure -- it passes through untouched, and
|
|
* once the executor has run, falling back to it would run the whole thing
|
|
AGAIN -- a wrapper that fails afterwards returns what already ran.
|
|
"""
|
|
import traceback
|
|
|
|
def guarded(executor, *args, **kwargs):
|
|
if key in dropped:
|
|
return executor(*args, **kwargs)
|
|
seen = _Executor(executor)
|
|
try:
|
|
return fn(seen, *args, **kwargs)
|
|
except Exception as exc: # noqa: BLE001
|
|
if exc is seen.raised:
|
|
raise
|
|
if dropped.record(key, exc):
|
|
from . import log
|
|
log.warning(
|
|
key,
|
|
"failed during sampling and is now OFF for the rest of this run; "
|
|
"the run continues without it\n"
|
|
+ "".join(traceback.format_exception(exc)).rstrip())
|
|
return seen.result if seen.done else executor(*args, **kwargs)
|
|
|
|
marker = getattr(fn, TAG, None)
|
|
if marker is not None:
|
|
guarded.__dict__[TAG] = marker
|
|
return guarded
|
|
|
|
|
|
# The neutral result for each hook ComfyUI offers, by the method that installs
|
|
# it. Core knows ComfyUI's own shapes here -- not FunPack's features -- because
|
|
# "what this hook returns when it does nothing" is a fact about ComfyUI.
|
|
def _neutral_wrapper(executor, *args, **kwargs):
|
|
return executor(*args, **kwargs)
|
|
|
|
|
|
NEUTRAL = {
|
|
"add_wrapper_with_key": _neutral_wrapper,
|
|
"add_callback_with_key": lambda *a, **k: None,
|
|
"set_model_sampler_pre_cfg_function": lambda args: args["conds_out"],
|
|
"set_model_sampler_post_cfg_function": lambda args: args["denoised"],
|
|
"set_model_sampler_cfg_function": lambda args: args["input"] - (
|
|
args["uncond_denoised"] + (args["cond_denoised"] - args["uncond_denoised"]) * args["cond_scale"]),
|
|
"set_model_unet_function_wrapper": lambda apply_fn, args: apply_fn(args["input"], args["timestep"],
|
|
**args.get("c", {})),
|
|
"set_model_denoise_mask_function": lambda sigma, denoise_mask, extra_options=None: denoise_mask,
|
|
}
|
|
|
|
# Which positional argument holds the callable, for methods that take more than one.
|
|
CALLABLE_AT = {"add_wrapper_with_key": 2, "add_callback_with_key": 2}
|
|
|
|
|
|
class GuardedPatcher:
|
|
"""A ModelPatcher that guards every hook installed through it.
|
|
|
|
A module is handed one of these instead of the real patcher, so it does not
|
|
opt in to being safe -- a module that could forget to is a module that can
|
|
break the foundation. Everything else forwards untouched.
|
|
"""
|
|
|
|
def __init__(self, patcher, key: str, dropped: Dropped, guarding: bool = True):
|
|
object.__setattr__(self, "_patcher", patcher)
|
|
object.__setattr__(self, "_key", key)
|
|
object.__setattr__(self, "_dropped", dropped)
|
|
object.__setattr__(self, "_guarding", guarding)
|
|
object.__setattr__(self, "unguarded", [])
|
|
|
|
def __getattr__(self, name):
|
|
target = getattr(self._patcher, name)
|
|
if not callable(target) or not name.startswith(("set_model_", "add_wrapper", "add_callback")):
|
|
return target
|
|
|
|
if not self._guarding:
|
|
# Full control: install exactly what the module wrote, and let it
|
|
# fail where it fails. Someone asking for the raw behaviour wants
|
|
# the real traceback at the real step, not a tidy recovery.
|
|
return target
|
|
|
|
neutral = NEUTRAL.get(name)
|
|
if neutral is None:
|
|
# Honest about the limit: an unknown hook shape has no known neutral
|
|
# result, so it is installed as-is and named. Silently pretending it
|
|
# was guarded would be worse than saying it was not.
|
|
self.unguarded.append(name)
|
|
return target
|
|
|
|
index = CALLABLE_AT.get(name, 0)
|
|
|
|
def install(*args, **kwargs):
|
|
args = list(args)
|
|
if len(args) > index and callable(args[index]):
|
|
args[index] = (guard_wrapper(args[index], self._key, self._dropped)
|
|
if name == "add_wrapper_with_key"
|
|
else guard(args[index], self._key, neutral, self._dropped))
|
|
return target(*args, **kwargs)
|
|
|
|
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)
|