47 lines
1.3 KiB
Python
47 lines
1.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Tiny MLX RIFE frame-interpolation smoke test."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import time
|
|
|
|
import numpy as np
|
|
|
|
from fastvideo.mlx_runtime.rife_interp import interpolate, load_model
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(
|
|
description="MLX RIFE 4.25 frame interpolation smoke test."
|
|
)
|
|
parser.add_argument(
|
|
"--self-test",
|
|
action="store_true",
|
|
help="Run a tiny two-frame interpolation test.",
|
|
)
|
|
args = parser.parse_args()
|
|
if not args.self_test:
|
|
raise SystemExit("Nothing to do; pass --self-test")
|
|
|
|
frame0 = np.zeros((64, 96, 3), dtype=np.uint8)
|
|
frame1 = np.zeros((64, 96, 3), dtype=np.uint8)
|
|
frame1[:, :, 0] = 255
|
|
start = time.perf_counter()
|
|
model = load_model()
|
|
load_s = time.perf_counter() - start
|
|
start = time.perf_counter()
|
|
frames = interpolate([frame0, frame1], factor=2, model=model)
|
|
interp_s = time.perf_counter() - start
|
|
assert len(frames) == 3
|
|
assert frames[1].shape == frame0.shape
|
|
assert frames[1].dtype == np.uint8
|
|
print(
|
|
"MLX RIFE self-test passed: "
|
|
f"load_s={load_s:.3f} interp_s={interp_s:.3f} shape={frames[1].shape}"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|