File size: 2,124 Bytes
e488ce7 | 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 | #!/usr/bin/env python3
"""H3 capability probe + collision report. Run from the ComfyUI root:
python custom_nodes/ComfyUI-MAINodes/tests/test_h3_capabilities.py
No GPU, no weights. Checks that the probe reads the installed core (every
key resolves to a bool or a named state, never raises), that a fake
ModelPatcher with a foreign double-block replacement produces exactly one
collision warning naming the foreign owner, and that a clean model
produces none.
"""
import os, sys, types
sys.path.insert(0, os.getcwd())
import comfy.options; comfy.options.enable_args_parsing()
sys.path.insert(0, os.path.join(os.getcwd(), "custom_nodes"))
from importlib import import_module
caps_mod = import_module("ComfyUI-MAINodes.h3_capabilities")
caps = caps_mod.probe_core()
for k in ("per_token_masks", "clip_guide", "audio_guide", "keyframe_plus_refs",
"tokenizer_special_tokens", "visual_cond_noise_aug", "audio_cond_noise_aug",
"tae_h3_decoder", "length_rounds_up", "arbitrary_keyframe_position"):
assert k in caps, k
assert caps[k] != "unknown", (k, "source unreadable")
print("probe:", {k: v for k, v in caps.items() if k not in ("comfy_root",)})
# a fake patcher with a foreign replacement on one block
class Fake:
def __init__(self, foreign):
self.model_options = {"transformer_options": {"patches_replace": {"dit": {}}}}
self.object_patches = {}
if foreign:
src = "def f(*a, **k):\n return a[0]\n"
ns = {}
code = compile(src, "/x/custom_nodes/SomeOtherPack/nodes.py", "exec")
exec(code, ns)
self.model_options["transformer_options"]["patches_replace"]["dit"][("double_block", 3)] = ns["f"]
rep = caps_mod.block_patch_report(Fake(True))
warns = caps_mod.collision_warnings(rep)
assert len(warns) == 1 and "SomeOtherPack" in warns[0], warns
clean = caps_mod.collision_warnings(caps_mod.block_patch_report(Fake(False)))
assert [w for w in clean if "double_block" in w] == [], clean
text = caps_mod.format_report(caps, rep)
assert "WARNINGS" in text and "SomeOtherPack" in text
print(text)
print("PASS")
|