135 lines
5.1 KiB
Python
135 lines
5.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""CPU-only contracts for spatial fast mode + flag composition."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from examples.inference.basic.mlx_wan22_generate import DEFAULT_HEIGHT, DEFAULT_NUM_FRAMES, DEFAULT_WIDTH
|
|
from fastvideo.mlx_runtime.fast_spatial import (
|
|
DEFAULT_FAST_SPATIAL_SHARPEN,
|
|
apply_fast_spatial_upsample,
|
|
plan_fast_spatial,
|
|
resolve_spatial_mode,
|
|
)
|
|
from fastvideo.mlx_runtime.frame_upsample import upsample_frames
|
|
|
|
|
|
def _frames(count: int, height: int, width: int, seed: int = 0) -> list[np.ndarray]:
|
|
rng = np.random.default_rng(seed)
|
|
return [rng.integers(0, 256, (height, width, 3), dtype=np.uint8) for _ in range(count)]
|
|
|
|
|
|
def test_resolve_spatial_mode_priority() -> None:
|
|
assert resolve_spatial_mode(refine=False, fast_spatial=False) == "off"
|
|
assert resolve_spatial_mode(refine=False, fast_spatial=True) == "fast_spatial"
|
|
# Refine is the quality path and wins when both are requested.
|
|
assert resolve_spatial_mode(refine=True, fast_spatial=True) == "refine"
|
|
assert resolve_spatial_mode(refine=True, fast_spatial=False) == "refine"
|
|
|
|
|
|
def test_plan_fast_spatial_matches_refine_geometry() -> None:
|
|
spatial = plan_fast_spatial(
|
|
height=480,
|
|
width=832,
|
|
num_frames=81,
|
|
spatial_scale=2,
|
|
vae_spatial_compression=8,
|
|
)
|
|
assert spatial.enabled
|
|
assert spatial.stage1_height == 240
|
|
assert spatial.stage1_width == 416
|
|
assert spatial.target_height == 480
|
|
assert spatial.target_width == 832
|
|
assert spatial.plan.stage1_latent_height == 30
|
|
assert spatial.plan.stage2_latent_width == 104
|
|
assert spatial.sharpen == DEFAULT_FAST_SPATIAL_SHARPEN
|
|
|
|
|
|
def test_plan_fast_spatial_disabled() -> None:
|
|
spatial = plan_fast_spatial(height=480, width=832, num_frames=81, enabled=False)
|
|
assert not spatial.enabled
|
|
assert spatial.scale == 1
|
|
|
|
|
|
def test_plan_fast_spatial_accepts_wan22_defaults() -> None:
|
|
assert (DEFAULT_HEIGHT, DEFAULT_WIDTH, DEFAULT_NUM_FRAMES) == (448, 832, 121)
|
|
spatial = plan_fast_spatial(
|
|
height=DEFAULT_HEIGHT,
|
|
width=DEFAULT_WIDTH,
|
|
num_frames=DEFAULT_NUM_FRAMES,
|
|
spatial_scale=2,
|
|
vae_spatial_compression=16,
|
|
vae_temporal_compression=4,
|
|
patch_size=(1, 2, 2),
|
|
)
|
|
assert (spatial.plan.stage1_latent_height, spatial.plan.stage1_latent_width) == (14, 26)
|
|
assert (spatial.plan.stage2_latent_height, spatial.plan.stage2_latent_width) == (28, 52)
|
|
|
|
|
|
def test_apply_fast_spatial_upsample_resizes_decoded_frames() -> None:
|
|
"""The upsample is pixel-space: it takes decoded frames, not latents.
|
|
|
|
Interpolating latents instead is what made spatial fast mode incoherent —
|
|
a blended Wan latent is off the decoder's manifold, so the decode came back
|
|
as the right silhouette under a smeared veil.
|
|
"""
|
|
spatial = plan_fast_spatial(
|
|
height=480,
|
|
width=832,
|
|
num_frames=81,
|
|
spatial_scale=2,
|
|
vae_spatial_compression=8,
|
|
upsample_mode="lanczos",
|
|
sharpen=0.0,
|
|
)
|
|
frames = _frames(3, 240, 416)
|
|
out = apply_fast_spatial_upsample(frames, spatial)
|
|
|
|
assert len(out) == 3
|
|
assert all(frame.shape == (480, 832, 3) for frame in out)
|
|
assert all(frame.dtype == np.uint8 for frame in out)
|
|
expected = upsample_frames(frames, width=832, height=480, mode="lanczos", sharpen=0.0)
|
|
for got, want in zip(out, expected, strict=True):
|
|
np.testing.assert_array_equal(got, want)
|
|
# Inputs are not mutated.
|
|
assert all(frame.shape == (240, 416, 3) for frame in frames)
|
|
|
|
|
|
def test_apply_fast_spatial_upsample_applies_sharpen() -> None:
|
|
frames = _frames(1, 240, 416)
|
|
soft = plan_fast_spatial(height=480, width=832, num_frames=81, spatial_scale=2, sharpen=0.0)
|
|
crisp = plan_fast_spatial(height=480, width=832, num_frames=81, spatial_scale=2, sharpen=0.8)
|
|
assert not np.array_equal(
|
|
apply_fast_spatial_upsample(frames, soft)[0],
|
|
apply_fast_spatial_upsample(frames, crisp)[0],
|
|
)
|
|
|
|
|
|
def test_apply_fast_spatial_noop_when_disabled() -> None:
|
|
frames = _frames(2, 32, 32, seed=1)
|
|
spatial = plan_fast_spatial(height=32, width=32, num_frames=5, enabled=False, patch_size=(1, 1, 1))
|
|
out = apply_fast_spatial_upsample(frames, spatial)
|
|
for got, want in zip(out, frames, strict=True):
|
|
np.testing.assert_array_equal(got, want)
|
|
|
|
|
|
def test_plan_fast_spatial_rejects_bad_mode() -> None:
|
|
with pytest.raises(ValueError, match="upsample mode"):
|
|
plan_fast_spatial(height=480, width=832, num_frames=81, upsample_mode="bicubic")
|
|
|
|
|
|
def test_plan_fast_spatial_rejects_negative_sharpen() -> None:
|
|
with pytest.raises(ValueError, match="sharpen"):
|
|
plan_fast_spatial(height=480, width=832, num_frames=81, sharpen=-0.1)
|
|
|
|
|
|
def test_geometry_errors_name_fast_spatial_not_refine() -> None:
|
|
"""Both modes share the resolution splitter; the error must name the caller.
|
|
|
|
832 is not divisible by 3, so scale 3 is rejected at plan time.
|
|
"""
|
|
with pytest.raises(ValueError, match=r"^fast-spatial requires"):
|
|
plan_fast_spatial(height=480, width=832, num_frames=81, spatial_scale=3)
|