digital-garbageandClaude Opus 5.5
78e5459422
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
64 lines
2.6 KiB
Python
64 lines
2.6 KiB
Python
"""Taste attention on a real H3 forward, and the rotation it relies on."""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _needs_comfy(comfyui):
|
|
"""Imports comfy."""
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def root(tmp_path, monkeypatch):
|
|
from modules.system.taste import store
|
|
monkeypatch.setattr(store, "ROOT", tmp_path / "taste")
|
|
|
|
|
|
def _load(tiny, **values):
|
|
from modules.sampling.modifiers.nodes import FunPackLoadModifiers
|
|
return FunPackLoadModifiers.execute(
|
|
tiny.patcher, {"taste": {"key": "fox"}, "q_steer": {"enabled": True, **values}}).result
|
|
|
|
|
|
def test_the_rotation_matches_the_models_own_and_undoes_exactly():
|
|
from comfy_kitchen.backends.eager.rope import apply_rope_split_half1
|
|
from comfy.ldm.minimax.model import rope_rotation_table
|
|
from modules.models.minimax_h3.rope import query_rotation
|
|
seq, heads, dim, rot = 6, 2, 8, 6
|
|
table = rope_rotation_table(torch.randn(seq, rot), torch.float32)
|
|
x = torch.randn(heads, seq, dim)
|
|
rotate = query_rotation(table, seq)
|
|
ours = rotate(x, slice(0, seq))
|
|
theirs = torch.cat([apply_rope_split_half1(x[..., :rot].unsqueeze(0).transpose(1, 2),
|
|
table).transpose(1, 2)[0], x[..., rot:]], -1)
|
|
assert torch.allclose(ours, theirs, atol=1e-5)
|
|
assert torch.allclose(rotate(ours, slice(0, seq), inverse=True), x, atol=1e-5)
|
|
assert query_rotation(table, seq + 1) is None
|
|
|
|
|
|
def test_captures_every_block_on_the_last_step(tiny_h3):
|
|
from modules.system.taste import store
|
|
patched, status = _load(tiny_h3, blocks="3")
|
|
assert "q_steer: learning at every block" in status
|
|
tiny_h3.sample(patched)
|
|
pending = torch.load(store.ROOT / "fox" / "q_steer.pending.pt")["rows"]
|
|
attn = tiny_h3.patcher.model.diffusion_model.blocks[0].attn
|
|
assert sorted(pending) == ["0", "1", "2", "3"]
|
|
assert pending["2"].shape == (attn.heads * attn.head_dim,)
|
|
|
|
|
|
def test_a_learned_direction_moves_the_picture_and_never_the_sound(tiny_h3):
|
|
from modules.system.taste import store
|
|
attn = tiny_h3.patcher.model.diffusion_model.blocks[0].attn
|
|
liked = torch.zeros(attn.heads * attn.head_dim)
|
|
liked[::attn.head_dim] = 1.0
|
|
for i, (vec, rating) in enumerate([(liked, "liked"), (-liked, "disliked")] * 2):
|
|
store.capture("fox", "q_steer", {3: vec}, prompt_id=f"t{i}")
|
|
store.rate(f"t{i}", rating)
|
|
base = tiny_h3.run()
|
|
patched, _status = _load(tiny_h3, blocks="3", strength=1.0)
|
|
video, audio = tiny_h3.sample(patched)
|
|
assert not torch.allclose(video, base[0])
|
|
assert torch.equal(audio, base[1]) # only picture queries moved
|