74 lines
2.3 KiB
Python
74 lines
2.3 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
|
|
|
|
"""Characterize atomic Comfy attn1 replacement installation."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
import pytest
|
|
import torch
|
|
from comfy.model_patcher import ModelPatcher
|
|
from torch import nn
|
|
|
|
from simple_syrup.runtime.model_attention_replacement_mutations import (
|
|
ModelAttn1ReplacementsMutation,
|
|
)
|
|
|
|
|
|
def test_attn1_replacements_install_every_declared_key_atomically() -> None:
|
|
"""Install one callback at each exact standard-UNet self-attention key."""
|
|
|
|
model = _patcher()
|
|
|
|
def replacement(*args: object) -> object:
|
|
"""Return the first callback argument for one inert fixture."""
|
|
|
|
return args[0]
|
|
|
|
ModelAttn1ReplacementsMutation(
|
|
(("input", 4, 0), ("middle", 0, 0)),
|
|
replacement,
|
|
).apply(model)
|
|
|
|
installed = model.model_options["transformer_options"]["patches_replace"]["attn1"]
|
|
assert installed == {
|
|
("input", 4, 0): replacement,
|
|
("middle", 0, 0): replacement,
|
|
}
|
|
|
|
|
|
def test_attn1_replacements_reject_broad_or_exact_collisions_before_mutation() -> None:
|
|
"""Preserve existing replacement ownership without partial installation."""
|
|
|
|
for existing_key in (("input", 4), ("input", 4, 0)):
|
|
model = _patcher()
|
|
|
|
def existing(*args: object) -> object:
|
|
"""Return the first argument through the foreign owner."""
|
|
|
|
return args[0]
|
|
|
|
replacements: dict[str, Any] = {"attn1": {existing_key: existing}}
|
|
model.model_options["transformer_options"]["patches_replace"] = replacements
|
|
before = replacements["attn1"].copy()
|
|
|
|
with pytest.raises(ValueError, match="already installed"):
|
|
ModelAttn1ReplacementsMutation(
|
|
(("input", 4, 0), ("middle", 0, 0)),
|
|
lambda *args: args[0],
|
|
).apply(model)
|
|
|
|
assert replacements["attn1"] == before
|
|
|
|
|
|
def _patcher() -> ModelPatcher:
|
|
"""Return one real CPU patcher with mutable Comfy transformer options."""
|
|
|
|
root = nn.Module()
|
|
root.diffusion_model = nn.Identity()
|
|
device = torch.device("cpu")
|
|
return ModelPatcher(root, load_device=device, offload_device=device)
|