CauchyInno
5752b3b51f
Implement native ComfyUI Anima caching with documented sampler presets, example workflows, and validation coverage.
308 lines
10 KiB
Python
308 lines
10 KiB
Python
import math
|
|
import sys
|
|
import types
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
|
|
from nodes import ( # noqa: E402
|
|
ANIMA_COEFFICIENTS,
|
|
_STEP_INFO_KEY,
|
|
AnimaTeaCache,
|
|
AnimaTeaCacheRuntime,
|
|
CacheDecision,
|
|
poly1d,
|
|
)
|
|
|
|
|
|
def step_options(index, total=5, labels=None, uuids=None):
|
|
options = {
|
|
_STEP_INFO_KEY: {
|
|
"index": index,
|
|
"total": total,
|
|
"percent": index / total,
|
|
"schedule_key": (total,),
|
|
}
|
|
}
|
|
if labels is not None:
|
|
options["cond_or_uncond"] = labels
|
|
if uuids is not None:
|
|
options["uuids"] = uuids
|
|
return options
|
|
|
|
|
|
def runtime(threshold=0.05, start=0.0, end=1.0):
|
|
return AnimaTeaCacheRuntime(threshold, start, end, torch.device("cpu"))
|
|
|
|
|
|
def seed_residual(cache_runtime, decision, tensor):
|
|
cache_runtime.commit_residual(decision, torch.full_like(tensor, 3.0))
|
|
|
|
|
|
def test_poly1d_is_quartic_with_five_coefficients():
|
|
value = 0.125
|
|
expected = sum(
|
|
coefficient * value ** (len(ANIMA_COEFFICIENTS) - index - 1)
|
|
for index, coefficient in enumerate(ANIMA_COEFFICIENTS)
|
|
)
|
|
assert math.isclose(poly1d(ANIMA_COEFFICIENTS, value), expected)
|
|
|
|
|
|
def test_first_and_last_steps_force_full():
|
|
cache_runtime = runtime(threshold=100.0)
|
|
tensor = torch.ones(1, 1, 1, 1, 2)
|
|
|
|
first = cache_runtime.decide(tensor, step_options(0))
|
|
assert first.should_calc
|
|
seed_residual(cache_runtime, first, tensor)
|
|
|
|
middle = cache_runtime.decide(tensor * 1.001, step_options(1))
|
|
assert not middle.should_calc
|
|
|
|
last = cache_runtime.decide(tensor * 1.002, step_options(4))
|
|
assert last.should_calc
|
|
|
|
|
|
def test_threshold_refreshes_residual():
|
|
cache_runtime = runtime(threshold=0.01)
|
|
tensor = torch.ones(1, 1, 1, 1, 2)
|
|
first = cache_runtime.decide(tensor, step_options(0))
|
|
seed_residual(cache_runtime, first, tensor)
|
|
|
|
decision = cache_runtime.decide(tensor * 2, step_options(1))
|
|
assert decision.should_calc
|
|
assert cache_runtime.entries[0].accumulated_distance == 0.0
|
|
|
|
|
|
def test_cfg_entries_are_separate_and_reused_in_batch_order():
|
|
cache_runtime = runtime(threshold=100.0)
|
|
tensor = torch.stack((torch.ones(2), torch.ones(2) * 2)).reshape(2, 1, 1, 1, 2)
|
|
options = step_options(0, labels=[1, 0], uuids=["negative", "positive"])
|
|
first = cache_runtime.decide(tensor, options)
|
|
assert first.keys == [(1, "negative", 0), (0, "positive", 0)]
|
|
cache_runtime.commit_residual(
|
|
first,
|
|
torch.stack((torch.ones(2) * 3, torch.ones(2) * 4)).reshape(2, 1, 1, 1, 2),
|
|
)
|
|
|
|
options = step_options(1, labels=[1, 0], uuids=["negative", "positive"])
|
|
second = cache_runtime.decide(tensor * 1.001, options)
|
|
assert not second.should_calc
|
|
reused = cache_runtime.reuse_residual(second, torch.zeros_like(tensor))
|
|
assert torch.equal(reused[0], torch.full_like(reused[0], 3.0))
|
|
assert torch.equal(reused[1], torch.full_like(reused[1], 4.0))
|
|
|
|
|
|
def test_reentering_step_zero_resets_interrupted_run():
|
|
cache_runtime = runtime(threshold=100.0)
|
|
tensor = torch.ones(1, 1, 1, 1, 2)
|
|
first = cache_runtime.decide(tensor, step_options(0))
|
|
seed_residual(cache_runtime, first, tensor)
|
|
cache_runtime.decide(tensor * 1.001, step_options(1))
|
|
|
|
restarted = cache_runtime.decide(tensor, step_options(0))
|
|
assert restarted.should_calc
|
|
assert cache_runtime.entries[0].previous_residual is None
|
|
|
|
|
|
def test_new_schedule_and_shape_change_reset_cache():
|
|
cache_runtime = runtime(threshold=100.0)
|
|
tensor = torch.ones(1, 1, 1, 1, 2)
|
|
first = cache_runtime.decide(tensor, step_options(0))
|
|
seed_residual(cache_runtime, first, tensor)
|
|
|
|
new_schedule = step_options(1)
|
|
new_schedule[_STEP_INFO_KEY]["schedule_key"] = (99,)
|
|
assert cache_runtime.decide(tensor, new_schedule).should_calc
|
|
|
|
cache_runtime.commit_residual(
|
|
CacheDecision(True, True, [0], [slice(0, 1)]),
|
|
tensor,
|
|
)
|
|
larger = torch.ones(1, 1, 1, 2, 2)
|
|
assert cache_runtime.decide(larger, step_options(2)).should_calc
|
|
|
|
|
|
def test_outside_percent_range_forces_full():
|
|
cache_runtime = runtime(threshold=100.0, start=0.25, end=0.75)
|
|
tensor = torch.ones(1, 1, 1, 1, 2)
|
|
first = cache_runtime.decide(tensor, step_options(0))
|
|
seed_residual(cache_runtime, first, tensor)
|
|
|
|
assert not cache_runtime.decide(tensor * 1.001, step_options(2)).should_calc
|
|
assert cache_runtime.decide(tensor * 1.002, step_options(4)).should_calc
|
|
|
|
|
|
def test_missing_step_metadata_disables_cache():
|
|
cache_runtime = runtime()
|
|
decision = cache_runtime.decide(torch.ones(1, 1, 1, 1, 2), {})
|
|
assert decision.should_calc
|
|
assert not decision.cacheable
|
|
assert cache_runtime.entries == {}
|
|
|
|
|
|
def test_build_step_info_matches_current_sigma():
|
|
cache_runtime = runtime()
|
|
sigmas = torch.tensor([1.0, 0.7, 0.4, 0.1, 0.0])
|
|
info = cache_runtime._build_step_info(torch.tensor([0.69]), sigmas)
|
|
assert info["index"] == 1
|
|
assert info["total"] == 4
|
|
assert math.isclose(info["percent"], 1 / 4)
|
|
|
|
|
|
def test_first_block_signal_matches_norm1_formula():
|
|
class Block:
|
|
use_adaln_lora = True
|
|
layer_norm_self_attn = torch.nn.Identity()
|
|
|
|
@staticmethod
|
|
def adaln_modulation_self_attn(embedding):
|
|
return torch.cat((embedding, embedding * 2, embedding * 3), dim=-1)
|
|
|
|
class Model:
|
|
blocks = [Block()]
|
|
|
|
cache_runtime = runtime()
|
|
hidden = torch.ones(1, 1, 1, 1, 2)
|
|
embedding = torch.tensor([[[1.0, 2.0]]])
|
|
adaln_lora = torch.ones(1, 1, 6)
|
|
extra = torch.full_like(hidden, 2.0)
|
|
|
|
result = cache_runtime._first_block_modulated_input(
|
|
Model(), hidden, embedding, adaln_lora, extra
|
|
)
|
|
shift = embedding + 1
|
|
scale = embedding * 2 + 1
|
|
expected = (hidden + extra) * (1 + scale[:, :, None, None, :]) + shift[
|
|
:, :, None, None, :
|
|
]
|
|
assert torch.equal(result, expected)
|
|
|
|
|
|
def test_diffusion_wrapper_preserves_later_wrappers_and_replaces_original():
|
|
cache_runtime = runtime()
|
|
calls = []
|
|
cache_runtime.forward = lambda model, *args, **kwargs: calls.append(model) or "teacache"
|
|
|
|
class Executor:
|
|
class_obj = object()
|
|
original = lambda: "original"
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
return self.original(*args, **kwargs)
|
|
|
|
executor = Executor()
|
|
result = cache_runtime.diffusion_model_wrapper(
|
|
executor,
|
|
transformer_options={
|
|
"sigmas": torch.tensor([1.0]),
|
|
"sample_sigmas": torch.tensor([1.0, 0.0]),
|
|
},
|
|
)
|
|
assert result == "teacache"
|
|
assert calls == [executor.class_obj]
|
|
|
|
|
|
def test_node_registers_diffusion_wrapper_in_model_options():
|
|
class FakeAnima:
|
|
pass
|
|
|
|
class FakeModelPatcher:
|
|
def __init__(self):
|
|
self.diffusion_model = FakeAnima()
|
|
self.model_options = {}
|
|
|
|
def get_model_object(self, _name):
|
|
return self.diffusion_model
|
|
|
|
def clone(self):
|
|
clone = FakeModelPatcher()
|
|
clone.diffusion_model = self.diffusion_model
|
|
clone.model_options = {
|
|
key: value.copy() if isinstance(value, dict) else value
|
|
for key, value in self.model_options.items()
|
|
}
|
|
return clone
|
|
|
|
def add_wrapper_with_key(wrapper_type, key, wrapper, options, is_model_options=False):
|
|
assert is_model_options
|
|
transformer_options = options.setdefault("transformer_options", {})
|
|
transformer_options.setdefault("wrappers", {}).setdefault(wrapper_type, {}).setdefault(
|
|
key, []
|
|
).append(wrapper)
|
|
|
|
modules = {
|
|
"comfy": types.ModuleType("comfy"),
|
|
"comfy.ldm": types.ModuleType("comfy.ldm"),
|
|
"comfy.ldm.anima": types.ModuleType("comfy.ldm.anima"),
|
|
"comfy.ldm.anima.model": types.ModuleType("comfy.ldm.anima.model"),
|
|
"comfy.model_management": types.ModuleType("comfy.model_management"),
|
|
"comfy.patcher_extension": types.ModuleType("comfy.patcher_extension"),
|
|
}
|
|
modules["comfy.ldm.anima.model"].Anima = FakeAnima
|
|
modules["comfy.model_management"].get_torch_device = lambda: torch.device("cpu")
|
|
modules["comfy.patcher_extension"].WrappersMP = types.SimpleNamespace(
|
|
DIFFUSION_MODEL="diffusion_model"
|
|
)
|
|
modules["comfy.patcher_extension"].add_wrapper_with_key = add_wrapper_with_key
|
|
modules["comfy"].model_management = modules["comfy.model_management"]
|
|
modules["comfy"].patcher_extension = modules["comfy.patcher_extension"]
|
|
previous = {name: sys.modules.get(name) for name in modules}
|
|
sys.modules.update(modules)
|
|
try:
|
|
patched = AnimaTeaCache().apply_teacache(
|
|
FakeModelPatcher(), 0.05, 0.0, 1.0, "cpu"
|
|
)[0]
|
|
wrappers = patched.model_options["transformer_options"]["wrappers"]
|
|
assert len(wrappers["diffusion_model"]["anima_teacache"]) == 1
|
|
finally:
|
|
for name, module in previous.items():
|
|
if module is None:
|
|
sys.modules.pop(name, None)
|
|
else:
|
|
sys.modules[name] = module
|
|
|
|
|
|
def test_non_anima_model_is_rejected():
|
|
class FakeModelPatcher:
|
|
@staticmethod
|
|
def get_model_object(_name):
|
|
return object()
|
|
|
|
class FakeAnima:
|
|
pass
|
|
|
|
comfy_module = types.ModuleType("comfy")
|
|
ldm_module = types.ModuleType("comfy.ldm")
|
|
anima_package = types.ModuleType("comfy.ldm.anima")
|
|
anima_model_module = types.ModuleType("comfy.ldm.anima.model")
|
|
anima_model_module.Anima = FakeAnima
|
|
previous = {
|
|
name: sys.modules.get(name)
|
|
for name in ("comfy", "comfy.ldm", "comfy.ldm.anima", "comfy.ldm.anima.model")
|
|
}
|
|
sys.modules.update(
|
|
{
|
|
"comfy": comfy_module,
|
|
"comfy.ldm": ldm_module,
|
|
"comfy.ldm.anima": anima_package,
|
|
"comfy.ldm.anima.model": anima_model_module,
|
|
}
|
|
)
|
|
try:
|
|
try:
|
|
AnimaTeaCache().apply_teacache(FakeModelPatcher(), 0.05, 0.0, 1.0, "cpu")
|
|
except TypeError as error:
|
|
assert "only supports" in str(error)
|
|
else:
|
|
raise AssertionError("Expected non-Anima model rejection")
|
|
finally:
|
|
for name, module in previous.items():
|
|
if module is None:
|
|
sys.modules.pop(name, None)
|
|
else:
|
|
sys.modules[name] = module
|