366 lines
16 KiB
Python
366 lines
16 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import patch
|
|
import wave
|
|
|
|
import torch
|
|
|
|
from comfy_test_stubs import install_comfy_stubs
|
|
|
|
install_comfy_stubs()
|
|
|
|
from ltx_reframe import (
|
|
LTXReframeCustomAudioLoader,
|
|
LTXReframeRectilinearCorrection,
|
|
LTXReframeTargetFit,
|
|
NODE_CLASS_MAPPINGS,
|
|
build_reframe_masks,
|
|
compute_reframe_placement,
|
|
ensure_audio_track,
|
|
fit_audio_track,
|
|
prepare_reframe_canvas,
|
|
rectilinear_correct_images,
|
|
resolve_frame_range,
|
|
snap_dimension,
|
|
)
|
|
|
|
|
|
class LTXReframeTests(unittest.TestCase):
|
|
@staticmethod
|
|
def _write_test_wav(path: Path, *, sample_rate: int = 8000, samples: int = 80):
|
|
with wave.open(str(path), "wb") as handle:
|
|
handle.setnchannels(1)
|
|
handle.setsampwidth(2)
|
|
handle.setframerate(sample_rate)
|
|
handle.writeframes(b"\x00\x00" * samples)
|
|
|
|
def test_custom_audio_loader_is_registered_and_accepts_an_empty_lazy_branch(self):
|
|
self.assertIs(
|
|
NODE_CLASS_MAPPINGS["LTXReframeCustomAudioLoader"],
|
|
LTXReframeCustomAudioLoader,
|
|
)
|
|
self.assertTrue(LTXReframeCustomAudioLoader.VALIDATE_INPUTS(""))
|
|
with self.assertRaisesRegex(ValueError, "no custom audio file is selected"):
|
|
LTXReframeCustomAudioLoader().load_audio("")
|
|
|
|
def test_custom_audio_loader_lists_audio_uploads_with_a_blank_default(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
with (
|
|
patch("ltx_reframe.folder_paths.get_input_directory", return_value=directory),
|
|
patch(
|
|
"ltx_reframe.folder_paths.filter_files_content_types",
|
|
return_value=["song.wav"],
|
|
),
|
|
):
|
|
spec = LTXReframeCustomAudioLoader.INPUT_TYPES()
|
|
|
|
options, config = spec["required"]["audio_file"]
|
|
self.assertEqual(options, ["", "song.wav"])
|
|
self.assertTrue(config["audio_upload"])
|
|
|
|
def test_custom_audio_loader_decodes_and_fingerprints_an_input_file(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
audio_path = Path(directory) / "clip.wav"
|
|
self._write_test_wav(audio_path)
|
|
with (
|
|
patch("ltx_reframe.folder_paths.get_input_directory", return_value=directory),
|
|
patch(
|
|
"ltx_reframe.folder_paths.annotated_filepath",
|
|
side_effect=lambda name: (name, None),
|
|
),
|
|
):
|
|
audio, = LTXReframeCustomAudioLoader().load_audio("clip.wav")
|
|
first_fingerprint = LTXReframeCustomAudioLoader.IS_CHANGED("clip.wav")
|
|
second_fingerprint = LTXReframeCustomAudioLoader.IS_CHANGED("clip.wav")
|
|
with audio_path.open("ab") as handle:
|
|
handle.write(b"changed")
|
|
changed_fingerprint = LTXReframeCustomAudioLoader.IS_CHANGED("clip.wav")
|
|
|
|
self.assertEqual(audio["sample_rate"], 8000)
|
|
self.assertEqual(tuple(audio["waveform"].shape), (1, 1, 80))
|
|
self.assertEqual(first_fingerprint, second_fingerprint)
|
|
self.assertNotEqual(first_fingerprint, changed_fingerprint)
|
|
self.assertEqual(len(first_fingerprint), 64)
|
|
|
|
def test_custom_audio_loader_reports_a_missing_input_file(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
with (
|
|
patch("ltx_reframe.folder_paths.get_input_directory", return_value=directory),
|
|
patch(
|
|
"ltx_reframe.folder_paths.annotated_filepath",
|
|
side_effect=lambda name: (name, None),
|
|
),
|
|
):
|
|
with self.assertRaisesRegex(ValueError, "file was not found"):
|
|
LTXReframeCustomAudioLoader().load_audio("missing.wav")
|
|
|
|
def test_custom_audio_loader_rejects_paths_outside_comfyui_input(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
with (
|
|
patch("ltx_reframe.folder_paths.get_input_directory", return_value=directory),
|
|
patch(
|
|
"ltx_reframe.folder_paths.annotated_filepath",
|
|
return_value=("../outside.wav", None),
|
|
),
|
|
):
|
|
with self.assertRaisesRegex(ValueError, "input directory"):
|
|
LTXReframeCustomAudioLoader().load_audio("../outside.wav")
|
|
|
|
def test_custom_audio_example_workflow_has_consistent_links_and_safe_defaults(self):
|
|
workflow_path = (
|
|
Path(__file__).resolve().parents[1]
|
|
/ "examples"
|
|
/ "gemmaREFRAME_Experimental_audio_on-off.json"
|
|
)
|
|
workflow = json.loads(workflow_path.read_text(encoding="utf-8"))
|
|
nodes = {node["id"]: node for node in workflow["nodes"]}
|
|
links = {link[0]: link for link in workflow["links"]}
|
|
|
|
self.assertEqual(len(nodes), len(workflow["nodes"]))
|
|
self.assertEqual(len(links), len(workflow["links"]))
|
|
for link_id, (_, source_id, source_slot, target_id, target_slot, _type) in links.items():
|
|
self.assertIn(source_id, nodes)
|
|
self.assertIn(target_id, nodes)
|
|
self.assertIn(link_id, nodes[source_id]["outputs"][source_slot].get("links") or [])
|
|
self.assertEqual(nodes[target_id]["inputs"][target_slot].get("link"), link_id)
|
|
|
|
custom_loader = nodes[9086]
|
|
selector = nodes[9112]
|
|
reframe = nodes[5553]
|
|
stage_one = nodes[9055]
|
|
create_video = nodes[5539]
|
|
preview = nodes[9088]
|
|
|
|
self.assertEqual(custom_loader["type"], "LTXReframeCustomAudioLoader")
|
|
self.assertEqual(custom_loader["widgets_values"], [""])
|
|
self.assertEqual(selector["type"], "LazySwitchKJ")
|
|
self.assertFalse(selector["widgets_values"][0])
|
|
self.assertEqual(set(selector["outputs"][0]["links"]), {14666, 14818, 14918})
|
|
self.assertEqual(reframe["outputs"][4]["links"], [14990, 14991])
|
|
self.assertEqual(stage_one["inputs"][10]["link"], 14818)
|
|
self.assertEqual(preview["inputs"][0]["link"], 14918)
|
|
self.assertEqual(create_video["inputs"][1]["link"], 14999)
|
|
self.assertEqual(create_video["inputs"][2]["link"], 14666)
|
|
|
|
def test_existing_audio_is_preserved(self):
|
|
audio = {
|
|
"waveform": torch.ones((1, 1, 100), dtype=torch.float32),
|
|
"sample_rate": 16000,
|
|
}
|
|
|
|
result, synthesized = ensure_audio_track(audio, frame_count=24, fps=24.0)
|
|
|
|
self.assertIs(result, audio)
|
|
self.assertFalse(synthesized)
|
|
|
|
def test_audio_less_video_gets_duration_matched_stereo_silence(self):
|
|
audio, synthesized = ensure_audio_track(None, frame_count=121, fps=24.0)
|
|
|
|
self.assertTrue(synthesized)
|
|
self.assertEqual(audio["sample_rate"], 44100)
|
|
self.assertEqual(tuple(audio["waveform"].shape), (1, 2, 222338))
|
|
self.assertEqual(torch.count_nonzero(audio["waveform"]).item(), 0)
|
|
|
|
def test_silent_audio_fallback_rejects_invalid_fps(self):
|
|
with self.assertRaisesRegex(ValueError, "valid video frame rate"):
|
|
ensure_audio_track(None, frame_count=121, fps=0.0)
|
|
|
|
def test_audio_is_cropped_and_padded_to_selected_duration(self):
|
|
cropped, synthesized = fit_audio_track(
|
|
{"waveform": torch.ones((1, 2, 2000)), "sample_rate": 1000},
|
|
frame_count=25,
|
|
fps=25.0,
|
|
)
|
|
padded, _ = fit_audio_track(
|
|
{"waveform": torch.ones((1, 2, 500)), "sample_rate": 1000},
|
|
frame_count=25,
|
|
fps=25.0,
|
|
)
|
|
|
|
self.assertFalse(synthesized)
|
|
self.assertEqual(cropped["waveform"].shape[-1], 1000)
|
|
self.assertEqual(padded["waveform"].shape[-1], 1000)
|
|
self.assertEqual(torch.count_nonzero(padded["waveform"][..., 500:]).item(), 0)
|
|
|
|
def test_full_video_range_is_not_aligned_or_trimmed(self):
|
|
selected = resolve_frame_range(70, 25.0, False, 8, 16)
|
|
|
|
self.assertFalse(selected.range_enabled)
|
|
self.assertEqual(selected.effective_first_frame, 0)
|
|
self.assertEqual(selected.effective_last_frame, 69)
|
|
self.assertEqual(selected.selected_frame_count, 70)
|
|
|
|
def test_valid_sixty_five_frame_range_is_preserved(self):
|
|
selected = resolve_frame_range(241, 25.0, True, 0, 64)
|
|
|
|
self.assertEqual(selected.effective_first_frame, 0)
|
|
self.assertEqual(selected.effective_last_frame, 64)
|
|
self.assertEqual(selected.selected_frame_count, 65)
|
|
|
|
def test_range_alignment_preserves_first_frame_and_trims_end(self):
|
|
zero_based = resolve_frame_range(241, 25.0, True, 0, 70)
|
|
offset = resolve_frame_range(241, 25.0, True, 100, 199)
|
|
|
|
self.assertEqual((zero_based.effective_first_frame, zero_based.effective_last_frame), (0, 64))
|
|
self.assertEqual(zero_based.selected_frame_count, 65)
|
|
self.assertEqual((offset.effective_first_frame, offset.effective_last_frame), (100, 196))
|
|
self.assertEqual(offset.selected_frame_count, 97)
|
|
|
|
def test_last_frame_minus_one_resolves_to_source_end_then_aligns(self):
|
|
selected = resolve_frame_range(100, 25.0, True, 10, -1)
|
|
|
|
self.assertEqual(selected.requested_last_frame, 99)
|
|
self.assertEqual(selected.effective_last_frame, 98)
|
|
self.assertEqual(selected.selected_frame_count, 89)
|
|
|
|
def test_invalid_frame_ranges_are_rejected(self):
|
|
with self.assertRaisesRegex(ValueError, "greater than or equal"):
|
|
resolve_frame_range(100, 25.0, True, 30, 20)
|
|
with self.assertRaisesRegex(ValueError, "outside"):
|
|
resolve_frame_range(100, 25.0, True, 0, 100)
|
|
with self.assertRaisesRegex(ValueError, "at least nine"):
|
|
resolve_frame_range(100, 25.0, True, 0, 7)
|
|
|
|
def test_single_frame_range_is_allowed(self):
|
|
selected = resolve_frame_range(100, 25.0, True, 42, 42)
|
|
|
|
self.assertEqual(selected.selected_frame_count, 1)
|
|
self.assertEqual(selected.effective_last_frame, 42)
|
|
|
|
def test_dimensions_snap_to_ltx_multiple(self):
|
|
self.assertEqual(snap_dimension(1919), 1920)
|
|
self.assertEqual(snap_dimension(1050), 1024)
|
|
self.assertEqual(snap_dimension(1), 256)
|
|
|
|
def test_portrait_canvas_centers_aspect_locked_landscape_source(self):
|
|
placement = compute_reframe_placement(1920, 1080, 1088, 1920, 0.5, 0.5, 0.7)
|
|
self.assertEqual(placement.target_width, 1088)
|
|
self.assertEqual(placement.target_height, 1920)
|
|
self.assertLess(abs(placement.box_width / placement.box_height - 16 / 9), 0.01)
|
|
self.assertLessEqual(abs(placement.box_x * 2 + placement.box_width - placement.target_width), 1)
|
|
self.assertLessEqual(abs(placement.box_y * 2 + placement.box_height - placement.target_height), 1)
|
|
|
|
def test_canvas_and_masks_preserve_source_region(self):
|
|
frames = torch.ones((2, 8, 16, 3), dtype=torch.float32)
|
|
placement = compute_reframe_placement(16, 8, 256, 256, 0.5, 0.5, 0.5)
|
|
canvas = prepare_reframe_canvas(frames, placement)
|
|
outpaint, source = build_reframe_masks(placement, feather=0)
|
|
|
|
self.assertEqual(tuple(canvas.shape), (2, 256, 256, 3))
|
|
self.assertEqual(tuple(outpaint.shape), (1, 256, 256))
|
|
self.assertEqual(tuple(source.shape), (1, 256, 256))
|
|
self.assertTrue(torch.allclose(outpaint + source, torch.ones_like(source)))
|
|
self.assertEqual(canvas[:, placement.box_y, placement.box_x].mean().item(), 1.0)
|
|
self.assertEqual(canvas[:, 0, 0].sum().item(), 0.0)
|
|
|
|
def test_feather_ramps_only_outside_source_box(self):
|
|
placement = compute_reframe_placement(16, 9, 256, 256, 0.5, 0.5, 0.5)
|
|
outpaint, source = build_reframe_masks(placement, feather=16)
|
|
inside_y = placement.box_y + placement.box_height // 2
|
|
inside_x = placement.box_x + placement.box_width // 2
|
|
self.assertEqual(outpaint[0, inside_y, inside_x].item(), 0.0)
|
|
self.assertEqual(source[0, inside_y, inside_x].item(), 1.0)
|
|
self.assertEqual(outpaint[0, 0, 0].item(), 1.0)
|
|
|
|
def test_zoom_in_can_extend_source_beyond_target(self):
|
|
placement = compute_reframe_placement(1920, 1080, 1088, 1920, 0.5, 0.5, 2.0)
|
|
|
|
self.assertEqual(placement.box_width, 2176)
|
|
self.assertEqual(placement.box_height, 1224)
|
|
self.assertEqual(placement.box_x, -544)
|
|
self.assertEqual(placement.box_y, 348)
|
|
self.assertEqual(placement.source_scale, 2.0)
|
|
|
|
def test_zoomed_source_is_cropped_safely_to_target_canvas(self):
|
|
frames = torch.ones((2, 8, 16, 3), dtype=torch.float32)
|
|
placement = compute_reframe_placement(16, 8, 256, 256, 0.5, 0.5, 4.0)
|
|
canvas = prepare_reframe_canvas(frames, placement)
|
|
outpaint, source = build_reframe_masks(placement, feather=32)
|
|
|
|
self.assertEqual(tuple(canvas.shape), (2, 256, 256, 3))
|
|
self.assertTrue(torch.allclose(canvas, torch.ones_like(canvas)))
|
|
self.assertTrue(torch.allclose(source, torch.ones_like(source)))
|
|
self.assertTrue(torch.allclose(outpaint, torch.zeros_like(outpaint)))
|
|
|
|
def test_target_fit_returns_exact_dimensions(self):
|
|
images = torch.rand((2, 96, 160, 3), dtype=torch.float32)
|
|
|
|
fitted, = LTXReframeTargetFit().fit(images, target_width=128, target_height=128)
|
|
|
|
self.assertEqual(tuple(fitted.shape), (2, 128, 128, 3))
|
|
self.assertGreaterEqual(fitted.min().item(), 0.0)
|
|
self.assertLessEqual(fitted.max().item(), 1.0)
|
|
|
|
def test_rectilinear_correction_is_temporally_deterministic(self):
|
|
horizontal = torch.linspace(0.0, 1.0, 65).view(1, 1, 65, 1)
|
|
frame = horizontal.expand(1, 33, 65, 3)
|
|
images = frame.expand(2, -1, -1, -1).clone()
|
|
|
|
corrected, correction_mask = rectilinear_correct_images(
|
|
images,
|
|
correction_strength=0.1,
|
|
center_x=0.5,
|
|
center_y=0.5,
|
|
protect_source=0.0,
|
|
protection_feather=0,
|
|
interpolation="bilinear",
|
|
chunk_size=1,
|
|
)
|
|
|
|
self.assertEqual(tuple(corrected.shape), tuple(images.shape))
|
|
self.assertEqual(tuple(correction_mask.shape), (1, 33, 65))
|
|
self.assertTrue(torch.allclose(corrected[0], corrected[1]))
|
|
self.assertGreater(corrected[0, 16, 0, 0].item(), images[0, 16, 0, 0].item())
|
|
self.assertLess(corrected[0, 16, -1, 0].item(), images[0, 16, -1, 0].item())
|
|
self.assertAlmostEqual(
|
|
corrected[0, 16, 32, 0].item(),
|
|
images[0, 16, 32, 0].item(),
|
|
places=5,
|
|
)
|
|
|
|
def test_rectilinear_source_mask_can_fully_protect_original_pixels(self):
|
|
images = torch.rand((2, 24, 40, 3), dtype=torch.float32)
|
|
source_mask = torch.ones((1, 24, 40), dtype=torch.float32)
|
|
|
|
corrected, correction_mask = rectilinear_correct_images(
|
|
images,
|
|
correction_strength=0.15,
|
|
center_x=0.5,
|
|
center_y=0.5,
|
|
protect_source=1.0,
|
|
protection_feather=0,
|
|
interpolation="bicubic",
|
|
chunk_size=2,
|
|
source_mask=source_mask,
|
|
)
|
|
|
|
self.assertTrue(torch.allclose(corrected, images))
|
|
self.assertEqual(torch.count_nonzero(correction_mask).item(), 0)
|
|
|
|
def test_rectilinear_node_bypass_is_pixel_exact(self):
|
|
images = torch.rand((3, 16, 20, 3), dtype=torch.float32)
|
|
|
|
bypassed, correction_mask = LTXReframeRectilinearCorrection().correct(
|
|
images,
|
|
enabled=False,
|
|
correction_strength=0.05,
|
|
center_x=0.5,
|
|
center_y=0.5,
|
|
protect_source=1.0,
|
|
protection_feather=96,
|
|
interpolation="bicubic",
|
|
chunk_size=8,
|
|
)
|
|
|
|
self.assertIs(bypassed, images)
|
|
self.assertEqual(tuple(correction_mask.shape), (1, 16, 20))
|
|
self.assertEqual(torch.count_nonzero(correction_mask).item(), 0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|