5d515ee617
Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
172 lines
7.1 KiB
Python
172 lines
7.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Unit tests for the float64 schedule memoization in fastvideo.models.utils."""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from fastvideo.models.utils import (
|
|
_FLOAT64_SCHEDULE_ATTR,
|
|
get_float64_schedule,
|
|
pred_noise_to_pred_video,
|
|
pred_noise_to_x_bound,
|
|
)
|
|
|
|
|
|
class FakeScheduler:
|
|
"""Minimal stand-in exposing the sigma and timestep tables the helper reads."""
|
|
|
|
def __init__(self, num_steps: int = 4) -> None:
|
|
self.set_timesteps(num_steps)
|
|
|
|
def set_timesteps(self, num_steps: int) -> None:
|
|
"""Rebuild the schedule tables the way a real scheduler does per call."""
|
|
self.timesteps = torch.linspace(1000.0, 0.0, num_steps)
|
|
self.sigmas = self.timesteps / 1000.0
|
|
|
|
|
|
def _reference_schedule(scheduler: FakeScheduler,
|
|
device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Convert the schedule the way the call sites did before memoization."""
|
|
return (scheduler.sigmas.double().to(device),
|
|
scheduler.timesteps.double().to(device))
|
|
|
|
|
|
def test_repeated_call_returns_cached_tensors():
|
|
"""A second call for the same schedule reuses the same tensor objects."""
|
|
scheduler = FakeScheduler()
|
|
device = torch.device("cpu")
|
|
sigmas_first, timesteps_first = get_float64_schedule(scheduler, device)
|
|
sigmas_second, timesteps_second = get_float64_schedule(scheduler, device)
|
|
assert sigmas_first is sigmas_second
|
|
assert timesteps_first is timesteps_second
|
|
|
|
|
|
def test_cached_values_match_direct_conversion():
|
|
"""Memoized tables are bitwise-equal to converting the schedule directly."""
|
|
scheduler = FakeScheduler()
|
|
device = torch.device("cpu")
|
|
sigmas, timesteps = get_float64_schedule(scheduler, device)
|
|
sigmas_reference, timesteps_reference = _reference_schedule(scheduler, device)
|
|
assert torch.equal(sigmas, sigmas_reference)
|
|
assert torch.equal(timesteps, timesteps_reference)
|
|
|
|
|
|
def test_tables_are_float64():
|
|
"""Both tables are promoted to float64 regardless of the source dtype."""
|
|
scheduler = FakeScheduler()
|
|
scheduler.sigmas = scheduler.sigmas.to(torch.float32)
|
|
scheduler.timesteps = scheduler.timesteps.to(torch.float32)
|
|
sigmas, timesteps = get_float64_schedule(scheduler, torch.device("cpu"))
|
|
assert sigmas.dtype == torch.float64
|
|
assert timesteps.dtype == torch.float64
|
|
|
|
|
|
def test_set_timesteps_invalidates_cache():
|
|
"""Swapping the schedule tables rebuilds them instead of serving stale ones."""
|
|
scheduler = FakeScheduler(num_steps=4)
|
|
device = torch.device("cpu")
|
|
sigmas_before, _ = get_float64_schedule(scheduler, device)
|
|
|
|
scheduler.set_timesteps(8)
|
|
sigmas_after, timesteps_after = get_float64_schedule(scheduler, device)
|
|
|
|
assert sigmas_after is not sigmas_before
|
|
assert sigmas_after.shape[0] == 8
|
|
sigmas_reference, timesteps_reference = _reference_schedule(scheduler, device)
|
|
assert torch.equal(sigmas_after, sigmas_reference)
|
|
assert torch.equal(timesteps_after, timesteps_reference)
|
|
|
|
|
|
def test_in_place_mutation_invalidates_cache():
|
|
"""Writing into the tables in place rebuilds them instead of serving stale ones."""
|
|
scheduler = FakeScheduler(num_steps=4)
|
|
device = torch.device("cpu")
|
|
sigmas_before, _ = get_float64_schedule(scheduler, device)
|
|
|
|
scheduler.sigmas[-1] = scheduler.sigmas[-2] # same tensor object, new values
|
|
sigmas_after, _ = get_float64_schedule(scheduler, device)
|
|
|
|
assert sigmas_after is not sigmas_before
|
|
sigmas_reference, _ = _reference_schedule(scheduler, device)
|
|
assert torch.equal(sigmas_after, sigmas_reference)
|
|
|
|
|
|
def test_cache_is_stored_on_the_scheduler():
|
|
"""The memo lives on the scheduler, so separate schedulers stay independent."""
|
|
first, second = FakeScheduler(), FakeScheduler()
|
|
device = torch.device("cpu")
|
|
assert getattr(first, _FLOAT64_SCHEDULE_ATTR, None) is None
|
|
|
|
first_sigmas, _ = get_float64_schedule(first, device)
|
|
assert getattr(first, _FLOAT64_SCHEDULE_ATTR, None) is not None
|
|
assert getattr(second, _FLOAT64_SCHEDULE_ATTR, None) is None
|
|
|
|
second_sigmas, _ = get_float64_schedule(second, device)
|
|
assert second_sigmas is not first_sigmas
|
|
|
|
|
|
def _reference_pred_noise_to_pred_video(pred_noise: torch.Tensor,
|
|
noise_input_latent: torch.Tensor,
|
|
timestep: torch.Tensor,
|
|
scheduler: FakeScheduler) -> torch.Tensor:
|
|
"""Reproduce the pre-memoization conversion for a numerical parity check."""
|
|
dtype = pred_noise.dtype
|
|
device = pred_noise.device
|
|
pred_noise = pred_noise.double().to(device)
|
|
noise_input_latent = noise_input_latent.double().to(device)
|
|
sigmas, timesteps = _reference_schedule(scheduler, device)
|
|
timestep_id = torch.argmin(
|
|
(timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
|
|
sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
|
|
return (noise_input_latent - sigma_t * pred_noise).to(dtype)
|
|
|
|
|
|
@pytest.mark.parametrize("num_steps", [2, 4, 8])
|
|
def test_pred_noise_to_pred_video_matches_reference(num_steps):
|
|
"""Memoizing the schedule leaves the converted latents bitwise unchanged."""
|
|
torch.manual_seed(0)
|
|
scheduler = FakeScheduler(num_steps=num_steps)
|
|
timestep = scheduler.timesteps[:min(3, num_steps)].clone()
|
|
batch_size = timestep.numel()
|
|
pred_noise = torch.randn(batch_size, 4, 8, 8)
|
|
noise_input_latent = torch.randn(batch_size, 4, 8, 8)
|
|
|
|
expected = _reference_pred_noise_to_pred_video(pred_noise, noise_input_latent,
|
|
timestep, scheduler)
|
|
actual = pred_noise_to_pred_video(pred_noise, noise_input_latent, timestep,
|
|
scheduler)
|
|
assert torch.equal(actual, expected)
|
|
|
|
|
|
def test_pred_noise_to_pred_video_is_stable_across_repeated_calls():
|
|
"""Reusing a cached schedule across steps keeps results identical."""
|
|
torch.manual_seed(0)
|
|
scheduler = FakeScheduler()
|
|
pred_noise = torch.randn(2, 4, 8, 8)
|
|
noise_input_latent = torch.randn(2, 4, 8, 8)
|
|
timestep = scheduler.timesteps[:2].clone()
|
|
|
|
first = pred_noise_to_pred_video(pred_noise, noise_input_latent, timestep,
|
|
scheduler)
|
|
second = pred_noise_to_pred_video(pred_noise, noise_input_latent, timestep,
|
|
scheduler)
|
|
assert torch.equal(first, second)
|
|
|
|
|
|
def test_pred_noise_to_x_bound_uses_the_same_schedule():
|
|
"""The boundary variant reads the memoized tables without altering them."""
|
|
torch.manual_seed(0)
|
|
scheduler = FakeScheduler()
|
|
pred_noise = torch.randn(2, 4, 8, 8)
|
|
noise_input_latent = torch.randn(2, 4, 8, 8)
|
|
timestep = scheduler.timesteps[:2].clone()
|
|
boundary_timestep = torch.full_like(timestep, float(scheduler.timesteps[-1]))
|
|
|
|
sigmas_before, _ = get_float64_schedule(scheduler, pred_noise.device)
|
|
result = pred_noise_to_x_bound(pred_noise, noise_input_latent, timestep,
|
|
boundary_timestep, scheduler)
|
|
sigmas_after, _ = get_float64_schedule(scheduler, pred_noise.device)
|
|
|
|
assert sigmas_after is sigmas_before
|
|
assert result.shape == pred_noise.shape
|