130 lines
4.3 KiB
Python
130 lines
4.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
|
|
|
|
"""Verify read-only effective model graph resolution through object patches."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import comfy.model_patcher
|
|
import pytest
|
|
import torch
|
|
from torch import nn
|
|
|
|
from simple_syrup.runtime.regional_lora.effective_model_graph import (
|
|
EffectiveModelGraphError,
|
|
EffectiveModelGraphResolver,
|
|
)
|
|
|
|
|
|
def test_resolver_reads_base_graph_without_mutation() -> None:
|
|
"""Return exact nested base objects when no active patch owns the path."""
|
|
|
|
patcher = _patcher()
|
|
layer = patcher.model.diffusion_model.layer
|
|
before = _snapshot(patcher)
|
|
|
|
result = EffectiveModelGraphResolver().resolve(
|
|
patcher,
|
|
"diffusion_model.layer.weight",
|
|
)
|
|
|
|
assert result.value is layer.weight
|
|
assert result.active_object_patch_paths == ()
|
|
assert _snapshot(patcher) == before
|
|
|
|
|
|
def test_resolver_traverses_one_unapplied_parent_object_patch() -> None:
|
|
"""Observe a pending replacement's children without applying the patch."""
|
|
|
|
patcher = _patcher()
|
|
original = patcher.model.diffusion_model.layer
|
|
replacement = nn.Linear(4, 6, bias=False)
|
|
patcher.add_object_patch("diffusion_model.layer", replacement)
|
|
|
|
module = EffectiveModelGraphResolver().resolve(
|
|
patcher,
|
|
"diffusion_model.layer",
|
|
)
|
|
parameter = EffectiveModelGraphResolver().resolve(
|
|
patcher,
|
|
"diffusion_model.layer.weight",
|
|
)
|
|
|
|
assert module.value is replacement
|
|
assert parameter.value is replacement.weight
|
|
assert module.active_object_patch_paths == ("diffusion_model.layer",)
|
|
assert parameter.active_object_patch_paths == ("diffusion_model.layer",)
|
|
assert patcher.model.diffusion_model.layer is original
|
|
|
|
|
|
def test_resolver_uses_exact_parameter_patch() -> None:
|
|
"""Observe exact pending parameter ownership independently of its module."""
|
|
|
|
patcher = _patcher()
|
|
replacement = nn.Parameter(torch.ones((6, 4)), requires_grad=False)
|
|
patcher.add_object_patch("diffusion_model.layer.weight", replacement)
|
|
|
|
result = EffectiveModelGraphResolver().resolve(
|
|
patcher,
|
|
"diffusion_model.layer.weight",
|
|
)
|
|
|
|
assert result.value is replacement
|
|
assert result.active_object_patch_paths == ("diffusion_model.layer.weight",)
|
|
|
|
|
|
def test_resolver_rejects_overlapping_parent_and_child_patches() -> None:
|
|
"""Reject insertion-order-dependent effective graph ownership."""
|
|
|
|
patcher = _patcher()
|
|
patcher.add_object_patch("diffusion_model.layer", nn.Linear(4, 6, bias=False))
|
|
patcher.add_object_patch(
|
|
"diffusion_model.layer.weight",
|
|
nn.Parameter(torch.ones((6, 4)), requires_grad=False),
|
|
)
|
|
|
|
with pytest.raises(EffectiveModelGraphError, match="overlapping object patch"):
|
|
EffectiveModelGraphResolver().resolve(
|
|
patcher,
|
|
"diffusion_model.layer.weight",
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("path", ["", ".layer", "layer.", "layer..weight"])
|
|
def test_resolver_rejects_invalid_or_missing_paths(path: str) -> None:
|
|
"""Fail clearly for malformed and absent effective graph paths."""
|
|
|
|
patcher = _patcher()
|
|
with pytest.raises((ValueError, EffectiveModelGraphError)):
|
|
EffectiveModelGraphResolver().resolve(patcher, path)
|
|
with pytest.raises(EffectiveModelGraphError, match="does not exist"):
|
|
EffectiveModelGraphResolver().resolve(patcher, "diffusion_model.missing")
|
|
|
|
|
|
def _patcher() -> comfy.model_patcher.ModelPatcher:
|
|
"""Return one real Comfy patcher around a minimal nested graph."""
|
|
|
|
model = nn.Module()
|
|
model.diffusion_model = nn.Module()
|
|
model.diffusion_model.layer = nn.Linear(4, 6, bias=False)
|
|
return comfy.model_patcher.ModelPatcher(
|
|
model,
|
|
torch.device("cpu"),
|
|
torch.device("cpu"),
|
|
)
|
|
|
|
|
|
def _snapshot(patcher: comfy.model_patcher.ModelPatcher) -> SimpleNamespace:
|
|
"""Snapshot graph and patch mapping identity without copying tensors."""
|
|
|
|
return SimpleNamespace(
|
|
model=patcher.model,
|
|
layer=patcher.model.diffusion_model.layer,
|
|
weight=patcher.model.diffusion_model.layer.weight,
|
|
object_patches=dict(patcher.object_patches),
|
|
backups=dict(patcher.object_patches_backup),
|
|
)
|