450 lines
13 KiB
Python
450 lines
13 KiB
Python
# SimpleSyrup - workflow-focused ComfyUI extensions for image generation
|
|
# Copyright (C) 2026 Artificial Sweetener and contributors
|
|
# SPDX-License-Identifier: AGPL-3.0-or-later
|
|
|
|
"""Tests for KSampler Extras scheduler runtime helpers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from collections.abc import Sequence
|
|
|
|
import comfy.samplers
|
|
import pytest
|
|
import torch
|
|
|
|
from simple_syrup.runtime import sampling_reference_schedules, sampling_schedulers
|
|
from simple_syrup.runtime.sampling_schedulers import (
|
|
calculate_sigmas,
|
|
)
|
|
|
|
|
|
class FakeModel:
|
|
"""Provide the model-sampling object expected by scheduler helpers."""
|
|
|
|
def __init__(self) -> None:
|
|
"""Create a fake model with a stable model_sampling object."""
|
|
|
|
self.model_sampling = object()
|
|
|
|
def get_model_object(self, name: str) -> object:
|
|
"""Return the requested fake model object."""
|
|
|
|
assert name == "model_sampling"
|
|
return self.model_sampling
|
|
|
|
|
|
class FakeLatentFormat:
|
|
"""Expose the spatial compression used to recover image dimensions."""
|
|
|
|
def __init__(self, spacial_downscale_ratio: int) -> None:
|
|
"""Store one deterministic latent-to-image scale."""
|
|
|
|
self.spacial_downscale_ratio = spacial_downscale_ratio
|
|
|
|
|
|
class FakeModelWithLatentFormat(FakeModel):
|
|
"""Provide model-sampling and latent-format objects for Flux2 tests."""
|
|
|
|
def __init__(self, spacial_downscale_ratio: int) -> None:
|
|
"""Create a fake model with the requested spatial compression."""
|
|
|
|
super().__init__()
|
|
self.latent_format = FakeLatentFormat(spacial_downscale_ratio)
|
|
|
|
def get_model_object(self, name: str) -> object:
|
|
"""Return the requested fake model object."""
|
|
|
|
if name == "latent_format":
|
|
return self.latent_format
|
|
return super().get_model_object(name)
|
|
|
|
|
|
class FakeDiscreteModelSampling:
|
|
"""Provide k-diffusion-style discrete sigma conversion for tests."""
|
|
|
|
def __init__(self, sigmas: Sequence[float]) -> None:
|
|
"""Create a fake discrete model sampling object."""
|
|
|
|
self.sigmas = torch.tensor(sigmas, dtype=torch.float32)
|
|
self.log_sigmas = self.sigmas.log()
|
|
|
|
def sigma(self, timestep: torch.Tensor) -> torch.Tensor:
|
|
"""Convert fractional timesteps to sigmas with log-linear interpolation."""
|
|
|
|
timestep = torch.clamp(
|
|
timestep.float().to(self.log_sigmas.device),
|
|
min=0,
|
|
max=(len(self.sigmas) - 1),
|
|
)
|
|
low_index = timestep.floor().long()
|
|
high_index = timestep.ceil().long()
|
|
weight = timestep.frac()
|
|
log_sigma = (1 - weight) * self.log_sigmas[
|
|
low_index
|
|
] + weight * self.log_sigmas[high_index]
|
|
return log_sigma.exp().to(timestep.device)
|
|
|
|
|
|
class FakeDiscreteModel:
|
|
"""Provide a discrete model_sampling object for automatic_a1111 tests."""
|
|
|
|
def __init__(self) -> None:
|
|
"""Create a fake model with k-diffusion-style sigmas."""
|
|
|
|
self.model_sampling = FakeDiscreteModelSampling((0.1, 0.3, 1.0))
|
|
|
|
def get_model_object(self, name: str) -> FakeDiscreteModelSampling:
|
|
"""Return the requested fake model sampling object."""
|
|
|
|
assert name == "model_sampling"
|
|
return self.model_sampling
|
|
|
|
|
|
def assert_sigmas_close(actual: torch.Tensor, expected: list[float]) -> None:
|
|
"""Assert that calculated sigmas match fixed reference values."""
|
|
|
|
assert torch.allclose(
|
|
actual.cpu(),
|
|
torch.tensor(expected, dtype=torch.float32),
|
|
atol=1e-5,
|
|
rtol=1e-5,
|
|
)
|
|
|
|
|
|
def reference_extra_sigmas(
|
|
scheduler_name: str,
|
|
sampler_name: str,
|
|
steps: int,
|
|
denoise: float,
|
|
) -> torch.Tensor:
|
|
"""Calculate expected extra scheduler sigmas with KSampler semantics."""
|
|
|
|
if denoise <= 0.0:
|
|
return torch.FloatTensor([])
|
|
|
|
schedule_steps = steps if denoise > 0.9999 else int(steps / denoise)
|
|
if sampler_name in comfy.samplers.KSampler.DISCARD_PENULTIMATE_SIGMA_SAMPLERS:
|
|
schedule_steps += 1
|
|
|
|
sigmas = reference_full_extra_schedule(scheduler_name, schedule_steps)
|
|
if sampler_name in comfy.samplers.KSampler.DISCARD_PENULTIMATE_SIGMA_SAMPLERS:
|
|
sigmas = torch.cat([sigmas[:-2], sigmas[-1:]])
|
|
|
|
if denoise <= 0.9999:
|
|
sigmas = sigmas[-(steps + 1) :]
|
|
return sigmas
|
|
|
|
|
|
def reference_full_extra_schedule(scheduler_name: str, steps: int) -> torch.Tensor:
|
|
"""Calculate expected full AYS/GITS formula output for tests."""
|
|
|
|
if scheduler_name == "AYS SD1":
|
|
return reference_ays_schedule("SD1", steps)
|
|
if scheduler_name == "AYS SDXL":
|
|
return reference_ays_schedule("SDXL", steps)
|
|
if scheduler_name == "GITS":
|
|
return reference_gits_schedule(steps)
|
|
if scheduler_name == "automatic_a1111":
|
|
return reference_automatic_a1111_schedule(FakeDiscreteModel(), steps)
|
|
raise ValueError(f"Unsupported reference scheduler '{scheduler_name}'.")
|
|
|
|
|
|
def reference_ays_schedule(model_type: str, steps: int) -> torch.Tensor:
|
|
"""Calculate full AYS schedule with Comfy Extras formula semantics."""
|
|
|
|
sigmas = list(sampling_reference_schedules.AYS_NOISE_LEVELS[model_type])
|
|
if (steps + 1) != len(sigmas):
|
|
sigmas = reference_loglinear_interpolate(sigmas, steps + 1)
|
|
sigmas[-1] = 0.0
|
|
return torch.FloatTensor(sigmas)
|
|
|
|
|
|
def reference_gits_schedule(steps: int) -> torch.Tensor:
|
|
"""Calculate full GITS schedule for the default coefficient."""
|
|
|
|
if steps <= 20:
|
|
sigmas = list(sampling_reference_schedules.GITS_DEFAULT_NOISE_LEVELS[steps - 2])
|
|
else:
|
|
sigmas = reference_loglinear_interpolate(
|
|
sampling_reference_schedules.GITS_DEFAULT_NOISE_LEVELS[-1],
|
|
steps + 1,
|
|
)
|
|
sigmas[-1] = 0.0
|
|
return torch.FloatTensor(sigmas)
|
|
|
|
|
|
def reference_automatic_a1111_schedule(
|
|
model: FakeDiscreteModel,
|
|
steps: int,
|
|
) -> torch.Tensor:
|
|
"""Calculate k-diffusion DiscreteSchedule.get_sigmas-style output."""
|
|
|
|
model_sampling = model.model_sampling
|
|
timesteps = torch.linspace(
|
|
len(model_sampling.sigmas) - 1,
|
|
0,
|
|
steps,
|
|
device=model_sampling.sigmas.device,
|
|
)
|
|
sigmas = model_sampling.sigma(timesteps)
|
|
return torch.cat([sigmas, sigmas.new_zeros([1])]).cpu()
|
|
|
|
|
|
def reference_loglinear_interpolate(
|
|
sigmas: Sequence[float],
|
|
num_steps: int,
|
|
) -> list[float]:
|
|
"""Interpolate reference sigma values in log space."""
|
|
|
|
reversed_logs = [math.log(value) for value in reversed(sigmas)]
|
|
source_max = len(reversed_logs) - 1
|
|
target_max = num_steps - 1
|
|
interpolated: list[float] = []
|
|
|
|
for target_index in range(num_steps):
|
|
source_position = target_index * source_max / target_max
|
|
left_index = math.floor(source_position)
|
|
right_index = min(left_index + 1, source_max)
|
|
fraction = source_position - left_index
|
|
left_value = reversed_logs[left_index]
|
|
right_value = reversed_logs[right_index]
|
|
interpolated.append(
|
|
math.exp(left_value + (right_value - left_value) * fraction)
|
|
)
|
|
|
|
return list(reversed(interpolated))
|
|
|
|
|
|
def test_ays_sd1_full_schedule_matches_reference_values() -> None:
|
|
"""AYS SD1 full schedules match fixed Comfy Extras reference values."""
|
|
|
|
model = FakeModel()
|
|
|
|
assert_sigmas_close(
|
|
calculate_sigmas(model, "AYS SD1", "euler", 10, 1.0),
|
|
[
|
|
14.61464119,
|
|
6.474576,
|
|
3.86367464,
|
|
2.69461513,
|
|
1.88419211,
|
|
1.39438045,
|
|
0.96425837,
|
|
0.65236861,
|
|
0.39774564,
|
|
0.15152326,
|
|
0.0,
|
|
],
|
|
)
|
|
assert_sigmas_close(
|
|
calculate_sigmas(model, "AYS SD1", "euler", 20, 1.0),
|
|
[
|
|
14.61464119,
|
|
9.72746658,
|
|
6.474576,
|
|
5.00156546,
|
|
3.86367464,
|
|
3.22662616,
|
|
2.69461513,
|
|
2.25325823,
|
|
1.88419211,
|
|
1.62088883,
|
|
1.39438045,
|
|
1.15954435,
|
|
0.96425837,
|
|
0.79312789,
|
|
0.65236861,
|
|
0.50938863,
|
|
0.39774564,
|
|
0.24549484,
|
|
0.15152326,
|
|
0.06647934,
|
|
0.0,
|
|
],
|
|
)
|
|
|
|
|
|
def test_ays_sdxl_full_schedule_matches_reference_values() -> None:
|
|
"""AYS SDXL full schedules match fixed Comfy Extras reference values."""
|
|
|
|
model = FakeModel()
|
|
|
|
assert_sigmas_close(
|
|
calculate_sigmas(model, "AYS SDXL", "euler", 10, 1.0),
|
|
[
|
|
14.61464119,
|
|
6.31844854,
|
|
3.76817894,
|
|
2.18114805,
|
|
1.34052444,
|
|
0.86207211,
|
|
0.55506933,
|
|
0.37985408,
|
|
0.23323642,
|
|
0.11141882,
|
|
0.0,
|
|
],
|
|
)
|
|
assert_sigmas_close(
|
|
calculate_sigmas(model, "AYS SDXL", "euler", 20, 1.0),
|
|
[
|
|
14.61464119,
|
|
9.60946751,
|
|
6.31844854,
|
|
4.87945127,
|
|
3.76817894,
|
|
2.86687231,
|
|
2.18114805,
|
|
1.70993638,
|
|
1.34052444,
|
|
1.07500172,
|
|
0.86207211,
|
|
0.69174403,
|
|
0.55506933,
|
|
0.45917898,
|
|
0.37985408,
|
|
0.29765046,
|
|
0.23323642,
|
|
0.16120461,
|
|
0.11141882,
|
|
0.05700676,
|
|
0.0,
|
|
],
|
|
)
|
|
|
|
|
|
def test_gits_full_schedule_matches_reference_values() -> None:
|
|
"""GITS full schedules match fixed reference values at the default coefficient."""
|
|
|
|
model = FakeModel()
|
|
|
|
assert_sigmas_close(
|
|
calculate_sigmas(model, "GITS", "euler", 10, 1.0),
|
|
[
|
|
14.61464119,
|
|
5.85520077,
|
|
2.84484982,
|
|
1.67050016,
|
|
1.08895338,
|
|
0.74807048,
|
|
0.50118381,
|
|
0.32104823,
|
|
0.19894916,
|
|
0.09824532,
|
|
0.0,
|
|
],
|
|
)
|
|
assert_sigmas_close(
|
|
calculate_sigmas(model, "GITS", "euler", 20, 1.0),
|
|
[
|
|
14.61464119,
|
|
7.49001646,
|
|
4.65472794,
|
|
3.07277966,
|
|
2.19988537,
|
|
1.61558151,
|
|
1.24153244,
|
|
0.95350921,
|
|
0.74807048,
|
|
0.59516323,
|
|
0.50118381,
|
|
0.41087446,
|
|
0.34370604,
|
|
0.29807833,
|
|
0.25053367,
|
|
0.22545385,
|
|
0.19894916,
|
|
0.17026083,
|
|
0.13792117,
|
|
0.09824532,
|
|
0.0,
|
|
],
|
|
)
|
|
assert_sigmas_close(
|
|
calculate_sigmas(model, "GITS", "euler", 30, 1.0),
|
|
[
|
|
14.61464119,
|
|
9.35946941,
|
|
6.39175272,
|
|
4.65472794,
|
|
3.52900577,
|
|
2.74887061,
|
|
2.19988537,
|
|
1.7906853,
|
|
1.47980762,
|
|
1.24153244,
|
|
1.04120433,
|
|
0.87942225,
|
|
0.74807048,
|
|
0.64230049,
|
|
0.56202602,
|
|
0.50118381,
|
|
0.43900734,
|
|
0.38714039,
|
|
0.34370604,
|
|
0.31257147,
|
|
0.28130382,
|
|
0.25053367,
|
|
0.23352164,
|
|
0.21624818,
|
|
0.19894916,
|
|
0.17933175,
|
|
0.1587158,
|
|
0.13792117,
|
|
0.11000647,
|
|
0.0655402,
|
|
0.0,
|
|
],
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("scheduler_name", ["AYS SD1", "AYS SDXL", "GITS"])
|
|
@pytest.mark.parametrize("sampler_name", ["euler", "dpm_2"])
|
|
@pytest.mark.parametrize("steps", [5, 10, 20, 30])
|
|
@pytest.mark.parametrize("denoise", [0.25, 0.5, 0.8, 1.0])
|
|
def test_extra_schedulers_follow_ksampler_denoise_semantics(
|
|
scheduler_name: str,
|
|
sampler_name: str,
|
|
steps: int,
|
|
denoise: float,
|
|
) -> None:
|
|
"""Extra schedulers use KSampler partial-denoise and sigma cleanup rules."""
|
|
|
|
actual = calculate_sigmas(
|
|
FakeModel(),
|
|
scheduler_name,
|
|
sampler_name,
|
|
steps,
|
|
denoise,
|
|
)
|
|
expected = reference_extra_sigmas(scheduler_name, sampler_name, steps, denoise)
|
|
|
|
assert actual.shape == expected.shape
|
|
assert torch.allclose(actual, expected, atol=1e-5, rtol=1e-5)
|
|
if denoise < 1.0:
|
|
assert actual.shape == (steps + 1,)
|
|
|
|
|
|
def test_reported_ays_sd1_partial_denoise_regression() -> None:
|
|
"""AYS SD1 with partial denoise returns a KSampler-length sigma schedule."""
|
|
|
|
actual = calculate_sigmas(FakeModel(), "AYS SD1", "euler", 5, 0.25)
|
|
expected = reference_extra_sigmas("AYS SD1", "euler", 5, 0.25)
|
|
|
|
assert actual.shape == (6,)
|
|
assert torch.allclose(actual, expected, atol=1e-5, rtol=1e-5)
|
|
|
|
|
|
def test_gits_scope_uses_default_coefficient() -> None:
|
|
"""The simple node contract exposes GITS at its default coefficient only."""
|
|
|
|
assert sampling_schedulers.GITS_DEFAULT_COEFF == 1.20
|
|
|
|
|
|
def test_beta57_scope_uses_res4lyf_preset_parameters() -> None:
|
|
"""The beta57 scheduler contract exposes RES4LYF's fixed beta preset."""
|
|
|
|
assert sampling_schedulers.BETA57_ALPHA == 0.5
|
|
assert sampling_schedulers.BETA57_BETA == 0.7
|