89 lines
2.8 KiB
Python
89 lines
2.8 KiB
Python
from __future__ import annotations
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import torch
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
if str(ROOT) not in sys.path:
|
|
sys.path.insert(0, str(ROOT))
|
|
|
|
from comfyui_spectrum_wan.config import SpectrumWanConfig
|
|
from comfyui_spectrum_wan.handlers import resolve_handler
|
|
from comfyui_spectrum_wan.runtime import SpectrumWanRuntime
|
|
|
|
|
|
class DummyModel:
|
|
model_name = "wan2.2_t2v_high_noise_14B_fp8_scaled.safetensors"
|
|
|
|
|
|
def _runtime(debug: bool) -> SpectrumWanRuntime:
|
|
cfg = SpectrumWanConfig(
|
|
backend="auto",
|
|
blend_weight=1.0,
|
|
degree=4,
|
|
ridge_lambda=0.1,
|
|
window_size=2.0,
|
|
flex_window=0.75,
|
|
warmup_steps=2,
|
|
history_size=8,
|
|
debug=debug,
|
|
).validated()
|
|
handler = resolve_handler(cfg.backend, DummyModel())
|
|
return SpectrumWanRuntime(cfg, handler)
|
|
|
|
|
|
def test_begin_step_debug_logs_to_stderr(capsys) -> None:
|
|
runtime = _runtime(debug=True)
|
|
sample_sigmas = torch.linspace(1.0, 0.0, 6)
|
|
transformer_options = {"sample_sigmas": sample_sigmas, "cond_or_uncond": [0, 1]}
|
|
|
|
runtime.begin_step(transformer_options, torch.tensor([sample_sigmas[0]]))
|
|
|
|
captured = capsys.readouterr()
|
|
assert captured.out == ""
|
|
assert "[Spectrum WAN]" in captured.err
|
|
assert "backend=wan22_high_noise" in captured.err
|
|
assert "phase=wan22_high_noise" in captured.err
|
|
assert "step=0" in captured.err
|
|
assert "global_step=0" in captured.err
|
|
assert "actual_forward=True" in captured.err
|
|
|
|
|
|
def test_begin_step_debug_false_stays_silent(capsys) -> None:
|
|
runtime = _runtime(debug=False)
|
|
sample_sigmas = torch.linspace(1.0, 0.0, 6)
|
|
transformer_options = {"sample_sigmas": sample_sigmas, "cond_or_uncond": [0, 1]}
|
|
|
|
runtime.begin_step(transformer_options, torch.tensor([sample_sigmas[0]]))
|
|
|
|
captured = capsys.readouterr()
|
|
assert captured.out == ""
|
|
assert captured.err == ""
|
|
|
|
|
|
def test_begin_step_debug_prefers_unready_bias_shift_over_ready_local_forecaster(capsys) -> None:
|
|
runtime = _runtime(debug=True)
|
|
runtime.cfg.warmup_steps = 0
|
|
sample_sigmas = torch.linspace(1.0, 0.0, 6)
|
|
transformer_options = {"sample_sigmas": sample_sigmas, "cond_or_uncond": [0, 1]}
|
|
|
|
stream = runtime._stream(transformer_options)
|
|
assert stream.forecaster is not None
|
|
stream.forecaster.update(0, torch.ones((1, 2, 2)))
|
|
stream.forecaster.update(1, torch.ones((1, 2, 2)) * 2)
|
|
assert stream.forecaster.ready()
|
|
|
|
stream.bias_shift_predictor = SimpleNamespace(
|
|
ready=lambda: False,
|
|
global_step=lambda step_idx: step_idx,
|
|
)
|
|
|
|
runtime.begin_step(transformer_options, torch.tensor([sample_sigmas[0]]))
|
|
|
|
captured = capsys.readouterr()
|
|
assert "actual_forward=True" in captured.err
|
|
assert "forecast_ready=False" in captured.err
|