386 lines
16 KiB
Python
386 lines
16 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
import sys
|
|
from types import SimpleNamespace
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
import torch
|
|
|
|
ROOT = Path(__file__).resolve().parent.parent
|
|
sys.path.insert(0, str(ROOT))
|
|
if "nodes" in sys.modules and not hasattr(sys.modules["nodes"], "__path__"):
|
|
del sys.modules["nodes"]
|
|
|
|
from nodes._interpolate_helpers import (
|
|
BLEND_METHOD,
|
|
FLOW_METHOD,
|
|
FrameJob,
|
|
_execute_jobs,
|
|
_pair_grouped_batches,
|
|
adjacent_frame_differences,
|
|
apply_scene_cuts,
|
|
backward_warp,
|
|
detect_scene_cuts,
|
|
interpolate_frames,
|
|
is_passthrough,
|
|
method_uses_optical_flow,
|
|
output_frame_count,
|
|
plan_frame_jobs,
|
|
raft_checkpoint_path,
|
|
weighted_blend,
|
|
)
|
|
|
|
CPU = torch.device("cpu")
|
|
|
|
|
|
def solid_batch(values: list[float]) -> torch.Tensor:
|
|
"""A BHWC batch where frame i is filled with values[i]."""
|
|
batch = torch.zeros((len(values), 4, 4, 3), dtype=torch.float32)
|
|
for index, value in enumerate(values):
|
|
batch[index] = value
|
|
return batch
|
|
|
|
|
|
class PlannerTests(unittest.TestCase):
|
|
def test_24_to_30_positions(self):
|
|
jobs = plan_frame_jobs(24, 24.0, 30.0)
|
|
self.assertEqual(len(jobs), 30)
|
|
# Every fifth output frame lands exactly on a source frame.
|
|
for job in jobs:
|
|
position = job.output_index * 24.0 / 30.0
|
|
if job.output_index % 5 == 0:
|
|
self.assertTrue(job.is_copy)
|
|
self.assertEqual(job.src_a, round(position))
|
|
elif job.output_index < 29:
|
|
self.assertFalse(job.is_copy)
|
|
self.assertEqual(job.src_a, int(position))
|
|
self.assertEqual(job.src_b, job.src_a + 1)
|
|
self.assertAlmostEqual(job.t, position - int(position), places=6)
|
|
# The tail position (23.2) is past the last pair and holds frame 23.
|
|
self.assertEqual(jobs[29], FrameJob(29, 23, 23, 0.0))
|
|
|
|
def test_30_to_24_positions(self):
|
|
jobs = plan_frame_jobs(30, 30.0, 24.0)
|
|
self.assertEqual(len(jobs), 24)
|
|
self.assertEqual(jobs[0], FrameJob(0, 0, 0, 0.0))
|
|
self.assertEqual(jobs[1].src_a, 1)
|
|
self.assertEqual(jobs[1].src_b, 2)
|
|
self.assertAlmostEqual(jobs[1].t, 0.25, places=6)
|
|
for job in jobs:
|
|
if job.output_index % 4 == 0:
|
|
self.assertTrue(job.is_copy)
|
|
self.assertEqual(job.src_a, job.output_index * 5 // 4)
|
|
else:
|
|
self.assertFalse(job.is_copy)
|
|
|
|
def test_single_frame_source_holds_it(self):
|
|
jobs = plan_frame_jobs(1, 1.0, 4.0)
|
|
self.assertEqual(len(jobs), 4)
|
|
for job in jobs:
|
|
self.assertEqual(job, FrameJob(job.output_index, 0, 0, 0.0))
|
|
|
|
def test_two_frames_upsampled_holds_the_tail(self):
|
|
jobs = plan_frame_jobs(2, 1.0, 4.0)
|
|
self.assertEqual(len(jobs), 8)
|
|
self.assertEqual(jobs[0], FrameJob(0, 0, 0, 0.0))
|
|
self.assertEqual(jobs[2], FrameJob(2, 0, 1, 0.5))
|
|
# Position 1.0 lands on the last frame; everything after holds it.
|
|
for job in jobs[4:]:
|
|
self.assertEqual((job.src_a, job.src_b, job.t), (1, 1, 0.0))
|
|
|
|
def test_direct_copies_never_reach_the_interpolation_path(self):
|
|
for jobs in (plan_frame_jobs(24, 24.0, 30.0), plan_frame_jobs(30, 30.0, 24.0)):
|
|
for job in jobs:
|
|
if job.t == 0.0 or job.src_a == job.src_b:
|
|
self.assertTrue(job.is_copy)
|
|
else:
|
|
self.assertTrue(0.0 < job.t < 1.0)
|
|
|
|
def test_output_frame_count_keeps_duration(self):
|
|
self.assertEqual(output_frame_count(24, 24.0, 30.0), 30)
|
|
self.assertEqual(output_frame_count(30, 30.0, 24.0), 24)
|
|
self.assertEqual(output_frame_count(10, 24.0, 30.0), 12)
|
|
self.assertEqual(output_frame_count(1, 30.0, 30.0), 1)
|
|
|
|
def test_bad_inputs_are_rejected(self):
|
|
with self.assertRaisesRegex(ValueError, "empty IMAGE batch"):
|
|
plan_frame_jobs(0, 24.0, 30.0)
|
|
with self.assertRaisesRegex(ValueError, "above zero"):
|
|
plan_frame_jobs(10, 0.0, 30.0)
|
|
with self.assertRaisesRegex(ValueError, "above zero"):
|
|
plan_frame_jobs(10, 24.0, -1.0)
|
|
|
|
|
|
class SceneCutTests(unittest.TestCase):
|
|
def test_adjacent_differences_measure_mean_abs_change(self):
|
|
frames = solid_batch([0.0, 0.1, 0.9])
|
|
differences = adjacent_frame_differences(frames, chunk_size=1)
|
|
self.assertEqual(len(differences), 2)
|
|
self.assertAlmostEqual(differences[0], 0.1, places=5)
|
|
self.assertAlmostEqual(differences[1], 0.8, places=5)
|
|
|
|
def test_detect_scene_cuts_thresholds_and_zero_disables(self):
|
|
differences = [0.1, 0.8, 0.2]
|
|
self.assertEqual(detect_scene_cuts(differences, 0.35), {1})
|
|
self.assertEqual(detect_scene_cuts(differences, 0.0), set())
|
|
self.assertEqual(detect_scene_cuts(differences, 0.9), set())
|
|
|
|
def test_interpolations_across_a_cut_become_holds(self):
|
|
jobs = [
|
|
FrameJob(0, 0, 0, 0.0),
|
|
FrameJob(1, 0, 1, 0.5),
|
|
FrameJob(2, 1, 2, 0.5),
|
|
FrameJob(3, 2, 2, 0.0),
|
|
]
|
|
adjusted = apply_scene_cuts(jobs, {1})
|
|
self.assertEqual(adjusted[0], jobs[0])
|
|
self.assertEqual(adjusted[1], jobs[1])
|
|
self.assertEqual(adjusted[2], FrameJob(2, 1, 1, 0.0))
|
|
self.assertEqual(adjusted[3], jobs[3])
|
|
|
|
def test_no_cuts_returns_jobs_unchanged(self):
|
|
jobs = [FrameJob(0, 0, 1, 0.5)]
|
|
self.assertEqual(apply_scene_cuts(jobs, set()), jobs)
|
|
|
|
|
|
class BlendAndWarpTests(unittest.TestCase):
|
|
def test_weighted_blend_midpoint(self):
|
|
a = torch.zeros((1, 2, 2, 3))
|
|
b = torch.ones((1, 2, 2, 3))
|
|
mid = weighted_blend(a, b, torch.tensor(0.5))
|
|
self.assertTrue(torch.allclose(mid, torch.full_like(mid, 0.5)))
|
|
|
|
def test_weighted_blend_endpoints(self):
|
|
a = torch.full((1, 2, 2, 3), 0.25)
|
|
b = torch.full((1, 2, 2, 3), 0.75)
|
|
self.assertTrue(torch.allclose(weighted_blend(a, b, torch.tensor(0.0)), a))
|
|
self.assertTrue(torch.allclose(weighted_blend(a, b, torch.tensor(1.0)), b))
|
|
|
|
def test_backward_warp_zero_flow_is_identity(self):
|
|
images = torch.rand((2, 3, 6, 8))
|
|
flow = torch.zeros((2, 2, 6, 8))
|
|
self.assertTrue(torch.allclose(backward_warp(images, flow), images, atol=1e-5))
|
|
|
|
def test_backward_warp_integer_shift_samples_the_neighbor(self):
|
|
images = torch.zeros((1, 1, 1, 4))
|
|
images[0, 0, 0] = torch.tensor([0.0, 1.0, 2.0, 3.0])
|
|
flow = torch.zeros((1, 2, 1, 4))
|
|
flow[:, 0] = 1.0 # sample one pixel to the right, border-padded
|
|
warped = backward_warp(images, flow)
|
|
expected = torch.tensor([1.0, 2.0, 3.0, 3.0])
|
|
self.assertTrue(torch.allclose(warped[0, 0, 0], expected, atol=1e-5))
|
|
|
|
|
|
class InterpolateFramesTests(unittest.TestCase):
|
|
def test_passthrough_early_out_returns_the_same_tensor(self):
|
|
frames = solid_batch([0.1, 0.2])
|
|
result, fps = interpolate_frames(frames, 24.0, 24.005, BLEND_METHOD, 0.35, 8, CPU)
|
|
self.assertIs(result, frames)
|
|
self.assertEqual(fps, 24.0)
|
|
|
|
def test_blend_midpoints_and_copies(self):
|
|
frames = solid_batch([0.0, 1.0])
|
|
result, fps = interpolate_frames(frames, 12.0, 24.0, BLEND_METHOD, 0.0, 8, CPU)
|
|
self.assertEqual(tuple(result.shape), (4, 4, 4, 3))
|
|
self.assertEqual(fps, 24.0)
|
|
self.assertTrue(torch.allclose(result[0], frames[0]))
|
|
self.assertTrue(torch.allclose(result[1], torch.full_like(result[1], 0.5)))
|
|
self.assertTrue(torch.allclose(result[2], frames[1]))
|
|
self.assertTrue(torch.allclose(result[3], frames[1]))
|
|
|
|
def test_scene_cut_holds_instead_of_morphing(self):
|
|
# Frames 0/1 are a gentle ramp; frame 2 is a hard cut to white.
|
|
frames = solid_batch([0.0, 0.1, 1.0])
|
|
result, _fps = interpolate_frames(frames, 12.0, 24.0, BLEND_METHOD, 0.35, 8, CPU)
|
|
self.assertEqual(int(result.shape[0]), 6)
|
|
self.assertTrue(torch.allclose(result[1], torch.full_like(result[1], 0.05)))
|
|
# The frame between 1 and 2 crosses the cut and holds frame 1.
|
|
self.assertTrue(torch.allclose(result[3], frames[1]))
|
|
self.assertTrue(torch.allclose(result[4], frames[2]))
|
|
|
|
def test_threshold_zero_blends_across_the_cut(self):
|
|
frames = solid_batch([0.0, 0.1, 1.0])
|
|
result, _fps = interpolate_frames(frames, 12.0, 24.0, BLEND_METHOD, 0.0, 8, CPU)
|
|
self.assertTrue(torch.allclose(result[3], torch.full_like(result[3], 0.55)))
|
|
|
|
def test_actual_fps_reports_the_rounded_rate(self):
|
|
frames = solid_batch([float(i) / 10.0 for i in range(10)])
|
|
result, fps = interpolate_frames(frames, 24.0, 30.0, BLEND_METHOD, 0.0, 8, CPU)
|
|
self.assertEqual(int(result.shape[0]), 12)
|
|
self.assertAlmostEqual(fps, 28.8, places=5)
|
|
|
|
def test_small_batch_size_matches_a_single_batch(self):
|
|
frames = solid_batch([0.0, 0.25, 0.5, 0.75, 1.0])
|
|
chunked, _ = interpolate_frames(frames, 24.0, 30.0, BLEND_METHOD, 0.0, 2, CPU)
|
|
whole, _ = interpolate_frames(frames, 24.0, 30.0, BLEND_METHOD, 0.0, 64, CPU)
|
|
self.assertTrue(torch.allclose(chunked, whole))
|
|
|
|
def test_method_names_are_validated(self):
|
|
self.assertFalse(method_uses_optical_flow(BLEND_METHOD))
|
|
self.assertTrue(method_uses_optical_flow(FLOW_METHOD))
|
|
with self.assertRaisesRegex(ValueError, "does not know the method"):
|
|
method_uses_optical_flow("cubic")
|
|
|
|
def test_bad_batches_are_rejected(self):
|
|
with self.assertRaisesRegex(ValueError, "BHWC"):
|
|
interpolate_frames(torch.zeros((2, 2, 3)), 24.0, 30.0, BLEND_METHOD, 0.0, 8, CPU)
|
|
with self.assertRaisesRegex(ValueError, "empty IMAGE batch"):
|
|
interpolate_frames(
|
|
torch.zeros((0, 2, 2, 3)), 24.0, 30.0, BLEND_METHOD, 0.0, 8, CPU
|
|
)
|
|
|
|
def test_is_passthrough_boundary(self):
|
|
self.assertTrue(is_passthrough(24.0, 24.0))
|
|
self.assertTrue(is_passthrough(24.0, 24.009))
|
|
self.assertFalse(is_passthrough(24.0, 24.02))
|
|
|
|
|
|
class WeightCacheTests(unittest.TestCase):
|
|
def test_checkpoint_path_uses_the_weights_filename_without_downloading(self):
|
|
weights = SimpleNamespace(
|
|
url="https://download.example/models/raft_small-01064c6d.pth"
|
|
)
|
|
self.assertEqual(
|
|
raft_checkpoint_path(weights, "/tmp/torch-hub"),
|
|
Path("/tmp/torch-hub/checkpoints/raft_small-01064c6d.pth"),
|
|
)
|
|
|
|
def test_checkpoint_path_rejects_a_missing_filename(self):
|
|
with self.assertRaisesRegex(RuntimeError, "checkpoint URL"):
|
|
raft_checkpoint_path(SimpleNamespace(url=""), "/tmp/torch-hub")
|
|
|
|
|
|
class FrameInterpolateNodeTests(unittest.TestCase):
|
|
def test_node_returns_frames_and_fps(self):
|
|
from nodes.node_frame_interpolate import AusBossFrameInterpolate
|
|
|
|
frames = solid_batch([0.0, 1.0])
|
|
images, fps = AusBossFrameInterpolate().interpolate(
|
|
frames, 12.0, 24.0, BLEND_METHOD, 0.35, 8
|
|
)
|
|
self.assertEqual(int(images.shape[0]), 4)
|
|
self.assertEqual(fps, 24.0)
|
|
|
|
|
|
|
|
|
|
class CopyPathIsBatchedTests(unittest.TestCase):
|
|
"""Copied frames go over in batches like the blends do.
|
|
|
|
Gathering every copy at once was the one allocation batch_size could not
|
|
bound: a long clip at an integer multiple copies most of its output in a
|
|
single index_select, on the device the input already lives on.
|
|
"""
|
|
|
|
def frames(self, count: int) -> torch.Tensor:
|
|
ramp = torch.linspace(0.0, 1.0, count).view(-1, 1, 1, 1)
|
|
return ramp.expand(count, 4, 4, 3).clone()
|
|
|
|
def test_the_result_is_identical_at_every_batch_size(self):
|
|
frames = self.frames(12)
|
|
want = interpolate_frames(frames, 24.0, 48.0, BLEND_METHOD, 0.0, 64, CPU)[0]
|
|
for batch in (1, 2, 3, 5, 64):
|
|
got = interpolate_frames(frames, 24.0, 48.0, BLEND_METHOD, 0.0, batch, CPU)[0]
|
|
self.assertTrue(torch.equal(got, want), f"batch_size={batch}")
|
|
|
|
def test_a_copy_only_plan_is_still_chunked(self):
|
|
# 30 -> 15 fps is pure decimation: every job is a copy, so this path
|
|
# is the only one that runs.
|
|
frames = self.frames(12)
|
|
want = interpolate_frames(frames, 30.0, 15.0, BLEND_METHOD, 0.0, 64, CPU)[0]
|
|
for batch in (1, 4):
|
|
got = interpolate_frames(frames, 30.0, 15.0, BLEND_METHOD, 0.0, batch, CPU)[0]
|
|
self.assertTrue(torch.equal(got, want), f"batch_size={batch}")
|
|
|
|
def test_no_single_gather_exceeds_the_batch_size(self):
|
|
seen = []
|
|
real = torch.Tensor.index_select
|
|
|
|
def spy(self, dim, index):
|
|
seen.append(int(index.numel()))
|
|
return real(self, dim, index)
|
|
|
|
frames = self.frames(24)
|
|
with patch.object(torch.Tensor, "index_select", spy):
|
|
interpolate_frames(frames, 24.0, 48.0, BLEND_METHOD, 0.0, 4, CPU)
|
|
self.assertTrue(seen)
|
|
self.assertLessEqual(max(seen), 8) # <= batch_size sources per gather
|
|
|
|
|
|
class FlowSolvedOncePerPairTests(unittest.TestCase):
|
|
"""Every (a, b) source pair is flow-solved exactly once for the whole
|
|
plan, whatever batch_size is - batches take whole pair groups and a
|
|
one-batch cache covers a pair that a tiny batch_size splits."""
|
|
|
|
class ContentFlowModel:
|
|
"""Stub RAFT whose flow DEPENDS on its inputs.
|
|
|
|
A zero-flow stub cannot catch a pair-association bug (identity
|
|
warps look right whichever pair's flow is used); a flow derived
|
|
from the frames themselves makes a mixed-up cache visible in the
|
|
output. Also counts the rows it is asked to solve.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.rows = 0
|
|
|
|
def __call__(self, first, second):
|
|
self.rows += int(first.shape[0])
|
|
shift = (first - second).mean(dim=(1, 2, 3), keepdim=True)
|
|
flow = torch.zeros((first.shape[0], 2, first.shape[2], first.shape[3]))
|
|
return [flow + shift * 3.0]
|
|
|
|
def run_plan(self, batch_size):
|
|
torch.manual_seed(7)
|
|
frames = torch.rand((12, 16, 16, 3))
|
|
jobs = apply_scene_cuts(plan_frame_jobs(12, 24.0, 120.0), set())
|
|
model = self.ContentFlowModel()
|
|
out = _execute_jobs(frames, jobs, batch_size, CPU, model)
|
|
blends = [job for job in jobs if not job.is_copy]
|
|
pairs = {(job.src_a, job.src_b) for job in blends}
|
|
return out, model.rows, 2 * len(pairs), len(blends)
|
|
|
|
def test_solves_equal_unique_pairs_at_every_batch_size(self):
|
|
for batch_size in (1, 2, 3, 8, 64):
|
|
_out, rows, ideal, blend_count = self.run_plan(batch_size)
|
|
with self.subTest(batch_size=batch_size):
|
|
self.assertEqual(rows, ideal)
|
|
# And the old failure mode really was bigger than this.
|
|
self.assertLess(ideal, 2 * blend_count)
|
|
|
|
def test_outputs_match_across_batch_sizes(self):
|
|
# Tight allclose, not torch.equal: the stub's own mean() reduction
|
|
# jitters at the last float32 bit (~1e-7) with batch layout, exactly
|
|
# as a real conv stack would. A pair association bug shows up at the
|
|
# shift magnitude - four orders of magnitude above this tolerance.
|
|
reference, _rows, _ideal, _n = self.run_plan(64)
|
|
for batch_size in (1, 2, 3, 8):
|
|
out, _rows, _ideal, _n = self.run_plan(batch_size)
|
|
with self.subTest(batch_size=batch_size):
|
|
self.assertTrue(torch.allclose(out, reference, atol=1e-5, rtol=0.0))
|
|
self.assertLess(float((out - reference).abs().max()), 1e-5)
|
|
|
|
def test_pair_groups_fill_batches_without_splitting_small_groups(self):
|
|
jobs = apply_scene_cuts(plan_frame_jobs(12, 24.0, 120.0), set())
|
|
blends = [job for job in jobs if not job.is_copy]
|
|
batches = _pair_grouped_batches(blends, 8)
|
|
self.assertEqual(sum(len(batch) for batch in batches), len(blends))
|
|
for batch in batches[:-1]:
|
|
self.assertLessEqual(len(batch), 8)
|
|
# A batch never holds a fragment of a pair another batch also holds,
|
|
# unless that pair alone is bigger than the step.
|
|
seen = {}
|
|
for index, batch in enumerate(batches):
|
|
for job in batch:
|
|
seen.setdefault((job.src_a, job.src_b), set()).add(index)
|
|
for pair, indices in seen.items():
|
|
jobs_for_pair = sum(1 for job in blends if (job.src_a, job.src_b) == pair)
|
|
if jobs_for_pair <= 8:
|
|
self.assertEqual(len(indices), 1, pair)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|