alisson-anjosandClaude Opus 5.5
d01ca239ce
With save_to = mp4 file, chosen or forced by the RAM guard, the images and comparison outputs were blocked, so the result nodes stayed empty and looked broken. They now carry the video read back from the written files, scaled to fit the free RAM (the files keep the full size), and the join shows the files on its node. Bump BFSNodes version to 1.105.1. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1212 lines
64 KiB
Python
1212 lines
64 KiB
Python
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import types
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
import torch
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
_TMP = tempfile.mkdtemp(prefix="bfs_shotloop_test_")
|
|
|
|
# The module only needs a few folder_paths calls; stub them so the tests run without ComfyUI.
|
|
_fp = types.ModuleType("folder_paths")
|
|
_fp.get_temp_directory = lambda: _TMP
|
|
_fp.get_input_directory = lambda: _TMP
|
|
_fp.get_annotated_filepath = lambda name: str(Path(_TMP) / name)
|
|
sys.modules.setdefault("folder_paths", _fp)
|
|
|
|
|
|
class _Blocker:
|
|
def __init__(self, message):
|
|
self.message = message
|
|
|
|
|
|
_ce = types.ModuleType("comfy_execution")
|
|
_gu = types.ModuleType("comfy_execution.graph_utils")
|
|
_gu.ExecutionBlocker = _Blocker
|
|
sys.modules.setdefault("comfy_execution", _ce)
|
|
sys.modules.setdefault("comfy_execution.graph_utils", _gu)
|
|
|
|
SPEC = importlib.util.spec_from_file_location("bfs_shot_loop", ROOT / "bfs_shot_loop.py")
|
|
assert SPEC is not None and SPEC.loader is not None
|
|
SL = importlib.util.module_from_spec(SPEC)
|
|
SPEC.loader.exec_module(SL)
|
|
H3 = "H3 (17n+5)"
|
|
|
|
|
|
def _shot(i, start, end, cut=False, gen=None, count=1, queue=False, run_id=""):
|
|
length = end - start
|
|
return {"index": i, "count": count, "start": start, "end": end, "length": length,
|
|
"gen_length": gen or SL.snap_up(length, H3), "fps": 24.0, "cut_before": cut,
|
|
"queue": queue, "run_id": run_id}
|
|
|
|
|
|
class GridTest(unittest.TestCase):
|
|
def test_h3_grid(self):
|
|
self.assertEqual([SL.snap_up(x, H3) for x in (1, 5, 6, 22, 23, 107, 108)], [5, 5, 22, 22, 39, 107, 124])
|
|
self.assertEqual([SL.snap_down(x, H3) for x in (4, 5, 21, 22, 106, 107, 110)], [5, 5, 5, 22, 90, 107, 107])
|
|
|
|
def test_other_grids(self):
|
|
self.assertEqual(SL.snap_up(50, "LTX / Wan (8n+1)"), 57)
|
|
self.assertEqual(SL.snap_down(50, "LTX / Wan (8n+1)"), 49)
|
|
self.assertEqual(SL.snap_up(6, "Wan (4n+1)"), 9)
|
|
self.assertEqual(SL.snap_up(37, "any"), 37)
|
|
|
|
def test_generation_size_is_on_the_multiple(self):
|
|
w, h = SL.generation_size(1280, 720, 0.15, 32)
|
|
self.assertEqual((w % 32, h % 32), (0, 0))
|
|
self.assertGreater(w, h)
|
|
|
|
def test_planner_defaults_to_half_megapixel(self):
|
|
self.assertEqual(SL.DEFAULT_PLAN["megapixels"], 0.5)
|
|
|
|
|
|
class TextGenerateCompatibilityTest(unittest.TestCase):
|
|
def _generate(self, execute):
|
|
calls = []
|
|
|
|
def wrapped(*args, **kwargs):
|
|
result = execute(*args, **kwargs)
|
|
calls.append(kwargs)
|
|
return result
|
|
|
|
textgen = types.ModuleType("comfy_extras.nodes_textgen")
|
|
textgen.TextGenerate = type("TextGenerate", (), {"execute": staticmethod(wrapped)})
|
|
comfy_extras = types.ModuleType("comfy_extras")
|
|
comfy_extras.__path__ = []
|
|
modules = {"comfy_extras": comfy_extras, "comfy_extras.nodes_textgen": textgen}
|
|
with patch.dict(sys.modules, modules):
|
|
result = SL.vlm_generate(object(), "describe this", torch.zeros(1, 2, 2, 3), 32)
|
|
self.assertEqual(result, "generated")
|
|
return calls[0]
|
|
|
|
def test_old_textgenerate_signature_does_not_receive_mtp(self):
|
|
def execute(clip, prompt, max_length, sampling_mode, image=None):
|
|
return types.SimpleNamespace(args=("generated",))
|
|
|
|
kwargs = self._generate(execute)
|
|
self.assertNotIn("mtp", kwargs)
|
|
|
|
def test_new_textgenerate_signature_keeps_mtp_disabled(self):
|
|
def execute(clip, prompt, max_length, sampling_mode, image=None, mtp="auto"):
|
|
return types.SimpleNamespace(args=("generated",))
|
|
|
|
kwargs = self._generate(execute)
|
|
self.assertEqual(kwargs["mtp"], "off")
|
|
|
|
|
|
class PlanTest(unittest.TestCase):
|
|
def test_fixed_mode_splits_evenly_under_the_limit(self):
|
|
segs = SL.plan_segments(300, [], "fixed", 107, 24)
|
|
self.assertEqual([s["end"] - s["start"] for s in segs], [100, 100, 100])
|
|
|
|
def test_camera_cut_mode_splits_a_continuous_range_at_max_length(self):
|
|
segs = SL.plan_segments(360, [], "shots", 107, 1)
|
|
self.assertEqual(len(segs), 4)
|
|
self.assertTrue(all(s["end"] - s["start"] <= 107 for s in segs))
|
|
self.assertEqual(segs[-1]["end"], 360)
|
|
|
|
def test_shots_split_at_cuts_and_long_shots_are_divided(self):
|
|
segs = SL.plan_segments(300, [50, 80], "shots", 107, 10)
|
|
bounds = [(s["start"], s["end"]) for s in segs]
|
|
self.assertEqual(bounds[:2], [(0, 50), (50, 80)])
|
|
self.assertTrue(all(e - s <= 107 for s, e in bounds))
|
|
self.assertEqual(bounds[-1][1], 300)
|
|
self.assertTrue(segs[1]["cut_before"])
|
|
|
|
def test_short_shots_merge_into_a_neighbour_that_still_fits(self):
|
|
segs = SL.plan_segments(200, [10, 100], "shots", 107, 24)
|
|
self.assertEqual([(s["start"], s["end"]) for s in segs], [(0, 100), (100, 200)])
|
|
|
|
def test_limits_and_manual_bounds(self):
|
|
self.assertEqual(len(SL.plan_segments(1000, [], "fixed", 100, 10, max_parts=3)), 3)
|
|
self.assertEqual(SL.plan_segments(1000, [], "fixed", 100, 10, max_total=250)[-1]["end"], 250)
|
|
manual = SL.plan_segments(100, [], "manual", 107, 10, manual=[30, 70])
|
|
self.assertEqual([(s["start"], s["end"]) for s in manual], [(0, 30), (30, 70), (70, 100)])
|
|
|
|
def test_builtin_detector_finds_spikes_and_respects_the_gap(self):
|
|
score = [1.0] * 100
|
|
raw = [0.01] * 100
|
|
for i in (20, 22, 60):
|
|
score[i], raw[i] = 50.0, 0.5
|
|
self.assertEqual(SL.detect_cuts(score, raw, 0.5, 24.0), [20, 60])
|
|
|
|
|
|
class JoinTest(unittest.TestCase):
|
|
def _frames(self, n, value):
|
|
return torch.full((n, 8, 8, 3), float(value))
|
|
|
|
def test_trims_to_true_lengths_in_order(self):
|
|
shots = [_shot(1, 30, 50, cut=True, gen=22), _shot(0, 0, 30, gen=39)]
|
|
imgs = [self._frames(22, 0.8), self._frames(39, 0.2)]
|
|
video, audio, fps = SL.BFSShotJoin()._join_all(imgs, shots, [0])
|
|
self.assertEqual(video.shape[0], 50)
|
|
self.assertAlmostEqual(float(video[0, 0, 0, 0]), 0.2)
|
|
self.assertAlmostEqual(float(video[-1, 0, 0, 0]), 0.8)
|
|
self.assertEqual(fps, 24.0)
|
|
|
|
def test_crossfade_only_across_soft_joins(self):
|
|
imgs = [self._frames(39, 0.0), self._frames(39, 1.0)]
|
|
soft = SL.BFSShotJoin()._join_all(imgs, [_shot(0, 0, 30, gen=39), _shot(1, 30, 60, gen=39)], [4])[0]
|
|
hard = SL.BFSShotJoin()._join_all(imgs, [_shot(0, 0, 30, gen=39), _shot(1, 30, 60, cut=True, gen=39)], [4])[0]
|
|
self.assertGreater(float(soft[31, 0, 0, 0]), 0.0)
|
|
self.assertLess(float(soft[31, 0, 0, 0]), 1.0)
|
|
self.assertEqual(float(hard[31, 0, 0, 0]), 1.0)
|
|
|
|
def test_audio_is_trimmed_to_the_video(self):
|
|
audio = {"waveform": torch.zeros(1, 2, 48000 * 10), "sample_rate": 48000}
|
|
imgs = [self._frames(39, 0.5)]
|
|
_, a, _ = SL.BFSShotJoin()._join_all(imgs, [_shot(0, 0, 24, gen=39)], [0], [audio])
|
|
self.assertEqual(a["waveform"].shape[-1], 48000)
|
|
|
|
def test_queue_mode_blocks_until_every_shot_is_stored(self):
|
|
join = SL.BFSShotJoin()
|
|
shots = [_shot(0, 0, 20, count=2, queue=True, run_id="t1"), _shot(1, 20, 45, cut=True, count=2, queue=True, run_id="t1")]
|
|
first = join.join([self._frames(22, 0.1)], [shots[0]], [0])
|
|
self.assertIsInstance(first[0], _Blocker)
|
|
video, _, _, _, _ = join.join([self._frames(39, 0.9)], [shots[1]], [0])
|
|
self.assertEqual(video.shape[0], 45)
|
|
self.assertAlmostEqual(float(video[0, 0, 0, 0]), 0.1, places=2)
|
|
|
|
|
|
class FilterTest(unittest.TestCase):
|
|
STATS = {"persons": 1, "person_area": 0.2, "person_frames": 6, "faces": 1, "brightness": 0.4, "motion": 0.02, "sampled": 6}
|
|
|
|
def test_no_filter_keeps_everything(self):
|
|
self.assertEqual(SL.skip_reason(self.STATS, 50, dict(SL.DEFAULT_FILTERS)), "")
|
|
self.assertFalse(SL.filters_active(dict(SL.DEFAULT_FILTERS)))
|
|
|
|
def test_each_filter(self):
|
|
f = dict(SL.DEFAULT_FILTERS)
|
|
self.assertEqual(SL.skip_reason(dict(self.STATS, persons=0), 50, dict(f, person=True)), "no person")
|
|
self.assertIn("smaller", SL.skip_reason(dict(self.STATS, person_area=0.01), 50, dict(f, person=True, min_person_area=0.03)))
|
|
self.assertIn("more than", SL.skip_reason(dict(self.STATS, persons=5), 50, dict(f, max_persons=2)))
|
|
self.assertEqual(SL.skip_reason(dict(self.STATS, faces=0), 50, dict(f, face=True)), "no face")
|
|
self.assertEqual(SL.skip_reason(dict(self.STATS, brightness=0.02), 50, dict(f, skip_dark=True)), "dark / fade")
|
|
self.assertEqual(SL.skip_reason(dict(self.STATS, motion=0.0), 50, dict(f, skip_static=True)), "static")
|
|
self.assertIn("shorter", SL.skip_reason(self.STATS, 5, dict(f, min_frames=10)))
|
|
|
|
|
|
class TimelineJoinTest(unittest.TestCase):
|
|
def _tl(self, fill):
|
|
return {"path": "", "fps": 24.0, "width": 8, "height": 8, "fill": fill, "audio": None, "n": 60,
|
|
"src": list(range(60)),
|
|
"segs": [{"start": 0, "end": 20, "run": True, "cut_before": False, "run_index": 0},
|
|
{"start": 20, "end": 40, "run": False, "cut_before": True, "run_index": -1},
|
|
{"start": 40, "end": 60, "run": True, "cut_before": True, "run_index": 1}]}
|
|
|
|
def test_drop_removes_skipped_shots(self):
|
|
shots = [_shot(0, 0, 20, gen=22), _shot(1, 40, 60, cut=True, gen=22)]
|
|
imgs = [torch.zeros(22, 8, 8, 3), torch.ones(22, 8, 8, 3)]
|
|
video = SL.BFSShotJoin()._join_all(imgs, shots, [0], None, self._tl("drop"))[0]
|
|
self.assertEqual(video.shape[0], 40)
|
|
self.assertEqual(float(video[25, 0, 0, 0]), 1.0)
|
|
|
|
def test_original_fill_reads_the_source_frames(self):
|
|
shots = [_shot(0, 0, 20, gen=22), _shot(1, 40, 60, cut=True, gen=22)]
|
|
imgs = [torch.zeros(22, 8, 8, 3), torch.ones(22, 8, 8, 3)]
|
|
orig = SL._read_frames
|
|
SL._read_frames = lambda path, idx, size: [__import__("numpy").full((8, 8, 3), 128, "uint8") for _ in idx]
|
|
try:
|
|
video = SL.BFSShotJoin()._join_all(imgs, shots, [0], None, self._tl("original"))[0]
|
|
finally:
|
|
SL._read_frames = orig
|
|
self.assertEqual(video.shape[0], 60)
|
|
self.assertAlmostEqual(float(video[30, 0, 0, 0]), 128 / 255, places=3)
|
|
|
|
def test_skipped_shot_keeps_its_sound(self):
|
|
# 60 frames at 24 fps, sample rate 240 -> 10 samples per frame; each sample holds its timeline frame number
|
|
sr = 240
|
|
timeline_audio = {"waveform": torch.arange(600, dtype=torch.float32).div(10).floor().view(1, 1, 600), "sample_rate": sr}
|
|
runs_only = {"waveform": torch.full((1, 1, 400), -1.0), "sample_rate": sr, "bfs_runs_only": True}
|
|
tl = dict(self._tl("original"), audio=timeline_audio)
|
|
shots = [_shot(0, 0, 20, gen=22), _shot(1, 40, 60, cut=True, gen=22)]
|
|
imgs = [torch.zeros(22, 8, 8, 3), torch.ones(22, 8, 8, 3)]
|
|
orig = SL._read_frames
|
|
SL._read_frames = lambda path, idx, size: [__import__("numpy").full((8, 8, 3), 128, "uint8") for _ in idx]
|
|
try:
|
|
_, a, _ = SL.BFSShotJoin()._join_all(imgs, shots, [0], [runs_only], tl)
|
|
# a soundtrack from elsewhere is still used as given
|
|
_, other, _ = SL.BFSShotJoin()._join_all(imgs, shots, [0], [dict(runs_only, bfs_runs_only=False)], tl)
|
|
finally:
|
|
SL._read_frames = orig
|
|
wf = a["waveform"][0, 0]
|
|
self.assertEqual(wf.shape[0], 600)
|
|
self.assertEqual([float(wf[i * 10]) for i in (0, 25, 30, 45, 59)], [0.0, 25.0, 30.0, 45.0, 59.0]) # in sync, skipped shot not silent
|
|
self.assertNotIn("bfs_runs_only", a)
|
|
self.assertEqual(float(other["waveform"][0, 0, 0]), -1.0)
|
|
|
|
|
|
class CastTest(unittest.TestCase):
|
|
# person 0 on frames 0-47 (big face), person 1 on 48-95, both small on 96-119
|
|
CAST = {"step": 6, "people": [{"id": 0}, {"id": 1}], "samples":
|
|
[{"f": f, "faces": [{"pid": 0, "area": 0.1}]} for f in range(0, 48, 6)]
|
|
+ [{"f": f, "faces": [{"pid": 1, "area": 0.1}]} for f in range(48, 96, 6)]
|
|
+ [{"f": f, "faces": []} for f in range(96, 120, 6)]}
|
|
|
|
def test_shot_people_and_main(self):
|
|
self.assertEqual(SL.shot_people(self.CAST, 0, 60), {"people": [0, 1], "main": 0})
|
|
self.assertEqual(SL.shot_people(self.CAST, 96, 120), {"people": [], "main": -1})
|
|
|
|
def test_change_points_need_a_long_enough_run(self):
|
|
self.assertEqual(SL.person_change_points(self.CAST, 0, 120, 24), [48])
|
|
self.assertEqual(SL.person_change_points(self.CAST, 0, 60, 24), [])
|
|
|
|
def test_plan_splits_assigns_and_skips_by_person(self):
|
|
path = str(Path(_TMP) / "cast.mp4"); Path(path).write_bytes(b"x")
|
|
a = {"fps": 24.0, "n": 120, "score": [0.0] * 120, "raw": [0.0] * 120}
|
|
SL._CAST_CACHE[SL._cast_key(path, a)] = self.CAST
|
|
orig = SL.find_cuts
|
|
SL.find_cuts = lambda *args, **kw: ([], "test")
|
|
try:
|
|
plan = dict(SL.DEFAULT_PLAN, max_s=10, min_s=0.5, cast={"0": {"ref": "a.png"}}, cast_split=True, cast_only=True)
|
|
segs = SL.apply_filters(plan, a, path, SL.resolve_plan(plan, a, path))
|
|
finally:
|
|
SL.find_cuts = orig
|
|
self.assertEqual([(s["start"], s["end"]) for s in segs], [(0, 48), (48, 120)])
|
|
self.assertEqual(segs[0]["ref"], "a.png")
|
|
self.assertTrue(segs[0]["run"])
|
|
self.assertEqual(segs[1]["skip_reason"], "no linked person")
|
|
|
|
|
|
class ComparisonTest(unittest.TestCase):
|
|
def test_join_returns_a_labelled_side_by_side(self):
|
|
shots = [dict(_shot(0, 0, 20, gen=22), frames=torch.zeros(22, 32, 48, 3), ref=torch.ones(1, 40, 30, 3),
|
|
ref2=None, prompt="a test prompt"),
|
|
dict(_shot(1, 20, 40, cut=True, gen=22), frames=torch.zeros(22, 32, 48, 3), ref=None, ref2=None,
|
|
prompt="")]
|
|
imgs = [torch.full((22, 32, 48, 3), 0.5), torch.full((22, 32, 48, 3), 0.7)]
|
|
video, _, _, comp, _ = SL.BFSShotJoin().join(imgs, shots, [0], comparison=[True], label=["steps 20"])
|
|
self.assertEqual(video.shape[0], 40)
|
|
self.assertEqual(comp.shape[0], 40)
|
|
self.assertEqual(comp.shape[2] % 16, 0)
|
|
self.assertEqual(comp.shape[1] % 16, 0)
|
|
self.assertGreaterEqual(comp.shape[2], 48 * 2 + 32)
|
|
self.assertGreater(comp.shape[1], 32)
|
|
off = SL.BFSShotJoin().join(imgs, shots, [0])[3]
|
|
self.assertEqual(off.shape[0], 1)
|
|
|
|
|
|
class ContinuityTest(unittest.TestCase):
|
|
def test_pick_frame_ignores_the_overlap(self):
|
|
frames = torch.arange(30, dtype=torch.float32).view(30, 1, 1, 1).expand(30, 2, 2, 3)
|
|
self.assertEqual(float(SL.pick_frame(frames, 20, "first")[0, 0, 0, 0]), 0.0)
|
|
self.assertEqual(float(SL.pick_frame(frames, 20, "middle")[0, 0, 0, 0]), 10.0)
|
|
self.assertEqual(float(SL.pick_frame(frames, 20, "last")[0, 0, 0, 0]), 19.0)
|
|
|
|
def test_chain_uses_the_previous_result_only_when_asked(self):
|
|
frames = torch.arange(30, dtype=torch.float32).view(30, 1, 1, 1).expand(30, 2, 2, 3)
|
|
SL.remember_result({"index": 0, "count": 2}, frames)
|
|
shot = {"index": 1, "count": 2, "chain": "reference", "chain_frame": "last", "prev_length": 20}
|
|
self.assertEqual(float(SL.chain_image(shot)[0, 0, 0, 0]), 19.0)
|
|
self.assertIsNone(SL.chain_image(dict(shot, chain="off")))
|
|
self.assertIsNone(SL.chain_image(dict(shot, index=0)))
|
|
self.assertIsNone(SL.chain_image(dict(shot, index=3, count=4))) # not the shot right after
|
|
stored = torch.ones(1, 2, 2, 3)
|
|
self.assertIs(SL.chain_image(dict(shot, chain_image=stored)), stored)
|
|
|
|
def test_plan_carries_the_per_shot_setting(self):
|
|
a = {"fps": 24.0, "n": 100, "score": [0.0] * 100, "raw": [0.0] * 100}
|
|
plan = dict(SL.DEFAULT_PLAN, mode="fixed", max_s=2, bounds=[50],
|
|
segs=[{}, {"chain": "first frame", "chain_frame": "middle"}])
|
|
segs = SL.resolve_plan(plan, a)
|
|
self.assertEqual((segs[0]["chain"], segs[1]["chain"], segs[1]["chain_frame"]), ("head", "first frame", "middle")) # head is the default
|
|
|
|
|
|
class MaskCropTest(unittest.TestCase):
|
|
def test_spec_takes_target_from_shot_and_settings_from_global(self):
|
|
spec = SL.mask_spec({"text": "person", "padding": 0.9}, {"padding": 0.3, "text": "ignored"})
|
|
self.assertEqual((spec["text"], spec["padding"], spec["fill_holes"]), ("person", 0.3, True))
|
|
|
|
def test_shape_fills_holes_and_holds_in_time(self):
|
|
m = torch.zeros(5, 9, 9, dtype=torch.uint8)
|
|
m[2, 2:7, 2:7] = 1
|
|
m[2, 4, 4] = 0 # a hole
|
|
out = SL.shape_mask(m, dict(SL.DEFAULT_MASK, temporal_expand=1))
|
|
self.assertEqual(int(out[2, 4, 4]), 1)
|
|
self.assertEqual(int(out[1, 3, 3]), 1) # held one frame back
|
|
self.assertEqual(int(out[0, 3, 3]), 0)
|
|
blocks = SL.shape_mask(m, dict(SL.DEFAULT_MASK, temporal_expand=0, blockify=4))
|
|
self.assertEqual(tuple(blocks.shape), (5, 9, 9))
|
|
|
|
def test_box_is_the_union_plus_padding(self):
|
|
m = torch.zeros(2, 10, 10, dtype=torch.uint8)
|
|
m[0, 2:4, 2:4] = 1
|
|
m[1, 6:8, 6:8] = 1
|
|
self.assertEqual(SL.crop_box(m, 0.0), [0.2, 0.2, 0.8, 0.8])
|
|
self.assertIsNone(SL.crop_box(torch.zeros(2, 10, 10), 0.1))
|
|
|
|
def test_uncrop_pastes_only_inside_the_mask(self):
|
|
full = torch.zeros(3, 20, 20, 3)
|
|
mask = torch.zeros(3, 20, 20)
|
|
mask[:, 5:15, 5:15] = 1
|
|
shot = {"full_frames": full, "crop": {"box": [0.25, 0.25, 0.75, 0.75], "mask": mask, "paste": "mask",
|
|
"expand": 0, "feather": 0}}
|
|
out = SL.uncrop(torch.ones(3, 8, 8, 3), shot)
|
|
self.assertEqual(tuple(out.shape), (3, 20, 20, 3))
|
|
self.assertEqual(float(out[0, 10, 10, 0]), 1.0)
|
|
self.assertEqual(float(out[0, 1, 1, 0]), 0.0)
|
|
|
|
|
|
class DuetPanelCropTest(unittest.TestCase):
|
|
def test_join_side_cuts_the_panel_off(self):
|
|
shot = {"width": 64, "height": 32, "panel": {"position": "left", "h": 2, "w": 4, "strip_h": 0, "strip_w": 4}}
|
|
canvas = torch.zeros(3, 32, 128, 3)
|
|
canvas[:, :, 64:] = 1.0
|
|
out = SL.crop_panel(canvas, shot)
|
|
self.assertEqual(tuple(out.shape), (3, 32, 64, 3))
|
|
self.assertEqual(float(out.min()), 1.0)
|
|
self.assertIs(SL.crop_panel(out, shot), out) # already the video size
|
|
self.assertIs(SL.crop_panel(canvas, {"width": 64}), canvas) # no panel
|
|
|
|
|
|
class SilentAudioTest(unittest.TestCase):
|
|
def test_join_without_audio_returns_a_silent_track(self):
|
|
imgs = [torch.full((39, 8, 8, 3), 0.5)]
|
|
video, a, fps = SL.BFSShotJoin()._join_all(imgs, [_shot(0, 0, 24, gen=39)], [0])
|
|
self.assertEqual(tuple(a["waveform"].shape[:2]), (1, 2))
|
|
self.assertEqual(a["waveform"].shape[-1], 44100) # 24 frames at 24 fps = 1 s
|
|
self.assertEqual(float(a["waveform"].abs().max()), 0.0)
|
|
self.assertFalse(SL._usable({"waveform": torch.zeros(1, 2, 0), "sample_rate": 44100}))
|
|
|
|
|
|
class RepackTest(unittest.TestCase):
|
|
def test_replaces_only_connected_pieces_and_fits_the_guide(self):
|
|
shot = dict(_shot(0, 0, 20, gen=22), width=8, height=8, frames=torch.zeros(22, 8, 8, 3),
|
|
ref=torch.zeros(1, 4, 4, 3), ref2=None, prompt="old", audio=None)
|
|
ref = torch.ones(1, 6, 6, 3)
|
|
out = SL.BFSShotRepack().repack(shot, ref_image=ref)[0]
|
|
self.assertTrue(torch.equal(out["ref"], ref))
|
|
self.assertEqual(out["prompt"], "old")
|
|
self.assertIs(out["frames"], shot["frames"])
|
|
out = SL.BFSShotRepack().repack(shot, guide_frames=torch.ones(30, 16, 16, 3), prompt="new")[0]
|
|
self.assertEqual(tuple(out["frames"].shape), (22, 8, 8, 3))
|
|
self.assertEqual(out["prompt"], "new")
|
|
self.assertEqual(out["length"], 20)
|
|
|
|
|
|
class ConditioningWriterTest(unittest.TestCase):
|
|
"""The optional prompt writer of BFS Shot H3 Conditioning (template path, no models)."""
|
|
|
|
def setUp(self):
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
import bfs_h3_side_panel as SP
|
|
self.SP, self.seen = SP, {}
|
|
seen = self.seen
|
|
|
|
class R2V:
|
|
@staticmethod
|
|
def execute(**kw):
|
|
seen["prompt"] = kw["prompt"]
|
|
return types.SimpleNamespace(args=("pos", "lat"))
|
|
|
|
mod = types.ModuleType("comfy_extras.nodes_minimax_h3")
|
|
mod.MiniMaxH3ReferenceToVideo, mod.MiniMaxH3AddGuide = R2V, R2V
|
|
self._old = {k: sys.modules.get(k) for k in ("comfy_extras", "comfy_extras.nodes_minimax_h3")}
|
|
sys.modules["comfy_extras"] = types.ModuleType("comfy_extras")
|
|
sys.modules["comfy_extras.nodes_minimax_h3"] = mod
|
|
self._apply = SP.BFSH3SidePanel.apply
|
|
SP.BFSH3SidePanel.apply = lambda self, pos, lat, *a, **k: (pos, lat, {"position": "left"}, None, None)
|
|
self._patch = SP.patch_model_rope
|
|
SP.patch_model_rope = lambda m, info, gap: "patched"
|
|
|
|
def tearDown(self):
|
|
self.SP.BFSH3SidePanel.apply, self.SP.patch_model_rope = self._apply, self._patch
|
|
for k, v in self._old.items():
|
|
if v is None:
|
|
sys.modules.pop(k, None)
|
|
else:
|
|
sys.modules[k] = v
|
|
|
|
def _shot(self):
|
|
return dict(prompt="my own {layout} prompt", width=448, height=800, gen_length=22,
|
|
frames=torch.zeros(22, 8, 8, 3), ref=torch.zeros(1, 4, 4, 3), ref2=None, audio=None)
|
|
|
|
def _run(self, **kw):
|
|
c = SL.BFSShotH3Conditioning()
|
|
return c.condition(self._shot(), "clip", "vae", "none", True, "none", "match", **kw)
|
|
|
|
def test_planner_prompt_is_default_and_layout_filled(self):
|
|
out = self._run()
|
|
self.assertEqual(out[3], "my own prompt")
|
|
out = self._run(duet="canvas")
|
|
self.assertIn("the left half shows the original performance", out[3])
|
|
self.assertEqual(self.seen["prompt"], out[3])
|
|
|
|
def test_task_template_follows_the_duet_mode(self):
|
|
out = self._run(duet="canvas", task="character swap")
|
|
self.assertIn("split screen", out[3])
|
|
self.assertIn("<Picture 1>", out[3])
|
|
out = self._run(duet="shifted RoPE", model="m", task="character swap")
|
|
self.assertNotIn("split screen", out[3])
|
|
self.assertEqual(out[2], "patched")
|
|
|
|
def test_setting_picture_is_the_last_reference_with_static(self):
|
|
seen = {}
|
|
orig = sys.modules["comfy_extras.nodes_minimax_h3"].MiniMaxH3ReferenceToVideo.execute
|
|
|
|
def capture(**kw):
|
|
seen.update(kw)
|
|
return orig(**kw)
|
|
sys.modules["comfy_extras.nodes_minimax_h3"].MiniMaxH3ReferenceToVideo.execute = staticmethod(capture)
|
|
shot = self._shot()
|
|
shot["frames"] = torch.full((22, 8, 8, 3), 0.5)
|
|
mask = torch.zeros(22, 8, 8)
|
|
mask[:, 2:4, 2:4] = 1
|
|
out = SL.BFSShotH3Conditioning().condition(shot, "clip", "vae", "none", True, "none", "match", duet="canvas",
|
|
task="character swap", setting_ref="on (generation size)",
|
|
setting_mask=mask)
|
|
refs = list(seen["ref_images"].values())
|
|
self.assertEqual(len(refs), 2)
|
|
pic = refs[-1][0]
|
|
self.assertTrue(((pic[3, 3] == 0) | (pic[3, 3] == 1)).all()) # static inside the (grown) mask
|
|
self.assertTrue(torch.allclose(pic[7, 7], torch.tensor(0.5))) # the place outside it
|
|
self.assertIn("<Picture 2> shows the setting", out[3])
|
|
self.assertEqual(SL.add_setting("a {setting} b", 3, True), "a <Picture 3> b")
|
|
|
|
def test_task_needs_duet(self):
|
|
with self.assertRaises(ValueError):
|
|
self._run(task="style", instruction="anime")
|
|
|
|
|
|
class SimpleConditioningTest(unittest.TestCase):
|
|
def test_exposes_only_the_common_controls(self):
|
|
inputs = SL.BFSShotH3ConditioningSimple.INPUT_TYPES()
|
|
names = set(inputs["required"]) | set(inputs["optional"])
|
|
self.assertIn("with_audio", names)
|
|
self.assertIn("inpaint", names)
|
|
self.assertNotIn("duet", names)
|
|
self.assertNotIn("model", names)
|
|
self.assertNotIn("rope_gap", names)
|
|
self.assertNotIn("task", names)
|
|
|
|
def test_forwards_safe_advanced_defaults_and_hides_model_output(self):
|
|
seen = {}
|
|
original = SL.BFSShotH3Conditioning.condition
|
|
|
|
def fake(_self, *args, **kwargs):
|
|
seen["args"], seen["kwargs"] = args, kwargs
|
|
return "positive", "latent", "unchanged model", "prompt"
|
|
|
|
SL.BFSShotH3Conditioning.condition = fake
|
|
try:
|
|
out = SL.BFSShotH3ConditioningSimple().condition(
|
|
"shot", "clip", "vae", SL.BFSShotH3ConditioningSimple.GUIDE_MODES[0], True, "match",
|
|
audio_vae="audio vae", with_audio=True)
|
|
finally:
|
|
SL.BFSShotH3Conditioning.condition = original
|
|
|
|
self.assertEqual(out, ("positive", "latent", "prompt"))
|
|
self.assertEqual(seen["args"][5], "none")
|
|
self.assertTrue(seen["kwargs"]["with_audio"])
|
|
self.assertEqual(seen["kwargs"]["inpaint"], SL.INPAINT_MODES[0])
|
|
|
|
|
|
class TargetTest(unittest.TestCase):
|
|
def test_target_placeholder(self):
|
|
self.assertEqual(SL.fill_target("Replace {target} with <Subject 1>", "the man in a red coat"),
|
|
"Replace the man in a red coat with <Subject 1>")
|
|
self.assertEqual(SL.fill_target("Replace {target} with <Subject 1>", ""), "Replace the person with <Subject 1>")
|
|
self.assertEqual(SL.fill_target("no placeholder", "x"), "no placeholder")
|
|
|
|
|
|
class MaskOverlayTest(unittest.TestCase):
|
|
def test_overlay_marks_mask_and_box(self):
|
|
o = torch.zeros(3, 32, 32, 3)
|
|
m = torch.zeros(3, 32, 32); m[:, 8:16, 8:16] = 1
|
|
out = SL._mask_overlay(o, {"mask": m, "box": [0.1, 0.1, 0.9, 0.9]}, 3)
|
|
self.assertGreater(float(out[0, 12, 12, 0]), 0.3) # red inside the mask
|
|
self.assertEqual(float(out[0, 20, 20].sum()), 0.0) # untouched outside mask and box
|
|
self.assertGreater(float(out[0, 3, 16, 1]), 0.5) # yellow box edge
|
|
|
|
|
|
class StitchFinishTest(unittest.TestCase):
|
|
def _shot(self):
|
|
full = torch.full((4, 64, 64, 3), 0.5)
|
|
mask = torch.zeros(4, 64, 64); mask[:, 20:44, 24:40] = 1
|
|
return {"full_frames": full, "crop": {"box": [0.125, 0.125, 0.875, 0.875], "mask": mask, "paste": "mask",
|
|
"expand": 0, "feather": 2}}
|
|
|
|
def test_ring_colour_match_removes_drift_and_keeps_subject_contrast(self):
|
|
shot = self._shot()
|
|
gen = torch.full((4, 48, 48, 3), 0.62) # the model brightened the whole crop (+0.12)
|
|
gen[:, 12:36, 16:32] = 0.2 # a darker new subject
|
|
plain = SL.uncrop(gen, shot)
|
|
fixed = SL.uncrop(gen, shot, match_colors=1.0)
|
|
# around the subject the drift is gone, inside it the new subject stays darker than the scene
|
|
self.assertGreater(abs(float(plain[0, 21, 23].mean()) - 0.5), abs(float(fixed[0, 21, 23].mean()) - 0.5))
|
|
self.assertLess(float(fixed[0, 30, 30].mean()), 0.45)
|
|
|
|
def test_edge_hardness(self):
|
|
a = torch.tensor([[[0.1, 0.5, 0.9]]])
|
|
h = SL._alpha_hardness(a, 1.0)
|
|
self.assertEqual([round(float(x), 2) for x in h[0, 0]], [0.0, 0.5, 1.0])
|
|
self.assertTrue(torch.equal(SL._alpha_hardness(a, 0.0), a))
|
|
|
|
|
|
class InpaintInCropTest(unittest.TestCase):
|
|
def test_latent_starts_from_the_shot_and_only_the_mask_is_generated(self):
|
|
sys.path.insert(0, str(ROOT))
|
|
nested = sys.modules.get("comfy.nested_tensor")
|
|
if nested is None:
|
|
class _N:
|
|
is_nested = True
|
|
|
|
def __init__(self, t):
|
|
self.tensors = list(t)
|
|
comfy = sys.modules.setdefault("comfy", types.ModuleType("comfy"))
|
|
nested = types.ModuleType("comfy.nested_tensor"); nested.NestedTensor = _N
|
|
comfy.nested_tensor = nested
|
|
sys.modules["comfy.nested_tensor"] = nested
|
|
sys.modules.setdefault("comfy.utils", types.ModuleType("comfy.utils"))
|
|
|
|
class VAE:
|
|
def encode(self, px):
|
|
t = ((px.shape[0] - 5) // 17) * 5 + 2
|
|
x = px.mean(-1)[None, None]
|
|
return torch.nn.functional.interpolate(x, size=(t, px.shape[1] // 16, px.shape[2] // 16)).repeat(1, 24, 1, 1, 1)
|
|
|
|
lat = {"samples": nested.NestedTensor((torch.zeros(1, 24, 7, 8, 8), torch.zeros(1, 32, 2, 37)))}
|
|
frames = torch.full((22, 128, 128, 3), 0.4)
|
|
mask = torch.zeros(22, 128, 128); mask[:, 48:80, 48:80] = 1
|
|
pos, out = SL.inpaint_latent(lat, [], VAE(), frames, mask)
|
|
video = out["samples"].tensors[0]
|
|
self.assertAlmostEqual(float(video.mean()), 0.4, places=4) # starts from the shot's own frames
|
|
vm, am = out["noise_mask"].tensors
|
|
self.assertEqual(tuple(vm.shape), (1, 1, 7, 8, 8))
|
|
self.assertEqual(float(vm[0, 0, 3, 4, 4]), 1.0) # the person: generated
|
|
self.assertEqual(float(vm[0, 0, 3, 0, 0]), 0.0) # the rest of the crop: kept
|
|
self.assertEqual(float(vm[0, 0, 3, 2, 4]), 1.0) # grown by one latent cell
|
|
self.assertTrue(bool((am == 1).all())) # audio fully generated
|
|
_, out2 = SL.inpaint_latent(lat, [], VAE(), frames, mask, strength=0.8)
|
|
vm2 = out2["noise_mask"].tensors[0]
|
|
self.assertAlmostEqual(float(vm2[0, 0, 3, 4, 4]), 0.8, places=5) # partly regenerated inside the mask
|
|
self.assertEqual(float(vm2[0, 0, 3, 0, 0]), 0.0) # still kept outside
|
|
|
|
|
|
def test_generation_mask_with_and_without_crop(self):
|
|
m = torch.zeros(5, 64, 64); m[:, 20:40, 20:40] = 1
|
|
cropped = SL.generation_mask({"crop": {"crop_mask": m, "expand": 4}, "frames": torch.zeros(5, 64, 64, 3)})
|
|
self.assertEqual(float(cropped[0, 17, 30]), 1.0) # grown by expand
|
|
self.assertEqual(float(cropped[0, 10, 30]), 0.0)
|
|
with self.assertRaises(ValueError):
|
|
SL.generation_mask({"crop": None, "mask_src": None, "frames": torch.zeros(5, 64, 64, 3)})
|
|
orig = SL.shot_mask
|
|
SL.shot_mask = lambda *a: {"masks": torch.nn.functional.interpolate(m[:, None], size=(32, 32))[:, 0].to(torch.uint8)}
|
|
try:
|
|
full = SL.generation_mask({"crop": None, "frames": torch.zeros(5, 64, 64, 3),
|
|
"mask_src": {"path": "", "analysis": {}, "start": 0, "length": 5,
|
|
"spec": {"expand": 0}}})
|
|
finally:
|
|
SL.shot_mask = orig
|
|
self.assertEqual(tuple(full.shape), (5, 64, 64))
|
|
self.assertTrue(torch.equal(full, m))
|
|
|
|
|
|
def test_comparison_overlay_for_mask_only_shots(self):
|
|
m = torch.zeros(5, 32, 32); m[:, 8:24, 8:24] = 1
|
|
ms = {"path": __file__, "analysis": {"fps": 24.0}, "start": 0, "length": 5, "spec": dict(SL.DEFAULT_MASK, expand=0)}
|
|
shot = {"crop": None, "frames": torch.zeros(5, 64, 64, 3), "mask_src": ms, "inpaint": True}
|
|
orig = SL.shot_mask
|
|
SL.shot_mask = lambda *a: {"masks": m.to(torch.uint8)}
|
|
try:
|
|
ov = SL.overlay_of(shot)
|
|
self.assertEqual(tuple(ov["mask"].shape), (5, 64, 64))
|
|
self.assertIsNone(SL.overlay_of(dict(shot, inpaint=False))) # not used and not segmented yet: no SAM 3
|
|
self.assertIsNone(SL.overlay_of({"crop": None, "frames": shot["frames"]}))
|
|
finally:
|
|
SL.shot_mask = orig
|
|
o = SL._mask_overlay(torch.full((5, 32, 32, 3), 0.5), ov, 5)
|
|
self.assertGreater(float(o[0, 16, 16, 0]), float(o[0, 16, 16, 1])) # red inside the mask
|
|
self.assertAlmostEqual(float(o[0, 2, 2, 0]), 0.5) # untouched outside
|
|
|
|
|
|
class ExternalMaskTest(unittest.TestCase):
|
|
def _video(self, name, n=48, fps=24.0, w=64, h=36, box=(10, 5, 30, 25)):
|
|
import cv2
|
|
import numpy as np
|
|
path = str(Path(_TMP) / name)
|
|
vw = cv2.VideoWriter(path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
|
|
for i in range(n):
|
|
f = np.zeros((h, w, 3), np.uint8)
|
|
x0, y0, x1, y1 = box
|
|
f[y0:y1, x0 + i % 4:x1 + i % 4] = 255
|
|
vw.write(f)
|
|
vw.release()
|
|
return path
|
|
|
|
def setUp(self):
|
|
self.an = {"fps": 24.0, "n": 48, "n_src": 48, "fps_src": 24.0, "width": 64, "height": 36}
|
|
|
|
def test_precedence(self):
|
|
p = {"mask_video": "global.mp4"}
|
|
self.assertEqual(SL.plan_spec(p, {"video": "shot.mp4", "text": "man"})["ext"]["scope"], "shot")
|
|
self.assertNotIn("ext", SL.plan_spec(p, {"text": "man"})) # the shot's own SAM 3 text wins
|
|
self.assertEqual(SL.plan_spec(p, {})["ext"], {"file": "global.mp4", "scope": "video"})
|
|
SL._EXT_MASK["v.mp4"] = torch.ones(48, 36, 64)
|
|
try:
|
|
self.assertEqual(SL.plan_spec(p, {}, "v.mp4")["ext"], {"tensor": "v.mp4"}) # the node input beats the panel file
|
|
finally:
|
|
SL._EXT_MASK.pop("v.mp4")
|
|
self.assertFalse(SL.spec_has_mask(SL.plan_spec({}, {})))
|
|
|
|
def test_file_and_tensor_masks_replace_sam(self):
|
|
path = self._video("roto_whole.mp4")
|
|
self._video("roto_shot.mp4", n=12)
|
|
orig = SL.segment_frames
|
|
SL.segment_frames = lambda *a, **k: self.fail("SAM 3 must not run with an external mask")
|
|
try:
|
|
r = SL.shot_mask(path, self.an, 24, 12, SL.plan_spec({"mask_video": "roto_whole.mp4"}, {}))
|
|
self.assertEqual(r["masks"].shape[0], 12)
|
|
self.assertTrue(r["box"] is not None)
|
|
cov = float(r["masks"].float().mean())
|
|
self.assertTrue(0.1 < cov < 0.4, cov)
|
|
r2 = SL.shot_mask(path, self.an, 24, 12, SL.plan_spec({}, {"video": "roto_shot.mp4"}))
|
|
self.assertEqual(r2["masks"].shape[0], 12)
|
|
t = torch.zeros(48, 36, 64); t[:, 0:18, :] = 1 # top half
|
|
SL._EXT_MASK[path] = t
|
|
r3 = SL.shot_mask(path, self.an, 0, 10, SL.plan_spec({}, {}, path))
|
|
m = r3["masks"].float()
|
|
self.assertGreater(float(m[:, :8].mean()), 0.9)
|
|
self.assertLess(float(m[:, -8:].mean()), 0.1)
|
|
finally:
|
|
SL.segment_frames = orig
|
|
SL._EXT_MASK.pop(path, None)
|
|
|
|
|
|
class FramePasteTest(unittest.TestCase):
|
|
def test_paste_back_keeps_the_background_and_takes_the_new_outline(self):
|
|
base = torch.full((5, 64, 64, 3), 0.2) # the original frames
|
|
result = torch.full((5, 64, 64, 3), 0.9) # generated: everything changed
|
|
old = torch.zeros(5, 64, 64); old[:, 20:44, 20:36] = 1 # the original person
|
|
new = torch.zeros(5, 64, 64); new[:, 16:48, 18:44] = 1 # the new one is bigger
|
|
ms = {"path": "", "analysis": {}, "start": 0, "length": 5, "spec": dict(SL.DEFAULT_MASK, expand=0, feather=0)}
|
|
shot = {"index": 0, "frames": base, "mask_src": ms, "paste": True}
|
|
o1, o2 = SL.shot_mask, SL.segment_frames
|
|
SL.shot_mask = lambda *a: {"masks": old.to(torch.uint8)}
|
|
SL.segment_frames = lambda imgs, sp: torch.nn.functional.interpolate(new[:, None], size=imgs.shape[1:3])[:, 0]
|
|
try:
|
|
out = SL.paste_back(result, shot, True)
|
|
only_old = SL.paste_back(result, shot, False)
|
|
finally:
|
|
SL.shot_mask, SL.segment_frames = o1, o2
|
|
self.assertAlmostEqual(float(out[0, 2, 2, 0]), 0.2, places=4) # background: the original
|
|
self.assertAlmostEqual(float(out[0, 30, 28, 0]), 0.9, places=4) # the person: generated
|
|
self.assertAlmostEqual(float(out[0, 30, 41, 0]), 0.9, places=4) # the bigger new outline is pasted too
|
|
self.assertAlmostEqual(float(only_old[0, 30, 41, 0]), 0.2, places=4)
|
|
|
|
|
|
class MaskGuideTest(unittest.TestCase):
|
|
def test_masked_only_greys_everything_outside_the_mask(self):
|
|
fr = torch.full((4, 32, 32, 3), 0.9)
|
|
m = torch.zeros(4, 16, 16); m[:, 4:12, 4:12] = 1
|
|
g = SL.masked_only(fr, m)
|
|
self.assertEqual(tuple(g.shape), (4, 32, 32, 3))
|
|
self.assertAlmostEqual(float(g[0, 16, 16, 0]), 0.9) # the masked region
|
|
self.assertAlmostEqual(float(g[0, 2, 2, 0]), 0.5) # grey elsewhere
|
|
self.assertEqual(SL.MASK_GUIDES[0], "off")
|
|
fr2 = torch.rand(3, 64, 64, 3)
|
|
m2 = torch.zeros(3, 64, 64); m2[:, 10:50, 10:50] = 1
|
|
old_det = dict(SL._DET); SL._DET["pose"] = None # no YOLO pose model: a clear error
|
|
try:
|
|
with self.assertRaisesRegex(ValueError, "YOLO pose model"):
|
|
SL.masked_only(fr2, m2, look="pose (people)")
|
|
finally:
|
|
SL._DET.clear(); SL._DET.update(old_det)
|
|
for look in [x for x in SL.MASK_GUIDE_LOOKS if not x.startswith("pose")]:
|
|
g2 = SL.masked_only(fr2, m2, look=look)
|
|
self.assertEqual(tuple(g2.shape), (3, 64, 64, 3), look)
|
|
self.assertAlmostEqual(float(g2[0, 2, 2, 0]), 0.5, msg=look) # grey outside, every look
|
|
if look != "colour": # no colour inside
|
|
self.assertLess(float((g2[..., 0] - g2[..., 1]).abs().max()), 1e-5, look)
|
|
t = SL.add_mask_video("x subject_definitions: <Subject 1> is ...", 2)
|
|
self.assertIn("<Video 2> shows only the region", t)
|
|
self.assertEqual(SL.add_mask_video("follow {mask_video}", 1), "follow <Video 1>")
|
|
|
|
|
|
class ConditionFlowTest(unittest.TestCase):
|
|
"""BFS Shot H3 Conditioning with the native H3 nodes replaced by recorders."""
|
|
|
|
def setUp(self):
|
|
sys.path.insert(0, str(ROOT))
|
|
calls = self.calls = {"r2v": [], "guides": []}
|
|
|
|
class Out:
|
|
def __init__(self, *a):
|
|
self.args = a
|
|
|
|
class R2V:
|
|
@staticmethod
|
|
def execute(**kw):
|
|
calls["r2v"].append(kw)
|
|
N = sys.modules["comfy.nested_tensor"].NestedTensor
|
|
return Out([[torch.zeros(1), {}]], {"samples": N((torch.zeros(1, 24, 2, 4, 4), torch.zeros(1, 32, 2, 37)))})
|
|
|
|
class Guide:
|
|
@staticmethod
|
|
def execute(positive, latent, frame_idx, vae, image=None, **kw):
|
|
calls["guides"].append(image)
|
|
return Out(positive)
|
|
|
|
mod = types.ModuleType("comfy_extras.nodes_minimax_h3")
|
|
mod.MiniMaxH3ReferenceToVideo, mod.MiniMaxH3AddGuide = R2V, Guide
|
|
sys.modules.setdefault("comfy_extras", types.ModuleType("comfy_extras"))
|
|
self._old = sys.modules.get("comfy_extras.nodes_minimax_h3")
|
|
sys.modules["comfy_extras.nodes_minimax_h3"] = mod
|
|
if "comfy.nested_tensor" not in sys.modules:
|
|
class _N:
|
|
is_nested = True
|
|
|
|
def __init__(self, t):
|
|
self.tensors = list(t)
|
|
comfy = sys.modules.setdefault("comfy", types.ModuleType("comfy"))
|
|
nt = types.ModuleType("comfy.nested_tensor"); nt.NestedTensor = _N
|
|
comfy.nested_tensor = nt; sys.modules["comfy.nested_tensor"] = nt
|
|
|
|
def tearDown(self):
|
|
if self._old is not None:
|
|
sys.modules["comfy_extras.nodes_minimax_h3"] = self._old
|
|
else:
|
|
sys.modules.pop("comfy_extras.nodes_minimax_h3", None)
|
|
|
|
def _shot(self, **kw):
|
|
m = torch.zeros(22, 64, 64); m[:, 16:48, 16:40] = 1
|
|
shot = {"index": 0, "count": 1, "frames": torch.full((22, 64, 64, 3), 0.7), "ref": torch.zeros(1, 32, 32, 3),
|
|
"ref2": None, "prompt": "subject_definitions: <Subject 1> is ...", "width": 64, "height": 64,
|
|
"gen_length": 22, "audio": None, "crop": {"crop_mask": m, "expand": 0} if kw.get("crop") else None,
|
|
"inpaint": kw.get("inpaint", True)}
|
|
if not kw.get("crop"):
|
|
shot["mask_src"] = {"path": __file__, "analysis": {"fps": 24.0}, "start": 0, "length": 22, "spec": dict(SL.DEFAULT_MASK, expand=0)}
|
|
return shot, m
|
|
|
|
def run_cond(self, shot, m, **kw):
|
|
class VAE:
|
|
def encode(self, px):
|
|
return torch.zeros(1, 24, 2, px.shape[1] // 16, px.shape[2] // 16)
|
|
orig = SL.shot_mask
|
|
SL.shot_mask = lambda *a: {"masks": m.to(torch.uint8)}
|
|
try:
|
|
return SL.BFSShotH3Conditioning().condition(shot, None, VAE(), SL.BFSShotH3Conditioning.GUIDE_MODES[0], False,
|
|
"none", "match", **kw)
|
|
finally:
|
|
SL.shot_mask = orig
|
|
|
|
def test_reference_video_of_the_masked_region(self):
|
|
shot, m = self._shot()
|
|
self.run_cond(shot, m, inpaint=SL.INPAINT_MODES[0], mask_guide=SL.MASK_GUIDES[3], mask_ref_size="full",
|
|
mask_guide_look="colour")
|
|
kw = self.calls["r2v"][0]
|
|
v = kw["ref_videos"]["ref_video_1"]
|
|
self.assertAlmostEqual(float(v[0, 2, 2, 0]), 0.5) # grey outside the mask
|
|
self.assertAlmostEqual(float(v[0, 30, 30, 0]), 0.7) # the subject
|
|
self.assertIn("<Video 1> shows only the region", kw["prompt"])
|
|
self.assertEqual(len(self.calls["guides"]), 1) # the normal aligned guide is still there
|
|
self.calls["r2v"].clear()
|
|
self.run_cond(shot, m, inpaint=SL.INPAINT_MODES[0], mask_guide=SL.MASK_GUIDES[3], mask_ref_size="1/2")
|
|
self.assertEqual(tuple(self.calls["r2v"][0]["ref_videos"]["ref_video_1"].shape[1:3]), (32, 32))
|
|
|
|
def test_extra_and_instead_guides(self):
|
|
shot, m = self._shot()
|
|
self.run_cond(shot, m, inpaint=SL.INPAINT_MODES[0], mask_guide=SL.MASK_GUIDES[1])
|
|
self.assertEqual(len(self.calls["guides"]), 2)
|
|
self.calls["guides"].clear()
|
|
self.run_cond(shot, m, inpaint=SL.INPAINT_MODES[0], mask_guide=SL.MASK_GUIDES[2])
|
|
self.assertEqual(len(self.calls["guides"]), 1)
|
|
self.assertAlmostEqual(float(self.calls["guides"][0][0, 2, 2, 0]), 0.5) # it is the masked one
|
|
|
|
def test_ignored_outside_mask_only(self):
|
|
shot, m = self._shot(inpaint=False)
|
|
self.run_cond(shot, m, inpaint=SL.INPAINT_MODES[0], mask_guide=SL.MASK_GUIDES[3])
|
|
self.assertNotIn("ref_videos", self.calls["r2v"][0])
|
|
shot, m = self._shot(crop=True)
|
|
self.calls["r2v"].clear()
|
|
self.run_cond(shot, m, inpaint=SL.INPAINT_MODES[0], mask_guide=SL.MASK_GUIDES[3])
|
|
self.assertNotIn("ref_videos", self.calls["r2v"][0])
|
|
|
|
|
|
class DescribeShortTest(unittest.TestCase):
|
|
def test_short_details_fit_inside_the_prompt(self):
|
|
self.assertEqual(SL.tidy_details('"The green goblin in a torn leather vest."\nextra', {"describe_preset": "short"}),
|
|
"green goblin in a torn leather vest")
|
|
self.assertEqual(SL.tidy_details("A young woman.", {"describe_preset": "full body"}), "A young woman.")
|
|
self.assertIn("one English noun phrase", SL.describe_instruction({}))
|
|
|
|
def test_every_preset_removes_vlm_answer_labels(self):
|
|
for preset in SL.DESCRIBE_PRESETS:
|
|
with self.subTest(preset=preset):
|
|
self.assertEqual(SL.tidy_details("**Answer:** elderly woman with large glasses", {"describe_preset": preset}),
|
|
"elderly woman with large glasses")
|
|
|
|
def test_short_mode_honours_its_twenty_word_contract(self):
|
|
text = "Answer: The " + " ".join(f"word{i}" for i in range(30))
|
|
out = SL.tidy_details(text, {"describe_preset": "short"})
|
|
self.assertEqual(len(out.split()), 20)
|
|
self.assertFalse(out.lower().startswith(("answer", "the ")))
|
|
|
|
def test_long_modes_are_bounded_but_not_collapsed_to_short(self):
|
|
text = " ".join(f"word{i}" for i in range(170))
|
|
full = SL.tidy_details(text, {"describe_preset": "full body"})
|
|
face = SL.tidy_details(text, {"describe_preset": "head / face"})
|
|
self.assertEqual(len(full.split()), 150)
|
|
self.assertEqual(len(face.split()), 110)
|
|
self.assertTrue(full.endswith("."))
|
|
|
|
|
|
class FitPromptTest(unittest.TestCase):
|
|
P = ("body_swap: subject_definitions: <Subject 1> is the {details} shown in <Picture 1>, whose identity define the "
|
|
"character. The guide video provides the scene; only {target} is replaced: their face become <Subject 1>. "
|
|
"summary: Replace {target} in the guide video with <Subject 1>.")
|
|
|
|
def test_slot_is_the_sentence_with_a_blank(self):
|
|
d = SL.prompt_slot(self.P, "{details}")
|
|
self.assertIn("<Subject 1> is the ___ shown in <Picture 1>", d)
|
|
self.assertNotIn("{target}", d)
|
|
t = SL.prompt_slot(self.P, "{target}")
|
|
self.assertIn("only ___ is replaced", t)
|
|
self.assertEqual(SL.prompt_slot("no marker here", "{details}"), "")
|
|
|
|
def test_answer_matches_the_blank(self):
|
|
d = SL.prompt_slot(self.P, "{details}")
|
|
self.assertEqual(SL.tidy_slot("A golden retriever with a thick coat.", d), "golden retriever with a thick coat")
|
|
t = SL.prompt_slot(self.P, "{target}")
|
|
self.assertEqual(SL.tidy_slot("the young woman in a black corset.", t), "the young woman in a black corset")
|
|
self.assertIn("goes into the blank", SL.describe_instruction({"describe_preset": "short"}, d))
|
|
self.assertNotIn("goes into the blank", SL.describe_instruction({"describe_preset": "short", "fit_prompt": False}, d))
|
|
self.assertNotIn("goes into the blank", SL.describe_instruction({"describe_preset": "full body"}, d))
|
|
# the VLM echoing the sentence (the bug the user hit)
|
|
echo = ("humanoid robot is the ___ shown in <Picture 1>, whose identity define the character. The robot has a sleek "
|
|
"body with a metallic silver torso.")
|
|
self.assertEqual(SL.tidy_slot(echo, d), "humanoid robot")
|
|
self.assertEqual(SL.tidy_slot("sleek silver humanoid robot with black joints shown in <Picture 1>", d),
|
|
"sleek silver humanoid robot with black joints")
|
|
self.assertEqual(SL.tidy_slot("<Subject 1> is the green goblin in a vest", d), "green goblin in a vest")
|
|
|
|
def test_target_label_is_removed_even_without_a_prompt_slot(self):
|
|
old = SL.vlm_generate
|
|
SL.vlm_generate = lambda *args, **kwargs: "Answer: elderly woman with large tortoiseshell glasses"
|
|
try:
|
|
crop = __import__("numpy").zeros((8, 8, 3), dtype="uint8")
|
|
self.assertEqual(SL.describe_target(object(), crop),
|
|
"the elderly woman with large tortoiseshell glasses")
|
|
finally:
|
|
SL.vlm_generate = old
|
|
|
|
|
|
class RangeTest(unittest.TestCase):
|
|
def test_start_end_and_max_total_from_the_start(self):
|
|
self.assertEqual(SL.plan_range(1000, 24.0, 10, 20), (240, 480))
|
|
self.assertEqual(SL.plan_range(1000, 24.0, 10, 0), (240, 1000))
|
|
self.assertEqual(SL.plan_range(1000, 24.0, 10, 0, 5), (240, 360)) # max total counted from the start
|
|
self.assertEqual(SL.plan_range(1000, 24.0, 0, 0, 0, start_f=250, end_f=500), (250, 500)) # frames win
|
|
self.assertEqual(SL.plan_range(1000, 24.0, 0, 0, 0, start_f=250, end_f=0), (250, 1000))
|
|
segs = SL.plan_segments(480, [100, 300], "shots", 107, 10, start=240)
|
|
self.assertEqual(segs[0]["start"], 240)
|
|
self.assertEqual(segs[-1]["end"], 480)
|
|
self.assertTrue(all(240 <= s["start"] < s["end"] <= 480 for s in segs))
|
|
self.assertTrue(any(s["start"] == 300 for s in segs)) # the cut inside the range is kept
|
|
fixed = SL.plan_segments(480, [], "fixed", 107, 10, start=240)
|
|
self.assertEqual((fixed[0]["start"], fixed[-1]["end"]), (240, 480))
|
|
manual = SL.plan_segments(480, [], "manual", 107, 10, manual=[100, 350], start=240)
|
|
self.assertEqual([(s["start"], s["end"]) for s in manual], [(240, 350), (350, 480)])
|
|
|
|
|
|
class TrackModeTest(unittest.TestCase):
|
|
def test_track_is_used_by_shots_without_their_own_selection_and_sliced_per_shot(self):
|
|
p = {"track": {"points": [{"x": 0.5, "y": 0.5, "label": 1}], "frame": 120, "text": "woman"}, "fps": 24}
|
|
self.assertIn("track", SL.plan_spec(p, {}, "v.mp4")["ext"])
|
|
self.assertNotIn("ext", SL.plan_spec(p, {"text": "man"}, "v.mp4")) # the shot's own text wins
|
|
self.assertEqual(SL.plan_spec(p, {"video": "m.mp4"}, "v.mp4")["ext"]["scope"], "shot")
|
|
whole = torch.zeros(200, 36, 64, dtype=torch.uint8)
|
|
whole[100:150, 10:20, 10:20] = 1 # the subject only in frames 100-149
|
|
orig = SL.track_video
|
|
SL.track_video = lambda path, a, t: (whole, 50) # the track covers frames 50-249
|
|
try:
|
|
spec = SL.plan_spec(p, {}, "v.mp4")
|
|
an = {"n": 300, "fps": 24.0}
|
|
m = SL.ext_masks(spec["ext"], an, 140, 20, 64, 36) # shot 140-159
|
|
self.assertEqual(tuple(m.shape), (20, 36, 64))
|
|
self.assertFalse(bool(m[0].any())) # frame 140 = track index 90
|
|
self.assertTrue(bool(m[-1].any())) # frame 159 = index 109
|
|
finally:
|
|
SL.track_video = orig
|
|
|
|
|
|
class DuetPromptTest(unittest.TestCase):
|
|
def test_fields_are_found_and_the_vlm_fills_them(self):
|
|
tpl = ("subject_definitions:\n<Subject 1> is the woman from <Picture 1>: {skin, face}.\n<Subject 2> is the {room} "
|
|
"from {setting}.\n\nsummary:\nx {layout}")
|
|
self.assertEqual(SL.template_slots(tpl), ["skin, face", "room"])
|
|
asked = []
|
|
|
|
def gen(clip, q, images, max_tokens, temperature=0.0, seed=0):
|
|
asked.append(q)
|
|
if images is None:
|
|
return tpl.replace("{skin, face}", "pale skin, oval face").replace("{room}", "kitchen")
|
|
return '{"medium": "phone footage", "actions": "the person waves"}'
|
|
|
|
old = SL.vlm_generate, SL.vlm_describe
|
|
SL.vlm_generate, SL.vlm_describe = gen, lambda clip, ims, q, mt=320, temperature=0.0, seed=0: "the kitchen with white cabinets"
|
|
try:
|
|
fr = torch.rand(9, 32, 32, 3)
|
|
out = SL.duet_prompt(object(), tpl, fr, [torch.rand(1, 32, 32, 3)], "character swap", "", 512,
|
|
setting=torch.rand(1, 32, 32, 3))
|
|
self.assertIn("pale skin, oval face", out); self.assertIn("{setting}", out); self.assertIn("{layout}", out)
|
|
self.assertIn("the person waves", asked[-1])
|
|
self.assertEqual(SL.duet_prompt(None, tpl, fr, [], "character swap", ""), tpl.strip())
|
|
self.assertEqual(SL.duet_prompt(object(), "a full prompt", fr, [], "character swap", ""), "a full prompt")
|
|
w = SL.duet_prompt(object(), "", fr, [torch.rand(1, 32, 32, 3)], "character swap", "", 512,
|
|
setting=torch.rand(1, 32, 32, 3))
|
|
self.assertIn("<Subject 2> is the kitchen with white cabinets from {setting}.", w)
|
|
finally:
|
|
SL.vlm_generate, SL.vlm_describe = old
|
|
|
|
|
|
class ShotHistoryTest(unittest.TestCase):
|
|
def test_history_is_taken_out_of_the_shot(self):
|
|
h = torch.full((17, 8, 8, 3), 255, dtype=torch.uint8)
|
|
shot = {"chain": "history", "index": 1, "count": 2, "history": h, "history_src": h}
|
|
hist, src = SL.shot_history(shot)
|
|
self.assertEqual((tuple(hist.shape), float(hist.max()), "history" in shot), ((17, 8, 8, 3), 1.0, False))
|
|
self.assertEqual(SL.shot_history({"chain": "off", "index": 1}), (None, None))
|
|
|
|
def test_auto_loop_takes_the_previous_render(self):
|
|
SL._LAST_RESULT.clear()
|
|
SL._LAST_RESULT.update(index=0, count=2, frames=torch.rand(40, 8, 8, 3))
|
|
hist, _ = SL.shot_history({"chain": "history", "index": 1, "count": 2, "prev_length": 30, "chain_history": 34})
|
|
self.assertEqual(hist.shape[0], 22) # 30 real frames -> 22, the longest 17k+5 that fits
|
|
SL._LAST_RESULT.clear()
|
|
|
|
def test_history_lengths_are_whole_vae_clips(self):
|
|
# 17 frames encode to 2 tokens (the first 5 frames): only 17k+5 reaches the boundary
|
|
self.assertEqual([SL.history_frame_count(w, 200) for w in (5, 17, 22, 34, 39, 51, 56)], [5, 22, 22, 39, 39, 56, 56])
|
|
self.assertEqual([SL.history_frame_count(56, a) for a in (4, 5, 21, 22, 38, 40)], [0, 5, 5, 22, 22, 39])
|
|
|
|
|
|
class ShotHeadTest(unittest.TestCase):
|
|
"""Continuity 'head': the previous shot's end starts the next shot, pinned in the latent, and the join cuts it."""
|
|
|
|
@staticmethod
|
|
def _segs(chain="head", n=22, cut=False, gap=0):
|
|
return [{"start": 0, "end": 60, "chain": "off", "cut_before": False},
|
|
{"start": 60 + gap, "end": 120, "chain": chain, "chain_history": n, "cut_before": cut}]
|
|
|
|
def test_head_length(self):
|
|
self.assertEqual(SL.head_length(self._segs(), 1, H3), 22)
|
|
self.assertEqual(SL.head_length(self._segs(n=56), 1, H3), 39) # capped at 39
|
|
self.assertEqual(SL.head_length(self._segs(n=17), 1, H3), 22) # old plans: next valid length up
|
|
self.assertEqual(SL.head_length(self._segs(), 0, H3), 0) # first shot
|
|
self.assertEqual(SL.head_length(self._segs(cut=True), 1, H3), 0) # a camera cut
|
|
self.assertEqual(SL.head_length(self._segs(gap=5), 1, H3), 0) # a skipped shot in between
|
|
self.assertEqual(SL.head_length(self._segs(chain="history"), 1, H3), 0)
|
|
self.assertEqual(SL.head_length(self._segs(), 1, "LTX / Wan (8n+1)"), 0)
|
|
|
|
def test_join_cuts_the_head(self):
|
|
img = torch.arange(30, dtype=torch.float32).view(30, 1, 1, 1).expand(30, 2, 2, 3)
|
|
shot = {"head": 5, "gen_length": 30, "frames": img.clone(), "crop": {"mask": torch.zeros(30, 2, 2), "box": [0, 0, 1, 1]}}
|
|
out, sh = SL.drop_head(img, shot)
|
|
self.assertEqual((out.shape[0], float(out[0, 0, 0, 0]), sh["gen_length"], sh["frames"].shape[0],
|
|
sh["crop"]["mask"].shape[0], sh["head"], shot["frames"].shape[0]), (25, 5.0, 25, 25, 25, 0, 30))
|
|
self.assertIs(SL.drop_head(img, {"head": 0})[0], img)
|
|
|
|
def test_remember_result_skips_the_head(self):
|
|
SL._LAST_RESULT.clear()
|
|
SL.remember_result({"index": 0, "count": 2, "head": 5}, torch.zeros(30, 2, 2, 3))
|
|
self.assertEqual(SL._LAST_RESULT["frames"].shape[0], 25)
|
|
SL._LAST_RESULT.clear()
|
|
|
|
def test_pin_head_keeps_the_first_latent_steps(self):
|
|
class NT:
|
|
def __init__(self, t):
|
|
self.tensors, self.is_nested = list(t), True
|
|
nt = types.ModuleType("comfy.nested_tensor")
|
|
nt.NestedTensor = NT
|
|
old = {k: sys.modules.get(k) for k in ("comfy", "comfy.nested_tensor")}
|
|
comfy = types.ModuleType("comfy")
|
|
comfy.nested_tensor = nt
|
|
sys.modules["comfy"], sys.modules["comfy.nested_tensor"] = comfy, nt
|
|
try:
|
|
import bfs_h3_history as HH
|
|
|
|
class VAE:
|
|
def encode(self, px):
|
|
n = px.shape[0]
|
|
self.size = tuple(px.shape[1:3])
|
|
return torch.full((1, 24, 5 * -(-n // 17) - 3, px.shape[1] // 16, px.shape[2] // 16), 7.0)
|
|
vae = VAE()
|
|
latent = {"samples": NT((torch.zeros(1, 24, 12, 4, 4), torch.zeros(1, 8, 50)))}
|
|
out, n = HH.pin_head(latent, vae, torch.rand(30, 32, 32, 3)) # 30 frames -> the last 22 -> 7 steps
|
|
v, m = out["samples"].tensors[0], out["noise_mask"].tensors[0]
|
|
self.assertEqual((n, vae.size), (22, (64, 64)))
|
|
self.assertTrue(bool((v[:, :, :7] == 7).all()) and bool((v[:, :, 7:] == 0).all()))
|
|
self.assertTrue(bool((m[:, :, :7] == 0).all()) and bool((m[:, :, 7:] == 1).all()))
|
|
self.assertTrue(bool((out["noise_mask"].tensors[1] == 1).all()))
|
|
# an existing inpaint mask stays as it is after the head
|
|
inp = torch.full((1, 1, 12, 4, 4), 0.5)
|
|
out2, _ = HH.pin_head(dict(latent, noise_mask=NT((inp, torch.ones(1, 8, 50)))), vae, torch.rand(22, 64, 64, 3))
|
|
m2 = out2["noise_mask"].tensors[0]
|
|
self.assertTrue(bool((m2[:, :, :7] == 0).all()) and bool((m2[:, :, 7:] == 0.5).all()) and bool((inp == 0.5).all()))
|
|
finally:
|
|
for k, v in old.items():
|
|
if v is None:
|
|
sys.modules.pop(k, None)
|
|
else:
|
|
sys.modules[k] = v
|
|
|
|
|
|
class MemoryStatusTest(unittest.TestCase):
|
|
def test_shot_tokens(self):
|
|
# 107 H3 frames -> 7 clips -> 32 latent steps; 960x544 -> 60 x 34 cells
|
|
self.assertEqual(SL.shot_tokens(107, 960, 544), 32 * 34 * 60)
|
|
self.assertEqual(SL.shot_tokens(22, 32, 32), 7 * 2 * 2)
|
|
self.assertGreater(SL.shot_tokens(90, 960, 544), SL.shot_tokens(73, 960, 544)) # a head makes the shot bigger
|
|
self.assertEqual(SL.shot_tokens(81, 32, 32, "LTX / Wan (8n+1)"), 21 * 2 * 2)
|
|
|
|
def test_record_and_status(self):
|
|
import json as _json
|
|
path = Path(_TMP) / SL._VRAM_FILE
|
|
path.unlink(missing_ok=True)
|
|
cuda = types.SimpleNamespace(is_available=lambda: True, is_initialized=lambda: True, current_device=lambda: 0,
|
|
get_device_name=lambda d: "Test GPU", max_memory_allocated=lambda: 12 * 2 ** 30,
|
|
reset_peak_memory_stats=lambda: None, mem_get_info=lambda: (8 * 2 ** 30, 32 * 2 ** 30))
|
|
with patch.object(SL.torch, "cuda", cuda):
|
|
SL.vram_record(5000)
|
|
SL.vram_record(3000) # smaller: the biggest stays
|
|
st = SL.vram_status()
|
|
self.assertEqual((st["gpu"], st["max_tokens"], st["max_peak_gb"], st["runs"], st["total_gb"], st["free_gb"]),
|
|
("Test GPU", 5000, 12.0, 2, 32.0, 8.0))
|
|
self.assertIn("Test GPU", _json.loads(path.read_text()))
|
|
SL.vram_record(4000) # no CUDA in use: nothing recorded, nothing started
|
|
self.assertEqual(len(_json.loads(path.read_text())["Test GPU"]["samples"]), 2)
|
|
|
|
|
|
class ShotStoreTest(unittest.TestCase):
|
|
def test_storage_round_trip(self):
|
|
d = tempfile.mkdtemp(dir=_TMP)
|
|
y = torch.linspace(0, 1, 48).view(1, 48, 1, 1)
|
|
frames = (((torch.sin(torch.arange(12).view(12, 1, 1, 1) + 9 * y) + 1) / 2).expand(12, 48, 64, 3) * 255).round().byte()
|
|
for k, mode in enumerate(("lossless", "raw", "smaller")):
|
|
SL.store_shot(d, k, frames, {"index": k, "fps": 24.0}, mode)
|
|
back, meta = SL.load_shot(d, k)
|
|
self.assertEqual((tuple(back.shape), meta["index"]), ((12, 48, 64, 3), k))
|
|
diff = int((back.int() - frames.int()).abs().max())
|
|
if mode == "smaller":
|
|
self.assertLess(diff, 40)
|
|
else:
|
|
self.assertEqual(diff, 0, mode) # lossless and raw are exact
|
|
self.assertTrue(os.path.exists(os.path.join(d, "shot_0000.mkv")))
|
|
self.assertFalse(os.path.exists(os.path.join(d, "shot_0001.mkv"))) # raw: the frames are in the .pt
|
|
|
|
def test_status_cascades_through_continuity(self):
|
|
d = tempfile.mkdtemp(dir=_TMP)
|
|
for i in range(3):
|
|
open(os.path.join(d, f"shot_{i:04d}.pt"), "wb").close()
|
|
st = {"done": [0, 1, 2], "sig": {"0": "a", "1": "b", "2": "c"}, "gen": {"0": "g0", "1": "g1", "2": "g2"},
|
|
"from": {"1": "g0", "2": "g1"}, "cur": {"sigs": ["a", "b", "c"], "prev": [False, True, False]}}
|
|
self.assertEqual(SL.shot_status(st, d), [True, True, True])
|
|
st["gen"]["0"] = "new" # shot 1 was redone: shot 2 continues from it
|
|
self.assertEqual(SL.shot_status(st, d), [True, False, True]) # shot 3 does not continue (cut)
|
|
st["cur"]["sigs"][2] = "changed" # a setting of shot 3 changed
|
|
self.assertEqual(SL.shot_status(st, d), [True, False, False])
|
|
|
|
def test_join_can_write_an_mp4(self):
|
|
import cv2
|
|
d = tempfile.mkdtemp(dir=_TMP)
|
|
imgs = [torch.full((22, 32, 32, 3), 0.2), torch.full((39, 32, 32, 3), 0.8)]
|
|
shots = [_shot(0, 0, 20, gen=22), _shot(1, 20, 45, cut=True, gen=39)]
|
|
audio = {"waveform": torch.zeros(1, 2, 48000 * 3), "sample_rate": 48000}
|
|
with patch.object(SL, "_next_output_path", lambda prefix: os.path.join(d, prefix + "_00001.mp4")):
|
|
out = SL.BFSShotJoin().join(imgs, shots, [0], audio=[audio], save_to=[SL.SAVE_TO[1]],
|
|
filename_prefix=["test"])
|
|
res = out["result"]
|
|
self.assertEqual(res[0].shape[0], 45) # the output still carries the video (read back)
|
|
self.assertAlmostEqual(float(res[0][30].mean()), 0.8, places=1)
|
|
self.assertTrue(res[4].endswith("test_00001.mp4"))
|
|
self.assertEqual(out["ui"]["images"][0]["filename"], "test_00001.mp4") # shown on the node
|
|
cap = cv2.VideoCapture(res[4]); n = 0
|
|
while cap.read()[0]:
|
|
n += 1
|
|
self.assertEqual(n, 45)
|
|
|
|
|
|
class JoinRamGuardTest(unittest.TestCase):
|
|
def test_images_join_falls_back_to_mp4_when_ram_is_short(self):
|
|
d = tempfile.mkdtemp(dir=_TMP)
|
|
imgs = [torch.full((22, 32, 32, 3), 0.2), torch.full((39, 32, 32, 3), 0.8)]
|
|
shots = [_shot(0, 0, 20, gen=22), _shot(1, 20, 45, cut=True, gen=39)]
|
|
with patch.object(SL, "_ram_available", lambda: 1000), \
|
|
patch.object(SL, "_next_output_path", lambda prefix: os.path.join(d, prefix + "_00001.mp4")):
|
|
out = SL.BFSShotJoin().join(imgs, shots, [0], comparison=[True])
|
|
res = out["result"]
|
|
self.assertTrue(os.path.exists(res[4]) and os.path.exists(res[4][:-4] + "_comparison.mp4"))
|
|
self.assertEqual(res[0].shape[0], 45) # never an empty output node
|
|
self.assertTrue(res[0].shape[1] < 32 and res[3].shape[0] == 45) # scaled down to the (tiny) free RAM
|
|
normal = SL.BFSShotJoin().join(imgs, shots, [0]) # enough RAM: images as before
|
|
self.assertEqual((normal[0].shape[0], normal[4]), (45, ""))
|
|
|
|
def test_join_bytes(self):
|
|
self.assertEqual(SL._join_bytes(10, 100, 200, False, 360), 10 * 100 * 200 * 12)
|
|
self.assertGreater(SL._join_bytes(10, 100, 200, True, 1024), SL._join_bytes(10, 100, 200, True, 50))
|
|
|
|
|
|
class StoredComparisonTest(unittest.TestCase):
|
|
def test_originals_go_to_a_video_not_the_pt(self):
|
|
rid = "cmp_" + str(id(self))
|
|
y = torch.linspace(0, 1, 64).view(1, 64, 1, 1)
|
|
src = [((torch.sin(torch.arange(n).view(n, 1, 1, 1) + 7 * y) + 1) / 2).expand(n, 64, 64, 3).clone() for n in (22, 39)]
|
|
shots = [dict(_shot(0, 0, 20, count=2, queue=True, run_id=rid, gen=22), frames=src[0]),
|
|
dict(_shot(1, 20, 45, cut=True, count=2, queue=True, run_id=rid, gen=39), frames=src[1])]
|
|
SL.BFSShotJoin().join([src[0] * 0.5], [shots[0]], [0], comparison=[True], comparison_height=[64])
|
|
out = SL.BFSShotJoin().join([src[1] * 0.5], [shots[1]], [0], comparison=[True], comparison_height=[64])
|
|
d = SL.run_dir(rid)
|
|
self.assertLess(os.path.getsize(os.path.join(d, "shot_0000.pt")), 50_000) # info only
|
|
self.assertTrue(os.path.exists(os.path.join(d, "shot_0000_orig.mkv")))
|
|
video, comp = out[0], out[3]
|
|
self.assertEqual(video.shape[0], 45)
|
|
self.assertEqual(comp.shape[0], 45)
|
|
# the original column (left) shows the source, not the dark filler of a missing original
|
|
self.assertGreater(float(comp[:, comp.shape[1] // 2, 5:60].mean()), 0.2)
|
|
|
|
|
|
class ColorAnchorTest(unittest.TestCase):
|
|
def _shot(self, sat):
|
|
src = torch.full((12, 48, 48, 3), 0.5)
|
|
res = src.clone()
|
|
res[:, 12:36, 12:36] = torch.tensor([0.5 - sat, 0.5 + sat, 0.5 - sat]) # a green subject
|
|
res[:, 12:36, 12:36] += torch.linspace(-0.05, 0.05, 24).view(1, 24, 1, 1)
|
|
return res, src
|
|
|
|
def test_subject_chroma_goes_back_to_the_anchor(self):
|
|
r1, s1 = self._shot(0.15)
|
|
anchor = SL.color_stats(r1, s1)
|
|
r4, s4 = self._shot(0.30) # 4 shots later: twice as saturated
|
|
out, before = SL.anchor_colors(r4, s4, anchor)
|
|
after = SL.color_stats(out, s4)
|
|
self.assertGreater(abs(before["mu"][0] - anchor["mu"][0]), 5)
|
|
self.assertLess(abs(after["mu"][0] - anchor["mu"][0]), 1.5)
|
|
self.assertLess(abs(after["mu"][1] - anchor["mu"][1]), 1.5)
|
|
self.assertAlmostEqual(float(out[0, 2, 2, 0]), 0.5, places=2) # the background is untouched
|
|
self.assertIsNone(SL.color_stats(s1, s1)) # no subject: nothing to anchor
|
|
|
|
|
|
class FillDetailsTest(unittest.TestCase):
|
|
def test_render_node_fills_details_with_its_vlm(self):
|
|
old = SL.vlm_describe
|
|
SL.vlm_describe = lambda clip, ims, q, mt=320, temperature=0.0, seed=0: "The elderly man with white hair in a light blue t-shirt."
|
|
try:
|
|
out = SL.fill_details_now("<Subject 1> is the {details} shown in <Picture 1>.", object(), [torch.rand(1, 8, 8, 3)])
|
|
self.assertEqual(out, "<Subject 1> is the elderly man with white hair in a light blue t-shirt shown in <Picture 1>.")
|
|
self.assertIn("character", SL.fill_details_now("is the {details} shown", None, []))
|
|
self.assertEqual(SL.fill_details_now("no slot", None, []), "no slot")
|
|
finally:
|
|
SL.vlm_describe = old
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|