55 lines
2.6 KiB
Python
55 lines
2.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Compare actual Wan schedules to native Wan outputs and independent sigma arithmetic."""
|
|
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import numpy.typing as npt
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
|
|
FASTWAN_DMD_SCHEDULER_SHIFT,
|
|
build_wan_scheduler,
|
|
)
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
|
|
|
|
|
def test_wan_unipc_matches_captured_native_full_schedule() -> None:
|
|
# The fixture pins the official Wan source and records how to regenerate it.
|
|
golden = json.loads((Path(__file__).parent / "fixtures/native_unipc_shift12.json").read_text())
|
|
scheduler = build_wan_scheduler("unipc", golden["runtime_shift"])
|
|
for steps in (golden["steps"], 20, golden["steps"]):
|
|
scheduler.set_timesteps(steps, device="cpu", shift=golden["runtime_shift"])
|
|
if steps == golden["steps"]:
|
|
assert scheduler.timesteps.tolist() == golden["timesteps"]
|
|
torch.testing.assert_close(scheduler.sigmas, torch.tensor(golden["sigmas"]), rtol=0, atol=0)
|
|
|
|
|
|
@pytest.mark.parametrize("shift", [1.0, 3.0, 5.0, 12.0])
|
|
@pytest.mark.parametrize("steps", [20, 40])
|
|
def test_wan_unipc_shifts_unmodified_training_endpoints(shift: float, steps: int) -> None:
|
|
# Native Wan: float32 training sigma_max=(1-1/1000), sigma_min=0;
|
|
# interpolate in float64, apply the runtime shift, then cast sigma to FP32.
|
|
unshifted = np.linspace(float(np.float32(0.999)), 0.0, steps + 1)[:-1]
|
|
shifted = shift * unshifted / (1.0 + (shift - 1.0) * unshifted)
|
|
expected_timesteps = (shifted * 1000).astype(np.int64)
|
|
expected_sigmas: npt.NDArray[np.float32] = np.append(shifted, 0.0).astype(np.float32)
|
|
scheduler = build_wan_scheduler("unipc", shift)
|
|
scheduler.set_timesteps(steps, device="cpu", shift=shift)
|
|
assert scheduler.config.shift == scheduler.config["shift"] == 1.0
|
|
torch.testing.assert_close(scheduler.timesteps, torch.from_numpy(expected_timesteps), rtol=0, atol=0)
|
|
torch.testing.assert_close(scheduler.sigmas, torch.from_numpy(expected_sigmas), rtol=0, atol=0)
|
|
|
|
|
|
def test_wan_euler_keeps_dmd_shift_semantics() -> None:
|
|
scheduler = build_wan_scheduler("euler", FASTWAN_DMD_SCHEDULER_SHIFT)
|
|
scheduler.set_timesteps(3, device="cpu")
|
|
# The unchanged Euler branch starts at sigma=1, with shift=8 for DMD.
|
|
expected = torch.tensor([1.0, 16.0 / 17.0, 0.8, 0.0])
|
|
torch.testing.assert_close(scheduler.sigmas, expected, rtol=0, atol=1e-7)
|
|
torch.testing.assert_close(scheduler.timesteps, expected[:-1] * 1000, rtol=0, atol=1e-4)
|