Claude
3adebb731f
Old project/workflow names and browser settings migrate on load; /funpack redirects to /fun. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016acRss69eH1235eRdFTZHY
95 lines
3.5 KiB
Python
95 lines
3.5 KiB
Python
"""Block repeat on a real H3 forward."""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _needs_comfy(comfyui):
|
|
"""Imports comfy."""
|
|
|
|
|
|
def _load(tiny, patcher=None, **values):
|
|
from modules.sampling.modifiers.nodes import FunUILoadModifiers
|
|
return FunUILoadModifiers.execute(
|
|
patcher or tiny.patcher, {"block_repeat": {"enabled": True, **values}}).result
|
|
|
|
|
|
def _count_block_calls(tiny, patched, **run):
|
|
"""How many times each real block ran during one forward."""
|
|
counts = {}
|
|
handles = []
|
|
for i, block in enumerate(tiny.patcher.model.diffusion_model.blocks):
|
|
handles.append(block.register_forward_pre_hook(
|
|
lambda m, a, i=i: counts.__setitem__(i, counts.get(i, 0) + 1)))
|
|
try:
|
|
tiny.run(patched, **run)
|
|
finally:
|
|
for h in handles:
|
|
h.remove()
|
|
return counts
|
|
|
|
|
|
def test_each_runs_every_named_block_times_plus_one(tiny_h3):
|
|
patched, status = _load(tiny_h3, blocks="1-2", times=2)
|
|
assert "block_repeat: blocks 1,2 run 3x each" in status
|
|
assert _count_block_calls(tiny_h3, patched) == {0: 1, 1: 3, 2: 3, 3: 1}
|
|
|
|
|
|
def test_span_loops_the_range_and_runs_each_block_once_per_loop(tiny_h3):
|
|
patched, status = _load(tiny_h3, blocks="1-2", shape="span", times=1)
|
|
assert "blocks 1-2 loop 2x as one span" in status
|
|
assert _count_block_calls(tiny_h3, patched) == {0: 1, 1: 2, 2: 2, 3: 1}
|
|
|
|
|
|
def test_span_output_equals_running_the_span_twice_by_hand(tiny_h3):
|
|
"""The loop really feeds block 2's output back into block 1."""
|
|
patched, _ = _load(tiny_h3, blocks="1-2", shape="span")
|
|
looped = tiny_h3.run(patched)
|
|
# By hand: blocks 0,1,2,1,2,3 in order.
|
|
from core import dit_hooks
|
|
manual = tiny_h3.patcher.clone()
|
|
real = tiny_h3.patcher.model.diffusion_model.blocks
|
|
|
|
def twice(args, extra):
|
|
h = args["img"]
|
|
for i in (1, 2, 1, 2):
|
|
h = real[i](h, args["t_emb"], args["mod_segments"], args["rope_freqs"],
|
|
transformer_options=args["transformer_options"])
|
|
return {"img": h}
|
|
|
|
dit_hooks.add_block_hook(manual, "t", 1, twice)
|
|
dit_hooks.add_block_hook(manual, "t", 2, lambda a, e: {"img": a["img"]})
|
|
assert torch.allclose(looped[0], tiny_h3.run(manual)[0])
|
|
|
|
|
|
def test_a_scattered_list_is_refused_for_span(tiny_h3):
|
|
_p, status = _load(tiny_h3, blocks="0,2", shape="span")
|
|
assert "block_repeat: failed to install" in status
|
|
|
|
|
|
def test_step_window_repeats_only_on_the_last_steps(tiny_h3):
|
|
patched, _ = _load(tiny_h3, blocks="1", steps=1)
|
|
sched = torch.tensor([1.0, 0.7, 0.3, 0.0])
|
|
early = _count_block_calls(tiny_h3, patched, sigma=0.7, sigmas=sched)
|
|
last = _count_block_calls(tiny_h3, patched, sigma=0.3, sigmas=sched)
|
|
assert early[1] == 1 and last[1] == 2
|
|
|
|
|
|
def test_video_only_leaves_audio_as_a_single_pass_would(tiny_h3):
|
|
"""Repeating the LAST block with video only: the audio branch must match a
|
|
run with no repeat at all, the video branch must not."""
|
|
base = tiny_h3.run()
|
|
patched, _ = _load(tiny_h3, blocks="3", video_only=True)
|
|
out = tiny_h3.run(patched)
|
|
assert torch.allclose(out[1], base[1], atol=1e-6), "audio got the repeat"
|
|
assert not torch.allclose(out[0], base[0]), "video did not"
|
|
|
|
|
|
def test_video_only_span_leaves_audio_as_a_single_pass_would(tiny_h3):
|
|
base = tiny_h3.run()
|
|
patched, _ = _load(tiny_h3, blocks="2-3", shape="span", video_only=True)
|
|
out = tiny_h3.run(patched)
|
|
assert torch.allclose(out[1], base[1], atol=1e-6), "audio got the loop"
|
|
assert not torch.allclose(out[0], base[0])
|