Claude
6f6e3a730e
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Tw9fjW75aYB8UyVRZo6fCr
72 lines
3.6 KiB
Python
72 lines
3.6 KiB
Python
import pytest
|
|
import torch
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _needs_comfy(comfyui):
|
|
"""Imports comfy."""
|
|
|
|
|
|
def _load(tiny, **v):
|
|
from comfy.patcher_extension import WrappersMP
|
|
from modules.sampling.modifiers.nodes import FunUILoadModifiers
|
|
patched, status = FunUILoadModifiers.execute(
|
|
tiny.patcher, {"sharpness": {"enabled": True, **v}}).result
|
|
wrap = [w for ws in patched.wrappers[WrappersMP.APPLY_MODEL].values() for w in ws][-1]
|
|
return patched, wrap, status
|
|
|
|
|
|
def _to(i, steps=4):
|
|
s = torch.linspace(1.0, 0.0, steps + 1)
|
|
return {"sample_sigmas": s, "sigmas": s[i:i + 1]}
|
|
|
|
|
|
def test_off_installs_nothing(tiny_h3):
|
|
from comfy.patcher_extension import WrappersMP
|
|
from modules.sampling.modifiers.nodes import FunUILoadModifiers
|
|
patched, _ = FunUILoadModifiers.execute(tiny_h3.patcher, {"sharpness": {"enabled": False}}).result
|
|
assert not patched.wrappers.get(WrappersMP.APPLY_MODEL)
|
|
|
|
|
|
def test_a_plain_latent_is_pushed_away_from_the_previous_prediction_in_the_last_part_only(tiny_h3):
|
|
_p, wrap, status = _load(tiny_h3, amount=0.5, last_part=0.5)
|
|
assert "sharpness" in status
|
|
x = torch.randn(1, 4, 3, 8, 8)
|
|
outs = [x * f for f in (0.9, 0.5, 0.2, 0.1)]
|
|
ts = torch.linspace(1.0, 0.0, 5)
|
|
got = [wrap(lambda *a, o=outs[i], **kw: o, x, ts[i:i + 1], None, None, None, _to(i)) for i in range(4)]
|
|
assert torch.equal(got[0], outs[0]) and torch.equal(got[1], outs[1]) # outside the last half
|
|
kept = list(got) # the boost is measured against the sharpened answer before
|
|
for i in (2, 3):
|
|
assert torch.allclose(got[i], outs[i] + 0.25 * (outs[i] - kept[i - 1]), atol=1e-5)
|
|
|
|
|
|
def test_a_packed_model_is_edited_on_its_answer_including_the_last_step_and_never_the_sound(tiny_h3):
|
|
from conftest import packed_av, unpacked
|
|
_p, wrap, _ = _load(tiny_h3, amount=1.0, last_part=1.0)
|
|
x0, shapes = packed_av(torch.randn(1, 4, 3, 8, 8), torch.randn(1, 8, 5))
|
|
answers = [x0 * f for f in (0.9, 0.5, 0.2, 0.1)]
|
|
ts = torch.linspace(1.0, 0.0, 5)
|
|
got = [wrap(lambda *a, o=answers[i], **k: o, x0, ts[i:i + 1], None, None, None, _to(i), latent_shapes=shapes)
|
|
for i in range(4)]
|
|
assert torch.equal(got[0], answers[0]) # nothing to compare on step 0
|
|
for i in (1, 3):
|
|
assert not torch.allclose(unpacked(got[i], shapes)[0], unpacked(answers[i], shapes)[0])
|
|
assert torch.equal(unpacked(got[i], shapes)[1], unpacked(answers[i], shapes)[1]) # sound untouched
|
|
|
|
|
|
def test_a_hand_over_to_a_bigger_picture_is_never_compared_with_the_smaller_one_before_it(tiny_h3):
|
|
"""Two stages: the step index counts the whole schedule, and the first call at the full size has nothing of its
|
|
own size to compare with, so it is not sharpened against the first stage's smaller answer."""
|
|
_p, wrap, _ = _load(tiny_h3, amount=1.0, last_part=1.0)
|
|
ts = torch.linspace(1.0, 0.0, 5)
|
|
small = torch.randn(1, 4, 3, 4, 4)
|
|
big = torch.randn(1, 4, 3, 8, 8)
|
|
first = wrap(lambda *a, **k: small, small, ts[0:1], None, None, None, _to(0))
|
|
second = wrap(lambda *a, **k: small * 0.5, small, ts[1:2], None, None, None, _to(1))
|
|
assert torch.equal(first, small) and not torch.equal(second, small * 0.5) # same size: compared
|
|
handed = wrap(lambda *a, **k: big, big, ts[2:3], None, None, None, _to(2))
|
|
assert torch.equal(handed, big) # new size: left as the model said
|
|
after = wrap(lambda *a, **k: big * 0.5, big, ts[3:4], None, None, None, _to(3))
|
|
assert not torch.equal(after, big * 0.5) # and sharpening goes on in stage 2
|