from __future__ import annotations import os from pathlib import Path import sys import tempfile import time import unittest import unittest.mock import av import numpy as np from PIL import Image 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._media_helpers import decode_video_frame, video_metadata from nodes._transform_engine import ( scale_to_megapixels, TransformSpec, stable_file_fingerprint, transform_pil, transform_tensor_batch, ) def solid(width=12, height=8, color=(30, 80, 140, 255)): return Image.new("RGBA", (width, height), color) class TransformEngineTests(unittest.TestCase): def test_identity_preserves_dimensions_and_pixels(self): source = solid() output, mask, geometry = transform_pil(source, TransformSpec()) self.assertEqual(output.size, source.size) self.assertEqual(geometry.output_width, 12) self.assertEqual(np.asarray(mask).max(), 0) self.assertTrue(np.all(np.asarray(output) == np.array([30, 80, 140]))) def test_source_alpha_is_part_of_mask(self): source = solid(5, 5) source.putpixel((2, 2), (255, 0, 0, 0)) _, mask, _ = transform_pil(source, TransformSpec()) self.assertEqual(mask.getpixel((2, 2)), 255) self.assertEqual(mask.getpixel((0, 0)), 0) def test_positive_and_negative_rotation_boundaries(self): for angle in (-90, 90, 179.9, -179.9): with self.subTest(angle=angle): output, mask, geometry = transform_pil(solid(11, 7), TransformSpec(rotation_degrees=angle)) self.assertEqual(output.size, mask.size) self.assertGreaterEqual(geometry.rotated_width, 7) self.assertGreaterEqual(geometry.rotated_height, 7) if abs(angle) % 90 > 0.01: self.assertGreater(np.asarray(mask).max(), 0) def test_crop_clamps_to_rotated_source(self): output, _, geometry = transform_pil( solid(20, 10), TransformSpec(crop_x=18, crop_y=9, crop_width=999, crop_height=999) ) self.assertEqual(output.size, (2, 1)) self.assertEqual((geometry.crop_x, geometry.crop_y), (18, 9)) def test_ratio_is_enforced_in_backend(self): output, _, _ = transform_pil( solid(100, 100), TransformSpec(crop_aspect_ratio="16:9") ) self.assertEqual(output.size, (100, 56)) def test_every_padding_side_changes_output_and_mask(self): cases = { "pad_left": ((7, 4), (0, 1)), "pad_top": ((5, 6), (1, 0)), "pad_right": ((7, 4), (6, 1)), "pad_bottom": ((5, 6), (1, 5)), } for field, (expected_size, generated_pixel) in cases.items(): with self.subTest(field=field): output, mask, _ = transform_pil(solid(5, 4), TransformSpec(**{field: 2})) self.assertEqual(output.size, expected_size) self.assertEqual(mask.getpixel(generated_pixel), 255) def test_fill_color_is_used_for_padding_and_rotation_voids(self): output, _, _ = transform_pil( solid(4, 3, (255, 0, 0, 255)), TransformSpec(rotation_degrees=45, pad_left=2, fill_color="#123456"), ) self.assertEqual(output.getpixel((0, 0)), (0x12, 0x34, 0x56)) def test_feather_keeps_generated_pixels_and_softens_boundary(self): _, mask, _ = transform_pil(solid(20, 20), TransformSpec(pad_left=10, feather=3)) values = np.asarray(mask) self.assertEqual(values[:, :5].min(), 255) self.assertTrue(np.any((values[:, 10:15] > 0) & (values[:, 10:15] < 255))) # No seam: the ramp is still near full strength on the first kept # column (the old max-with-blur approach dropped to ~127 here). self.assertGreaterEqual(int(values[10, 10]), 200) # And it decays moving further into kept content. self.assertLess(int(values[10, 16]), int(values[10, 11])) def test_feather_fades_image_into_fill(self): output, _, _ = transform_pil( solid(20, 20, (255, 255, 255, 255)), TransformSpec(pad_left=10, feather=3, fill_color="#000000"), ) pixels = np.asarray(output) row = pixels[10] self.assertLessEqual(int(row[9].max()), 40) # padding stays fill color self.assertTrue(40 < int(row[11][0]) < 240) # visible blend band self.assertGreaterEqual(int(row[25].min()), 250) # interior untouched def test_canvas_multiple_rounds_only_right_and_bottom(self): output, mask, geometry = transform_pil( solid(101, 99), TransformSpec(canvas_multiple=8) ) self.assertEqual(output.size, (104, 104)) self.assertEqual((geometry.pad_left, geometry.pad_top), (0, 0)) self.assertEqual((geometry.pad_right, geometry.pad_bottom), (3, 5)) self.assertEqual(mask.getpixel((103, 103)), 255) def test_tensor_batch_preserves_bhwc_and_bhw(self): source = torch.zeros((3, 9, 7, 3), dtype=torch.float32) output, mask, _ = transform_tensor_batch(source, TransformSpec(pad_bottom=1)) self.assertEqual(tuple(output.shape), (3, 10, 7, 3)) self.assertEqual(tuple(mask.shape), (3, 10, 7)) self.assertEqual(output.dtype, source.dtype) def test_invalid_shape_and_oversized_padding_are_actionable(self): with self.assertRaisesRegex(ValueError, "expected BHWC"): transform_tensor_batch(torch.zeros((4, 4, 3)), TransformSpec()) with self.assertRaisesRegex(ValueError, "pad_left"): transform_pil(solid(), TransformSpec(pad_left=40000)) def test_fingerprint_changes_with_inputs_and_file_state(self): with tempfile.TemporaryDirectory() as directory: path = Path(directory, "source.bin") path.write_bytes(b"one") first = stable_file_fingerprint(path, {"crop": 0}) second = stable_file_fingerprint(path, {"crop": 1}) self.assertNotEqual(first, second) time.sleep(0.01) path.write_bytes(b"two-two") os.utime(path, None) self.assertNotEqual(first, stable_file_fingerprint(path, {"crop": 0})) class VideoDecodeTests(unittest.TestCase): def setUp(self): self._tempdir = tempfile.TemporaryDirectory() self.directory = Path(self._tempdir.name) def tearDown(self): # Persistent scrub sessions hold the file open (locked on Windows); # release them before the temporary directory is removed. from nodes import _media_helpers _media_helpers.close_scrub_sessions() self._tempdir.cleanup() def _write_video(self, path: Path, frames: int = 5, step: int = 45): container = av.open(str(path), mode="w") stream = container.add_stream("mpeg4", rate=5) stream.width = 16 stream.height = 16 stream.pix_fmt = "yuv420p" for index in range(frames): array = np.zeros((16, 16, 3), dtype=np.uint8) array[..., 0] = (index * step) % 256 frame = av.VideoFrame.from_ndarray(array, format="rgb24") for packet in stream.encode(frame): container.mux(packet) for packet in stream.encode(): container.mux(packet) container.close() def test_first_middle_last_frame_decoding(self): path = Path(self.directory, "frames.mp4") self._write_video(path) metadata = video_metadata(path) self.assertEqual(metadata["width"], 16) self.assertGreaterEqual(metadata["frame_count"], 5) reds = [] for index in (0, 2, 4): frame, actual_index, _ = decode_video_frame(path, "frame index", index, 0.0) self.assertEqual(actual_index, index) reds.append(float(np.asarray(frame)[..., 0].mean())) self.assertLess(reds[0], reds[1]) self.assertLess(reds[1], reds[2]) def test_keyframe_seek_agrees_with_sequential_scan(self): # 64 frames at GOP defaults spans several keyframes, so mid and late # targets exercise the container.seek fast path against ground truth. from nodes._media_helpers import _decode_sequential path = Path(self.directory, "long.mp4") self._write_video(path, frames=64, step=4) fps = float(video_metadata(path)["fps"] or 0.0) self.assertGreater(fps, 0) for target in (0, 7, 33, 63): with self.subTest(target=target): fast_frame, fast_index, fast_time = decode_video_frame( path, "frame index", target, 0.0 ) slow_frame, slow_index, slow_time = _decode_sequential(path, target, fps) self.assertEqual(fast_index, slow_index) self.assertAlmostEqual(fast_time, slow_time, places=4) self.assertTrue( np.array_equal(np.asarray(fast_frame), np.asarray(slow_frame)) ) def test_time_seek_mode_lands_on_matching_frame(self): path = Path(self.directory, "timed.mp4") self._write_video(path, frames=25, step=10) # 5 fps: 2.0 seconds is exactly frame 10. _, actual_index, actual_time = decode_video_frame(path, "time seconds", 0, 2.0) self.assertEqual(actual_index, 10) self.assertAlmostEqual(actual_time, 2.0, places=3) def test_out_of_range_index_clamps_to_last_frame(self): path = Path(self.directory, "short.mp4") self._write_video(path, frames=6, step=40) _, actual_index, _ = decode_video_frame(path, "frame index", 999999, 0.0) self.assertEqual(actual_index, 5) def test_session_forward_and_backward_decode_match_fresh_decode(self): from nodes import _media_helpers path = Path(self.directory, "session.mp4") self._write_video(path, frames=40, step=6) fps = float(video_metadata(path)["fps"] or 0.0) decode_video_frame(path, "frame index", 5, 0.0) # Forward within the gap uses the persistent session's decoder # state; backward forces a re-seek on the same open container. for target in (9, 2): with self.subTest(target=target): fast = decode_video_frame(path, "frame index", target, 0.0) fresh = _media_helpers._decode_sequential(path, target, fps) self.assertEqual(fast[1], fresh[1]) self.assertTrue(np.array_equal(np.asarray(fast[0]), np.asarray(fresh[0]))) def test_repeating_a_frame_request_returns_that_frame(self): # A scrub that settles on the frame it already showed asks for the # same index twice. The session used to decode "forward" to it and # hand back the frame after - the picture jumped one frame on every # release. path = Path(self.directory, "repeat.mp4") self._write_video(path, frames=40, step=6) first = decode_video_frame(path, "frame index", 12, 0.0) again = decode_video_frame(path, "frame index", 12, 0.0) self.assertEqual((first[1], again[1]), (12, 12)) self.assertTrue(np.array_equal(np.asarray(first[0]), np.asarray(again[0]))) following = decode_video_frame(path, "frame index", 13, 0.0) self.assertEqual(following[1], 13) self.assertFalse(np.array_equal(np.asarray(first[0]), np.asarray(following[0]))) def test_preview_cache_serves_a_revisited_frame_without_decoding(self): from nodes import _media_helpers path = Path(self.directory, "revisit.mp4") self._write_video(path, frames=20, step=12) _media_helpers.clear_preview_cache() body, index, moment = _media_helpers.cached_preview(path, "frame index", 7, 0.0, 640, 640) self.assertEqual(index, 7) self.assertGreater(len(body), 0) with unittest.mock.patch.object( _media_helpers, "decode_video_frame", side_effect=AssertionError("decoded again") ): hit = _media_helpers.cached_preview(path, "frame index", 7, 0.0, 640, 640) # Time mode naming the same frame is the same cache entry. fps = float(video_metadata(path)["fps"]) by_time = _media_helpers.cached_preview(path, "time seconds", 0, 7 / fps, 640, 640) self.assertEqual(hit, (body, index, moment)) self.assertEqual(by_time, (body, index, moment)) # Another size is another entry, decoded on its own. with _media_helpers._PREVIEW_CACHE_LOCK: entries = len(_media_helpers._PREVIEW_CACHE) _media_helpers.cached_preview(path, "frame index", 7, 0.0, 320, 320) with _media_helpers._PREVIEW_CACHE_LOCK: self.assertEqual(len(_media_helpers._PREVIEW_CACHE), entries + 1) _media_helpers.clear_preview_cache() def test_preview_cache_forgets_a_rewritten_file(self): from nodes import _media_helpers path = Path(self.directory, "rewrite.mp4") self._write_video(path, frames=10, step=20) _media_helpers.clear_preview_cache() before = _media_helpers.cached_preview(path, "frame index", 3, 0.0, 640, 640) _media_helpers.close_scrub_sessions() self._write_video(path, frames=10, step=1) os.utime(path, (time.time() + 5, time.time() + 5)) after = _media_helpers.cached_preview(path, "frame index", 3, 0.0, 640, 640) self.assertNotEqual(before[0], after[0]) _media_helpers.clear_preview_cache() def test_storyboard_tiles_spread_over_the_clip(self): from nodes import _media_helpers path = Path(self.directory, "spread.mp4") # 40 frames at 5 fps: 8 s, so 16 tiles half a second apart. self._write_video(path, frames=40, step=6) key = _media_helpers._file_key(path) with _media_helpers._STORYBOARDS_LOCK: _media_helpers._STORYBOARDS[key] = {"status": "building"} _media_helpers._build_storyboard(path, key) ready = _media_helpers.storyboard_payload(path) self.assertEqual(ready["status"], "ready") self.assertEqual(ready["count"], 16) for step, moment in enumerate(ready["times"]): self.assertAlmostEqual(moment, step * 0.5, delta=0.11) def test_storyboard_seek_path_reaches_its_targets(self): from nodes import _media_helpers path = Path(self.directory, "seekboard.mp4") self._write_video(path, frames=40, step=6) key = _media_helpers._file_key(path) with _media_helpers._STORYBOARDS_LOCK: _media_helpers._STORYBOARDS[key] = {"status": "building"} # Force the long-clip path: keyframe seek plus a capped forward decode. with unittest.mock.patch.object(_media_helpers, "_STORYBOARD_SEQUENTIAL_FRAMES", 0): _media_helpers._build_storyboard(path, key) ready = _media_helpers.storyboard_payload(path) self.assertEqual(ready["status"], "ready") self.assertGreaterEqual(ready["count"], 12) self.assertGreater(ready["times"][-1], 6.0) self.assertEqual(ready["times"], sorted(ready["times"])) def test_storyboard_builds_ready_payload(self): import base64 from io import BytesIO from nodes import _media_helpers path = Path(self.directory, "board.mp4") self._write_video(path, frames=30, step=8) key = _media_helpers._file_key(path) # Seed "building" so storyboard_payload does not spawn its # background thread; the build below is fully deterministic. with _media_helpers._STORYBOARDS_LOCK: _media_helpers._STORYBOARDS[key] = {"status": "building"} self.assertEqual(_media_helpers.storyboard_payload(path)["status"], "building") _media_helpers._build_storyboard(path, key) ready = _media_helpers.storyboard_payload(path) self.assertEqual(ready["status"], "ready") self.assertGreaterEqual(ready["count"], 1) self.assertEqual(ready["times"], sorted(ready["times"])) sprite_bytes = base64.b64decode(ready["sprite"].split(",", 1)[1]) with Image.open(BytesIO(sprite_bytes)) as sprite: self.assertEqual( sprite.size, (ready["tile_width"] * ready["count"], ready["tile_height"]) ) def test_metadata_cache_hits_and_invalidates_on_file_change(self): from nodes import _media_helpers path = Path(self.directory, "cached.mp4") self._write_video(path, frames=4) first = _media_helpers.cached_video_metadata(path) with unittest.mock.patch.object( _media_helpers, "video_metadata", side_effect=AssertionError("cache miss") ): # Same file state: served from cache, video_metadata untouched. self.assertEqual(_media_helpers.cached_video_metadata(path), first) self._write_video(path, frames=8) os.utime(path, None) refreshed = _media_helpers.cached_video_metadata(path) self.assertGreaterEqual(int(refreshed["frame_count"]), 8) class LocalPreviewGateTests(unittest.TestCase): def test_nothing_outside_comfyuis_folders_is_readable(self): from nodes._media_helpers import local_preview_allowed self.assertFalse(local_preview_allowed(str(Path.home() / "video.mp4"))) self.assertFalse(local_preview_allowed("")) self.assertFalse(local_preview_allowed(r"\\host\share\video.mp4")) self.assertFalse(local_preview_allowed("//host/share/video.mp4")) def test_comfyuis_folders_are_readable(self): from nodes import _media_helpers with tempfile.TemporaryDirectory() as folder: root = Path(folder).resolve() with unittest.mock.patch.object(_media_helpers, "_comfy_managed_roots", lambda: [root]): self.assertTrue(_media_helpers.local_preview_allowed(str(root / "sub" / "clip.mp4"))) self.assertFalse(_media_helpers.local_preview_allowed(str(root.parent / "clip.mp4"))) def test_the_retired_environment_switch_opens_nothing(self): from nodes._media_helpers import local_preview_allowed with tempfile.TemporaryDirectory() as folder, unittest.mock.patch.dict( os.environ, {"AUSBOSS_TRANSFORM_LOCAL_PREVIEW": "1"} ): self.assertFalse(local_preview_allowed(str(Path(folder) / "video.mp4"))) def test_queued_runs_and_previews_explain_the_rule(self): from nodes import _media_helpers with tempfile.TemporaryDirectory() as folder: clip = Path(folder) / "clip.mp4" clip.write_bytes(b"") with self.assertRaises(ValueError) as caught: _media_helpers.resolve_video_path("local path", "", str(clip)) self.assertIn("input, output or temp", str(caught.exception)) # The preview routes pass it through: it names no server path. self.assertEqual(_media_helpers._safe_route_error(caught.exception), str(caught.exception)) class TestScaleToMegapixels(unittest.TestCase): def test_square_megapixel_is_1024(self): self.assertEqual(scale_to_megapixels(1024, 1024, 1.0, 1), (1024, 1024)) self.assertEqual(scale_to_megapixels(512, 512, 4.0, 1), (2048, 2048)) def test_hd_at_one_megapixel_on_64_grid(self): # The everyday case: 1920x1080 to ~1MP on a VAE-friendly grid. self.assertEqual(scale_to_megapixels(1920, 1080, 1.0, 64), (1344, 768)) def test_steps_round_each_dimension_independently(self): self.assertEqual(scale_to_megapixels(1000, 707, 1.0, 8), (1216, 864)) def test_never_below_one_step(self): width, height = scale_to_megapixels(100, 100, 0.01, 64) self.assertEqual((width, height), (128, 128)) if __name__ == "__main__": unittest.main()