110 lines
5.1 KiB
Python
110 lines
5.1 KiB
Python
import pytest
|
|
|
|
from modules.output.interpolate import framegen
|
|
from modules.output.render import upscale
|
|
from modules.output.render.tests.clips import comfy, make_clip, needs_ffmpeg # noqa: F401
|
|
|
|
torch = pytest.importorskip("torch")
|
|
|
|
|
|
def test_scaled_fps_keeps_the_fraction():
|
|
assert upscale.scaled_fps("24/1", 2) == "48/1"
|
|
assert upscale.scaled_fps("30000/1001", 2) == "60000/1001"
|
|
assert upscale.scaled_fps("25", 3) == "75/1"
|
|
|
|
|
|
@needs_ffmpeg
|
|
def test_interpolation_stream_adds_frames_and_keeps_length_and_sound(comfy):
|
|
for mult in (2, 3, 4): # the length is source x factor at every factor, not short by (m-1)/m of a frame
|
|
src, dst = str(comfy.output / f"a{mult}.mp4"), str(comfy.output / f"fi{mult}.mp4")
|
|
make_clip(src, seconds=1.0, size="64x48", rate=10, audio=True)
|
|
last = [None]
|
|
|
|
def between(x, mult=mult, last=last): # stands in for the effect: (mult-1) in-betweens per pair, carried across chunks
|
|
out = []
|
|
for f in x:
|
|
if last[0] is not None:
|
|
out += [(last[0] * (mult - i) + f * i) / mult for i in range(1, mult)]
|
|
out.append(f)
|
|
last[0] = f
|
|
return torch.stack(out)
|
|
|
|
assert upscale.upscale_file(src, dst, between, fps_factor=mult) == 10 * mult
|
|
w, h, fps, sound = upscale.probe(dst)
|
|
assert (w, h, fps, sound) == (64, 48, f"{10 * mult}/1", True)
|
|
|
|
|
|
def test_frame_generation_options_are_checked_by_name():
|
|
assert framegen.options("{}") == framegen.DEFAULTS
|
|
assert framegen.options({"multiplier": 4, "mode": "HIGH"})["multiplier"] == 4
|
|
for bad in ({"mode": "ULTRA"}, {"multiplier": 1}, {"multiplier": 5}, {"multiplier": 2.5}, {"multiplier": True}, {"tenbit": 1}):
|
|
with pytest.raises(ValueError):
|
|
framegen.options(bad)
|
|
with pytest.raises(ValueError):
|
|
framegen.options("not json")
|
|
|
|
|
|
def test_frame_generation_without_the_add_on_says_how_to_get_it(monkeypatch):
|
|
monkeypatch.setitem(__import__("sys").modules, "nvvfx", None)
|
|
with pytest.raises(RuntimeError, match="Settings ▸ Modules"):
|
|
framegen.make(framegen.options("{}"))
|
|
|
|
|
|
def test_frame_generation_feeds_pairs_in_order_across_chunks(monkeypatch):
|
|
"""A stand-in nvvfx (CPU) records the calls: every pair once, in order, indexes 1..N-1, the seam between chunks included."""
|
|
import sys
|
|
import types
|
|
calls, loads = [], []
|
|
|
|
class Effect:
|
|
class Mode: LOW = 0; MEDIUM = 1; HIGH = 2
|
|
class ImageEncoding: RGB8 = 0; RGB10A2 = 1
|
|
def __init__(self, mode, **kw): self.mode = mode
|
|
def load(self): loads.append(1)
|
|
def run(self, a, b, frame_index=1):
|
|
calls.append((float(a.mean()), float(b.mean()), frame_index))
|
|
return types.SimpleNamespace(image=(a + b) / 2)
|
|
|
|
from core import nvidia
|
|
monkeypatch.setattr(nvidia, "_held", {})
|
|
monkeypatch.setitem(sys.modules, "nvvfx", types.SimpleNamespace(VideoFrameGeneration=Effect))
|
|
from core import addons
|
|
monkeypatch.setattr(addons, "enabled", lambda _id: True)
|
|
monkeypatch.setattr(torch.Tensor, "to", lambda self, *a, **k: self) # no GPU here: "cuda" is this device
|
|
monkeypatch.setattr(torch, "from_dlpack", lambda t: t)
|
|
fn = framegen.make(framegen.options({"multiplier": 3}))
|
|
frames = torch.stack([torch.full((2, 2, 3), v) for v in (0.0, 0.3, 0.6, 0.9)])
|
|
a, b = fn(frames[:2]), fn(frames[2:])
|
|
assert a.shape[0] == 4 and b.shape[0] == 6 # 1 + 3, then 3 + 3: the seam pair got its frames too
|
|
assert [round(c[1], 2) for c in calls[::2]] == [0.3, 0.6, 0.9] and [c[2] for c in calls] == [1, 2] * 3
|
|
assert calls[2][0] == pytest.approx(0.3) and calls[2][1] == pytest.approx(0.6) # the seam pair was interpolated
|
|
again = framegen.make(framegen.options({"multiplier": 3})) # a second job: NVIDIA's model stays loaded
|
|
again(frames[:2])
|
|
assert len(loads) == 1
|
|
|
|
|
|
@needs_ffmpeg
|
|
def test_the_job_writes_a_doubled_file_beside_the_original(comfy, monkeypatch):
|
|
import sys
|
|
import types
|
|
from modules.output.interpolate import nodes
|
|
make_clip(str(comfy.output / "a.mp4"), seconds=1.0, size="64x48", rate=10, audio=True)
|
|
monkeypatch.setitem(sys.modules, "folder_paths", types.SimpleNamespace(get_output_directory=lambda: str(comfy.output)))
|
|
monkeypatch.setattr(nodes.video_pipe, "comfy_dir", lambda kind: str(comfy.output))
|
|
last = [None]
|
|
|
|
def between(x):
|
|
out = []
|
|
for f in x:
|
|
if last[0] is not None:
|
|
out.append((last[0] + f) / 2)
|
|
out.append(f)
|
|
last[0] = f
|
|
return torch.stack(out)
|
|
monkeypatch.setattr(nodes.framegen, "make", lambda o: between)
|
|
got = nodes.interpolate({"filename": "a.mp4", "subfolder": "", "type": "output"}, {})["media"]
|
|
assert got == {"filename": "a_x2.mp4", "subfolder": "funui_interpolated", "type": "output"}
|
|
assert upscale.probe(str(comfy.output / "funui_interpolated" / "a_x2.mp4")) == (64, 48, "20/1", True)
|
|
with pytest.raises(nodes.InterpolateError, match="not in the output folder"):
|
|
nodes.interpolate({"filename": "../x.mp4"}, {})
|