159 lines
5.0 KiB
Python
159 lines
5.0 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
|
|
|
|
"""Verify benchmark-only LoRA execution observation nodes."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from tools.attention_coupling_benchmark.comfy_probe.lora_execution_probe import (
|
|
InstrumentLoraModelV3,
|
|
ReadLoraMetricsV3,
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class _Hook:
|
|
"""Provide mutable schedule strength for transition observation."""
|
|
|
|
hook_ref: str = "hook-ref"
|
|
_strength_model: float = 0.75
|
|
strength: float = 1.0
|
|
|
|
@property
|
|
def strength_model(self) -> float:
|
|
"""Return the effective scheduled model strength."""
|
|
|
|
return self._strength_model * self.strength
|
|
|
|
|
|
@dataclass
|
|
class _HookGroup:
|
|
"""Expose ordered fake hooks."""
|
|
|
|
hooks: list[_Hook]
|
|
|
|
|
|
class _FakeModel:
|
|
"""Provide the ModelPatcher surface observed by the LoRA probe."""
|
|
|
|
def __init__(self) -> None:
|
|
"""Initialize one static patch and one registered hook."""
|
|
|
|
self.model_options: dict[str, Any] = {}
|
|
self.patches = {"diffusion_model.a.weight": [(0.75, object(), 1.0, None, None)]}
|
|
self.hook = _Hook()
|
|
self.current_hooks = _HookGroup([self.hook])
|
|
self.hook_patches = {
|
|
self.hook.hook_ref: {"diffusion_model.a.weight": [object()]}
|
|
}
|
|
self.wrapper: Any = None
|
|
self.callbacks: dict[str, Any] = {}
|
|
|
|
def clone(self) -> _FakeModel:
|
|
"""Return an independent patcher shell sharing the fake hook."""
|
|
|
|
cloned = _FakeModel()
|
|
cloned.patches = self.patches
|
|
cloned.hook = self.hook
|
|
cloned.current_hooks = self.current_hooks
|
|
cloned.hook_patches = self.hook_patches
|
|
return cloned
|
|
|
|
def set_model_unet_function_wrapper(self, wrapper: Any) -> None:
|
|
"""Record the installed observation wrapper."""
|
|
|
|
self.wrapper = wrapper
|
|
|
|
def add_callback_with_key(self, call_type: str, key: str, callback: Any) -> None:
|
|
"""Record one keyed ModelPatcher callback by event type."""
|
|
|
|
del key
|
|
self.callbacks[call_type] = callback
|
|
|
|
def trigger_registered_hooks(self) -> None:
|
|
"""Simulate Comfy registering hooks on the runtime clone."""
|
|
|
|
self.callbacks["on_register_all_hook_patches"](
|
|
self, self.current_hooks, {}, {}, self.current_hooks
|
|
)
|
|
|
|
def trigger_applied_hooks(self) -> None:
|
|
"""Simulate Comfy applying the current hook schedule."""
|
|
|
|
self.callbacks["on_apply_hooks"](self, self.current_hooks)
|
|
|
|
|
|
def test_probe_records_patch_order_schedule_transitions_and_outputs(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Capture exact evidence while returning the delegated tensor unchanged."""
|
|
|
|
monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
|
|
instrumented = InstrumentLoraModelV3.execute(
|
|
_FakeModel(), "lora-run", '["adapter-a"]', True
|
|
)
|
|
(model,) = instrumented.result
|
|
tensor = torch.zeros((1, 2, 2, 2), dtype=torch.float16)
|
|
arguments = {
|
|
"input": tensor,
|
|
"timestep": torch.tensor([1.0]),
|
|
"c": {},
|
|
}
|
|
model.trigger_registered_hooks()
|
|
model.trigger_applied_hooks()
|
|
|
|
def apply_model(
|
|
input_x: torch.Tensor,
|
|
timestep: torch.Tensor,
|
|
**conditioning: object,
|
|
) -> torch.Tensor:
|
|
"""Return a recognizable denoiser output."""
|
|
|
|
del timestep, conditioning
|
|
return input_x + 1
|
|
|
|
assert torch.equal(model.wrapper(apply_model, arguments), tensor + 1)
|
|
model.hook.strength = 0.5
|
|
model.trigger_applied_hooks()
|
|
arguments["timestep"] = torch.tensor([0.5])
|
|
assert torch.equal(model.wrapper(apply_model, arguments), tensor + 1)
|
|
|
|
measured = ReadLoraMetricsV3.execute({"samples": tensor}, "lora-run")
|
|
metrics = measured.ui["lora_benchmark_metrics"][0]
|
|
|
|
assert metrics["model_call_count"] == 2
|
|
assert metrics["static_patches"]["target_count"] == 1
|
|
assert metrics["hook_patches"][0]["identity"] == "adapter-a"
|
|
assert [
|
|
item["effective_strengths"] for item in metrics["schedule_observations"]
|
|
] == [
|
|
[0.75],
|
|
[0.375],
|
|
]
|
|
assert len(metrics["denoiser_outputs"]) == 2
|
|
assert json.loads(measured.result[1])["run_id"] == "lora-run"
|
|
|
|
|
|
def test_probe_schemas_are_dev_only_and_identity_input_fails_closed() -> None:
|
|
"""Expose dedicated node contracts and reject duplicate identities."""
|
|
|
|
instrument = InstrumentLoraModelV3.define_schema()
|
|
metrics = ReadLoraMetricsV3.define_schema()
|
|
assert instrument.node_id == "SimpleSyrupBenchmark.InstrumentLoraModel"
|
|
assert instrument.is_dev_only is True
|
|
assert metrics.node_id == "SimpleSyrupBenchmark.ReadLoraMetrics"
|
|
assert metrics.is_output_node is True
|
|
|
|
with pytest.raises(ValueError, match="unique"):
|
|
InstrumentLoraModelV3.execute(
|
|
_FakeModel(), "bad-run", '["adapter-a", "adapter-a"]', False
|
|
)
|