File size: 7,544 Bytes
bc50ea0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 | """Lazy AV denoise-mask payload compatibility for H3 masked continuation.
Normal guide continuation never installs this wrapper. Native ComfyUI support
always wins; the wrapper only adds the two AV mask conditions when the running
MiniMaxH3 implementation does not already expose PR #15375-equivalent payload
handling.
Adapted from seitanism/ComfyUI-H3-Motion-Context-MultiRef (GPL-3.0).
"""
from __future__ import annotations
import functools
import inspect
import logging
import types
import torch
import comfy.conds
import comfy.model_base as model_base
import comfy.utils as utils
_LOG = logging.getLogger("minimax_h3_context_loop.masked_prefix")
# Version 4 matches the final merged #15375 helper-based payload path. Version
# 3 predates merge commit c676536 and rounds before token-grid pooling; it must
# not wrap or overwrite the native merged result.
_MARKER = "_h3_existing_video_av_mask_payload_compat_v4"
_LEGACY_MARKERS = (
"_h3_existing_video_av_mask_payload_compat_v2",
"_h3_existing_video_av_mask_payload_compat_v3",
)
# Kept as an internal compatibility alias for focused third-party tests that
# marked the immediately preceding wrapper directly.
_LEGACY_MARKER = _LEGACY_MARKERS[-1]
def _walk_wrapped(fn):
seen = set()
while fn is not None and id(fn) not in seen:
seen.add(id(fn))
yield fn
fn = getattr(fn, "__wrapped__", None)
def _is_ours(fn):
code = getattr(fn, "__code__", None)
return bool(
getattr(fn, _MARKER, False)
and code is not None
and code.co_name == "wrapper"
and code.co_filename == __file__
)
def _is_compatible_wrapper(fn):
code = getattr(fn, "__code__", None)
return bool(
getattr(fn, _MARKER, False)
and code is not None
and code.co_name == "wrapper"
)
def _is_legacy_wrapper(fn):
code = getattr(fn, "__code__", None)
return bool(
any(getattr(fn, marker, False) for marker in _LEGACY_MARKERS)
and code is not None
and code.co_name == "wrapper"
)
def _signature_has(fn, *names):
try:
params = inspect.signature(fn).parameters
except (TypeError, ValueError):
return False
return all(name in params for name in names)
def _code_strings(code):
"""Yield strings/names from a function and nested code objects."""
if not isinstance(code, types.CodeType):
return
yield from code.co_names
def walk_const(value):
if isinstance(value, str):
yield value
elif isinstance(value, types.CodeType):
yield from _code_strings(value)
elif isinstance(value, (tuple, list, set, frozenset)):
for item in value:
yield from walk_const(item)
for value in code.co_consts:
yield from walk_const(value)
def _function_mentions_native_payload(fn):
for item in _walk_wrapped(fn):
if _is_ours(item) or _is_legacy_wrapper(item):
continue
strings = set(_code_strings(getattr(item, "__code__", None)) or ())
if "denoise_mask" in strings and "audio_denoise_mask" in strings:
return True
return False
def _function_calls_native_helper(fn):
for item in _walk_wrapped(fn):
if _is_ours(item) or _is_legacy_wrapper(item):
continue
strings = set(_code_strings(getattr(item, "__code__", None)) or ())
if "denoise_mask" in strings and "_denoise_mask_conds" in strings:
return True
return False
def _native_h3_mask_hooks(cls, fn):
if cls is None:
return False
scale = cls.__dict__.get("scale_latent_inpaint")
merged_helpers = all(callable(cls.__dict__.get(name)) for name in (
"_pool_masks_to_token_grid",
"_token_grid_masks",
"_denoise_mask_values",
"_denoise_mask_conds",
))
merged_path = bool(
merged_helpers
and _function_calls_native_helper(fn)
)
return bool(
merged_path
and callable(scale)
and _signature_has(scale, "x", "denoise_mask")
)
def _native_av_mask_payload(cls, fn):
return bool(fn and _native_h3_mask_hooks(cls, fn))
def capability_status():
cls = getattr(model_base, "MiniMaxH3", None)
fn = getattr(cls, "extra_conds", None) if cls is not None else None
direct = bool(fn and _function_mentions_native_payload(fn))
native_hooks = _native_h3_mask_hooks(cls, fn)
native = bool(fn and native_hooks)
return {
"available": fn is not None,
"native_av_mask_payload": native,
"native_payload_direct": direct,
"native_h3_mask_hooks": native_hooks,
"wrapper_present": bool(
fn and any(_is_compatible_wrapper(item)
for item in _walk_wrapped(fn))),
}
def _add_av_mask_conditions(owner, out, kwargs):
if not isinstance(out, dict):
return
denoise_mask = kwargs.get("denoise_mask")
latent_shapes = kwargs.get("latent_shapes")
if denoise_mask is None or latent_shapes is None or len(latent_shapes) < 2:
return
native_helper = getattr(owner, "_denoise_mask_conds", None)
if callable(native_helper):
out.update(native_helper(denoise_mask, latent_shapes))
return
# Last-resort compatibility when this module is used without the engine
# gate. The normal pack path installs the exact helper set first. Ceil to
# the merged 8-bit token strength rather than the pre-merge nearest value.
masks = utils.unpack_latents(denoise_mask, latent_shapes)
if len(masks) < 2:
return
pool = getattr(owner, "_pool_masks_to_token_grid", None)
if callable(pool):
masks = pool(masks)
masks = [torch.ceil(mask * 256.0) / 256.0 for mask in masks]
if torch.amin(masks[0]).item() < 1.0 - 1e-3:
out["denoise_mask"] = comfy.conds.CONDRegular(
masks[0][:1, :1].clone())
if torch.amin(masks[1]).item() < 1.0 - 1e-3:
out["audio_denoise_mask"] = comfy.conds.CONDRegular(
masks[1][:1].amax(dim=1, keepdim=True))
def _make_wrapper(base):
@functools.wraps(base, updated=())
def wrapper(self, **kwargs):
out = base(self, **kwargs)
_add_av_mask_conditions(self, out, kwargs)
return out
setattr(wrapper, _MARKER, True)
return wrapper
def ensure_av_mask_payload_compat():
cls = getattr(model_base, "MiniMaxH3", None)
if cls is None or not hasattr(cls, "extra_conds"):
raise RuntimeError(
"h3_masked_prefix: MiniMaxH3.extra_conds not found.")
current = cls.extra_conds
original = current
# Strip only obsolete top-level versions before probing native support.
# Otherwise their wrapped native function can make the chain look healthy
# while the outer v3 wrapper overwrites the merged token-grid conditions.
while _is_legacy_wrapper(current) and callable(
getattr(current, "__wrapped__", None)):
current = current.__wrapped__
if current is not original:
cls.extra_conds = current
if _native_av_mask_payload(cls, current):
return True
if any(_is_compatible_wrapper(item) for item in _walk_wrapped(current)):
return True
# Unknown third-party wrappers remain in the chain; the v4 layer only
# replaces the two mask conditions after their base result is produced.
cls.extra_conds = _make_wrapper(current)
_LOG.info(
"h3_masked_prefix: PR #15375 AV-mask payload compatibility enabled")
return True
|