501 lines
20 KiB
Python
501 lines
20 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",
|
|
"numpy",
|
|
"PIL",
|
|
"PIL.Image",
|
|
"folder_paths",
|
|
"dumas_image_nodes",
|
|
"dumas_h3_longvideos",
|
|
)
|
|
}
|
|
|
|
fake_torch = types.SimpleNamespace(
|
|
cuda=types.SimpleNamespace(OutOfMemoryError=RuntimeError),
|
|
float32="float32",
|
|
)
|
|
fake_numpy = types.SimpleNamespace(
|
|
clip=lambda array, _low, _high: array,
|
|
uint8="uint8",
|
|
)
|
|
fake_pil_image_module = types.SimpleNamespace(fromarray=lambda _array: None)
|
|
fake_pil_module = types.SimpleNamespace(Image=fake_pil_image_module)
|
|
fake_folder_paths = types.SimpleNamespace(
|
|
get_temp_directory=lambda: "/tmp",
|
|
get_output_directory=lambda: "/tmp",
|
|
get_save_image_path=lambda prefix, _out, _width, _height: ("/tmp", prefix, 1, "", prefix),
|
|
)
|
|
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["numpy"] = fake_numpy
|
|
sys.modules["PIL"] = fake_pil_module
|
|
sys.modules["PIL.Image"] = fake_pil_image_module
|
|
sys.modules["folder_paths"] = fake_folder_paths
|
|
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.image_module = importlib.import_module("dumas_image_nodes")
|
|
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_per_shot_directive_helpers_parse_new_controls(self):
|
|
beat = (
|
|
"ref_mode: every shot + handoff ref\n"
|
|
"ref_noise_aug: 0.87\n"
|
|
"continuity: keyframe carry\n"
|
|
"The courier waits under the sign."
|
|
)
|
|
|
|
self.assertEqual(
|
|
self.module.beat_ref_mode_directive(beat),
|
|
"every shot + handoff ref",
|
|
)
|
|
self.assertEqual(self.module.beat_ref_noise_aug_directive(beat), 0.87)
|
|
self.assertEqual(
|
|
self.module.beat_continuity_directive(beat),
|
|
"keyframe carry",
|
|
)
|
|
|
|
def test_distribute_generations_canonicalizes_per_shot_audio_and_anchor_directives(self):
|
|
generations = self.module.distribute_generations(
|
|
"",
|
|
[
|
|
"anchor_add: harsh sodium spill, wet asphalt reflections\n"
|
|
"soundscape: distant traffic hiss, loose sign rattle\n"
|
|
"music: low pulsing synth tension\n"
|
|
"continuity: hard cut\n"
|
|
"ref_mode: every shot\n"
|
|
"ref_noise_aug: 0.88\n"
|
|
"A courier waits under the streetlight."
|
|
],
|
|
"global rain",
|
|
"global score",
|
|
)
|
|
|
|
block = generations[0]
|
|
self.assertIn("harsh sodium spill, wet asphalt reflections", block)
|
|
self.assertIn("overall_soundscape: distant traffic hiss, loose sign rattle", block)
|
|
self.assertIn("non_diegetic_music: low pulsing synth tension", block)
|
|
self.assertNotIn("\nsoundscape:", block)
|
|
self.assertNotIn("\nmusic:", block)
|
|
self.assertNotIn("\ncontinuity:", block)
|
|
self.assertNotIn("\nref_mode:", block)
|
|
self.assertNotIn("\nref_noise_aug:", block)
|
|
self.assertNotIn("\nanchor_add:", block)
|
|
|
|
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,
|
|
{"kind": "character", "image": "img2", "name": "Jon"},
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
{"kind": "character", "image": "img7", "name": "Mara"},
|
|
None,
|
|
{"kind": "location", "image": "img9", "name": "Watchtower"},
|
|
]
|
|
|
|
text, references, 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(
|
|
[self.module._reference_image(ref) for ref in references],
|
|
["img2", "img7", "img9"],
|
|
)
|
|
self.assertEqual(dropped, [])
|
|
|
|
def test_resolve_tagged_refs_drops_unconnected_sparse_slots(self):
|
|
refs = [
|
|
None,
|
|
{"kind": "character", "image": "img2", "name": "Jon"},
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
{"kind": "character", "image": "img7", "name": "Mara"},
|
|
None,
|
|
None,
|
|
]
|
|
|
|
text, references, 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(
|
|
[self.module._reference_image(ref) for ref in references],
|
|
["img2", "img7"],
|
|
)
|
|
self.assertEqual(dropped, [4])
|
|
|
|
def test_shot_references_uses_all_connected_sparse_slots(self):
|
|
refs = [
|
|
None,
|
|
{"kind": "character", "image": "img2"},
|
|
None,
|
|
{"kind": "character", "image": "img4"},
|
|
None,
|
|
None,
|
|
{"kind": "location", "image": "img7"},
|
|
None,
|
|
None,
|
|
]
|
|
|
|
for mode, shot_index in (("auto ref2v", 0), ("first shot", 0), ("every shot", 3)):
|
|
self.assertEqual(
|
|
[self.module._reference_image(ref) for ref in self.module.shot_references(refs, mode, shot_index, 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_{index}", optional)
|
|
|
|
def test_node_appends_per_beat_list_outputs_without_reordering_existing_slots(self):
|
|
self.assertEqual(
|
|
self.module.H3LongVideos.RETURN_NAMES[-2:],
|
|
("beat_images", "beat_audio"),
|
|
)
|
|
self.assertEqual(
|
|
self.module.H3LongVideos.OUTPUT_IS_LIST[-2:],
|
|
(True, True),
|
|
)
|
|
|
|
def test_reference_context_matches_character_names_and_location_tags(self):
|
|
refs = [
|
|
{"kind": "character", "image": "img1", "name": "Mara", "description": "silver hair", "wardrobe": "red jacket"},
|
|
{"kind": "location", "image": "img2", "name": "Hangar", "description": "wet concrete floor"},
|
|
]
|
|
|
|
context = self.module._reference_context_for_text(
|
|
"[Generation 1] Mara crosses the room toward <Picture 2>.",
|
|
refs,
|
|
)
|
|
|
|
self.assertIn("Persistent appearance for Mara: silver hair.", context)
|
|
self.assertIn("Persistent wardrobe/style for Mara: red jacket.", context)
|
|
self.assertIn("Location context for Hangar: wet concrete floor.", context)
|
|
|
|
def test_reference_character_memory_uses_character_wardrobe_only(self):
|
|
refs = [
|
|
{"kind": "character", "image": "img1", "name": "Mara", "wardrobe": "red jacket, black boots"},
|
|
{"kind": "location", "image": "img2", "name": "Hangar", "description": "wet concrete floor", "wardrobe": "should be ignored"},
|
|
]
|
|
|
|
self.assertEqual(
|
|
self.module._reference_character_memory(refs),
|
|
"Mara = red jacket, black boots",
|
|
)
|
|
|
|
def test_ref_mode_defaults_are_ref2v_biased(self):
|
|
optional = self.module.H3LongVideos.INPUT_TYPES()["optional"]
|
|
|
|
self.assertEqual(optional["ref_mode"][1]["default"], "auto ref2v")
|
|
self.assertEqual(optional["ref_noise_aug"][1]["default"], 0.95)
|
|
|
|
def test_only_canonical_h3_long_videos_node_is_exposed(self):
|
|
self.assertEqual(
|
|
self.module.NODE_CLASS_MAPPINGS,
|
|
{"DumasH3LongVideos": self.module.H3LongVideos},
|
|
)
|
|
self.assertEqual(
|
|
self.module.NODE_DISPLAY_NAME_MAPPINGS,
|
|
{"DumasH3LongVideos": "Dumas H3 Long Videos (FL2VA + REF2VA)"},
|
|
)
|
|
|
|
def test_compose_persistent_does_not_expand_ambiguous_plural_to_full_cast(self):
|
|
active = self.module.parse_wardrobe(
|
|
"Maya = she, red jacket\n"
|
|
"Jon = he, navy overalls\n"
|
|
"Becca = she, green coat"
|
|
)
|
|
|
|
shot = self.module.compose_persistent(
|
|
"Both of them walk to the door.",
|
|
active,
|
|
"",
|
|
speaking=False,
|
|
)
|
|
|
|
self.assertEqual(shot, "Both of them walk to the door.")
|
|
|
|
def test_compose_persistent_keeps_two_person_plural_binding(self):
|
|
active = self.module.parse_wardrobe(
|
|
"Maya = she, red jacket\n"
|
|
"Jon = he, navy overalls"
|
|
)
|
|
|
|
shot = self.module.compose_persistent(
|
|
"They walk to the door.",
|
|
active,
|
|
"",
|
|
speaking=False,
|
|
)
|
|
|
|
self.assertIn("Maya (red jacket)", shot)
|
|
self.assertIn("Jon (navy overalls)", shot)
|
|
self.assertIn("They walk to the door.", shot)
|
|
|
|
def test_compose_persistent_all_three_characters_binds_full_cast(self):
|
|
active = self.module.parse_wardrobe(
|
|
"Maya = she, red jacket\n"
|
|
"Jon = he, navy overalls\n"
|
|
"Becca = she, green coat"
|
|
)
|
|
|
|
shot = self.module.compose_persistent(
|
|
"The three characters walk to the door.",
|
|
active,
|
|
"",
|
|
speaking=False,
|
|
)
|
|
|
|
self.assertIn("Maya (red jacket)", shot)
|
|
self.assertIn("Jon (navy overalls)", shot)
|
|
self.assertIn("Becca (green coat)", shot)
|
|
|
|
def test_plan_only_returns_joined_generations_on_script_socket_with_anchor_override(self):
|
|
module = self.module
|
|
node = module.H3LongVideos()
|
|
|
|
class _Clip:
|
|
def tokenize(self, text):
|
|
return text
|
|
|
|
def encode_from_tokens_scheduled(self, tokens):
|
|
return tokens
|
|
|
|
class _TorchStub:
|
|
@staticmethod
|
|
def zeros(shape):
|
|
return ("zeros", shape)
|
|
|
|
original_torch = module.torch
|
|
original_vram_gb = module.vram_gb
|
|
original_dit_resident_gb = module.dit_resident_gb
|
|
original_lora_overhead_gb = module.lora_overhead_gb
|
|
original_check_vae_wiring = module.check_vae_wiring
|
|
original_check_text_encoder = module.check_text_encoder
|
|
original_apply_h3_model_sampling = module.apply_h3_model_sampling
|
|
original_sla_pairing = module.sla_pairing
|
|
original_lora_hint_notes = module.lora_hint_notes
|
|
original_schedule_balance_note = module.schedule_balance_note
|
|
original_kernel_backend_note = module.kernel_backend_note
|
|
original_audio_scale_note = module.audio_scale_note
|
|
original_quant_accel_note = module.quant_accel_note
|
|
original_lora_active = module.lora_active
|
|
original_resolve_shot_frames = module.resolve_shot_frames
|
|
original_plan_beat_frames = module.plan_beat_frames
|
|
original_dialogue_fit_warnings = module.dialogue_fit_warnings
|
|
original_dialogue_filler_warnings = module.dialogue_filler_warnings
|
|
original_distribute_generations = module.distribute_generations
|
|
original_continuity_warnings = module.continuity_warnings
|
|
original_empty_av_latent = module._empty_av_latent
|
|
try:
|
|
module.torch = _TorchStub()
|
|
module.vram_gb = lambda: (0.0, 0.0)
|
|
module.dit_resident_gb = lambda _model: 0.0
|
|
module.lora_overhead_gb = lambda _model: 0.0
|
|
module.check_vae_wiring = lambda *_args, **_kwargs: None
|
|
module.check_text_encoder = lambda *_args, **_kwargs: None
|
|
module.apply_h3_model_sampling = lambda model, *_args, **_kwargs: (model, "")
|
|
module.sla_pairing = lambda *_args, **_kwargs: ("", False, "")
|
|
module.lora_hint_notes = lambda *_args, **_kwargs: []
|
|
module.schedule_balance_note = lambda *_args, **_kwargs: ""
|
|
module.kernel_backend_note = lambda *_args, **_kwargs: ""
|
|
module.audio_scale_note = lambda *_args, **_kwargs: ""
|
|
module.quant_accel_note = lambda *_args, **_kwargs: ""
|
|
module.lora_active = lambda _model: False
|
|
module.resolve_shot_frames = lambda *_args, **_kwargs: (73, "")
|
|
module.plan_beat_frames = lambda beats, fps, budget, per_beat=True: ([73] * len(beats), [])
|
|
module.dialogue_fit_warnings = lambda *_args, **_kwargs: []
|
|
module.dialogue_filler_warnings = lambda *_args, **_kwargs: []
|
|
module.distribute_generations = lambda anchor, beats, *_args, **_kwargs: [
|
|
f"[Generation 1] {anchor}. {beats[0]}",
|
|
f"[Generation 2] {anchor}. {beats[1]}{module.ANATOMY_STATE}",
|
|
]
|
|
module.continuity_warnings = lambda _gens: []
|
|
module._empty_av_latent = lambda *_args, **_kwargs: ({"samples": "latent"},)
|
|
|
|
result = node.run(
|
|
model=object(),
|
|
clip=_Clip(),
|
|
vae=object(),
|
|
audio_vae=object(),
|
|
prompt="Francine stands alone.\n\nFrancine and Frankie walk together.",
|
|
resolution="16:9",
|
|
steps=20,
|
|
cfg=1.0,
|
|
sampler_name="res_multistep",
|
|
scheduler="simple",
|
|
seed=1,
|
|
anchor_override="editorial room, soft practical lighting",
|
|
character_memory="Francine = white top\nFrankie = black jacket",
|
|
plan_only=True,
|
|
)
|
|
|
|
self.assertIn("# anatomy_guard: injected on shot(s) 2", result[3])
|
|
self.assertIn(
|
|
"[Generation 1] editorial room, soft practical lighting. Francine stands alone.",
|
|
result[3],
|
|
)
|
|
self.assertIn(
|
|
"[Generation 2] editorial room, soft practical lighting. Francine and Frankie walk together.",
|
|
result[3],
|
|
)
|
|
self.assertIn("2 beat(s)", result[2])
|
|
self.assertIn("2 shot(s)", result[2])
|
|
self.assertIn("ANATOMY -- guard injected on shot(s) 2", result[2])
|
|
self.assertEqual(result[-2:], ([], []))
|
|
finally:
|
|
module.torch = original_torch
|
|
module.vram_gb = original_vram_gb
|
|
module.dit_resident_gb = original_dit_resident_gb
|
|
module.lora_overhead_gb = original_lora_overhead_gb
|
|
module.check_vae_wiring = original_check_vae_wiring
|
|
module.check_text_encoder = original_check_text_encoder
|
|
module.apply_h3_model_sampling = original_apply_h3_model_sampling
|
|
module.sla_pairing = original_sla_pairing
|
|
module.lora_hint_notes = original_lora_hint_notes
|
|
module.schedule_balance_note = original_schedule_balance_note
|
|
module.kernel_backend_note = original_kernel_backend_note
|
|
module.audio_scale_note = original_audio_scale_note
|
|
module.quant_accel_note = original_quant_accel_note
|
|
module.lora_active = original_lora_active
|
|
module.resolve_shot_frames = original_resolve_shot_frames
|
|
module.plan_beat_frames = original_plan_beat_frames
|
|
module.dialogue_fit_warnings = original_dialogue_fit_warnings
|
|
module.dialogue_filler_warnings = original_dialogue_filler_warnings
|
|
module.distribute_generations = original_distribute_generations
|
|
module.continuity_warnings = original_continuity_warnings
|
|
module._empty_av_latent = original_empty_av_latent
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|