Files
olivv-cs-ComfyUI-FunPack/modules/sampling/attention_temperature/tests/test_attention_temperature.py
T
digital-garbageandClaude Opus 5.5 c3ad0dfd79 feat(v5): randomizer and block repeat modules for H3, on shared chainable block hooks
Block repeat's video-only mode now really keeps audio single-pass: H3 blocks
write in place, so v4 restored audio from an already-repeated stream.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-09-27 17:29:42 +03:00

76 lines
2.9 KiB
Python

"""The randomizer, through the real modifier loader onto a real H3 forward."""
import pytest
import torch
@pytest.fixture(autouse=True)
def _needs_comfy(comfyui):
"""Imports comfy."""
def _load(tiny, **values):
from modules.sampling.modifiers.nodes import FunPackLoadModifiers
settings = {"attention_temperature": {"enabled": True, **values}}
return FunPackLoadModifiers.execute(tiny.patcher, settings).result
def test_it_changes_the_output(tiny_h3):
base = tiny_h3.run()
patched, status = _load(tiny_h3, amount=-0.5, blocks="1-2")
assert "attention_temperature: 0.5x (sharper) at blocks 1-2" in status
out = tiny_h3.run(patched)
assert not torch.allclose(out[0], base[0]), "temperature ran and changed nothing"
def test_blocks_outside_the_list_see_the_original_query(tiny_h3):
from core import dit_hooks
patched, _ = _load(tiny_h3, amount=1.0, blocks="3")
seen = []
def spy(func, q, k, v, *a, **kw):
seen.append(float(q.abs().mean()))
return func(q, k, v, *a, **kw)
# Installed on top, so it sees q BEFORE the temperature override divides it,
# and a second spy underneath sees it after.
under = []
def spy_under(func, q, k, v, *a, **kw):
under.append(float(q.abs().mean()))
return func(q, k, v, *a, **kw)
probe = patched.clone()
dit_hooks.add_attention_override(probe, "test.spy", spy)
tiny_h3.run(probe)
# Rebuild with the spy UNDER the temperature: clone the base, spy first.
low = tiny_h3.patcher.clone()
dit_hooks.add_attention_override(low, "test.spy", spy_under)
from modules.sampling import attention_temperature as m
m.install(low, {"enabled": True, "amount": 1.0, "blocks": "3"}, key="funpack.attention_temperature")
tiny_h3.run(low)
# call order: refiner, blocks 0..3 -> only the last is halved
ratios = [u / s for u, s in zip(under, seen)]
assert [round(r, 4) for r in ratios] == [1.0, 1.0, 1.0, 1.0, 0.5]
def test_off_or_zero_installs_nothing(tiny_h3):
from modules.sampling.modifiers.nodes import FunPackLoadModifiers
_p, status = FunPackLoadModifiers.execute(
tiny_h3.patcher, {"attention_temperature": {"enabled": True, "amount": 0.0}}).result
assert "attention_temperature" not in status.splitlines()[0]
def test_a_block_list_with_nothing_usable_is_refused_not_silently_inert(tiny_h3):
_p, status = _load(tiny_h3, amount=0.5, blocks="40-49")
assert "attention_temperature: failed to install" in status
def test_loading_twice_does_not_double_the_effect(tiny_h3):
once, _ = _load(tiny_h3, amount=-0.5, blocks="1")
from modules.sampling.modifiers.nodes import FunPackLoadModifiers
twice, status = FunPackLoadModifiers.execute(
once, {"attention_temperature": {"enabled": True, "amount": -0.5, "blocks": "1"}}).result
assert torch.allclose(tiny_h3.run(once)[0], tiny_h3.run(twice)[0])