130 lines
3.3 KiB
Python
130 lines
3.3 KiB
Python
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import sys
|
|
import types
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
|
|
def _install_comfy_stubs() -> None:
|
|
comfy = types.ModuleType("comfy")
|
|
conds = types.ModuleType("comfy.conds")
|
|
|
|
class CONDList(list):
|
|
pass
|
|
|
|
conds.CONDList = CONDList
|
|
comfy.conds = conds
|
|
sys.modules["comfy"] = comfy
|
|
sys.modules["comfy.conds"] = conds
|
|
|
|
|
|
def _load_nodes_module():
|
|
_install_comfy_stubs()
|
|
path = Path(__file__).resolve().parents[1] / "nodes.py"
|
|
spec = importlib.util.spec_from_file_location("comfyui_anima_edit_reference_nodes", path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
assert spec.loader is not None
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
class DummyInnerModel:
|
|
def __init__(self):
|
|
self.processed = []
|
|
|
|
def process_latent_in(self, latent):
|
|
self.processed.append(latent)
|
|
return latent + 100
|
|
|
|
def extra_conds(self, **kwargs):
|
|
return {"base": "cond"}
|
|
|
|
|
|
class DummyPatcher:
|
|
def __init__(self):
|
|
self.model = DummyInnerModel()
|
|
self.model_options = {}
|
|
self.wrapper = None
|
|
|
|
def clone(self):
|
|
clone = DummyPatcher()
|
|
clone.model = self.model
|
|
clone.model_options = dict(self.model_options)
|
|
return clone
|
|
|
|
def add_object_patch(self, name, value):
|
|
setattr(self.model, name, value)
|
|
|
|
def set_model_unet_function_wrapper(self, wrapper):
|
|
self.wrapper = wrapper
|
|
|
|
|
|
class DummyModelApply:
|
|
def __init__(self):
|
|
self.seen_input = None
|
|
|
|
def __call__(self, x, timestep, **kwargs):
|
|
self.seen_input = x
|
|
return torch.full_like(x, 7)
|
|
|
|
|
|
def test_node_mapping_uses_namespaced_key():
|
|
nodes = _load_nodes_module()
|
|
|
|
assert list(nodes.NODE_CLASS_MAPPINGS) == ["AnimaEditLora"]
|
|
assert (
|
|
nodes.NODE_DISPLAY_NAME_MAPPINGS["AnimaEditLora"]
|
|
== "Anima Edit LoRA"
|
|
)
|
|
|
|
|
|
def test_extra_conds_converts_reference_latents_to_ref_latents_cond_list():
|
|
nodes = _load_nodes_module()
|
|
patcher = DummyPatcher()
|
|
|
|
(patched,) = nodes.AnimaEditLora().patch(patcher)
|
|
|
|
ref_a = torch.ones(1, 4, 8, 8)
|
|
ref_b = torch.ones(1, 4, 1, 8, 8) * 2
|
|
out = patched.model.extra_conds(reference_latents=[ref_a, ref_b])
|
|
|
|
assert out["base"] == "cond"
|
|
assert len(out["ref_latents"]) == 2
|
|
assert torch.equal(out["ref_latents"][0], ref_a + 100)
|
|
assert torch.equal(out["ref_latents"][1], ref_b + 100)
|
|
|
|
|
|
def test_wrapper_appends_multi_reference_latents_on_time_axis_and_crops_output():
|
|
nodes = _load_nodes_module()
|
|
patcher = DummyPatcher()
|
|
(patched,) = nodes.AnimaEditLora().patch(patcher)
|
|
|
|
x = torch.zeros(2, 4, 1, 8, 8)
|
|
ref_a = torch.ones(1, 4, 8, 8)
|
|
ref_b = torch.ones(1, 4, 2, 8, 8) * 2
|
|
model_apply = DummyModelApply()
|
|
|
|
out = patched.wrapper(
|
|
model_apply,
|
|
{
|
|
"input": x,
|
|
"timestep": torch.tensor([1.0, 1.0]),
|
|
"c": {"ref_latents": [ref_a, ref_b]},
|
|
},
|
|
)
|
|
|
|
assert model_apply.seen_input.shape == (2, 4, 4, 8, 8)
|
|
assert torch.equal(model_apply.seen_input[:, :, :1], x)
|
|
assert torch.equal(
|
|
model_apply.seen_input[:, :, 1:2],
|
|
ref_a.unsqueeze(2).expand(2, -1, -1, -1, -1),
|
|
)
|
|
assert torch.equal(
|
|
model_apply.seen_input[:, :, 2:4],
|
|
ref_b.expand(2, -1, -1, -1, -1),
|
|
)
|
|
assert out.shape == x.shape
|