33 lines
1.4 KiB
Python
33 lines
1.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Unit tests for Wan2.2 warped DMD schedule construction."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
|
|
from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
|
|
from fastvideo.mlx_runtime.wan22_sample import build_wan22_dmd_schedule
|
|
|
|
|
|
def test_warped_dmd_timesteps_differ_from_raw_indices() -> None:
|
|
schedule, warped = build_wan22_dmd_schedule([1000, 757, 522], flow_shift=5.0, warp_denoising_step=True)
|
|
_, raw = build_wan22_dmd_schedule([1000, 757, 522], flow_shift=5.0, warp_denoising_step=False)
|
|
assert raw == [1000.0, 757.0, 522.0]
|
|
# Warping maps step indices into the continuous flow-match schedule.
|
|
assert warped[0] == 1000.0
|
|
assert warped[1] != 757.0
|
|
assert warped[2] != 522.0
|
|
# Sigmas are monotone-ish decreasing along the schedule.
|
|
sigmas = [schedule.sigma_for(t) for t in warped]
|
|
assert sigmas[0] >= sigmas[1] >= sigmas[2]
|
|
|
|
|
|
def test_warped_dmd_timesteps_match_training_scheduler() -> None:
|
|
steps = torch.tensor([1000, 757, 522], dtype=torch.long)
|
|
scheduler = FlowMatchEulerDiscreteScheduler(shift=5.0)
|
|
expected = torch.cat((scheduler.timesteps, torch.tensor([0.0])))[1000 - steps]
|
|
|
|
_, warped = build_wan22_dmd_schedule(steps.tolist(), flow_shift=5.0, warp_denoising_step=True)
|
|
|
|
torch.testing.assert_close(torch.tensor(warped), expected)
|