85 lines
3.1 KiB
Python
85 lines
3.1 KiB
Python
import math
|
|
|
|
import torch
|
|
|
|
from comfyui_spectrum.config import SpectrumConfig
|
|
from comfyui_spectrum.forecast import ChebyshevSpectrumForecaster
|
|
from comfyui_spectrum.runtime import SpectrumRuntime
|
|
|
|
|
|
def test_config_and_forecaster() -> None:
|
|
assert SpectrumConfig().validate().backend == "anima"
|
|
|
|
forecaster = ChebyshevSpectrumForecaster(degree=1, ridge_lambda=0.1, max_history=16)
|
|
forecaster.update(-1.0, torch.tensor([10.0]))
|
|
forecaster.update(-7.0 / 9.0, torch.tensor([9.0]))
|
|
prediction = forecaster.predict(time_coord=1.0, blend_weight=0.0)
|
|
assert torch.allclose(prediction, torch.tensor([1.0]), atol=1e-5)
|
|
|
|
|
|
def test_runtime_forecast_step() -> None:
|
|
cfg = SpectrumConfig(
|
|
degree=1,
|
|
ridge_lambda=0.1,
|
|
window_size=2.0,
|
|
flex_window=0.75,
|
|
warmup_steps=2,
|
|
tail_actual_steps=0,
|
|
max_history=16,
|
|
).validate()
|
|
runtime = SpectrumRuntime(cfg)
|
|
sigmas = torch.linspace(1.0, 0.0, 9)
|
|
run_id = runtime.start_run(sigmas, "sample_euler", supports_solver_steps=True)
|
|
|
|
for step_id in range(2):
|
|
decision = runtime.begin_solver_step(run_id, step_id, runtime.time_coord_for_step(step_id), 8)
|
|
assert decision["actual_forward"] is True
|
|
call_id = runtime.register_model_hook_call(run_id, step_id, expected_shape=(1, 2, 2, 4, 8))
|
|
runtime.observe_actual_feature(
|
|
run_id,
|
|
step_id,
|
|
torch.full((1, 2, 2, 4, 8), float(10 + step_id)),
|
|
call_id=call_id,
|
|
)
|
|
runtime.finalize_solver_step(run_id, step_id, used_forecast=False)
|
|
|
|
decision = runtime.begin_solver_step(run_id, 2, runtime.time_coord_for_step(2), 8)
|
|
assert decision["actual_forward"] is False
|
|
call_id = runtime.register_model_hook_call(run_id, 2, expected_shape=(1, 2, 2, 4, 8))
|
|
prediction = runtime.predict_feature(run_id, 2, expected_shape=(1, 2, 2, 4, 8), call_id=call_id)
|
|
assert prediction is not None
|
|
assert prediction.shape == (1, 2, 2, 4, 8)
|
|
runtime.finalize_solver_step(run_id, 2, used_forecast=True)
|
|
assert runtime.stats.forecasted_count == 1
|
|
assert runtime.stats.forecast_disabled is False
|
|
runtime.end_run(run_id)
|
|
|
|
|
|
def test_nonuniform_schedule_coordinates() -> None:
|
|
runtime = SpectrumRuntime(SpectrumConfig(degree=1, max_history=8).validate())
|
|
run_id = runtime.start_run(torch.tensor([10.0, 9.0, 1.0, 0.0]), "sample_euler", supports_solver_steps=True)
|
|
coords = [runtime.time_coord_for_step(i) for i in range(3)]
|
|
assert math.isclose(coords[0], -1.0)
|
|
assert math.isclose(coords[2], 1.0)
|
|
assert math.isclose(coords[1], -7.0 / 9.0)
|
|
runtime.end_run(run_id)
|
|
|
|
|
|
def test_er_sde_supports_solver_step_tracking() -> None:
|
|
runtime = SpectrumRuntime(SpectrumConfig(max_history=8).validate())
|
|
run_id = runtime.start_run(torch.linspace(1.0, 0.0, 9), "sample_er_sde", supports_solver_steps=True)
|
|
assert runtime.stats.forecast_disabled is False
|
|
runtime.end_run(run_id)
|
|
|
|
|
|
def main() -> None:
|
|
test_config_and_forecaster()
|
|
test_runtime_forecast_step()
|
|
test_nonuniform_schedule_coordinates()
|
|
test_er_sde_supports_solver_step_tracking()
|
|
print("ok")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|