697 lines
31 KiB
Python
697 lines
31 KiB
Python
"""Tests for the SPA ComfyUI node (schema + patching behaviour).
|
|
|
|
The functional (patching) tests exercise ``apply_spa_to_model`` directly rather
|
|
than loading the extension's ``__init__.py``. Loading ``__init__.py`` transitively
|
|
imports the Qwen2D VAE patch / PixelRush / FreeScale nodes, which require the
|
|
*real* ``comfy`` package (``comfy.model_management``, ``comfy.ops``,
|
|
``comfy.ldm...``) that is not available under pytest's mocked ``comfy_api``.
|
|
``apply_spa_to_model`` itself only needs ``comfy_api.latest.io`` (mocked), so it
|
|
is the right unit to verify end-to-end patching for every supported model type.
|
|
The node *wiring* (schema, inputs, registration) is covered separately by the
|
|
schema text-checks below, which read ``__init__.py`` directly.
|
|
"""
|
|
import pathlib
|
|
import re
|
|
import types
|
|
|
|
import pytest
|
|
|
|
_INIT = pathlib.Path(__file__).parent.parent / "__init__.py"
|
|
|
|
|
|
class _MockModel:
|
|
"""Minimal stand-in for comfy.model_patcher.ModelPatcher.
|
|
|
|
Self-contained so the SPA node test does not depend on the (real-comfy
|
|
incompatible) ``mock_flux_model`` conftest fixture.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.model = types.SimpleNamespace()
|
|
self.model.diffusion_model = types.SimpleNamespace()
|
|
self._object_patches = {}
|
|
|
|
def _copy_dm(self, src):
|
|
dst = types.SimpleNamespace()
|
|
for k, v in vars(src).items():
|
|
setattr(dst, k, v)
|
|
return dst
|
|
|
|
def clone(self):
|
|
new = _MockModel()
|
|
new.model.diffusion_model = self._copy_dm(self.model.diffusion_model)
|
|
new._object_patches = dict(self._object_patches)
|
|
return new
|
|
|
|
def add_object_patch(self, path, obj):
|
|
self._object_patches[path] = obj
|
|
|
|
def set_model_unet_function_wrapper(self, fn):
|
|
self._unet_wrapper = fn
|
|
|
|
|
|
def _make_flux_mock():
|
|
m = _MockModel()
|
|
m.model.diffusion_model.pe_embedder = types.SimpleNamespace(
|
|
theta=10000, axes_dim=[16, 56, 56]
|
|
)
|
|
return m
|
|
|
|
|
|
def _make_qwen_mock():
|
|
m = _MockModel()
|
|
m.model.diffusion_model.pe_embedder = types.SimpleNamespace(
|
|
theta=10000, axes_dim=[16, 56, 56]
|
|
)
|
|
return m
|
|
|
|
|
|
def _make_zimage_mock():
|
|
m = _MockModel()
|
|
m.model.diffusion_model.rope_embedder = types.SimpleNamespace(
|
|
theta=10000, axes_dim=[16, 56, 56]
|
|
)
|
|
return m
|
|
|
|
|
|
def _make_nunchaku_mock():
|
|
m = _MockModel()
|
|
m.model.diffusion_model.model = types.SimpleNamespace(
|
|
pos_embed=types.SimpleNamespace(theta=10000, axes_dim=[16, 56, 56])
|
|
)
|
|
return m
|
|
|
|
|
|
def _make_anima_mock():
|
|
m = _MockModel()
|
|
# apply_spa_to_model computes Anima theta/axes_dim from these model attrs
|
|
# (it does not read them off the original embedder).
|
|
dm = m.model.diffusion_model
|
|
dm.model_channels = 1152
|
|
dm.num_heads = 16
|
|
dm.rope_h_extrapolation_ratio = 1.0
|
|
dm.rope_w_extrapolation_ratio = 1.0
|
|
dm.rope_t_extrapolation_ratio = 1.0
|
|
dm.pos_embedder = types.SimpleNamespace(
|
|
theta=[10000.0, 10000.0, 10000.0], axes_dim=[44, 42, 42]
|
|
)
|
|
return m
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestSpaNodeSchema:
|
|
def test_node_class_defined(self):
|
|
assert "class SPA" in _INIT.read_text(encoding="utf-8")
|
|
|
|
def test_node_registered(self):
|
|
content = _INIT.read_text(encoding="utf-8")
|
|
assert "return [DyPE_FLUX, SEGA, SPA" in content
|
|
|
|
def test_imports_apply_spa(self):
|
|
assert "apply_spa_to_model" in _INIT.read_text(encoding="utf-8")
|
|
|
|
def test_filter_error_scope_is_narrow(self):
|
|
"""The 'invalid spa_layer_filter' prefix must wrap ONLY filter-parse
|
|
failures — not every ValueError from apply_spa_to_model (2026-08-24:
|
|
the mutual-exclusion guard surfaced as 'invalid spa_layer_filter'
|
|
because the node re-wrapped all ValueErrors)."""
|
|
content = _INIT.read_text(encoding="utf-8")
|
|
start = content.index("class SPA(io.ComfyNode):")
|
|
end = content.index("class DyPEExtension")
|
|
spa_section = content[start:end]
|
|
# Pre-parses the filter in its own try/except.
|
|
assert "parse_layer_filter(spa_layer_filter)" in spa_section
|
|
assert "invalid spa_layer_filter" in spa_section
|
|
# The apply call itself is NOT inside a blanket try/except ValueError
|
|
# that would re-wrap unrelated errors with the filter message.
|
|
assert "except ValueError as exc:\n raise type(exc)" not in spa_section
|
|
|
|
def test_has_bundle_size_input(self):
|
|
content = _INIT.read_text(encoding="utf-8")
|
|
assert "bundle_size" in content
|
|
|
|
def test_has_enable_spa_input(self):
|
|
assert "enable_spa" in _INIT.read_text(encoding="utf-8")
|
|
|
|
def test_has_model_type_combo(self):
|
|
content = _INIT.read_text(encoding="utf-8")
|
|
for opt in ('"auto"', '"flux"', '"qwen"', '"anima"'):
|
|
assert opt in content
|
|
|
|
def test_category_and_output(self):
|
|
content = _INIT.read_text(encoding="utf-8")
|
|
assert "model_patches/position_encoding" in content
|
|
assert "io.Model.Output" in content
|
|
|
|
def test_has_spa_start_sigma_input(self):
|
|
"""T2.1: the node schema exposes the step-gating ``spa_start_sigma`` input."""
|
|
content = _INIT.read_text(encoding="utf-8")
|
|
assert '"spa_start_sigma"' in content
|
|
assert "io.Float.Input" in content
|
|
|
|
def test_no_method_input_on_spa_node(self):
|
|
"""Guard: the SPA node must NOT expose a ``method`` combo.
|
|
|
|
SPA always applies the model's native no-extrapolation RoPE
|
|
(``ntk_factor=1.0``) on the bundled coords, so the DyPE extrapolation
|
|
methods (ntk/yarn/vision_yarn/pi) are a no-op for SPA. The knob was
|
|
removed to stop users A/B-testing a dead input. Scoped to the SPA
|
|
class section only (the DyPE node legitimately keeps its own
|
|
``method`` input).
|
|
"""
|
|
content = _INIT.read_text(encoding="utf-8")
|
|
start = content.index("class SPA(io.ComfyNode):")
|
|
end = content.index("class DyPEExtension")
|
|
spa_section = content[start:end]
|
|
assert '"method"' not in spa_section, (
|
|
"SPA node still exposes a 'method' input — it is a no-op for SPA "
|
|
"and must stay removed")
|
|
# execute() must not accept a method parameter either.
|
|
assert "def execute(cls, model, width: int, height: int, model_type: str, enable_spa: bool" in spa_section
|
|
assert "method: str" not in spa_section
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestSpaNodePatching:
|
|
"""Functional tests that drive ``apply_spa_to_model`` end-to-end.
|
|
|
|
These mirror exactly what ``SPA.execute`` does (the node body is thin glue:
|
|
``bs = None if bundle_size<=0 else int(bundle_size)`` then
|
|
``apply_spa_to_model(model, model_type, width, height,
|
|
enable_spa=enable_spa, bundle_size=bs)``). The node no longer exposes a
|
|
``method`` input (SPA always uses the model's native no-extrapolation
|
|
RoPE); the tests below pass ``"ntk"`` positionally only because
|
|
``apply_spa_to_model`` keeps the parameter for the DyPE-base constructor
|
|
chain — it is a no-op for the SPA math.
|
|
"""
|
|
|
|
def test_patches_flux_with_auto_bundle_size(self):
|
|
from src.models.spa_flux import PosEmbedSPAFlux
|
|
from src.spa import apply_spa_to_model
|
|
|
|
out = apply_spa_to_model(_make_flux_mock(), "flux", 4096, 4096, "ntk", enable_spa=True)
|
|
embedder = out._object_patches["diffusion_model.pe_embedder"]
|
|
assert isinstance(embedder, PosEmbedSPAFlux)
|
|
# auto -> bundle_size == 0 (minimal in-distribution compression), paper-N
|
|
assert embedder.bundle_size == 0
|
|
assert embedder.enable_spa is True
|
|
|
|
def test_auto_default_is_zero(self):
|
|
from src.models.spa_flux import PosEmbedSPAFlux
|
|
from src.spa import apply_spa_to_model
|
|
|
|
# SPA is always active when enabled (auto=0 regardless of resolution);
|
|
# only an explicit bundle_size==1 yields a no-op passthrough.
|
|
out = apply_spa_to_model(_make_flux_mock(), "flux", 1024, 1024, "ntk", enable_spa=True)
|
|
embedder = out._object_patches["diffusion_model.pe_embedder"]
|
|
assert isinstance(embedder, PosEmbedSPAFlux)
|
|
assert embedder.bundle_size == 0
|
|
|
|
def test_explicit_n_stored_verbatim_and_pass_count(self):
|
|
import torch
|
|
|
|
from src.models.spa_flux import PosEmbedSPAFlux
|
|
from src.spa import (
|
|
SPA_MAX_PASSES,
|
|
apply_spa_to_model,
|
|
build_bundle_id_variants,
|
|
derive_bundle_s,
|
|
)
|
|
|
|
def _latent_ids(H, W):
|
|
ids = torch.zeros(1, H * W, 3)
|
|
ids[..., 1] = torch.arange(H).unsqueeze(1).expand(H, W).reshape(-1).float()
|
|
ids[..., 2] = torch.arange(W).unsqueeze(0).expand(H, W).reshape(-1).float()
|
|
return ids
|
|
|
|
# Paper-N semantics: an explicit knob (2..8) is stored verbatim.
|
|
out = apply_spa_to_model(_make_flux_mock(), "flux", 1024, 1024, "ntk", enable_spa=True, bundle_size=3)
|
|
embedder = out._object_patches["diffusion_model.pe_embedder"]
|
|
assert isinstance(embedder, PosEmbedSPAFlux)
|
|
assert embedder.bundle_size == 3 # stored verbatim (paper's N)
|
|
|
|
# derive_bundle_s is the single source of truth. At a 128x128 grid
|
|
# (max_pos=127 > trained_extent=64): auto -> s=ceil(127/79)=2 (3 passes);
|
|
# N=3 -> s=max(3,2)=3 (5 passes); N=5 -> s=5 (9 passes). All in-dist.
|
|
ids = _latent_ids(128, 128)
|
|
assert derive_bundle_s(127, 0) == 2
|
|
assert derive_bundle_s(127, 3) == 3
|
|
assert derive_bundle_s(127, 5) == 5
|
|
assert len(build_bundle_id_variants(ids, 0)) == 3
|
|
assert len(build_bundle_id_variants(ids, 3)) == 5
|
|
assert len(build_bundle_id_variants(ids, 5)) == 9
|
|
# Pass count never exceeds the cap.
|
|
for N in (0, 2, 3, 5, 8):
|
|
assert len(build_bundle_id_variants(ids, N)) <= SPA_MAX_PASSES
|
|
|
|
def test_legacy_group_num_values_migrate_to_auto(self, caplog):
|
|
"""Decision M1: legacy knob values >= 32 (old group_num semantics) are
|
|
migrated to auto (0) with a one-time WARNING, not stored verbatim."""
|
|
import logging
|
|
|
|
from src.models.spa_flux import PosEmbedSPAFlux
|
|
from src.spa import apply_spa_to_model
|
|
|
|
# Reset the one-time warning latch so caplog captures it deterministically.
|
|
apply_spa_to_model._spa_legacy_warned = False
|
|
with caplog.at_level(logging.WARNING, logger="ComfyUI-DyPE"):
|
|
out = apply_spa_to_model(_make_flux_mock(), "flux", 1024, 1024, "ntk",
|
|
enable_spa=True, bundle_size=80)
|
|
embedder = out._object_patches["diffusion_model.pe_embedder"]
|
|
assert isinstance(embedder, PosEmbedSPAFlux)
|
|
# Legacy 80 -> migrated to auto (0), NOT stored as 80.
|
|
assert embedder.bundle_size == 0
|
|
assert any("legacy" in r.message.lower() for r in caplog.records)
|
|
|
|
def test_explicit_bundle_size_one_is_off(self):
|
|
from src.models.spa_flux import PosEmbedSPAFlux
|
|
from src.spa import apply_spa_to_model
|
|
|
|
# bundle_size == 1 is the explicit OFF knob (true passthrough) and must NOT
|
|
# be clamped (the cost guard only applies to active, non-1 values).
|
|
out = apply_spa_to_model(_make_flux_mock(), "flux", 1024, 1024, "ntk", enable_spa=True, bundle_size=1)
|
|
embedder = out._object_patches["diffusion_model.pe_embedder"]
|
|
assert isinstance(embedder, PosEmbedSPAFlux)
|
|
assert embedder.bundle_size == 1
|
|
|
|
def test_disable_spa(self):
|
|
from src.models.spa_flux import PosEmbedSPAFlux
|
|
from src.spa import apply_spa_to_model
|
|
|
|
out = apply_spa_to_model(_make_flux_mock(), "flux", 4096, 4096, "ntk", enable_spa=False)
|
|
embedder = out._object_patches["diffusion_model.pe_embedder"]
|
|
assert isinstance(embedder, PosEmbedSPAFlux)
|
|
assert embedder.enable_spa is False
|
|
|
|
def test_spa_start_sigma_stored_on_patcher(self):
|
|
"""T2.1: ``spa_start_sigma`` flows from apply_spa_to_model onto the patcher."""
|
|
from src.spa import apply_spa_to_model
|
|
|
|
# Default: 1.0 (always active, backward compatible).
|
|
out = apply_spa_to_model(_make_flux_mock(), "flux", 4096, 4096, "ntk", enable_spa=True)
|
|
assert getattr(out, "_spa_start_sigma", None) == 1.0
|
|
|
|
# Explicit threshold is stored verbatim for the unet wrapper to read.
|
|
out2 = apply_spa_to_model(_make_flux_mock(), "flux", 4096, 4096, "ntk",
|
|
enable_spa=True, spa_start_sigma=0.7)
|
|
assert out2._spa_start_sigma == 0.7
|
|
|
|
def test_spa_steps_param(self):
|
|
"""T2.1: ``spa_steps`` (leading-step count gate) flows onto the patcher.
|
|
|
|
Default is 3 (HRDiT ``--spa_steps [3, 0]``); 0 = all steps (backward
|
|
compat); the step counter / last-sigma slots are initialized so the
|
|
unet wrapper's generation-boundary detection starts clean.
|
|
"""
|
|
from src.spa import apply_spa_to_model
|
|
|
|
# Default: 3 leading steps (HRDiT-faithful speed/quality tradeoff).
|
|
out = apply_spa_to_model(_make_flux_mock(), "flux", 4096, 4096, "ntk", enable_spa=True)
|
|
assert getattr(out, "_spa_steps", None) == 3
|
|
assert getattr(out, "_spa_step_counter", None) == 0
|
|
assert getattr(out, "_spa_last_sigma", "sentinel") is None
|
|
|
|
# Explicit value is stored verbatim.
|
|
out5 = apply_spa_to_model(_make_flux_mock(), "flux", 4096, 4096, "ntk",
|
|
enable_spa=True, spa_steps=5)
|
|
assert out5._spa_steps == 5
|
|
|
|
# 0 = all steps (backward compatible).
|
|
out0 = apply_spa_to_model(_make_flux_mock(), "flux", 4096, 4096, "ntk",
|
|
enable_spa=True, spa_steps=0)
|
|
assert out0._spa_steps == 0
|
|
|
|
# Negative values are clamped to 0 (all steps), never a broken gate.
|
|
outneg = apply_spa_to_model(_make_flux_mock(), "flux", 4096, 4096, "ntk",
|
|
enable_spa=True, spa_steps=-2)
|
|
assert outneg._spa_steps == 0
|
|
|
|
def test_qwen_and_zimage_patched(self):
|
|
from src.models.spa_qwen import PosEmbedSPAQwen
|
|
from src.models.spa_zimage import PosEmbedSPAZImage
|
|
from src.spa import apply_spa_to_model
|
|
|
|
out_q = apply_spa_to_model(_make_qwen_mock(), "qwen", 4096, 4096, "ntk", enable_spa=True)
|
|
assert isinstance(out_q._object_patches["diffusion_model.pe_embedder"], PosEmbedSPAQwen)
|
|
|
|
out_z = apply_spa_to_model(_make_zimage_mock(), "zimage", 4096, 4096, "ntk", enable_spa=True)
|
|
assert isinstance(out_z._object_patches["diffusion_model.rope_embedder"], PosEmbedSPAZImage)
|
|
|
|
def test_nunchaku_and_anima_patched(self):
|
|
from src.models.spa_anima import PosEmbedSPAAnima
|
|
from src.spa import apply_spa_to_model
|
|
|
|
# Nunchaku is unsupported (decision 4): SPA returns the model UNCHANGED and
|
|
# does NOT install an embedder patch for it.
|
|
out_n = apply_spa_to_model(_make_nunchaku_mock(), "nunchaku", 4096, 4096, "ntk", enable_spa=True)
|
|
assert "diffusion_model.model.pos_embed" not in out_n._object_patches
|
|
|
|
out_a = apply_spa_to_model(_make_anima_mock(), "anima", 4096, 4096, "ntk", enable_spa=True)
|
|
embedder = out_a._object_patches["diffusion_model.pos_embedder"]
|
|
assert isinstance(embedder, PosEmbedSPAAnima)
|
|
# Anima: axes_dim derived from model_channels/num_heads.
|
|
# head_dim = 1152//16 = 72; dim_h = (72//6)*2 = 24; dim_t = 72-2*24 = 24
|
|
assert embedder.axes_dim == [24, 24, 24]
|
|
assert embedder.thetas == [10000.0, 10000.0, 10000.0]
|
|
|
|
def test_auto_detects_flux_via_pe_embedder(self):
|
|
"""auto with a FLUX-like diffusion_model (has pe_embedder) selects flux."""
|
|
from src.models.spa_flux import PosEmbedSPAFlux
|
|
from src.spa import apply_spa_to_model
|
|
|
|
out = apply_spa_to_model(_make_flux_mock(), "auto", 4096, 4096, "ntk", enable_spa=True)
|
|
assert isinstance(out._object_patches["diffusion_model.pe_embedder"], PosEmbedSPAFlux)
|
|
assert out._object_patches["diffusion_model.pe_embedder"].bundle_size == 0
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestSpaDoubleApplyWarning:
|
|
"""W9.g (NTH-108): a GENUINE SPA re-apply warns; the HAP chain does not."""
|
|
|
|
def test_double_apply_warns(self, caplog):
|
|
"""Applying SPA onto a patcher that ALREADY has the hook AND an SPA
|
|
embedder emits the double-application WARNING."""
|
|
import logging
|
|
|
|
from src.models.spa_flux import PosEmbedSPAFlux
|
|
from src.spa import apply_spa_to_model
|
|
|
|
apply_spa_to_model( # first apply establishes the baseline state
|
|
_make_flux_mock(), "flux", 2048, 2048, "ntk",
|
|
enable_spa=True, bundle_size=3,
|
|
)
|
|
# Simulate a genuine re-apply: the SOURCE patcher carries BOTH the
|
|
# hook flag AND an SPA embedder AS THE LIVE ATTRIBUTE (the mock clone
|
|
# copies dm attributes, mirroring how get_model_object resolves the
|
|
# patched embedder in real ComfyUI).
|
|
second_src = _make_flux_mock()
|
|
second_src._spa_installed = True
|
|
second_src.model.diffusion_model.pe_embedder = PosEmbedSPAFlux(
|
|
theta=10000, axes_dim=[16, 56, 56], method="ntk",
|
|
base_resolution=1024, enable_spa=True, bundle_size=3,
|
|
)
|
|
with caplog.at_level(logging.WARNING, logger="ComfyUI-DyPE"):
|
|
apply_spa_to_model(
|
|
second_src, "flux", 2048, 2048, "ntk",
|
|
enable_spa=True, bundle_size=3,
|
|
)
|
|
assert any(
|
|
"applied twice" in r.message for r in caplog.records
|
|
), "expected the double-application WARNING"
|
|
|
|
def test_fresh_apply_no_warning(self, caplog):
|
|
"""A fresh apply (no prior hook) must NOT warn."""
|
|
import logging
|
|
|
|
from src.spa import apply_spa_to_model
|
|
|
|
with caplog.at_level(logging.WARNING, logger="ComfyUI-DyPE"):
|
|
apply_spa_to_model(
|
|
_make_flux_mock(), "flux", 2048, 2048, "ntk",
|
|
enable_spa=True, bundle_size=3,
|
|
)
|
|
assert not any("applied twice" in r.message for r in caplog.records)
|
|
|
|
def test_hap_chain_no_false_positive(self, caplog):
|
|
"""SPA-after-HAP: the source patcher HAS _spa_installed but its
|
|
embedder is NOT an SPA embedder -> no warning."""
|
|
import logging
|
|
|
|
from src.spa import apply_spa_to_model
|
|
|
|
src_patcher = _make_flux_mock()
|
|
src_patcher._spa_installed = True # HAP installed the shared hook
|
|
with caplog.at_level(logging.WARNING, logger="ComfyUI-DyPE"):
|
|
apply_spa_to_model(
|
|
src_patcher, "flux", 2048, 2048, "ntk",
|
|
enable_spa=True, bundle_size=3,
|
|
)
|
|
assert not any("applied twice" in r.message for r in caplog.records), (
|
|
"false positive: the legitimate HAP->SPA chain was flagged"
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# P3 — install policy / lifecycle (T-P3-3, T-P3-4)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
@pytest.mark.unit
|
|
class TestSpaInstallPolicy:
|
|
def test_nunchaku_guard_warns_and_skips_patch(self, caplog):
|
|
"""T-P3-3: SPA is unsupported on Nunchaku; log a warning and return m UNCHANGED.
|
|
|
|
No object patch is applied and the module-level ``optimized_attention`` hook
|
|
is never installed.
|
|
"""
|
|
import logging
|
|
|
|
import comfy.ldm.modules.attention as attn_mod
|
|
|
|
from src.spa import apply_spa_to_model
|
|
|
|
orig_attn = attn_mod.optimized_attention
|
|
with caplog.at_level(logging.WARNING, logger="ComfyUI-DyPE"):
|
|
out = apply_spa_to_model(
|
|
_make_nunchaku_mock(), "nunchaku", 4096, 4096, "ntk", enable_spa=True
|
|
)
|
|
# The model is returned unchanged: no embedder patch added.
|
|
assert "diffusion_model.model.pos_embed" not in out._object_patches
|
|
# The optimized_attention hook was NOT installed.
|
|
assert attn_mod.optimized_attention is orig_attn
|
|
# A clear warning was emitted mentioning Nunchaku.
|
|
assert any("Nunchaku" in r.message for r in caplog.records)
|
|
|
|
def test_bundle_size_one_installs_no_hook(self):
|
|
"""T-P3-4: bundle_size==1 -> no attention hook; module-level optimized_attention untouched."""
|
|
import comfy.ldm.modules.attention as attn_mod
|
|
|
|
from src.models.spa_flux import PosEmbedSPAFlux
|
|
from src.spa import apply_spa_to_model
|
|
|
|
orig_attn = attn_mod.optimized_attention
|
|
out = apply_spa_to_model(
|
|
_make_flux_mock(), "flux", 4096, 4096, "ntk", enable_spa=True, bundle_size=1
|
|
)
|
|
# Embedder IS replaced (base RoPE passthrough), but the hook must not be installed.
|
|
assert "diffusion_model.pe_embedder" in out._object_patches
|
|
assert isinstance(out._object_patches["diffusion_model.pe_embedder"], PosEmbedSPAFlux)
|
|
assert out._object_patches["diffusion_model.pe_embedder"].bundle_size == 1
|
|
# The module-level optimized_attention remains the original (untouched).
|
|
assert attn_mod.optimized_attention is orig_attn
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# P5 — composition / mutual exclusivity (T-P5-1, T-P5-2, T-P5-3, T-P5-4)
|
|
# ---------------------------------------------------------------------------
|
|
#
|
|
# Remediation decision 6: SPA and DyPE/SEGA are mutually exclusive in v1. The
|
|
# guards live in two places and must BOTH reject:
|
|
# * ``src.spa._spa_ensure_no_incompatible_embedder`` — rejects a DyPE/SEGA
|
|
# embedder when SPA is being applied.
|
|
# * ``src.patch_utils._dype_sega_reject_spa`` — rejects an SPA embedder when
|
|
# DyPE/SEGA is being applied.
|
|
#
|
|
# The mock's ``add_object_patch`` only records into ``_object_patches`` (it does
|
|
# NOT mutate the live embedder attribute), so to exercise the guard we install
|
|
# the *incompatible* embedder instance directly onto the live
|
|
# ``m.model.diffusion_model.pe_embedder`` — exactly what the guards read.
|
|
|
|
_SRC = pathlib.Path(__file__).parent.parent / "src"
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestSpaComposition:
|
|
# --- SPA rejects DyPE / SEGA (apply_spa_to_model on an incompatible model) ---
|
|
def test_spa_rejects_dype_embedder(self):
|
|
"""T-P5-1: applying SPA onto a DyPE-patched FLUX model raises ValueError."""
|
|
from src.models.flux import PosEmbedFlux
|
|
from src.spa import apply_spa_to_model
|
|
|
|
m = _make_flux_mock()
|
|
# Simulate a model already patched by DyPE (live embedder, what the guard reads).
|
|
m.model.diffusion_model.pe_embedder = PosEmbedFlux(10000, [16, 56, 56])
|
|
with pytest.raises(ValueError, match="mutually exclusive"):
|
|
apply_spa_to_model(m, "flux", 4096, 4096, "ntk", enable_spa=True)
|
|
|
|
def test_spa_rejects_sega_embedder(self):
|
|
"""T-P5-4: applying SPA onto a SEGA-patched FLUX model raises ValueError."""
|
|
from src.models.sega_flux import SegAPosEmbedFlux
|
|
from src.spa import apply_spa_to_model
|
|
|
|
m = _make_flux_mock()
|
|
# Simulate a model already patched by SEGA.
|
|
m.model.diffusion_model.pe_embedder = SegAPosEmbedFlux(10000, [16, 56, 56])
|
|
with pytest.raises(ValueError, match="mutually exclusive"):
|
|
apply_spa_to_model(m, "flux", 4096, 4096, "ntk", enable_spa=True)
|
|
|
|
# --- DyPE / SEGA reject SPA (apply_dype/apply_sega on an SPA-patched model) ---
|
|
def test_dype_rejects_spa_embedder(self):
|
|
"""T-P5-2: applying DyPE onto an SPA-patched FLUX model raises ValueError."""
|
|
from src.models.spa_flux import PosEmbedSPAFlux
|
|
from src.patch_utils import apply_dype_to_model
|
|
|
|
m = _make_flux_mock()
|
|
# Simulate a model already patched by SPA (live embedder, what the guard reads).
|
|
m.model.diffusion_model.pe_embedder = PosEmbedSPAFlux(10000, [16, 56, 56])
|
|
with pytest.raises(ValueError, match="mutually exclusive"):
|
|
apply_dype_to_model(m, "flux", 4096, 4096, "ntk", False, False, 2.0, 2.0, 0.5, 1.15)
|
|
|
|
def test_sega_rejects_spa_embedder(self):
|
|
"""T-P5-4: applying SEGA onto an SPA-patched FLUX model raises ValueError."""
|
|
from src.models.spa_flux import PosEmbedSPAFlux
|
|
from src.patch_utils import apply_sega_to_model
|
|
|
|
m = _make_flux_mock()
|
|
m.model.diffusion_model.pe_embedder = PosEmbedSPAFlux(10000, [16, 56, 56])
|
|
with pytest.raises(ValueError, match="mutually exclusive"):
|
|
apply_sega_to_model(m, "flux", 4096, 4096)
|
|
|
|
# --- Improved guard message + node error scoping (2026-08-24) ------------
|
|
|
|
def test_guard_message_names_embedder_class_and_recovery(self):
|
|
"""The SEGA->SPA guard names the offending embedder class, the method
|
|
kind (SEGA/DyPE), and the recovery path (reload)."""
|
|
from src.models.sega_anima import SegAPosEmbedAnima
|
|
from src.spa import _spa_ensure_no_incompatible_embedder
|
|
|
|
embedder = SegAPosEmbedAnima(theta=[10000.0, 10000.0, 10000.0],
|
|
axes_dim=[16, 56, 56])
|
|
with pytest.raises(ValueError) as exc_info:
|
|
_spa_ensure_no_incompatible_embedder(embedder)
|
|
msg = str(exc_info.value)
|
|
assert "SegAPosEmbedAnima" in msg # embedder class named
|
|
assert "SEGA" in msg # method kind named
|
|
assert "reload" in msg.lower() # recovery path given
|
|
assert "mutually exclusive" in msg # original semantics kept
|
|
|
|
def test_guard_message_dype_kind(self):
|
|
"""A DyPE (non-SEGA) embedder reports kind 'DyPE'."""
|
|
from src.models.flux import PosEmbedFlux
|
|
from src.spa import _spa_ensure_no_incompatible_embedder
|
|
|
|
embedder = PosEmbedFlux(10000, [16, 56, 56])
|
|
with pytest.raises(ValueError) as exc_info:
|
|
_spa_ensure_no_incompatible_embedder(embedder)
|
|
msg = str(exc_info.value)
|
|
assert "DyPE" in msg
|
|
assert "PosEmbedFlux" in msg
|
|
|
|
def test_parse_layer_filter_idempotent_on_frozenset(self):
|
|
"""parse_layer_filter passes an already-parsed frozenset through
|
|
unchanged (the node pre-parses; apply_spa_to_model re-parses)."""
|
|
from src.spa import parse_layer_filter
|
|
|
|
fs = frozenset({0, 1, 5})
|
|
assert parse_layer_filter(fs) is fs
|
|
assert parse_layer_filter(None) is None
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestSpaNoLegacyAveraging:
|
|
def test_no_legacy_rope_matrix_averaging(self):
|
|
"""T-P5-3: the root-cause bug (average RoPE *matrices* then one softmax) is gone.
|
|
|
|
The fix runs N attention passes and averages the *attention outputs*
|
|
(``torch.stack(outs, dim=0).mean(dim=0)`` in ``spa_attn.py``). Assert that:
|
|
* ``spa.py`` never averages any tensor (the embedder/base path is matrix-free),
|
|
* the only stacked-then-meaned object in the attention path is the *outputs*
|
|
(``outs``), never the RoPE embeddings (``embs``) — the legacy pattern is
|
|
absent from the code that performs attention,
|
|
* every SPA model adapter explicitly documents that the legacy embedding-
|
|
averaging path was removed (a positive signal, not executable code).
|
|
"""
|
|
spa_src = _SRC.joinpath("spa.py").read_text(encoding="utf-8")
|
|
# The base embedder path must never reduce tensors via .mean(...)
|
|
assert ".mean(" not in spa_src, "spa.py must not average any tensor (RoPE math)"
|
|
|
|
attn_src = _SRC.joinpath("spa_attn.py").read_text(encoding="utf-8")
|
|
# Faithful mechanism: average ATTENTION OUTPUTS.
|
|
assert "torch.stack(outs, dim=0).mean(dim=0)" in attn_src
|
|
# Legacy bug: average RoPE rotation MATRICES before attention — must be absent
|
|
# from the module that actually performs the averaging.
|
|
assert "torch.stack(embs" not in attn_src
|
|
|
|
# Every SPA model adapter must document that the legacy embedding-averaging
|
|
# path is removed (the only `torch.stack(embs` mentions live in such docstrings).
|
|
documented_removal = False
|
|
for f in _SRC.glob("models/spa_*.py"):
|
|
src = f.read_text(encoding="utf-8")
|
|
if "torch.stack(embs" in src:
|
|
# It must be framed as "the legacy ... path is removed", not live code.
|
|
assert "removed" in src, f"{f.name}: legacy embedding averaging not documented as removed"
|
|
documented_removal = True
|
|
assert documented_removal, "expected at least one SPA adapter to document legacy-path removal"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# P6 — documentation assertion (T-P6-1)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
@pytest.mark.unit
|
|
class TestSpaDocs:
|
|
def test_readme_documents_faithful_mechanism_and_constraints(self):
|
|
"""T-P6-1: README documents the output-averaging mechanism, mutual exclusivity,
|
|
and Nunchaku non-support (the three corrections made during remediation)."""
|
|
readme = pathlib.Path(__file__).parent.parent / "README.md"
|
|
content = readme.read_text(encoding="utf-8")
|
|
# Faithful mechanism: average attention OUTPUTS (not RoPE rotation matrices).
|
|
assert "averages the attention outputs" in content
|
|
# SPA and DyPE/SEGA are mutually exclusive in v1.
|
|
assert "mutually exclusive" in content
|
|
# Nunchaku is explicitly unsupported for SPA.
|
|
assert re.search(r"Nunchaku (is not supported|not supported)", content), (
|
|
"README must state that Nunchaku is not supported for SPA")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# P4 (2026-08-16 fix) — SPA-side carry-over of HRDiT state across clone()
|
|
# ---------------------------------------------------------------------------
|
|
#
|
|
# The real ModelPatcher.clone() drops custom attributes. apply_spa_to_model
|
|
# clones the incoming patcher, so any HRDiT state the source already carries
|
|
# (e.g. an HAP runtime installed by an upstream HAP node) must be carried onto
|
|
# the clone — otherwise HAP silently dies the moment an SPA node is chained
|
|
# after it. The HAP-side mirror of this test lives in test_hap_node.py
|
|
# (TestChainOrderIndependence); this class covers the SPA apply function's
|
|
# perspective.
|
|
|
|
@pytest.mark.unit
|
|
class TestSpaCarriesHrditState:
|
|
def test_spa_after_hap_carries_hap_ctx(self):
|
|
"""apply_spa_to_model carries an existing _hap_ctx across its clone()."""
|
|
from src.spa import apply_spa_to_model
|
|
|
|
m = _make_flux_mock()
|
|
sentinel = object()
|
|
m._hap_ctx = sentinel
|
|
|
|
out = apply_spa_to_model(m, "flux", 4096, 4096, "ntk", enable_spa=True)
|
|
assert out is not m
|
|
assert getattr(out, "_hap_ctx", None) is sentinel
|
|
|
|
def test_spa_carries_hrdit_consumers_and_state_ref(self):
|
|
"""apply_spa_to_model carries _hrdit_consumers and re-points _hrdit_state_ref."""
|
|
from src.spa import apply_spa_to_model
|
|
|
|
m = _make_flux_mock()
|
|
m._hrdit_consumers = {"spa"}
|
|
m._hrdit_state_ref = [m]
|
|
|
|
out = apply_spa_to_model(m, "flux", 4096, 4096, "ntk", enable_spa=True)
|
|
assert getattr(out, "_hrdit_consumers", None) == {"spa"}
|
|
# The shared state ref must now point at the newest clone so the
|
|
# already-installed unet wrapper reads the authoritative state.
|
|
assert getattr(out, "_hrdit_state_ref", [None])[0] is out
|
|
|
|
def test_spa_carry_noop_on_bare_model(self):
|
|
"""Carry-over must not raise when the source has no HRDiT attributes."""
|
|
from src.spa import apply_spa_to_model
|
|
|
|
out = apply_spa_to_model(
|
|
_make_flux_mock(), "flux", 4096, 4096, "ntk", enable_spa=True
|
|
)
|
|
# No HRDiT attrs were present on the source -> none invented on the clone
|
|
# (the install hook may legitimately add its own attrs afterwards).
|
|
assert not hasattr(out, "_hap_ctx")
|
|
|