235 lines
8.4 KiB
Python
235 lines
8.4 KiB
Python
# SimpleSyrup - workflow-focused ComfyUI extensions for image generation
|
|
# Copyright (C) 2026 Artificial Sweetener and contributors
|
|
# SPDX-License-Identifier: AGPL-3.0-or-later
|
|
|
|
"""Prove exact read-only discovery of the installed Anima patch surface."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from types import MethodType
|
|
|
|
import pytest
|
|
from anima_module_surface_fixtures import installed_meta_anima
|
|
from comfy.ldm.anima.model import Anima
|
|
from comfy.ldm.cosmos.predict2 import Attention, Block, GPT2FeedForward
|
|
from torch import nn
|
|
|
|
from simple_syrup.runtime.regional_lora.anima_module_surface import (
|
|
ANIMA_MODULE_SURFACE_DISCOVERY,
|
|
AnimaModuleSurfaceError,
|
|
)
|
|
from simple_syrup.runtime.regional_lora.anima_targets import (
|
|
ANIMA_BLOCK_COUNT,
|
|
AnimaLoraTargetFamily,
|
|
anima_lora_target_name,
|
|
expected_anima_lora_features,
|
|
)
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def installed_anima() -> nn.Module:
|
|
"""Build the real installed Anima graph with allocation-free meta weights."""
|
|
|
|
model = installed_meta_anima()
|
|
if not isinstance(model, nn.Module):
|
|
raise AssertionError("Installed Anima must be a torch module")
|
|
return model
|
|
|
|
|
|
def test_discovers_all_installed_blocks_and_lora_targets_read_only(
|
|
installed_anima: nn.Module,
|
|
) -> None:
|
|
"""Bind all 448 canonical targets without changing graph or global methods."""
|
|
|
|
named_modules_before = tuple(installed_anima.named_modules())
|
|
state_keys_before = tuple(installed_anima.state_dict().keys())
|
|
class_forwards_before = (
|
|
Anima.forward,
|
|
Block.forward,
|
|
Attention.forward,
|
|
GPT2FeedForward.forward,
|
|
)
|
|
|
|
surface = ANIMA_MODULE_SURFACE_DISCOVERY.discover(installed_anima)
|
|
|
|
assert surface.diffusion_model is installed_anima
|
|
assert len(surface.blocks) == ANIMA_BLOCK_COUNT
|
|
assert len(surface.lora_targets) == 448
|
|
assert [target.target_name for target in surface.lora_targets] == [
|
|
anima_lora_target_name(block_index, family)
|
|
for block_index in range(ANIMA_BLOCK_COUNT)
|
|
for family in AnimaLoraTargetFamily
|
|
]
|
|
assert tuple(installed_anima.named_modules()) == named_modules_before
|
|
assert tuple(installed_anima.state_dict().keys()) == state_keys_before
|
|
assert (
|
|
Anima.forward,
|
|
Block.forward,
|
|
Attention.forward,
|
|
GPT2FeedForward.forward,
|
|
) == class_forwards_before
|
|
|
|
|
|
def test_descriptors_retain_exact_installed_owners_and_shapes(
|
|
installed_anima: nn.Module,
|
|
) -> None:
|
|
"""Retain every block owner and target module by exact object identity."""
|
|
|
|
surface = ANIMA_MODULE_SURFACE_DISCOVERY.discover(installed_anima)
|
|
|
|
for block_surface in surface.blocks:
|
|
block = block_surface.block
|
|
assert block_surface.self_attention is block.self_attn
|
|
assert block_surface.cross_attention is block.cross_attn
|
|
assert block_surface.mlp is block.mlp
|
|
assert block_surface.adaln_self_attention is block.adaln_modulation_self_attn
|
|
assert block_surface.adaln_cross_attention is block.adaln_modulation_cross_attn
|
|
assert block_surface.adaln_mlp is block.adaln_modulation_mlp
|
|
assert len(block_surface.lora_targets) == len(AnimaLoraTargetFamily)
|
|
for target in block_surface.lora_targets:
|
|
assert target.module is _resolve_target(block, target.family)
|
|
assert (target.input_features, target.output_features) == (
|
|
expected_anima_lora_features(target.family)
|
|
)
|
|
assert target.forward_signature
|
|
|
|
|
|
def test_rejects_foreign_diffusion_module_before_graph_inspection() -> None:
|
|
"""Keep unverified Cosmos variants outside the Anima patch contract."""
|
|
|
|
with pytest.raises(AnimaModuleSurfaceError, match="exact installed.*Anima"):
|
|
ANIMA_MODULE_SURFACE_DISCOVERY.discover(nn.Identity())
|
|
|
|
|
|
def test_aggregates_configuration_and_block_count_drift(
|
|
installed_anima: nn.Module,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Report every incompatible model-level setting in one pre-patch failure."""
|
|
|
|
surface = ANIMA_MODULE_SURFACE_DISCOVERY.discover(installed_anima)
|
|
monkeypatch.setattr(installed_anima, "model_channels", 1024)
|
|
monkeypatch.setattr(installed_anima, "num_heads", 8)
|
|
monkeypatch.setattr(installed_anima, "use_adaln_lora", False)
|
|
monkeypatch.setattr(
|
|
installed_anima,
|
|
"blocks",
|
|
nn.ModuleList([item.block for item in surface.blocks[:-1]]),
|
|
)
|
|
|
|
with pytest.raises(AnimaModuleSurfaceError) as captured:
|
|
ANIMA_MODULE_SURFACE_DISCOVERY.discover(installed_anima)
|
|
|
|
message = str(captured.value)
|
|
assert "diffusion_model.model_channels" in message
|
|
assert "diffusion_model.num_heads" in message
|
|
assert "diffusion_model.use_adaln_lora" in message
|
|
assert "expected 28 blocks, observed 27" in message
|
|
|
|
|
|
def test_aggregates_block_owner_adaln_and_target_shape_drift(
|
|
installed_anima: nn.Module,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Reject mixed block owners, incomplete AdaLN paths, and wrong linear shapes."""
|
|
|
|
surface = ANIMA_MODULE_SURFACE_DISCOVERY.discover(installed_anima)
|
|
monkeypatch.setattr(surface.blocks[0].block, "self_attn", nn.Identity())
|
|
monkeypatch.setattr(
|
|
surface.blocks[1].block,
|
|
"adaln_modulation_cross_attn",
|
|
nn.Sequential(nn.SiLU()),
|
|
)
|
|
monkeypatch.setattr(
|
|
surface.blocks[2].block.mlp,
|
|
"layer1",
|
|
nn.Linear(32, 64, bias=False, device="meta"),
|
|
)
|
|
monkeypatch.setattr(
|
|
surface.blocks[3].self_attention,
|
|
"k_proj",
|
|
surface.blocks[3].self_attention.q_proj,
|
|
)
|
|
|
|
with pytest.raises(AnimaModuleSurfaceError) as captured:
|
|
ANIMA_MODULE_SURFACE_DISCOVERY.discover(installed_anima)
|
|
|
|
message = str(captured.value)
|
|
assert "diffusion_model.blocks.0.self_attn" in message
|
|
assert "diffusion_model.blocks.1.adaln_modulation_cross_attn" in message
|
|
assert "diffusion_model.blocks.2.mlp.layer1" in message
|
|
assert "input=2048, output=8192" in message
|
|
assert "diffusion_model.blocks.3.self_attn.k_proj" in message
|
|
assert "aliases the distinct canonical target" in message
|
|
|
|
|
|
def test_aggregates_every_patch_relevant_forward_signature_drift(
|
|
installed_anima: nn.Module,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Reject forward drift at diffusion, block, attention, MLP, and target layers."""
|
|
|
|
surface = ANIMA_MODULE_SURFACE_DISCOVERY.discover(installed_anima)
|
|
monkeypatch.setattr(
|
|
installed_anima,
|
|
"forward",
|
|
MethodType(_incompatible_forward, installed_anima),
|
|
)
|
|
monkeypatch.setattr(
|
|
surface.blocks[0].block,
|
|
"forward",
|
|
MethodType(_incompatible_forward, surface.blocks[0].block),
|
|
)
|
|
monkeypatch.setattr(
|
|
surface.blocks[1].self_attention,
|
|
"forward",
|
|
MethodType(_incompatible_forward, surface.blocks[1].self_attention),
|
|
)
|
|
monkeypatch.setattr(
|
|
surface.blocks[2].mlp,
|
|
"forward",
|
|
MethodType(_incompatible_forward, surface.blocks[2].mlp),
|
|
)
|
|
monkeypatch.setattr(
|
|
surface.blocks[3].lora_targets[0].module,
|
|
"forward",
|
|
MethodType(_no_input_forward, surface.blocks[3].lora_targets[0].module),
|
|
)
|
|
|
|
with pytest.raises(AnimaModuleSurfaceError) as captured:
|
|
ANIMA_MODULE_SURFACE_DISCOVERY.discover(installed_anima)
|
|
|
|
message = str(captured.value)
|
|
assert "diffusion_model.forward" in message
|
|
assert "diffusion_model.blocks.0.forward" in message
|
|
assert "diffusion_model.blocks.1.self_attn.forward" in message
|
|
assert "diffusion_model.blocks.2.mlp.forward" in message
|
|
assert "diffusion_model.blocks.3.self_attn.q_proj.forward" in message
|
|
|
|
|
|
def _resolve_target(block: Block, family: AnimaLoraTargetFamily) -> nn.Module:
|
|
"""Resolve one expected target independently for descriptor identity checks."""
|
|
|
|
current: object = block
|
|
for part in family.value.split("."):
|
|
if part.isdecimal() and isinstance(current, nn.Sequential):
|
|
current = current[int(part)]
|
|
else:
|
|
current = getattr(current, part)
|
|
if not isinstance(current, nn.Module):
|
|
raise AssertionError("Anima target did not resolve to a module")
|
|
return current
|
|
|
|
|
|
def _incompatible_forward(self: object, unexpected: object) -> object:
|
|
"""Provide an intentionally incompatible bound forward for guard tests."""
|
|
|
|
del self
|
|
return unexpected
|
|
|
|
|
|
def _no_input_forward(self: object) -> None:
|
|
"""Provide an intentionally input-free target forward for guard tests."""
|
|
|
|
del self
|