ausboss
5b8cad57d7
Image Crop + Rotate + Pad, Load Image + Pad and Load Video refused a
picture named the way ComfyUI names a generated file
("photo.png [output]", "[temp]") with "The selected source is outside
ComfyUI's input folder." Core Load Image takes those names.
resolve_input_path now reads the name the same way the video nodes'
Outputs list already does (resolve_server_file): the tag picks the
input, output or temp folder, and the rest is a name below it. A whole
path, "~", "." and "..", and a network path are refused from the text
alone; a link that leads out of the folder is refused once resolved.
Both functions share that code. A whole path inside the input folder,
accepted before, is now refused too, as the security rule asks.
Load Video's player and the shared upstream preview stripped the tag
and asked /view for the input folder; they now keep the folder.
Tests cover input, [input], [output] and [temp] names through the
resolver and through ComfyUI's own validation next to core Load Image,
and the refused escapes. Help pages, SECURITY.md and CHANGELOG say so.
578 lines
24 KiB
Python
578 lines
24 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import inspect
|
|
import math
|
|
import os
|
|
import tempfile
|
|
import unittest
|
|
from fractions import Fraction
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
import sys
|
|
|
|
import numpy as np
|
|
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"]
|
|
|
|
import av
|
|
|
|
from nodes._video_load_helpers import (
|
|
LazyAudio,
|
|
core_trim_args,
|
|
core_trimmed_video,
|
|
decode_audio_range,
|
|
decode_video_range,
|
|
effective_load_args,
|
|
lazy_audio_range,
|
|
memory_budget_error,
|
|
output_size,
|
|
trim_window,
|
|
)
|
|
from nodes import _video_load_helpers, node_load_video
|
|
|
|
|
|
def core_video_api_available() -> bool:
|
|
"""True when ComfyUI's comfy_api video types import in this interpreter.
|
|
|
|
ComfyUI goes into CPU mode first. The decode loop checks for a cancel
|
|
through comfy.model_management, which claims the GPU as it is imported;
|
|
these tests need no card, so CPU mode keeps them off it and lets them run
|
|
with the GPU hidden (CUDA_VISIBLE_DEVICES="").
|
|
"""
|
|
try:
|
|
import comfy.cli_args
|
|
|
|
comfy.cli_args.args.cpu = True
|
|
import comfy_api.input_impl # noqa: F401
|
|
except Exception:
|
|
return False
|
|
return True
|
|
|
|
|
|
def ensure_core_video_api(case: unittest.TestCase | None = None) -> bool:
|
|
"""Make comfy_api importable when AUSBOSS_COMFY_ROOT points at a ComfyUI
|
|
checkout; appended (never prepended) so ComfyUI's top-level nodes.py can
|
|
never shadow this pack's nodes package.
|
|
|
|
``case`` registers the cleanup that takes the entry back out again. Left
|
|
on sys.path it changes the answer to "is ComfyUI importable?" for every
|
|
test that runs afterwards, which silently flips each `except ImportError`
|
|
fail-soft seam in the pack onto its other branch - the branch the offline
|
|
assertions are not written for.
|
|
"""
|
|
if core_video_api_available():
|
|
return True
|
|
root = os.environ.get("AUSBOSS_COMFY_ROOT", "")
|
|
if root and (Path(root) / "comfy_api").is_dir() and root not in sys.path:
|
|
sys.path.append(root)
|
|
if case is not None:
|
|
case.addClassCleanup(_drop_from_sys_path, root)
|
|
return core_video_api_available()
|
|
|
|
|
|
def _drop_from_sys_path(entry: str) -> None:
|
|
while entry in sys.path:
|
|
sys.path.remove(entry)
|
|
|
|
|
|
def run_node(result):
|
|
"""Run a node FUNCTION result that may be sync or a coroutine."""
|
|
return asyncio.run(result) if asyncio.iscoroutine(result) else result
|
|
|
|
|
|
# The timeline's windowSeconds (js/shared/timeline_math.mjs), transcribed:
|
|
# IN is the largest hundredth that still keeps `first`, OUT the smallest that
|
|
# still keeps `last` and drops the frame after it.
|
|
def frame_window_seconds(fps: float, count: int, first: int, last: int) -> tuple[float, float]:
|
|
epsilon = 1e-4
|
|
start = 0.0 if first == 0 else max(0.0, round(math.floor((first / fps + epsilon) / 0.01 + 1e-9) * 0.01, 6))
|
|
end = 0.0 if last >= count - 1 else round(math.ceil((last / fps + epsilon) / 0.01 - 1e-9) * 0.01, 6)
|
|
return start, end
|
|
|
|
|
|
class FakeInterrupt(BaseException):
|
|
"""Mirrors ComfyUI's InterruptProcessingException: not an Exception."""
|
|
|
|
|
|
FPS = 12
|
|
FRAMES = 24
|
|
WIDTH, HEIGHT = 64, 48
|
|
AUDIO_RATE = 32000
|
|
|
|
|
|
def write_test_video(path: Path, with_audio: bool) -> None:
|
|
with av.open(str(path), "w") as container:
|
|
video = container.add_stream("mpeg4", rate=FPS)
|
|
video.width, video.height = WIDTH, HEIGHT
|
|
video.pix_fmt = "yuv420p"
|
|
audio = container.add_stream("aac", rate=AUDIO_RATE) if with_audio else None
|
|
for index in range(FRAMES):
|
|
level = min(255, index * 10)
|
|
array = np.full((HEIGHT, WIDTH, 3), level, dtype=np.uint8)
|
|
frame = av.VideoFrame.from_ndarray(array, format="rgb24")
|
|
for packet in video.encode(frame):
|
|
container.mux(packet)
|
|
for packet in video.encode():
|
|
container.mux(packet)
|
|
if audio is not None:
|
|
samples = AUDIO_RATE * FRAMES // FPS
|
|
tone = (0.2 * np.sin(np.linspace(0, 880 * np.pi, samples))).astype(np.float32)
|
|
audio_frame = av.AudioFrame.from_ndarray(tone.reshape(1, -1), format="flt", layout="mono")
|
|
audio_frame.sample_rate = AUDIO_RATE
|
|
audio_frame.time_base = Fraction(1, AUDIO_RATE)
|
|
audio_frame.pts = 0
|
|
for packet in audio.encode(audio_frame):
|
|
container.mux(packet)
|
|
for packet in audio.encode():
|
|
container.mux(packet)
|
|
|
|
|
|
class VideoLoadHelperTests(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls._tmp = tempfile.TemporaryDirectory()
|
|
cls.video = Path(cls._tmp.name) / "clip.mp4"
|
|
cls.silent = Path(cls._tmp.name) / "silent.mp4"
|
|
write_test_video(cls.video, with_audio=True)
|
|
write_test_video(cls.silent, with_audio=False)
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
cls._tmp.cleanup()
|
|
|
|
def test_full_decode_shape_fps_and_brightness_ramp(self):
|
|
frames, fps = decode_video_range(self.video, 0.0, 0.0, 0, 0)
|
|
self.assertEqual(tuple(frames.shape), (FRAMES, HEIGHT, WIDTH, 3))
|
|
self.assertAlmostEqual(fps, FPS, places=3)
|
|
self.assertLess(float(frames[0].mean()), float(frames[-1].mean()))
|
|
|
|
def test_trim_window_selects_the_middle_second(self):
|
|
frames, fps = decode_video_range(self.video, 0.5, 1.5, 0, 0)
|
|
self.assertEqual(frames.shape[0], FPS)
|
|
# Frame at t=0.5 is index 6 (brightness 60/255); allow codec loss.
|
|
self.assertAlmostEqual(float(frames[0].mean()), 60 / 255, delta=0.08)
|
|
|
|
def test_every_nth_thins_the_batch(self):
|
|
frames, fps = decode_video_range(self.video, 0.0, 0.0, 0, 0, every_nth=2)
|
|
self.assertEqual(frames.shape[0], FRAMES // 2)
|
|
# The helper reports the source rate; the node divides it.
|
|
self.assertAlmostEqual(fps, FPS, places=3)
|
|
# Kept frames are the even source indices: brightness steps by 20/255.
|
|
self.assertAlmostEqual(
|
|
float(frames[1].mean()) - float(frames[0].mean()), 20 / 255, delta=0.05
|
|
)
|
|
|
|
def test_max_frames_caps_the_decode(self):
|
|
frames, _fps = decode_video_range(self.video, 0.0, 0.0, 0, 0, max_frames=5)
|
|
self.assertEqual(frames.shape[0], 5)
|
|
|
|
def test_every_nth_and_cap_compose(self):
|
|
frames, _fps = decode_video_range(
|
|
self.video, 0.0, 0.0, 0, 0, every_nth=3, max_frames=4
|
|
)
|
|
self.assertEqual(frames.shape[0], 4)
|
|
|
|
def test_single_custom_dimension_preserves_aspect(self):
|
|
frames, _fps = decode_video_range(self.video, 0.0, 0.0, 32, 0)
|
|
self.assertEqual((frames.shape[2], frames.shape[1]), (32, 24))
|
|
|
|
def test_audio_window_matches_duration(self):
|
|
result = decode_audio_range(self.video, 0.5, 1.5)
|
|
waveform, rate = result["waveform"], result["sample_rate"]
|
|
self.assertEqual(rate, AUDIO_RATE)
|
|
self.assertEqual(waveform.ndim, 3)
|
|
self.assertAlmostEqual(waveform.shape[-1] / rate, 1.0, delta=0.05)
|
|
self.assertGreater(float(waveform.abs().max()), 0.0)
|
|
|
|
def test_silent_fallback_for_video_without_audio_track(self):
|
|
result = decode_audio_range(self.silent, 0.0, 2.0)
|
|
self.assertEqual(float(result["waveform"].abs().max()), 0.0)
|
|
self.assertAlmostEqual(
|
|
result["waveform"].shape[-1] / result["sample_rate"], 2.0, delta=0.01
|
|
)
|
|
|
|
def test_output_size_rules(self):
|
|
self.assertEqual(output_size(640, 480, 0, 0), (640, 480))
|
|
self.assertEqual(output_size(640, 480, 320, 0), (320, 240))
|
|
self.assertEqual(output_size(640, 480, 0, 240), (320, 240))
|
|
self.assertEqual(output_size(640, 480, 100, 100), (100, 100))
|
|
|
|
def test_lazy_audio_defers_the_loader_until_first_key_read(self):
|
|
calls = []
|
|
|
|
def loader():
|
|
calls.append(1)
|
|
return {"waveform": torch.zeros((1, 1, 4)), "sample_rate": 22050}
|
|
|
|
audio = LazyAudio(loader)
|
|
self.assertEqual(calls, []) # construction must not decode
|
|
self.assertEqual(audio["sample_rate"], 22050)
|
|
self.assertEqual(len(calls), 1)
|
|
self.assertEqual(audio.get("sample_rate"), 22050)
|
|
self.assertEqual(set(audio), {"waveform", "sample_rate"})
|
|
self.assertEqual(len(audio), 2)
|
|
self.assertEqual(len(calls), 1) # cached: still a single decode
|
|
|
|
def test_lazy_audio_range_matches_the_eager_decode(self):
|
|
lazy = lazy_audio_range(self.video, 0.5, 1.5)
|
|
eager = decode_audio_range(self.video, 0.5, 1.5)
|
|
self.assertEqual(lazy["sample_rate"], eager["sample_rate"])
|
|
self.assertTrue(torch.equal(lazy["waveform"], eager["waveform"]))
|
|
|
|
def test_memory_budget_allows_batches_inside_the_budget(self):
|
|
# 24 frames at 64x48 = 24 * 48 * 64 * 12 bytes, far under 1 GB.
|
|
self.assertIsNone(memory_budget_error(24, 64, 48, 1_000_000_000))
|
|
|
|
def test_memory_budget_blocks_oversized_batches_with_a_clear_error(self):
|
|
# 3000 frames at 1920x1080 needs ~74.6 GB float32.
|
|
message = memory_budget_error(3000, 1920, 1080, 16_000_000_000)
|
|
self.assertIsNotNone(message)
|
|
self.assertIn("3000", message)
|
|
self.assertIn("74.6 GB", message)
|
|
self.assertIn("16.0 GB", message)
|
|
self.assertIn("custom_width", message)
|
|
self.assertIn("shorter", message)
|
|
|
|
def test_memory_budget_names_the_caller_and_only_its_inputs(self):
|
|
# Load Video's wording is unchanged; another caller names itself and
|
|
# suggests its own inputs instead of custom_width/custom_height.
|
|
self.assertEqual(
|
|
memory_budget_error(3000, 1920, 1080, 16_000_000_000),
|
|
"Load Video would need about 74.6 GB for 3000 frames at 1920x1080, "
|
|
"but only 16.0 GB of memory is available. Trim a shorter start/end "
|
|
"window or set custom_width/custom_height to shrink the frames.",
|
|
)
|
|
message = memory_budget_error(
|
|
3000, 1920, 1080, 16_000_000_000, source="Clip", advice="Raise every_nth."
|
|
)
|
|
self.assertEqual(
|
|
message,
|
|
"Clip would need about 74.6 GB for 3000 frames at 1920x1080, but only "
|
|
"16.0 GB of memory is available. Raise every_nth.",
|
|
)
|
|
|
|
def test_memory_budget_applies_the_safety_factor(self):
|
|
needed = 10 * 8 * 8 * 3 * 4
|
|
self.assertIsNone(memory_budget_error(10, 8, 8, needed * 2))
|
|
self.assertIsNotNone(memory_budget_error(10, 8, 8, needed)) # 0.8 * needed < needed
|
|
|
|
def test_memory_budget_fails_soft_without_availability_data(self):
|
|
self.assertIsNone(memory_budget_error(10**9, 4096, 4096, None))
|
|
self.assertIsNone(memory_budget_error(10**9, 4096, 4096, 0))
|
|
self.assertIsNone(memory_budget_error(0, 4096, 4096, 1))
|
|
self.assertIsNone(memory_budget_error(10, 0, 0, 1))
|
|
|
|
def test_trim_validation_errors(self):
|
|
with self.assertRaisesRegex(ValueError, "smaller than"):
|
|
trim_window(10.0, 5.0, 5.0)
|
|
with self.assertRaisesRegex(ValueError, "only"):
|
|
trim_window(2.0, 3.0, 0.0)
|
|
|
|
def test_trim_errors_name_the_calling_node(self):
|
|
with self.assertRaisesRegex(ValueError, "^Load Video needs start_seconds smaller"):
|
|
trim_window(10.0, 5.0, 5.0)
|
|
with self.assertRaisesRegex(ValueError, "^Clip needs start_seconds smaller"):
|
|
trim_window(10.0, 5.0, 5.0, "Clip")
|
|
with self.assertRaisesRegex(ValueError, r"^Clip starts at 3\.00s"):
|
|
trim_window(2.0, 3.0, 0.0, "Clip")
|
|
|
|
|
|
class TrimFrameWindowTests(unittest.TestCase):
|
|
"""The timeline stores hundredth-grid seconds for a frame window; the
|
|
decode must keep exactly those frames."""
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls._tmp = tempfile.TemporaryDirectory()
|
|
cls.video = Path(cls._tmp.name) / "window.mp4"
|
|
write_test_video(cls.video, with_audio=False)
|
|
whole, fps = decode_video_range(cls.video, 0.0, 0.0, 0, 0)
|
|
cls.whole, cls.fps = whole, fps
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
cls._tmp.cleanup()
|
|
|
|
def test_grid_seconds_keep_exactly_the_chosen_frames(self):
|
|
count = int(self.whole.shape[0])
|
|
self.assertEqual(count, FRAMES)
|
|
for first, last in [(0, 0), (0, 5), (1, 1), (3, 10), (7, 22), (11, 23), (23, 23), (5, 23)]:
|
|
with self.subTest(first=first, last=last):
|
|
start, end = frame_window_seconds(self.fps, count, first, last)
|
|
batch, _ = decode_video_range(self.video, start, end, 0, 0)
|
|
self.assertEqual(int(batch.shape[0]), last - first + 1)
|
|
self.assertTrue(torch.equal(batch[0], self.whole[first]))
|
|
self.assertTrue(torch.equal(batch[-1], self.whole[last]))
|
|
|
|
|
|
class EffectiveLoadArgsTests(unittest.TestCase):
|
|
def test_trim_mode_passes_through_sanitized(self):
|
|
self.assertEqual(effective_load_args(False, 2.5, 2, 10), (2.5, 2, 10))
|
|
self.assertEqual(effective_load_args(False, 0.0, 0, -3), (0.0, 1, 0))
|
|
|
|
def test_single_frame_overrides_to_a_one_frame_open_window(self):
|
|
self.assertEqual(effective_load_args(True, 7.0, 4, 99), (0.0, 1, 1))
|
|
|
|
|
|
class CoreTrimArgsTests(unittest.TestCase):
|
|
def test_zeroes_pass_through_as_no_trim(self):
|
|
self.assertEqual(core_trim_args(0.0, 0.0), (0.0, 0.0))
|
|
|
|
def test_start_and_end_become_start_and_duration(self):
|
|
self.assertEqual(core_trim_args(0.5, 1.5), (0.5, 1.0))
|
|
|
|
def test_end_zero_means_until_the_end(self):
|
|
self.assertEqual(core_trim_args(2.0, 0.0), (2.0, 0.0))
|
|
|
|
def test_negative_start_clamps_to_zero(self):
|
|
self.assertEqual(core_trim_args(-1.0, 3.0), (0.0, 3.0))
|
|
|
|
def test_returns_none_when_the_core_api_is_unavailable(self):
|
|
# Forced, not inferred from the interpreter: reading the real import
|
|
# state made this assertion depend on whether an earlier class had
|
|
# already put ComfyUI on sys.path, so it self-skipped under any
|
|
# runner that reordered the classes and covered nothing.
|
|
with patch.dict(sys.modules, {"comfy_api.input_impl": None}):
|
|
self.assertIsNone(core_trimmed_video(Path("/tmp/never.mp4"), 0.0, 0.0))
|
|
|
|
|
|
class CoreVideoAdapterIntegrationTests(unittest.TestCase):
|
|
"""Constructs the core VIDEO adapter against a tiny generated clip.
|
|
|
|
Runs only when comfy_api is importable — either natively or via
|
|
AUSBOSS_COMFY_ROOT pointing at a ComfyUI checkout."""
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
if not ensure_core_video_api(cls):
|
|
raise unittest.SkipTest(
|
|
"comfy_api is not importable; set AUSBOSS_COMFY_ROOT to a ComfyUI checkout"
|
|
)
|
|
cls._tmp = tempfile.TemporaryDirectory()
|
|
cls.video = Path(cls._tmp.name) / "clip.mp4"
|
|
write_test_video(cls.video, with_audio=True)
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
cls._tmp.cleanup()
|
|
|
|
def test_trimmed_adapter_reads_metadata_without_decoding_frames(self):
|
|
video = core_trimmed_video(self.video, 0.5, 1.5)
|
|
self.assertIsNotNone(video)
|
|
# get_components materializes every frame; blocking it proves the
|
|
# metadata reads stay decode-free.
|
|
with patch.object(
|
|
type(video), "get_components", side_effect=AssertionError("decoded all frames")
|
|
):
|
|
self.assertEqual(video.get_dimensions(), (WIDTH, HEIGHT))
|
|
self.assertAlmostEqual(video.get_duration(), 1.0, delta=0.15)
|
|
self.assertAlmostEqual(float(video.get_frame_rate()), FPS, delta=0.01)
|
|
|
|
def test_trimmed_adapter_components_cover_the_requested_window(self):
|
|
video = core_trimmed_video(self.video, 0.5, 1.5)
|
|
components = video.get_components()
|
|
self.assertAlmostEqual(int(components.images.shape[0]), FPS, delta=2)
|
|
self.assertEqual(
|
|
(int(components.images.shape[2]), int(components.images.shape[1])),
|
|
(WIDTH, HEIGHT),
|
|
)
|
|
self.assertAlmostEqual(float(components.frame_rate), FPS, delta=0.01)
|
|
|
|
def test_untrimmed_adapter_spans_the_whole_clip(self):
|
|
video = core_trimmed_video(self.video, 0.0, 0.0)
|
|
self.assertAlmostEqual(video.get_duration(), FRAMES / FPS, delta=0.15)
|
|
|
|
|
|
class LoadVideoNodeTests(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls._tmp = tempfile.TemporaryDirectory()
|
|
cls.video = Path(cls._tmp.name) / "clip.mp4"
|
|
write_test_video(cls.video, with_audio=True)
|
|
# The loader reads a file by its name in ComfyUI's input folder; the
|
|
# fixture's folder stands in for it.
|
|
from types import SimpleNamespace
|
|
|
|
from nodes import _media_helpers
|
|
|
|
root = cls._tmp.name
|
|
stand_in = SimpleNamespace(
|
|
get_input_directory=lambda: root,
|
|
get_annotated_filepath=lambda name: os.path.join(root, name),
|
|
)
|
|
cls._folders = patch.object(_media_helpers, "folder_paths", stand_in)
|
|
cls._folders.start()
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
cls._folders.stop()
|
|
cls._tmp.cleanup()
|
|
|
|
def test_video_output_is_appended_after_the_original_seven(self):
|
|
node = node_load_video.AusBossLoadVideo
|
|
self.assertEqual(
|
|
node.RETURN_TYPES,
|
|
("IMAGE", "AUDIO", "INT", "FLOAT", "INT", "INT", "FLOAT", "VIDEO"),
|
|
)
|
|
self.assertEqual(
|
|
node.RETURN_NAMES,
|
|
("frames", "audio", "frame_count", "fps", "width", "height", "duration", "video"),
|
|
)
|
|
self.assertEqual(len(node.OUTPUT_TOOLTIPS), len(node.RETURN_TYPES))
|
|
|
|
def test_node_returns_the_trimmed_window_plus_a_core_video(self):
|
|
result = run_node(
|
|
node_load_video.AusBossLoadVideo().load_video("clip.mp4", 0.5, 1.5, 0, 0)
|
|
)
|
|
self.assertEqual(len(result), 8)
|
|
frames, _audio, frame_count, fps, width, height, duration, core_video = result
|
|
self.assertEqual(frame_count, int(frames.shape[0]))
|
|
self.assertEqual((width, height), (WIDTH, HEIGHT))
|
|
self.assertAlmostEqual(fps, FPS, places=3)
|
|
self.assertAlmostEqual(duration, 1.0, delta=0.2)
|
|
if core_video_api_available():
|
|
self.assertAlmostEqual(core_video.get_duration(), 1.0, delta=0.15)
|
|
else:
|
|
self.assertIsNone(core_video)
|
|
|
|
def test_single_frame_returns_one_frame_at_the_in_time(self):
|
|
# end_seconds deliberately sits before start_seconds: a stale trim
|
|
# window left over from trim mode must not affect a single-frame load.
|
|
result = run_node(
|
|
node_load_video.AusBossLoadVideo().load_video(
|
|
"clip.mp4", 0.5, 0.2, 0, 0, single_frame=True
|
|
)
|
|
)
|
|
frames, _audio, frame_count, fps, _width, _height, duration, _core = result
|
|
self.assertEqual(frame_count, 1)
|
|
self.assertEqual(int(frames.shape[0]), 1)
|
|
# The brightness ramp puts frame t=0.5 (index 6) at 60/255.
|
|
self.assertAlmostEqual(float(frames[0].mean()), 60 / 255, delta=0.08)
|
|
self.assertAlmostEqual(fps, FPS, places=3)
|
|
self.assertAlmostEqual(duration, 1 / FPS, delta=0.02)
|
|
|
|
def test_validation_reads_only_the_source(self):
|
|
# ComfyUI files a failed check once per input VALIDATE_INPUTS names,
|
|
# so naming the trim bounds too turned one missing file into four
|
|
# errors. The window is checked when the loader runs, where a wired
|
|
# bound has its value and single_frame applies.
|
|
node = node_load_video.AusBossLoadVideo
|
|
self.assertEqual(list(inspect.signature(node.VALIDATE_INPUTS).parameters), ["video"])
|
|
with patch.object(node_load_video, "resolve_input_path", lambda _name: self.video):
|
|
self.assertIs(node.VALIDATE_INPUTS("clip.mp4"), True)
|
|
gone = ValueError("The selected source file no longer exists.")
|
|
with patch.object(node_load_video, "resolve_input_path", side_effect=gone):
|
|
self.assertEqual(node.VALIDATE_INPUTS("gone.mp4"), "Load Video: The selected source file no longer exists.")
|
|
|
|
def test_a_backwards_trim_window_stops_the_loader_when_it_runs(self):
|
|
with self.assertRaisesRegex(ValueError, "Load Video needs start_seconds smaller than end_seconds"):
|
|
run_node(node_load_video.AusBossLoadVideo().load_video("clip.mp4", 1.5, 0.5, 0, 0))
|
|
|
|
def test_the_node_function_is_a_coroutine(self):
|
|
node = node_load_video.AusBossLoadVideo
|
|
self.assertTrue(inspect.iscoroutinefunction(getattr(node, node.FUNCTION)))
|
|
|
|
def test_the_event_loop_keeps_running_while_the_decode_blocks(self):
|
|
ticks = 0
|
|
|
|
async def drive():
|
|
nonlocal ticks
|
|
|
|
async def heartbeat():
|
|
nonlocal ticks
|
|
while True:
|
|
ticks += 1
|
|
await asyncio.sleep(0)
|
|
|
|
beat = asyncio.ensure_future(heartbeat())
|
|
try:
|
|
return await node_load_video.AusBossLoadVideo().load_video(
|
|
"clip.mp4", 0.0, 0.0, 0, 0
|
|
)
|
|
finally:
|
|
beat.cancel()
|
|
|
|
with patch.object(node_load_video, "resolve_input_path", lambda _name: self.video):
|
|
result = asyncio.run(drive())
|
|
|
|
self.assertEqual(int(result[2]), FRAMES)
|
|
# A blocking decode inside the coroutine would have starved the
|
|
# heartbeat entirely; off the loop it keeps being scheduled.
|
|
self.assertGreater(ticks, 10)
|
|
|
|
|
|
class RecordingBar:
|
|
def __init__(self):
|
|
self.updates = []
|
|
|
|
def update_absolute(self, value, total=None, preview=None):
|
|
self.updates.append((value, total, preview))
|
|
|
|
|
|
class DecodeLoopTests(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls._tmp = tempfile.TemporaryDirectory()
|
|
cls.video = Path(cls._tmp.name) / "clip.mp4"
|
|
write_test_video(cls.video, with_audio=False)
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
cls._tmp.cleanup()
|
|
|
|
def test_every_decoded_frame_advances_the_progress_bar(self):
|
|
bar = RecordingBar()
|
|
with patch.object(_video_load_helpers, "frame_progress", lambda _total: bar):
|
|
frames, _fps = decode_video_range(self.video, 0.0, 0.0, 0, 0)
|
|
|
|
count = int(frames.shape[0])
|
|
self.assertEqual(count, FRAMES)
|
|
self.assertEqual([update[0] for update in bar.updates], list(range(1, count + 1)))
|
|
self.assertTrue(all(update[1] >= count for update in bar.updates))
|
|
# The in-node player already previews; progress carries no images.
|
|
self.assertTrue(all(update[2] is None for update in bar.updates))
|
|
|
|
def test_an_unknown_frame_total_skips_tracking_without_a_bar(self):
|
|
totals = []
|
|
with patch.object(
|
|
_video_load_helpers, "frame_progress", lambda total: totals.append(total) or None
|
|
):
|
|
frames, _fps = decode_video_range(self.video, 0.0, 0.0, 0, 0)
|
|
self.assertEqual(int(frames.shape[0]), FRAMES)
|
|
self.assertEqual(len(totals), 1)
|
|
|
|
def test_a_broken_progress_bar_never_breaks_the_decode(self):
|
|
class BrokenBar:
|
|
def update_absolute(self, *_args, **_kwargs):
|
|
raise RuntimeError("the websocket went away")
|
|
|
|
with patch.object(_video_load_helpers, "frame_progress", lambda _total: BrokenBar()):
|
|
frames, _fps = decode_video_range(self.video, 0.0, 0.0, 0, 0)
|
|
self.assertEqual(int(frames.shape[0]), FRAMES)
|
|
|
|
def test_the_decode_loop_aborts_promptly_on_an_interrupt(self):
|
|
calls = []
|
|
|
|
def interrupt_on_the_third_frame():
|
|
calls.append(1)
|
|
if len(calls) > 2:
|
|
raise FakeInterrupt()
|
|
|
|
with patch.object(
|
|
_video_load_helpers, "raise_if_interrupted", interrupt_on_the_third_frame
|
|
):
|
|
with self.assertRaises(FakeInterrupt):
|
|
decode_video_range(self.video, 0.0, 0.0, 0, 0)
|
|
|
|
# Stopped on the third frame instead of decoding all 24.
|
|
self.assertEqual(len(calls), 3)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|