183 lines
6.6 KiB
Python
183 lines
6.6 KiB
Python
import importlib
|
|
import sys
|
|
import types
|
|
import unittest
|
|
|
|
|
|
class DumasH3LongVideosHelperTests(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls._saved_modules = {
|
|
name: sys.modules.get(name)
|
|
for name in (
|
|
"torch",
|
|
"nodes",
|
|
"comfy",
|
|
"comfy.utils",
|
|
"comfy.samplers",
|
|
"comfy.nested_tensor",
|
|
"comfy.model_management",
|
|
"node_helpers",
|
|
"dumas_h3_longvideos",
|
|
)
|
|
}
|
|
|
|
fake_torch = types.SimpleNamespace(
|
|
cuda=types.SimpleNamespace(OutOfMemoryError=RuntimeError),
|
|
float32="float32",
|
|
)
|
|
fake_nodes = types.SimpleNamespace(common_ksampler=lambda *args, **kwargs: ({},))
|
|
fake_comfy_samplers = types.SimpleNamespace(
|
|
KSampler=types.SimpleNamespace(SAMPLERS=("res_multistep",), SCHEDULERS=("simple",))
|
|
)
|
|
fake_comfy_utils = types.SimpleNamespace(ProgressBar=lambda total: None)
|
|
fake_mm = types.SimpleNamespace(
|
|
current_loaded_models=[],
|
|
free_memory=lambda *args, **kwargs: None,
|
|
get_torch_device=lambda: "cpu",
|
|
soft_empty_cache=lambda *args, **kwargs: None,
|
|
unload_all_models=lambda *args, **kwargs: None,
|
|
get_free_memory=lambda *args, **kwargs: 0,
|
|
get_total_memory=lambda *args, **kwargs: 0,
|
|
)
|
|
fake_comfy = types.SimpleNamespace(
|
|
utils=fake_comfy_utils,
|
|
samplers=fake_comfy_samplers,
|
|
nested_tensor=types.SimpleNamespace(),
|
|
model_management=fake_mm,
|
|
)
|
|
|
|
sys.modules["torch"] = fake_torch
|
|
sys.modules["nodes"] = fake_nodes
|
|
sys.modules["comfy"] = fake_comfy
|
|
sys.modules["comfy.utils"] = fake_comfy_utils
|
|
sys.modules["comfy.samplers"] = fake_comfy_samplers
|
|
sys.modules["comfy.nested_tensor"] = fake_comfy.nested_tensor
|
|
sys.modules["comfy.model_management"] = fake_mm
|
|
sys.modules["node_helpers"] = types.SimpleNamespace()
|
|
|
|
cls.module = importlib.import_module("dumas_h3_longvideos")
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
for name, module in cls._saved_modules.items():
|
|
if module is None:
|
|
sys.modules.pop(name, None)
|
|
else:
|
|
sys.modules[name] = module
|
|
|
|
def test_extract_wardrobe_is_cached(self):
|
|
fn = self.module.extract_wardrobe
|
|
fn.cache_clear()
|
|
|
|
beat = "walks forward\nwardrobe: red jacket, grey shorts\nlooks back"
|
|
self.assertEqual(fn(beat), ("walks forward\nlooks back", "red jacket, grey shorts"))
|
|
self.assertEqual(fn(beat), ("walks forward\nlooks back", "red jacket, grey shorts"))
|
|
self.assertGreater(fn.cache_info().hits, 0)
|
|
|
|
def test_dialogue_helpers_keep_existing_outputs_and_cache(self):
|
|
spans_cache = self.module._dialogue_spans_cached
|
|
sec_fn = self.module.dialogue_seconds
|
|
words_fn = self.module.dialogue_words
|
|
|
|
spans_cache.cache_clear()
|
|
sec_fn.cache_clear()
|
|
words_fn.cache_clear()
|
|
|
|
beat = 'Mara says, "Open it now." Jon replies, "Do it."'
|
|
self.assertEqual(self.module.dialogue_spans(beat), [3, 2])
|
|
self.assertEqual(words_fn(beat), 5)
|
|
self.assertAlmostEqual(sec_fn(beat), 3.5)
|
|
self.assertAlmostEqual(sec_fn(beat, pad=False), 2.5)
|
|
|
|
self.module.dialogue_spans(beat)
|
|
sec_fn(beat)
|
|
words_fn(beat)
|
|
|
|
self.assertGreater(spans_cache.cache_info().hits, 0)
|
|
self.assertGreater(sec_fn.cache_info().hits, 0)
|
|
self.assertGreater(words_fn.cache_info().hits, 0)
|
|
|
|
def test_directive_and_estimate_helpers_are_cached(self):
|
|
directive_fn = self.module.beat_seconds_directive
|
|
estimate_fn = self.module.estimate_beat_seconds
|
|
action_fn = self.module.action_clauses
|
|
|
|
directive_fn.cache_clear()
|
|
estimate_fn.cache_clear()
|
|
action_fn.cache_clear()
|
|
|
|
beat = 'seconds: 7.5\nShe opens the hatch and climbs inside.'
|
|
self.assertEqual(directive_fn(beat), 7.5)
|
|
self.assertEqual(action_fn(beat), 2)
|
|
self.assertAlmostEqual(estimate_fn(beat), 7.0)
|
|
|
|
directive_fn(beat)
|
|
action_fn(beat)
|
|
estimate_fn(beat)
|
|
|
|
self.assertGreater(directive_fn.cache_info().hits, 0)
|
|
self.assertGreater(action_fn.cache_info().hits, 0)
|
|
self.assertGreater(estimate_fn.cache_info().hits, 0)
|
|
|
|
def test_has_speech_cache_respects_written_text_filter(self):
|
|
fn = self.module.has_speech
|
|
fn.cache_clear()
|
|
|
|
written = 'She reads the sign marked "EXIT" and keeps walking.'
|
|
spoken = 'She says, "Exit now." and points to the door.'
|
|
self.assertFalse(fn(written))
|
|
self.assertTrue(fn(spoken))
|
|
fn(written)
|
|
fn(spoken)
|
|
self.assertGreaterEqual(fn.cache_info().hits, 2)
|
|
|
|
def test_resolve_tagged_refs_preserves_sparse_socket_numbers(self):
|
|
refs = [None, "img2", None, None, None, None, "img7", None, "img9"]
|
|
|
|
text, images, dropped = self.module.resolve_tagged_refs(
|
|
"Mara <Picture 7> turns toward Jon <Picture 2> while <Picture 9> watches.",
|
|
refs,
|
|
)
|
|
|
|
self.assertEqual(
|
|
text,
|
|
"Mara <Picture 2> turns toward Jon <Picture 1> while <Picture 3> watches.",
|
|
)
|
|
self.assertEqual(images, ["img2", "img7", "img9"])
|
|
self.assertEqual(dropped, [])
|
|
|
|
def test_resolve_tagged_refs_drops_unconnected_sparse_slots(self):
|
|
refs = [None, "img2", None, None, None, None, "img7", None, None]
|
|
|
|
text, images, dropped = self.module.resolve_tagged_refs(
|
|
"Use <Picture 7>, skip <Picture 4>, keep <Picture 2>.",
|
|
refs,
|
|
)
|
|
|
|
self.assertEqual(text, "Use <Picture 2>, skip, keep <Picture 1>.")
|
|
self.assertEqual(images, ["img2", "img7"])
|
|
self.assertEqual(dropped, [4])
|
|
|
|
def test_shot_references_uses_all_connected_sparse_slots(self):
|
|
refs = [None, "img2", None, "img4", None, None, "img7", None, None]
|
|
|
|
self.assertEqual(
|
|
self.module.shot_references(refs, "first shot", 0, None),
|
|
["img2", "img4", "img7"],
|
|
)
|
|
self.assertEqual(
|
|
self.module.shot_references(refs, "every shot", 3, None),
|
|
["img2", "img4", "img7"],
|
|
)
|
|
|
|
def test_input_types_expose_nine_ref_slots(self):
|
|
optional = self.module.H3LongVideos.INPUT_TYPES()["optional"]
|
|
|
|
for index in range(1, 10):
|
|
self.assertIn(f"ref_image_{index}", optional)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|