458 lines
17 KiB
Python
458 lines
17 KiB
Python
"""Tests for model-patch ControlNets: Anima LLLite and Z-Image Fun (#365), and
|
|
Anima VACE routed to the ComfyUI-Advanced-ControlNet-Anima fork (#378).
|
|
|
|
Stubs the ComfyUI runtime (including ``comfy_extras.nodes_model_patch``) and
|
|
loads the sources under a synthetic package, matching test_preset_loader.py.
|
|
Runs ``FlakeStack.execute`` end to end and records what reaches the core
|
|
ControlNet / model-patch nodes.
|
|
|
|
Run directly (``python tests/test_model_patch_cn.py``) or via pytest.
|
|
"""
|
|
import importlib.machinery
|
|
import importlib.util
|
|
import json
|
|
import logging
|
|
import os
|
|
import sys
|
|
import types
|
|
|
|
|
|
PKG = "flakes_under_test_patch_cn"
|
|
|
|
FILES = {"controlnet": [], "model_patches": []}
|
|
MTIMES = {}
|
|
HEADERS = {}
|
|
NODE_MAPPINGS = {}
|
|
CALLS = {"patch_load": [], "lllite": [], "zimage": [], "cn_load": [], "cn_apply": [],
|
|
"acn_load": [], "acn_apply": []}
|
|
|
|
ACN_PKG = "fake_acn.adv_control"
|
|
|
|
|
|
class FakeModel(str):
|
|
"""A MODEL stand-in that compares by name; flow-model sigmas = 1 - percent."""
|
|
|
|
def get_model_object(self, name):
|
|
assert name == "model_sampling"
|
|
return types.SimpleNamespace(percent_to_sigma=lambda p: 1.0 - p)
|
|
|
|
|
|
class ACNLoader:
|
|
def load_controlnet(self, cnet, _tk_opt=None):
|
|
CALLS["acn_load"].append(cnet)
|
|
return (f"ACN({cnet})",)
|
|
|
|
|
|
class ACNApply:
|
|
def apply_controlnet(self, positive, negative, control_net, image, strength, start_percent, end_percent,
|
|
mask_optional=None, vae_optional=None):
|
|
CALLS["acn_apply"].append((control_net, image, strength, start_percent, end_percent, vae_optional))
|
|
return (positive + [["VACE", {}]], negative)
|
|
|
|
|
|
ACNLoader.__module__ = ACNApply.__module__ = f"{ACN_PKG}.nodes_main"
|
|
|
|
|
|
def _install_acn(with_vace):
|
|
for name in ("fake_acn", ACN_PKG):
|
|
mod = types.ModuleType(name)
|
|
mod.__path__ = []
|
|
sys.modules[name] = mod
|
|
sys.modules.pop(f"{ACN_PKG}.control_anima_vace", None)
|
|
if with_vace:
|
|
vace = types.ModuleType(f"{ACN_PKG}.control_anima_vace")
|
|
vace.__spec__ = importlib.machinery.ModuleSpec(vace.__name__, None)
|
|
sys.modules[vace.__name__] = vace
|
|
NODE_MAPPINGS["ACN_ControlNetLoaderAdvanced"] = ACNLoader
|
|
NODE_MAPPINGS["ACN_AdvancedControlNetApply_v2"] = ACNApply
|
|
|
|
|
|
def _stub_externals():
|
|
fp = types.ModuleType("folder_paths")
|
|
fp.get_filename_list = lambda cat: list(FILES.get(cat, []))
|
|
fp.get_full_path_or_raise = lambda cat, name: f"/fake/{cat}/{name}"
|
|
fp.get_input_directory = lambda: "/fake/input"
|
|
fp.base_path = "/fake"
|
|
sys.modules["folder_paths"] = fp
|
|
|
|
for name in ("numpy", "torch"):
|
|
if name not in sys.modules:
|
|
sys.modules[name] = types.ModuleType(name)
|
|
if "PIL" not in sys.modules:
|
|
pil = types.ModuleType("PIL")
|
|
pil.Image = types.SimpleNamespace(open=None)
|
|
pil.ImageOps = types.SimpleNamespace(exif_transpose=None)
|
|
sys.modules["PIL"] = pil
|
|
sys.modules["PIL.Image"] = pil.Image
|
|
sys.modules["PIL.ImageOps"] = pil.ImageOps
|
|
|
|
comfy = types.ModuleType("comfy")
|
|
comfy.sd = types.ModuleType("comfy.sd")
|
|
comfy.utils = types.ModuleType("comfy.utils")
|
|
sys.modules["comfy"] = comfy
|
|
sys.modules["comfy.sd"] = comfy.sd
|
|
sys.modules["comfy.utils"] = comfy.utils
|
|
comfy.utils.safetensors_header = lambda path: HEADERS.get(path)
|
|
# No scheduled LoRAs here, so carry_schedules sees no hooks.
|
|
comfy.hooks = types.ModuleType("comfy.hooks")
|
|
comfy.hooks.HookGroup = types.SimpleNamespace(combine_all_hooks=lambda groups: None)
|
|
comfy.hooks.set_hooks_for_conditioning = lambda cond, hooks, append_hooks=True: cond
|
|
sys.modules["comfy.hooks"] = comfy.hooks
|
|
sys.modules["node_helpers"] = types.ModuleType("node_helpers")
|
|
|
|
class CLIPTextEncode:
|
|
def encode(self, clip, text):
|
|
return ([[f"COND({text})", {}]],)
|
|
|
|
class EmptyLatentImage:
|
|
def generate(self, w, h, b):
|
|
return ({"samples": f"latent{w}x{h}"},)
|
|
|
|
class ControlNetLoader:
|
|
def load_controlnet(self, name):
|
|
CALLS["cn_load"].append(name)
|
|
return (f"CN({name})",)
|
|
|
|
class ControlNetApplyAdvanced:
|
|
def apply_controlnet(self, pos, neg, cn, image, strength, start, end):
|
|
CALLS["cn_apply"].append((cn, strength, start, end))
|
|
return pos + [["CTRL", {}]], neg
|
|
|
|
nodes = types.ModuleType("nodes")
|
|
nodes.CLIPTextEncode = CLIPTextEncode
|
|
nodes.EmptyLatentImage = EmptyLatentImage
|
|
nodes.ControlNetLoader = ControlNetLoader
|
|
nodes.ControlNetApplyAdvanced = ControlNetApplyAdvanced
|
|
nodes.LoraLoader = type("LoraLoader", (), {})
|
|
nodes.NODE_CLASS_MAPPINGS = NODE_MAPPINGS
|
|
sys.modules["nodes"] = nodes
|
|
|
|
class ModelPatchLoader:
|
|
def load_model_patch(self, name):
|
|
CALLS["patch_load"].append(name)
|
|
return (object(),)
|
|
|
|
class AnimaLLLiteApply:
|
|
def apply_patch(self, model, model_patch, image, strength, start_percent, end_percent):
|
|
CALLS["lllite"].append((model, model_patch, image, strength, start_percent, end_percent))
|
|
return (FakeModel(f"{model}+lllite"),)
|
|
|
|
class ZImageFunControlnet:
|
|
def diffsynth_controlnet(self, model, model_patch, vae, image=None, strength=1.0):
|
|
CALLS["zimage"].append((model, model_patch, vae, image, strength))
|
|
return (f"{model}+zfun",)
|
|
|
|
extras = types.ModuleType("comfy_extras")
|
|
nmp = types.ModuleType("comfy_extras.nodes_model_patch")
|
|
nmp.ModelPatchLoader = ModelPatchLoader
|
|
nmp.AnimaLLLiteApply = AnimaLLLiteApply
|
|
nmp.ZImageFunControlnet = ZImageFunControlnet
|
|
extras.nodes_model_patch = nmp
|
|
sys.modules["comfy_extras"] = extras
|
|
sys.modules["comfy_extras.nodes_model_patch"] = nmp
|
|
|
|
|
|
_stub_externals()
|
|
|
|
|
|
def _load(name):
|
|
root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
if PKG not in sys.modules:
|
|
pkg = types.ModuleType(PKG)
|
|
pkg.__path__ = [root]
|
|
sys.modules[PKG] = pkg
|
|
full = f"{PKG}.{name}"
|
|
if full in sys.modules:
|
|
return sys.modules[full]
|
|
spec = importlib.util.spec_from_file_location(full, os.path.join(root, f"{name}.py"))
|
|
mod = importlib.util.module_from_spec(spec)
|
|
sys.modules[full] = mod
|
|
spec.loader.exec_module(mod)
|
|
return mod
|
|
|
|
|
|
flake_io = _load("flake_io")
|
|
flake_compose = _load("flake_compose")
|
|
full_flake_node = _load("full_flake_node")
|
|
|
|
flake_compose._load_cn_image = lambda name: f"IMG({name})"
|
|
flake_compose._fit_cn_image = lambda image, w, h: f"{image}@{w}x{h}"
|
|
flake_compose.os = types.SimpleNamespace(path=types.SimpleNamespace(
|
|
getmtime=lambda p: MTIMES.get(p.rsplit("/", 1)[-1], 1.0)))
|
|
|
|
|
|
def _reset():
|
|
for v in CALLS.values():
|
|
v.clear()
|
|
for v in FILES.values():
|
|
v.clear()
|
|
MTIMES.clear()
|
|
HEADERS.clear()
|
|
NODE_MAPPINGS.clear()
|
|
flake_compose._CN_PATCH_CACHE.clear()
|
|
|
|
|
|
def _cn(cn_type="depth", **kw):
|
|
base = dict(type=cn_type, model_name="", image_name="pose_a", strength=0.8,
|
|
start_percent=0.1, end_percent=0.9)
|
|
base.update(kw)
|
|
return flake_io.ControlNetEntry(**base)
|
|
|
|
|
|
def _run(model_family, *flake_cns, upstream_family=None):
|
|
flakes = {f"f{i}": flake_io.Flake(name=f"f{i}", controlnets=list(cns)) for i, cns in enumerate(flake_cns)}
|
|
flake_io.resolve = lambda e: flakes[e["name"]]
|
|
flake_data = (
|
|
(FakeModel("BASE_MODEL"), "CLIP", "VAE"),
|
|
([["P", {}]], [["N", {}]], {"samples": "l"}, 1024, 1024, "", "", {"stems": [], "model_family": upstream_family}),
|
|
(30, 4.0, "euler", "simple", 1.0),
|
|
)
|
|
entries = [{"name": n} for n in flakes]
|
|
import json
|
|
return full_flake_node.FlakeStack().execute(model_family, flake_data, json.dumps(entries))[0]
|
|
|
|
|
|
def test_anima_depth_patches_model_with_lllite():
|
|
_reset()
|
|
FILES["model_patches"].append("anima/anima-lllite-depth-1.safetensors")
|
|
out = _run("Anima/Base", [_cn("depth")])
|
|
(model, _clip, _vae), gen, _ = out
|
|
assert model == "BASE_MODEL+lllite"
|
|
assert CALLS["patch_load"] == ["anima/anima-lllite-depth-1.safetensors"]
|
|
m, _patch, image, strength, start, end = CALLS["lllite"][0]
|
|
assert (m, image, strength, start, end) == ("BASE_MODEL", "IMG(pose_a)@1024x1024", 0.8, 0.1, 0.9)
|
|
# Conditioning is untouched — no SDXL ControlNet involved.
|
|
assert not CALLS["cn_load"] and not CALLS["cn_apply"]
|
|
assert gen[0] == [["COND()", {}]]
|
|
|
|
|
|
def test_multiple_anima_cns_chain_on_the_model():
|
|
_reset()
|
|
FILES["model_patches"] += ["anima/anima-lllite-depth-1.safetensors", "anima/anima-lllite-pose-1.safetensors"]
|
|
out = _run("Anima/Base", [_cn("depth")], [_cn("openpose", image_name="pose_b")])
|
|
assert out[0][0] == "BASE_MODEL+lllite+lllite"
|
|
assert [c[0] for c in CALLS["lllite"]] == ["BASE_MODEL", "BASE_MODEL+lllite"]
|
|
|
|
|
|
def test_zimage_passes_vae_and_strength():
|
|
_reset()
|
|
FILES["model_patches"].append("zimage/Z-Image-Turbo-Fun-Controlnet-Union-2.1.safetensors")
|
|
out = _run("ZImage/Turbo", [_cn("canny", start_percent=0.0, end_percent=1.0)])
|
|
assert out[0][0] == "BASE_MODEL+zfun"
|
|
m, _patch, vae, image, strength = CALLS["zimage"][0]
|
|
assert (m, vae, image, strength) == ("BASE_MODEL", "VAE", "IMG(pose_a)@1024x1024", 0.8)
|
|
|
|
|
|
def test_preset_family_wins_over_stale_widget():
|
|
# A reroute/subgraph leaves the widget on SDXL/Base; the family recorded by
|
|
# the preset loader must still route Anima CNs to LLLite (#382).
|
|
_reset()
|
|
FILES["model_patches"].append("anima/anima-lllite-depth-1.safetensors")
|
|
out = _run("SDXL/Base", [_cn("depth")], upstream_family="Anima/Base")
|
|
assert out[0][0] == "BASE_MODEL+lllite"
|
|
assert not CALLS["cn_load"] and not CALLS["cn_apply"]
|
|
|
|
|
|
def test_sdxl_path_unchanged():
|
|
_reset()
|
|
FILES["controlnet"].append("sdxl/controlnet_depth_sdxl.safetensors")
|
|
out = _run("SDXL/Illustrious", [_cn("depth")])
|
|
assert out[0][0] == "BASE_MODEL"
|
|
assert CALLS["cn_load"] == ["sdxl/controlnet_depth_sdxl.safetensors"]
|
|
assert CALLS["cn_apply"][0][1:] == (0.8, 0.1, 0.9)
|
|
assert not CALLS["patch_load"] and not CALLS["lllite"]
|
|
|
|
|
|
def test_zero_strength_and_missing_image_are_skipped():
|
|
_reset()
|
|
FILES["model_patches"].append("anima/anima-lllite-depth-1.safetensors")
|
|
out = _run("Anima/Base", [_cn("depth", strength=0), _cn("depth", image_name=" ")])
|
|
assert out[0][0] == "BASE_MODEL"
|
|
assert not CALLS["patch_load"]
|
|
|
|
|
|
def test_patch_reused_across_combo_runs():
|
|
_reset()
|
|
FILES["model_patches"].append("anima/anima-lllite-depth-1.safetensors")
|
|
_run("Anima/Base", [_cn("depth", image_name="pose_a")])
|
|
_run("Anima/Base", [_cn("depth", image_name="pose_b")])
|
|
_run("Anima/Base", [_cn("depth", image_name="pose_c")])
|
|
assert CALLS["patch_load"] == ["anima/anima-lllite-depth-1.safetensors"]
|
|
patches = {c[1] for c in CALLS["lllite"]}
|
|
assert len(patches) == 1
|
|
|
|
|
|
def test_patch_reloads_when_file_changes():
|
|
_reset()
|
|
name = "anima/anima-lllite-depth-1.safetensors"
|
|
flake_compose.load_cn_patch(name)
|
|
MTIMES["anima-lllite-depth-1.safetensors"] = 2.0
|
|
flake_compose.load_cn_patch(name)
|
|
assert CALLS["patch_load"] == [name, name]
|
|
|
|
|
|
def test_patch_cache_is_bounded_lru():
|
|
_reset()
|
|
names = [f"p{i}.safetensors" for i in range(flake_compose._CN_PATCH_CACHE_SIZE + 1)]
|
|
for n in names:
|
|
flake_compose.load_cn_patch(n)
|
|
assert names[0] not in flake_compose._CN_PATCH_CACHE
|
|
assert list(flake_compose._CN_PATCH_CACHE) == names[1:]
|
|
# Touching the oldest survivor makes it most recent.
|
|
flake_compose.load_cn_patch(names[1])
|
|
assert list(flake_compose._CN_PATCH_CACHE)[-1] == names[1]
|
|
|
|
|
|
VACE_FILE = "anima/Anima-ControlNet-VACE-Depth.safetensors"
|
|
|
|
|
|
def _add_vace_file(keys=("control_blocks.0.after_proj.weight", "control_blocks.0.attn.weight")):
|
|
FILES["controlnet"].append(VACE_FILE)
|
|
HEADERS[f"/fake/controlnet/{VACE_FILE}"] = json.dumps({k: {} for k in keys}).encode()
|
|
|
|
|
|
def test_anima_vace_routes_to_fork_with_vae():
|
|
_reset()
|
|
_install_acn(with_vace=True)
|
|
_add_vace_file()
|
|
FILES["model_patches"].append("anima/anima-lllite-depth-1.safetensors")
|
|
out = _run("Anima/Base", [_cn("depth", model_name="Anima-ControlNet-VACE-Depth")])
|
|
(model, _clip, _vae), gen, _ = out
|
|
assert CALLS["acn_load"] == [VACE_FILE]
|
|
assert CALLS["acn_apply"] == [(f"ACN({VACE_FILE})", "IMG(pose_a)@1024x1024", 0.8, 0.1, 0.9, "VAE")]
|
|
assert gen[0][-1] == ["VACE", {}]
|
|
# Conditioning-level control: the model is not patched.
|
|
assert model == "BASE_MODEL"
|
|
assert not CALLS["patch_load"] and not CALLS["lllite"]
|
|
|
|
|
|
def test_anima_lllite_stays_default_when_vace_file_present():
|
|
_reset()
|
|
_install_acn(with_vace=True)
|
|
_add_vace_file()
|
|
FILES["model_patches"].append("anima/anima-lllite-depth-1.safetensors")
|
|
out = _run("Anima/Base", [_cn("depth")])
|
|
assert out[0][0] == "BASE_MODEL+lllite"
|
|
assert not CALLS["acn_load"]
|
|
|
|
|
|
def test_anima_controlnet_file_without_vace_keys_uses_lllite():
|
|
_reset()
|
|
_install_acn(with_vace=True)
|
|
_add_vace_file(keys=("lllite_unet.weight",))
|
|
FILES["model_patches"].append("anima/anima-lllite-depth-1.safetensors")
|
|
out = _run("Anima/Base", [_cn("depth", model_name="Anima-ControlNet-VACE-Depth")])
|
|
assert out[0][0] == "BASE_MODEL+lllite"
|
|
assert not CALLS["acn_load"]
|
|
|
|
|
|
def _vace_error():
|
|
try:
|
|
_run("Anima/Base", [_cn("depth", model_name="Anima-ControlNet-VACE-Depth")])
|
|
except RuntimeError as exc:
|
|
return str(exc)
|
|
return None
|
|
|
|
|
|
def test_anima_vace_without_fork_errors():
|
|
_reset()
|
|
_add_vace_file()
|
|
assert "ComfyUI-Advanced-ControlNet-Anima" in (_vace_error() or "")
|
|
|
|
|
|
def test_anima_vace_with_upstream_acn_errors():
|
|
_reset()
|
|
_install_acn(with_vace=False)
|
|
_add_vace_file()
|
|
assert "ComfyUI-Advanced-ControlNet-Anima" in (_vace_error() or "")
|
|
assert not CALLS["acn_load"]
|
|
|
|
|
|
LLLITE_FILES = ["anima/anima-lllite-depth-1.safetensors", "anima/anima-lllite-any-test-like-v2.safetensors"]
|
|
|
|
|
|
def test_anima_explicit_lllite_file_wins_over_type():
|
|
_reset()
|
|
FILES["model_patches"] += LLLITE_FILES
|
|
out = _run("Anima/Base", [_cn("depth", model_name="anima/anima-lllite-any-test-like-v2.safetensors")])
|
|
assert out[0][0] == "BASE_MODEL+lllite"
|
|
assert CALLS["patch_load"] == ["anima/anima-lllite-any-test-like-v2.safetensors"]
|
|
|
|
|
|
def test_anima_explicit_vace_wins_over_type():
|
|
_reset()
|
|
_install_acn(with_vace=True)
|
|
_add_vace_file()
|
|
FILES["model_patches"] += LLLITE_FILES
|
|
out = _run("Anima/Base", [_cn("depth", model_name=VACE_FILE)])
|
|
assert out[0][0] == "BASE_MODEL"
|
|
assert CALLS["acn_load"] == [VACE_FILE]
|
|
assert not CALLS["patch_load"]
|
|
|
|
|
|
def test_anima_stale_model_falls_back_to_type():
|
|
_reset()
|
|
FILES["model_patches"] += LLLITE_FILES + ["zimage/Z-Image-Turbo-Fun-Controlnet-Union-2.1.safetensors"]
|
|
FILES["controlnet"].append("sdxl/controlnet_depth_sdxl.safetensors")
|
|
for stale in ("controlnet_depth_sdxl", "Z-Image-Turbo-Fun-Controlnet-Union-2.1", "anima/deleted.safetensors"):
|
|
CALLS["patch_load"].clear()
|
|
flake_compose._CN_PATCH_CACHE.clear()
|
|
out = _run("Anima/Base", [_cn("depth", model_name=stale)])
|
|
assert out[0][0] == "BASE_MODEL+lllite"
|
|
assert CALLS["patch_load"] == ["anima/anima-lllite-depth-1.safetensors"], stale
|
|
|
|
|
|
def test_explicit_model_ignored_with_type_outside_anima():
|
|
_reset()
|
|
FILES["controlnet"] += ["sdxl/controlnet_depth_sdxl.safetensors", "sdxl/other_depth.safetensors"]
|
|
_run("SDXL/Base", [_cn("depth", model_name="other_depth")])
|
|
assert CALLS["cn_load"] == ["sdxl/controlnet_depth_sdxl.safetensors"]
|
|
|
|
|
|
def test_anima_cn_files_lists_lllite_and_vace():
|
|
_reset()
|
|
_add_vace_file()
|
|
FILES["controlnet"] += ["anima/not-vace.safetensors", "sdxl/controlnet_depth_sdxl.safetensors"]
|
|
HEADERS["/fake/controlnet/anima/not-vace.safetensors"] = json.dumps({"x.weight": {}}).encode()
|
|
FILES["model_patches"] += LLLITE_FILES + ["zimage/Z-Image-Turbo-Fun-Controlnet-Union-2.1.safetensors"]
|
|
assert flake_compose.anima_cn_files() == [{"name": n, "kind": "lllite"} for n in LLLITE_FILES] + [
|
|
{"name": VACE_FILE, "kind": "vace"}]
|
|
|
|
|
|
def _logged_run(*args):
|
|
records = []
|
|
handler = logging.Handler()
|
|
handler.emit = records.append
|
|
logger = logging.getLogger()
|
|
level = logger.level
|
|
logger.addHandler(handler)
|
|
logger.setLevel(logging.INFO)
|
|
try:
|
|
_run(*args)
|
|
finally:
|
|
logger.removeHandler(handler)
|
|
logger.setLevel(level)
|
|
return [r.getMessage() for r in records if r.getMessage().startswith("[flakes] Anima")]
|
|
|
|
|
|
def test_anima_patch_logs_file_kind_and_sigma_range():
|
|
_reset()
|
|
_install_acn(with_vace=True)
|
|
_add_vace_file()
|
|
FILES["model_patches"] += LLLITE_FILES
|
|
lines = _logged_run("Anima/Base", [_cn("depth"), _cn("depth", model_name=VACE_FILE)])
|
|
assert lines == [
|
|
"[flakes] Anima lllite anima/anima-lllite-depth-1.safetensors image=pose_a fit=1024x1024 "
|
|
"strength=0.80 start=0.10 end=0.90 sigma=0.900..0.100",
|
|
f"[flakes] Anima vace {VACE_FILE} image=pose_a fit=1024x1024 "
|
|
"strength=0.80 start=0.10 end=0.90 sigma=0.900..0.100",
|
|
]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
for fname, fn in sorted(globals().items()):
|
|
if fname.startswith("test_") and callable(fn):
|
|
fn()
|
|
print(f"PASS {fname}")
|
|
print("\nALL PASSED")
|