Claude
0dc605f013
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Dy3QqNuYH9NSJgob83xcTX
49 lines
1.9 KiB
Python
49 lines
1.9 KiB
Python
import pytest
|
|
import torch
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _needs_comfy(comfyui):
|
|
"""Imports comfy."""
|
|
|
|
|
|
def _line(model, x, sigmas, name):
|
|
import comfy.k_diffusion.sampling as ks
|
|
return getattr(ks, f"sample_{name}")(model, x, sigmas)
|
|
|
|
|
|
def _model(x, sigma, **_):
|
|
return x * 0.5 + 0.1 * sigma.reshape(-1, *([1] * (x.dim() - 1))) # any smooth x0 prediction
|
|
|
|
|
|
def test_the_presets_are_ordinary_comfy_sampler_names():
|
|
import comfy.samplers
|
|
import modules.sampling.distilled_flow # noqa: F401
|
|
for name in ("dflow_ord1", "dflow_ord2", "dflow_ord2_ab2"):
|
|
assert name in comfy.samplers.SAMPLER_NAMES
|
|
assert comfy.samplers.sampler_object(name).sampler_function is not None
|
|
|
|
|
|
def test_order_one_is_euler():
|
|
import comfy.k_diffusion.sampling as ks
|
|
import modules.sampling.distilled_flow # noqa: F401
|
|
x, sigmas = torch.randn(1, 4, 2, 4, 4), torch.linspace(1.0, 0.0, 5)
|
|
assert torch.allclose(_line(_model, x, sigmas, "dflow_ord1"), ks.sample_euler(_model, x, sigmas), atol=1e-5)
|
|
|
|
|
|
def test_order_two_differs_and_ramp_differs_again():
|
|
import modules.sampling.distilled_flow # noqa: F401
|
|
x, sigmas = torch.randn(1, 4, 2, 4, 4), torch.linspace(1.0, 0.0, 5)
|
|
one, two, ramp = (_line(_model, x, sigmas, n) for n in ("dflow_ord1", "dflow_ord2", "dflow_ord2_ab2"))
|
|
assert not torch.allclose(one, two) and not torch.allclose(two, ramp)
|
|
|
|
|
|
def test_the_sound_stream_stays_first_order():
|
|
from comfy.nested_tensor import NestedTensor
|
|
import modules.sampling.distilled_flow # noqa: F401
|
|
video, audio = torch.randn(1, 4, 2, 4, 4), torch.randn(1, 8, 5)
|
|
sigmas = torch.linspace(1.0, 0.0, 5)
|
|
both = _line(lambda x, s, **k: x * 0.5, NestedTensor((video, audio)), sigmas, "dflow_ord2_ab2")
|
|
alone = _line(lambda x, s, **k: x * 0.5, audio, sigmas, "dflow_ord1")
|
|
assert torch.allclose(both.unbind()[1], alone, atol=1e-5)
|