1547 lines
66 KiB
Python
1547 lines
66 KiB
Python
import importlib
|
|
import inspect
|
|
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_latent_upscale",
|
|
"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_expand_beats_auto_preserves_multiline_paragraph_as_one_beat(self):
|
|
beats, note = self.module.expand_beats(
|
|
["wardrobe: Maya = red jacket\nMaya enters the room.\nShe sits at the table."],
|
|
"auto",
|
|
)
|
|
|
|
self.assertEqual(
|
|
beats,
|
|
["wardrobe: Maya = red jacket\nMaya enters the room.\nShe sits at the table."],
|
|
)
|
|
self.assertEqual(note, "")
|
|
|
|
def test_expand_beats_legacy_blank_line_value_falls_back_to_auto(self):
|
|
beats, note = self.module.expand_beats(
|
|
["Maya enters the room.\nShe sits at the table."],
|
|
"blank line",
|
|
)
|
|
|
|
self.assertEqual(
|
|
beats,
|
|
["Maya enters the room.\nShe sits at the table."],
|
|
)
|
|
self.assertEqual(note, "")
|
|
|
|
def test_expand_beats_each_line_still_splits_multiline_paragraphs(self):
|
|
beats, note = self.module.expand_beats(
|
|
["wardrobe: Maya = red jacket\nMaya enters the room.\nseconds: 8\nShe sits at the table."],
|
|
"each line",
|
|
)
|
|
|
|
self.assertEqual(
|
|
beats,
|
|
[
|
|
"wardrobe: Maya = red jacket\nMaya enters the room.",
|
|
"seconds: 8\nShe sits at the table.",
|
|
],
|
|
)
|
|
self.assertIn("beat_split 'each line' split 1 multi-line paragraph(s) into 2 beats", note)
|
|
|
|
def test_timing_summary_reports_retry_and_bucket_totals(self):
|
|
note = self.module._format_timing_note([
|
|
{
|
|
"shot": 1,
|
|
"total": 12.4,
|
|
"retry_elapsed": 1.2,
|
|
"attempts": 2,
|
|
"sample": 8.0,
|
|
"latent_upscale_sample": 0.5,
|
|
"decode_video": 2.1,
|
|
"decode_audio": 0.4,
|
|
"cleanup": 0.2,
|
|
},
|
|
{
|
|
"shot": 2,
|
|
"total": 7.6,
|
|
"retry_elapsed": 0.0,
|
|
"attempts": 1,
|
|
"sample": 6.5,
|
|
"decode_video": 0.5,
|
|
"decode_audio": 0.3,
|
|
"cleanup": 0.1,
|
|
},
|
|
])
|
|
|
|
self.assertIn("timing: 2 shot(s) total 20.0s", note)
|
|
self.assertIn("sample 14.5s", note)
|
|
self.assertIn("decode video 2.6s", note)
|
|
self.assertIn("decode audio 0.7s", note)
|
|
self.assertIn("cleanup 0.3s", note)
|
|
self.assertIn("retry elapsed 1.2s", note)
|
|
self.assertIn("latent upscale 0.5s", note)
|
|
self.assertIn("retries 1", note)
|
|
self.assertIn("slowest shot 1 12.4s", note)
|
|
|
|
def test_latent_upscale_refines_video_but_preserves_audio(self):
|
|
class FakeTensor:
|
|
def __init__(self, name):
|
|
self.name = name
|
|
|
|
def detach(self):
|
|
return self
|
|
|
|
def to(self, *args, **kwargs):
|
|
return self
|
|
|
|
class FakeNestedTensor:
|
|
def __init__(self, parts):
|
|
self._parts = tuple(parts)
|
|
self.is_nested = True
|
|
|
|
def unbind(self):
|
|
return self._parts
|
|
|
|
calls = []
|
|
first_out = {"samples": FakeNestedTensor((FakeTensor("v1"), FakeTensor("a1")))}
|
|
second_out = {"samples": FakeNestedTensor((FakeTensor("v2"), FakeTensor("a2")))}
|
|
|
|
original_common_ksampler = self.module.nodes.common_ksampler
|
|
original_build = self.module._build_shot_conditioning
|
|
original_evict = self.module._evict_all_but
|
|
original_decode_video = self.module._decode_video
|
|
original_decode_audio = self.module._decode_audio
|
|
original_cleanup = self.module._deep_cleanup
|
|
original_upscale = self.module._upscale_latent_video
|
|
original_copy_sample = self.module._copy_sample_latent
|
|
original_nested = getattr(self.module.comfy.nested_tensor, "NestedTensor", None)
|
|
try:
|
|
self.module.comfy.nested_tensor.NestedTensor = FakeNestedTensor
|
|
|
|
def common_ksampler(*args, **kwargs):
|
|
calls.append((args, kwargs))
|
|
return (first_out if len(calls) == 1 else second_out,)
|
|
|
|
self.module.nodes.common_ksampler = common_ksampler
|
|
self.module._build_shot_conditioning = lambda *_args, **_kwargs: (
|
|
"cond",
|
|
{"samples": FakeNestedTensor((FakeTensor("basev"), FakeTensor("basea")))},
|
|
)
|
|
self.module._evict_all_but = lambda *_args, **_kwargs: None
|
|
self.module._upscale_latent_video = lambda video, param: (FakeTensor("upv"), 8, 16)
|
|
self.module._copy_sample_latent = lambda sampled: sampled["samples"].unbind()
|
|
self.module._decode_video = lambda _vae, out_latent, *_args, **_kwargs: out_latent
|
|
self.module._decode_audio = lambda _vae, out_latent: out_latent
|
|
self.module._deep_cleanup = lambda: None
|
|
|
|
result = self.module.H3LongVideos()._render(
|
|
model=object(),
|
|
clip=types.SimpleNamespace(
|
|
tokenize=lambda text, **kwargs: text,
|
|
encode_from_tokens_scheduled=lambda tokens: tokens,
|
|
),
|
|
vae=object(),
|
|
audio_vae=object(),
|
|
negative="negative",
|
|
prompt="beat",
|
|
w=128,
|
|
h=64,
|
|
ln=24,
|
|
fps=24,
|
|
tiled=False,
|
|
sa=(123, 20, 1.0, "res_multistep", "simple", 1.0),
|
|
handoff=None,
|
|
latent_upscale_param={
|
|
"mode": "model",
|
|
"model_name": "upscale.safetensors",
|
|
"device": "cpu",
|
|
"precision": "fp16",
|
|
"sampler_name": "euler_ancestral",
|
|
"scheduler": "simple",
|
|
"steps": 2,
|
|
"denoise": 0.4,
|
|
"megapixels": 1.5,
|
|
},
|
|
)
|
|
|
|
self.assertEqual(len(calls), 2)
|
|
self.assertIsNot(calls[1][0][8], first_out)
|
|
self.assertEqual(calls[1][0][8]["samples"].unbind()[0].name, "upv")
|
|
self.assertEqual(calls[1][0][2], 2)
|
|
self.assertEqual(calls[1][0][4], "euler_ancestral")
|
|
self.assertEqual(calls[1][0][5], "simple")
|
|
self.assertAlmostEqual(calls[1][1]["denoise"], 0.4)
|
|
self.assertEqual(result[1], first_out)
|
|
self.assertEqual(result[2][0].name, "v2")
|
|
self.assertEqual(result[2][1].name, "a1")
|
|
self.assertEqual(result[0]["samples"].unbind()[0].name, "v2")
|
|
self.assertEqual(result[0]["samples"].unbind()[-1].name, "a1")
|
|
finally:
|
|
self.module.nodes.common_ksampler = original_common_ksampler
|
|
self.module._build_shot_conditioning = original_build
|
|
self.module._evict_all_but = original_evict
|
|
self.module._decode_video = original_decode_video
|
|
self.module._decode_audio = original_decode_audio
|
|
self.module._deep_cleanup = original_cleanup
|
|
self.module._upscale_latent_video = original_upscale
|
|
self.module._copy_sample_latent = original_copy_sample
|
|
if original_nested is None:
|
|
delattr(self.module.comfy.nested_tensor, "NestedTensor")
|
|
else:
|
|
self.module.comfy.nested_tensor.NestedTensor = original_nested
|
|
|
|
def test_latent_upscale_decodes_audio_before_video_and_cleans_up(self):
|
|
class FakeTensor:
|
|
def __init__(self, name):
|
|
self.name = name
|
|
|
|
def detach(self):
|
|
return self
|
|
|
|
def to(self, *args, **kwargs):
|
|
return self
|
|
|
|
class FakeNestedTensor:
|
|
def __init__(self, parts):
|
|
self._parts = tuple(parts)
|
|
self.is_nested = True
|
|
|
|
def unbind(self):
|
|
return self._parts
|
|
|
|
order = []
|
|
first_out = {"samples": FakeNestedTensor((FakeTensor("v1"), FakeTensor("a1")))}
|
|
second_out = {"samples": FakeNestedTensor((FakeTensor("v2"), FakeTensor("a2")))}
|
|
|
|
original_common_ksampler = self.module.nodes.common_ksampler
|
|
original_build = self.module._build_shot_conditioning
|
|
original_evict = self.module._evict_all_but
|
|
original_decode_video = self.module._decode_video
|
|
original_decode_audio = self.module._decode_audio
|
|
original_cleanup = self.module._deep_cleanup
|
|
original_upscale = self.module._upscale_latent_video
|
|
original_copy_sample = self.module._copy_sample_latent
|
|
original_unload = getattr(self.module.mm, "unload_model_and_clones", None)
|
|
original_unload_all = self.module.mm.unload_all_models
|
|
original_nested = getattr(self.module.comfy.nested_tensor, "NestedTensor", None)
|
|
try:
|
|
self.module.comfy.nested_tensor.NestedTensor = FakeNestedTensor
|
|
|
|
def common_ksampler(*args, **kwargs):
|
|
order.append("latent_upscale_sample" if len(order) else "sample")
|
|
return (first_out if len([x for x in order if x.endswith("sample")]) == 1 else second_out,)
|
|
|
|
def decode_audio(_vae, out_latent):
|
|
order.append("audio")
|
|
self.assertIs(out_latent, first_out)
|
|
return out_latent
|
|
|
|
def decode_video(_vae, out_latent, *_args, **_kwargs):
|
|
order.append("video")
|
|
self.assertIsNot(out_latent, first_out)
|
|
self.assertIs(out_latent["samples"].unbind()[0], second_out["samples"].unbind()[0])
|
|
return out_latent
|
|
|
|
def cleanup():
|
|
order.append("cleanup")
|
|
|
|
def unload_model_and_clones(*_args, **_kwargs):
|
|
order.append("unload_h3_failed")
|
|
raise RuntimeError("model wrapper does not expose clone metadata")
|
|
|
|
def unload_all_models(*_args, **_kwargs):
|
|
order.append("unload_all")
|
|
|
|
def upscale_latent_video(video, param):
|
|
order.append("upscale")
|
|
return FakeTensor("upv"), 8, 16
|
|
|
|
self.module.nodes.common_ksampler = common_ksampler
|
|
self.module._build_shot_conditioning = lambda *_args, **_kwargs: (
|
|
"cond",
|
|
{"samples": FakeNestedTensor((FakeTensor("basev"), FakeTensor("basea")))},
|
|
)
|
|
self.module._evict_all_but = lambda *_args, **_kwargs: None
|
|
self.module.mm.unload_model_and_clones = unload_model_and_clones
|
|
self.module.mm.unload_all_models = unload_all_models
|
|
self.module._upscale_latent_video = upscale_latent_video
|
|
self.module._copy_sample_latent = lambda sampled: sampled["samples"].unbind()
|
|
self.module._decode_video = decode_video
|
|
self.module._decode_audio = decode_audio
|
|
self.module._deep_cleanup = cleanup
|
|
|
|
self.module.H3LongVideos()._render(
|
|
model=object(),
|
|
clip=types.SimpleNamespace(
|
|
tokenize=lambda text, **kwargs: text,
|
|
encode_from_tokens_scheduled=lambda tokens: tokens,
|
|
),
|
|
vae=object(),
|
|
audio_vae=object(),
|
|
negative="negative",
|
|
prompt="beat",
|
|
w=128,
|
|
h=64,
|
|
ln=24,
|
|
fps=24,
|
|
tiled=False,
|
|
sa=(123, 20, 1.0, "res_multistep", "simple", 1.0),
|
|
handoff=None,
|
|
latent_upscale_param={
|
|
"mode": "model",
|
|
"model_name": "upscale.safetensors",
|
|
"device": "cuda",
|
|
"precision": "fp16",
|
|
"sampler_name": "euler_ancestral",
|
|
"scheduler": "simple",
|
|
"steps": 2,
|
|
"denoise": 0.4,
|
|
"megapixels": 1.5,
|
|
},
|
|
)
|
|
|
|
self.assertEqual(order[0], "sample")
|
|
self.assertLess(order.index("unload_all"), order.index("upscale"))
|
|
self.assertLess(order.index("upscale"), order.index("latent_upscale_sample"))
|
|
self.assertLess(order.index("audio"), order.index("video"))
|
|
self.assertEqual(order[-1], "cleanup")
|
|
finally:
|
|
self.module.nodes.common_ksampler = original_common_ksampler
|
|
self.module._build_shot_conditioning = original_build
|
|
self.module._evict_all_but = original_evict
|
|
self.module._decode_video = original_decode_video
|
|
self.module._decode_audio = original_decode_audio
|
|
self.module._deep_cleanup = original_cleanup
|
|
self.module._upscale_latent_video = original_upscale
|
|
self.module._copy_sample_latent = original_copy_sample
|
|
if original_unload is None:
|
|
delattr(self.module.mm, "unload_model_and_clones")
|
|
else:
|
|
self.module.mm.unload_model_and_clones = original_unload
|
|
self.module.mm.unload_all_models = original_unload_all
|
|
if original_nested is None:
|
|
delattr(self.module.comfy.nested_tensor, "NestedTensor")
|
|
else:
|
|
self.module.comfy.nested_tensor.NestedTensor = original_nested
|
|
|
|
def test_latent_upscale_off_skips_second_pass(self):
|
|
calls = []
|
|
original_common_ksampler = self.module.nodes.common_ksampler
|
|
original_build = self.module._build_shot_conditioning
|
|
original_evict = self.module._evict_all_but
|
|
original_decode_video = self.module._decode_video
|
|
original_decode_audio = self.module._decode_audio
|
|
original_cleanup = self.module._deep_cleanup
|
|
original_upscale = self.module._upscale_latent_video
|
|
original_copy_sample = self.module._copy_sample_latent
|
|
try:
|
|
self.module.nodes.common_ksampler = lambda *args, **kwargs: (calls.append((args, kwargs)) or {"samples": "latent"},)
|
|
self.module._build_shot_conditioning = lambda *_args, **_kwargs: ("cond", {"samples": "base"})
|
|
self.module._evict_all_but = lambda *_args, **_kwargs: None
|
|
self.module._decode_video = lambda _vae, out_latent, *_args, **_kwargs: out_latent
|
|
self.module._decode_audio = lambda _vae, out_latent: out_latent
|
|
self.module._deep_cleanup = lambda: None
|
|
self.module._upscale_latent_video = lambda *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError("should not run"))
|
|
self.module._copy_sample_latent = lambda sampled: sampled
|
|
|
|
self.module.H3LongVideos()._render(
|
|
model=object(),
|
|
clip=object(),
|
|
vae=object(),
|
|
audio_vae=object(),
|
|
negative="negative",
|
|
prompt="beat",
|
|
w=128,
|
|
h=64,
|
|
ln=24,
|
|
fps=24,
|
|
tiled=False,
|
|
sa=(123, 20, 1.0, "res_multistep", "simple", 1.0),
|
|
handoff=None,
|
|
latent_upscale_param={"mode": "off"},
|
|
)
|
|
|
|
self.assertEqual(len(calls), 1)
|
|
finally:
|
|
self.module.nodes.common_ksampler = original_common_ksampler
|
|
self.module._build_shot_conditioning = original_build
|
|
self.module._evict_all_but = original_evict
|
|
self.module._decode_video = original_decode_video
|
|
self.module._decode_audio = original_decode_audio
|
|
self.module._deep_cleanup = original_cleanup
|
|
self.module._upscale_latent_video = original_upscale
|
|
self.module._copy_sample_latent = original_copy_sample
|
|
|
|
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_resolve_prompt_refs_keeps_named_character_images_alongside_tagged_location(self):
|
|
refs = [
|
|
{"kind": "character", "image": "img1", "name": "Mara"},
|
|
{"kind": "character", "image": "img2", "name": "Jon"},
|
|
{"kind": "location", "image": "img3", "name": "Hangar"},
|
|
]
|
|
|
|
text, references, dropped = self.module.resolve_prompt_refs(
|
|
"Mara and Jon argue inside <Picture 3>.",
|
|
refs,
|
|
)
|
|
|
|
self.assertEqual(text, "Mara and Jon argue inside <Picture 1>.")
|
|
self.assertEqual(
|
|
[self.module._reference_image(ref) for ref in references],
|
|
["img3", "img1", "img2"],
|
|
)
|
|
self.assertEqual(dropped, [])
|
|
|
|
def test_resolve_prompt_refs_where_tagged_mode_stays_tag_only(self):
|
|
refs = [
|
|
{"kind": "character", "image": "img1", "name": "Mara"},
|
|
{"kind": "location", "image": "img2", "name": "Hangar"},
|
|
]
|
|
|
|
text, references, dropped = self.module.resolve_prompt_refs(
|
|
"Mara waits in the hangar near <Picture 2>.",
|
|
refs,
|
|
include_named=False,
|
|
)
|
|
|
|
self.assertEqual(text, "Mara waits in the hangar near <Picture 1>.")
|
|
self.assertEqual([self.module._reference_image(ref) for ref in references], ["img2"])
|
|
self.assertEqual(dropped, [])
|
|
|
|
def test_resolve_tag_driven_prompt_refs_keeps_untagged_shot_on_handoff(self):
|
|
refs = [
|
|
{"kind": "character", "image": "img1", "name": "Mara"},
|
|
{"kind": "location", "image": "img2", "name": "Hangar"},
|
|
]
|
|
|
|
text, references, dropped = self.module.resolve_tag_driven_prompt_refs(
|
|
"Mara waits in the hangar.",
|
|
refs,
|
|
)
|
|
|
|
self.assertEqual(text, "Mara waits in the hangar.")
|
|
self.assertEqual(references, [])
|
|
self.assertEqual(dropped, [])
|
|
|
|
def test_resolve_tag_driven_prompt_refs_keeps_named_refs_on_tagged_shot(self):
|
|
refs = [
|
|
{"kind": "character", "image": "img1", "name": "Mara"},
|
|
{"kind": "location", "image": "img2", "name": "Hangar"},
|
|
]
|
|
|
|
text, references, dropped = self.module.resolve_tag_driven_prompt_refs(
|
|
"Mara waits in the hangar near <Picture 2>.",
|
|
refs,
|
|
)
|
|
|
|
self.assertEqual(text, "Mara waits in the hangar near <Picture 1>.")
|
|
self.assertEqual(
|
|
[self.module._reference_image(ref) for ref in references],
|
|
["img2", "img1"],
|
|
)
|
|
self.assertEqual(dropped, [])
|
|
|
|
def test_resolve_shot_references_uses_named_characters_without_picture_tags(self):
|
|
refs = [
|
|
{"kind": "character", "image": "img1", "name": "Mara"},
|
|
{"kind": "character", "image": "img2", "name": "Jon"},
|
|
{"kind": "location", "image": "img3", "name": "Hangar"},
|
|
]
|
|
|
|
text, references, dropped, shot_tag_driven, mode_eff = self.module.resolve_shot_references(
|
|
"[Generation 1] Mara crosses the hangar.",
|
|
refs,
|
|
"auto ref2v",
|
|
0,
|
|
None,
|
|
)
|
|
|
|
self.assertEqual(text, "[Generation 1] Mara crosses the hangar.")
|
|
self.assertEqual([self.module._reference_image(ref) for ref in references], ["img1"])
|
|
self.assertEqual(dropped, [])
|
|
self.assertFalse(shot_tag_driven)
|
|
self.assertEqual(mode_eff, "auto ref2v")
|
|
|
|
def test_resolve_prompt_refs_prioritizes_characters_before_locations(self):
|
|
refs = [
|
|
{"kind": "location", "image": "img1", "name": "Hangar"},
|
|
{"kind": "character", "image": "img2", "name": "Mara"},
|
|
]
|
|
|
|
text, references, dropped = self.module.resolve_prompt_refs(
|
|
"[Generation 1] Mara waits in the hangar.",
|
|
refs,
|
|
)
|
|
|
|
self.assertEqual(text, "[Generation 1] Mara waits in the hangar.")
|
|
self.assertEqual([self.module._reference_image(ref) for ref in references], ["img2", "img1"])
|
|
self.assertEqual(dropped, [])
|
|
|
|
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_annotate_script_refs_handles_named_characters_and_locations(self):
|
|
refs = [
|
|
{"kind": "character", "image": "img1", "name": "Mara"},
|
|
{"kind": "location", "image": "img2", "name": "Hangar"},
|
|
]
|
|
|
|
report = self.module.annotate_script_refs(
|
|
["Mara waits in the Hangar.", "Nobody else is here."],
|
|
refs,
|
|
)
|
|
|
|
self.assertIn("# shot 1 refs: Picture 1 Mara (by name); Picture 2 Hangar (by name)", report)
|
|
self.assertIn("# shot 2 refs: none", report)
|
|
|
|
def test_annotate_script_debug_groups_each_prompt_with_its_beat_info(self):
|
|
refs = [
|
|
{"kind": "character", "image": "img1", "name": "Mara"},
|
|
{"kind": "location", "image": "img2", "name": "Hangar"},
|
|
]
|
|
|
|
report = self.module.annotate_script_debug(
|
|
["[Generation 1] Mara waits in the Hangar.", "[Generation 2] Nobody else is here."],
|
|
[1],
|
|
"auto",
|
|
refs,
|
|
)
|
|
|
|
self.assertIn("Prompt 1\n[Generation 1] Mara waits in the Hangar.", report)
|
|
self.assertIn("Beat 1 info\nAnatomy guard: injected into this prompt", report)
|
|
self.assertIn("References used: <Picture 1> Mara (matched by name); <Picture 2> Hangar (matched by name)", report)
|
|
self.assertIn("Prompt 2\n[Generation 2] Nobody else is here.", report)
|
|
self.assertIn("Beat 2 info\nAnatomy guard: not injected for this prompt\nReferences used: none", report)
|
|
|
|
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)
|
|
|
|
names = list(optional)
|
|
ref_positions = [names.index(f"ref_{index}") for index in range(1, 10)]
|
|
self.assertEqual(ref_positions, list(range(ref_positions[0], ref_positions[0] + 9)))
|
|
|
|
def test_input_types_do_not_expose_legacy_ref_image_aliases(self):
|
|
optional = self.module.H3LongVideos.INPUT_TYPES()["optional"]
|
|
|
|
for index in range(1, 10):
|
|
self.assertNotIn(f"ref_image_{index}", optional)
|
|
self.assertNotIn("per_beat_length", optional)
|
|
self.assertNotIn("cleanup_between_shots", optional)
|
|
self.assertNotIn("detail_pass", optional)
|
|
self.assertNotIn("detail_sampler_name", optional)
|
|
self.assertNotIn("detail_scheduler", optional)
|
|
self.assertNotIn("detail_steps", optional)
|
|
self.assertNotIn("detail_denoise", optional)
|
|
self.assertIn("latent_upscale_param", optional)
|
|
|
|
def test_shot_seconds_tooltip_describes_ceiling_behavior(self):
|
|
optional = self.module.H3LongVideos.INPUT_TYPES()["optional"]
|
|
tooltip = optional["shot_seconds"][1]["tooltip"]
|
|
|
|
self.assertIn("GLOBAL per-shot maximum", tooltip)
|
|
self.assertIn("A beat's own `seconds:` directive can still ask for less", tooltip)
|
|
self.assertIn("let the render fail instead of shrinking it", tooltip)
|
|
|
|
def test_resolve_shot_frames_honors_forced_request_over_budget(self):
|
|
original_estimate_shot_frames = self.module.estimate_shot_frames
|
|
try:
|
|
self.module.estimate_shot_frames = lambda *_args, **_kwargs: 73
|
|
frames, note = self.module.resolve_shot_frames(10.0, 24, 16.0, 8.0, 1.5)
|
|
|
|
self.assertEqual(frames, 243)
|
|
self.assertIn("honoring it", note)
|
|
finally:
|
|
self.module.estimate_shot_frames = original_estimate_shot_frames
|
|
|
|
def test_run_defaults_match_declared_ref_widget_defaults(self):
|
|
node = self.module.H3LongVideos()
|
|
optional = node.INPUT_TYPES()["optional"]
|
|
params = inspect.signature(node.run).parameters
|
|
|
|
self.assertEqual(params["ref_mode"].default, optional["ref_mode"][1]["default"])
|
|
self.assertEqual(params["ref_image_size"].default, optional["ref_image_size"][1]["default"])
|
|
self.assertEqual(params["ref_noise_aug"].default, optional["ref_noise_aug"][1]["default"])
|
|
|
|
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",
|
|
"aliases": ["Xtina"],
|
|
"description": "silver hair",
|
|
"wardrobe": "red jacket",
|
|
"general": "wears a long grey coat",
|
|
"facts": {
|
|
"gender": "female",
|
|
"age": "41",
|
|
"nationality": "English",
|
|
"occupation": "a detective",
|
|
"height_feet": "6",
|
|
"height_inches": "2",
|
|
"accent": "English",
|
|
},
|
|
},
|
|
{"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("Character facts for Mara:", context)
|
|
self.assertIn("also known as Xtina", context)
|
|
self.assertIn("female", context)
|
|
self.assertIn("41 years old", context)
|
|
self.assertIn("English", context)
|
|
self.assertIn("works as a detective", context)
|
|
self.assertIn("6 foot 2 tall", context)
|
|
self.assertIn("speaks with a English accent", context)
|
|
self.assertIn("Persistent appearance for Mara: silver hair.", context)
|
|
self.assertIn("Persistent wardrobe/style for Mara: red jacket.", context)
|
|
self.assertIn("Character notes for Mara: wears a long grey coat.", context)
|
|
|
|
def test_reference_context_matches_location_names_without_picture_tag(self):
|
|
refs = [
|
|
{"kind": "location", "image": "img2", "name": "Hangar", "aliases": ["loading bay"], "description": "wet concrete floor", "general": "cold industrial lighting"},
|
|
]
|
|
|
|
context = self.module._reference_context_for_text(
|
|
"[Generation 1] They argue in the hangar near the loading bay.",
|
|
refs,
|
|
)
|
|
|
|
self.assertIn("Location context for Hangar: wet concrete floor.", context)
|
|
self.assertIn("Location notes for Hangar: cold industrial lighting.", context)
|
|
|
|
def test_reference_context_uses_resolved_picture_numbers_for_per_beat_refs(self):
|
|
refs = [
|
|
{"kind": "character", "image": "img1", "name": "Bill", "description": "very tall", "facts": {"age": "25"}},
|
|
{"kind": "location", "image": "img2", "name": "Pub", "description": "warm wood bar"},
|
|
]
|
|
|
|
rewritten, resolved_refs, dropped = self.module.resolve_prompt_refs(
|
|
"[Generation 1] Bill leans on <Picture 2>.",
|
|
refs,
|
|
)
|
|
context = self.module._reference_context_for_text(
|
|
rewritten,
|
|
refs,
|
|
resolved_refs=resolved_refs,
|
|
)
|
|
|
|
self.assertEqual(dropped, [])
|
|
self.assertIn("Character facts for <Picture 2> Bill: 25 years old.", context)
|
|
self.assertIn("Persistent appearance for <Picture 2> Bill: very tall.", context)
|
|
self.assertIn("Location context for <Picture 1> Pub: warm wood bar.", context)
|
|
|
|
def test_reference_context_skips_ambiguous_name_matches(self):
|
|
refs = [
|
|
{"kind": "character", "image": "img1", "name": "Alex", "description": "short dark hair"},
|
|
{"kind": "character", "image": "img2", "name": "Alex", "description": "tall blond hair"},
|
|
]
|
|
|
|
context = self.module._reference_context_for_text(
|
|
"[Generation 1] Alex enters the room.",
|
|
refs,
|
|
)
|
|
|
|
self.assertEqual(context, "")
|
|
|
|
def test_run_uses_reference_slots_directly(self):
|
|
calls = {}
|
|
original_parse_resolution = self.module.parse_resolution
|
|
original_connected_refs = self.module._connected_refs
|
|
original_reference_character_memory = self.module._reference_character_memory
|
|
original_vram_gb = self.module.vram_gb
|
|
original_dit_resident_gb = self.module.dit_resident_gb
|
|
original_lora_overhead_gb = self.module.lora_overhead_gb
|
|
original_resolve_shot_frames = self.module.resolve_shot_frames
|
|
original_lora_active = self.module.lora_active
|
|
original_sla_pairing = self.module.sla_pairing
|
|
original_apply_h3_model_sampling = self.module.apply_h3_model_sampling
|
|
original_split_paragraphs = self.module.split_paragraphs
|
|
original_expand_beats = self.module.expand_beats
|
|
original_anchor_warnings = self.module.anchor_warnings
|
|
original_anchor_contributes_nothing = self.module.anchor_contributes_nothing
|
|
original_anchor_is_action_beat = self.module.anchor_is_action_beat
|
|
original_distribute_generations = self.module.distribute_generations
|
|
original_continuity_warnings = self.module.continuity_warnings
|
|
original_speech_flags = self.module.speech_flags
|
|
original_annotate_script_debug = self.module.annotate_script_debug
|
|
original_empty_av_latent = self.module._empty_av_latent
|
|
original_torch_zeros = getattr(self.module.torch, "zeros", None)
|
|
try:
|
|
self.module.torch.zeros = lambda shape: shape
|
|
self.module.parse_resolution = lambda _resolution: (640, 360)
|
|
self.module._connected_refs = lambda refs: [ref for ref in refs if ref is not None]
|
|
self.module._reference_character_memory = lambda refs: (calls.setdefault("refs", tuple(refs)), "")[1]
|
|
self.module.vram_gb = lambda: (0, 0)
|
|
self.module.dit_resident_gb = lambda _model: 0
|
|
self.module.lora_overhead_gb = lambda _model: 0
|
|
self.module.resolve_shot_frames = lambda *args, **kwargs: (53, "")
|
|
self.module.lora_active = lambda _model: False
|
|
self.module.sla_pairing = lambda *_args, **_kwargs: ("", False, "")
|
|
self.module.apply_h3_model_sampling = lambda model, *_args: (model, "")
|
|
self.module.split_paragraphs = lambda _prompt, _sep: ["Anchor.", "Beat."]
|
|
self.module.expand_beats = lambda beat_paras, _split: (list(beat_paras), "")
|
|
self.module.anchor_warnings = lambda _anchor: []
|
|
self.module.anchor_contributes_nothing = lambda *_args, **_kwargs: False
|
|
self.module.anchor_is_action_beat = lambda *_args, **_kwargs: False
|
|
self.module.distribute_generations = lambda _anchor, beats, *_args, **_kwargs: list(beats)
|
|
self.module.continuity_warnings = lambda _gens: []
|
|
self.module.speech_flags = lambda _beats: []
|
|
self.module.annotate_script_debug = lambda *_args, **_kwargs: "script"
|
|
self.module._empty_av_latent = lambda *_args, **_kwargs: ({"samples": "latent"}, 5)
|
|
|
|
clip = types.SimpleNamespace(
|
|
tokenize=lambda text, **kwargs: text,
|
|
encode_from_tokens_scheduled=lambda tokens: tokens,
|
|
)
|
|
|
|
result = self.module.H3LongVideos().run(
|
|
model=object(),
|
|
clip=clip,
|
|
vae=object(),
|
|
audio_vae=object(),
|
|
prompt="Anchor only.",
|
|
resolution="16:9",
|
|
steps=6,
|
|
cfg=1,
|
|
sampler_name="res_multistep",
|
|
scheduler="simple",
|
|
seed=1,
|
|
plan_only=True,
|
|
ref_1={"image": "live-1"},
|
|
ref_3={"image": "live-3"},
|
|
latent_upscale_param={
|
|
"mode": "interp",
|
|
"method": "bilinear",
|
|
"sampler_name": "euler_ancestral",
|
|
"scheduler": "simple",
|
|
"steps": 2,
|
|
"denoise": 0.2,
|
|
"megapixels": 1.5,
|
|
},
|
|
)
|
|
|
|
self.assertEqual(calls["refs"][0]["image"], "live-1")
|
|
self.assertIsNone(calls["refs"][1])
|
|
self.assertEqual(calls["refs"][2]["image"], "live-3")
|
|
self.assertEqual(result[2].count("ref2va: 2 reference image(s)"), 1)
|
|
self.assertIn("latent upscale:", result[2])
|
|
self.assertIn("euler_ancestral/simple", result[2])
|
|
finally:
|
|
self.module.parse_resolution = original_parse_resolution
|
|
self.module._connected_refs = original_connected_refs
|
|
self.module._reference_character_memory = original_reference_character_memory
|
|
self.module.vram_gb = original_vram_gb
|
|
self.module.dit_resident_gb = original_dit_resident_gb
|
|
self.module.lora_overhead_gb = original_lora_overhead_gb
|
|
self.module.resolve_shot_frames = original_resolve_shot_frames
|
|
self.module.lora_active = original_lora_active
|
|
self.module.sla_pairing = original_sla_pairing
|
|
self.module.apply_h3_model_sampling = original_apply_h3_model_sampling
|
|
self.module.split_paragraphs = original_split_paragraphs
|
|
self.module.expand_beats = original_expand_beats
|
|
self.module.anchor_warnings = original_anchor_warnings
|
|
self.module.anchor_contributes_nothing = original_anchor_contributes_nothing
|
|
self.module.anchor_is_action_beat = original_anchor_is_action_beat
|
|
self.module.distribute_generations = original_distribute_generations
|
|
self.module.continuity_warnings = original_continuity_warnings
|
|
self.module.speech_flags = original_speech_flags
|
|
self.module.annotate_script_debug = original_annotate_script_debug
|
|
self.module._empty_av_latent = original_empty_av_latent
|
|
if original_torch_zeros is None:
|
|
delattr(self.module.torch, "zeros")
|
|
else:
|
|
self.module.torch.zeros = original_torch_zeros
|
|
|
|
def test_reference_context_matches_tagged_character_without_name_in_text(self):
|
|
refs = [
|
|
{"kind": "character", "image": "img1", "name": "Mara", "description": "silver hair", "wardrobe": "red jacket"},
|
|
]
|
|
|
|
context = self.module._reference_context_for_text(
|
|
"[Generation 1] <Picture 1> walks into the room.",
|
|
refs,
|
|
)
|
|
|
|
self.assertIn("Persistent appearance for Mara: silver hair.", context)
|
|
self.assertIn("Persistent wardrobe/style for Mara: red jacket.", context)
|
|
|
|
def test_reference_context_can_skip_character_wardrobe_when_live_memory_is_explicit(self):
|
|
refs = [
|
|
{"kind": "character", "image": "img1", "name": "Mara", "description": "silver hair", "wardrobe": "red jacket"},
|
|
]
|
|
|
|
context = self.module._reference_context_for_text(
|
|
"[Generation 1] Mara walks into the room.",
|
|
refs,
|
|
include_character_wardrobe=False,
|
|
)
|
|
|
|
self.assertIn("Persistent appearance for Mara: silver hair.", context)
|
|
self.assertNotIn("Persistent wardrobe/style for Mara: red jacket.", context)
|
|
|
|
def test_resolve_prompt_refs_adds_named_location_refs(self):
|
|
refs = [
|
|
{"kind": "location", "image": "img2", "name": "Hangar", "description": "wet concrete floor"},
|
|
]
|
|
|
|
rewritten, matched, dropped = self.module.resolve_prompt_refs(
|
|
"[Generation 1] They wait in the hangar.",
|
|
refs,
|
|
)
|
|
|
|
self.assertEqual(rewritten, "[Generation 1] They wait in the hangar.")
|
|
self.assertEqual(dropped, [])
|
|
self.assertEqual(len(matched), 1)
|
|
self.assertEqual(matched[0]["name"], "Hangar")
|
|
|
|
def test_resolve_tagged_refs_drops_reference_without_image(self):
|
|
refs = [
|
|
{"kind": "character", "image": None, "name": "Mara"},
|
|
{"kind": "character", "image": "img2", "name": "Jon"},
|
|
]
|
|
|
|
rewritten, matched, dropped = self.module.resolve_tagged_refs(
|
|
"[Generation 1] <Picture 1> faces <Picture 2>.",
|
|
refs,
|
|
)
|
|
|
|
self.assertEqual(rewritten, "[Generation 1] faces <Picture 1>.")
|
|
self.assertEqual(dropped, [1])
|
|
self.assertEqual(len(matched), 1)
|
|
self.assertEqual(matched[0]["name"], "Jon")
|
|
|
|
def test_reference_context_injects_immediately_after_generation_label(self):
|
|
block = (
|
|
"[Generation 1] Classic sitcom lighting and staging. "
|
|
"Duke walks into the room."
|
|
)
|
|
context = "Character facts for Duke: female, 25 years old."
|
|
|
|
result = self.module._inject_reference_context(block, context)
|
|
|
|
self.assertEqual(
|
|
result,
|
|
"[Generation 1] Character facts for Duke: female, 25 years old. "
|
|
"Classic sitcom lighting and staging. "
|
|
"Duke walks into the room.",
|
|
)
|
|
|
|
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_latent_upscale_params_node_is_exposed(self):
|
|
latent = importlib.import_module("dumas_h3_latent_upscale")
|
|
required = latent.H3LatentUpscaleParams.INPUT_TYPES()["required"]
|
|
|
|
self.assertEqual(
|
|
latent.NODE_CLASS_MAPPINGS,
|
|
{"DumasH3LatentUpscaleParams": latent.H3LatentUpscaleParams},
|
|
)
|
|
self.assertEqual(
|
|
latent.NODE_DISPLAY_NAME_MAPPINGS,
|
|
{"DumasH3LatentUpscaleParams": "Dumas H3 Latent Upscale Params"},
|
|
)
|
|
self.assertEqual(required["sampler_name"][1]["default"], "euler_ancestral")
|
|
self.assertEqual(required["scheduler"][1]["default"], "simple")
|
|
self.assertEqual(required["steps"][1]["default"], 2)
|
|
self.assertEqual(required["denoise"][1]["default"], 0.2)
|
|
self.assertEqual(required["megapixels"][1]["default"], 1.0)
|
|
self.assertEqual(required["tile_width"][1]["default"], 512)
|
|
self.assertEqual(required["tile_height"][1]["default"], 512)
|
|
self.assertEqual(required["overlap"][1]["default"], 64)
|
|
self.assertEqual(required["fade_width"][1]["default"], 32)
|
|
self.assertEqual(required["fade_height"][1]["default"], 32)
|
|
self.assertEqual(required["overlap_mode"][1]["default"], "earlier")
|
|
self.assertEqual(required["overlap_blend"][1]["default"], "linear")
|
|
self.assertEqual(required["tile_size_mode"][1]["default"], "specific_size")
|
|
self.assertEqual(required["grid_rows"][1]["default"], 2)
|
|
self.assertEqual(required["grid_cols"][1]["default"], 2)
|
|
self.assertEqual(required["spatial_w_overlap"][1]["default"], 128)
|
|
self.assertEqual(required["spatial_h_overlap"][1]["default"], 128)
|
|
self.assertEqual(required["min_tile_size"][1]["default"], 256)
|
|
self.assertEqual(required["masked_area_noise"][1]["default"], 0.0)
|
|
self.assertFalse(required["brightness_match"][1]["default"])
|
|
self.assertEqual(required["dynamic_fade"][1]["default"], "off")
|
|
self.assertEqual(required["dynamic_fade_min"][1]["default"], 32)
|
|
self.assertEqual(required["chunk_length"][1]["default"], 17)
|
|
self.assertEqual(required["temporal_overlap"][1]["default"], 0)
|
|
self.assertFalse(required["resize_conditioning"][1]["default"])
|
|
self.assertEqual(required["anchor_strength"][1]["default"], 0.999)
|
|
|
|
def test_latent_upscale_mode_infers_legacy_model_payloads(self):
|
|
self.assertEqual(self.module._latent_upscale_mode({"model_name": "foo.safetensors"}), "model")
|
|
self.assertEqual(self.module._latent_upscale_mode({"method": "bilinear"}), "interp")
|
|
self.assertEqual(self.module._latent_upscale_mode({"mode": "model"}), "model")
|
|
self.assertEqual(self.module._latent_upscale_mode({}), "off")
|
|
|
|
def test_tag_oom_stage_marks_oom_exceptions(self):
|
|
exc = RuntimeError("CUDA out of memory")
|
|
tagged = self.module._tag_oom_stage(exc, "latent_upscale")
|
|
self.assertIs(tagged, exc)
|
|
self.assertEqual(getattr(tagged, "_h3_stage", ""), "latent_upscale")
|
|
|
|
def test_shrink_model_tile_param_reduces_tile_size(self):
|
|
latent = importlib.import_module("dumas_h3_latent_upscale")
|
|
smaller = latent._shrink_model_tile_param({
|
|
"tile_size_mode": "specific_size",
|
|
"tile_width": 512,
|
|
"tile_height": 512,
|
|
"overlap": 64,
|
|
"fade_width": 32,
|
|
"fade_height": 32,
|
|
})
|
|
self.assertIsNotNone(smaller)
|
|
self.assertEqual(smaller["tile_size_mode"], "rows_cols")
|
|
self.assertEqual(smaller["grid_rows"], 4)
|
|
self.assertEqual(smaller["grid_cols"], 4)
|
|
self.assertEqual(smaller["spatial_w_overlap"], 0)
|
|
self.assertEqual(smaller["spatial_h_overlap"], 0)
|
|
self.assertEqual(smaller["fade_width"], 0)
|
|
self.assertEqual(smaller["fade_height"], 0)
|
|
self.assertEqual(smaller["min_tile_size"], 32)
|
|
|
|
def test_shrink_model_tile_param_rows_cols_resets_overlap(self):
|
|
latent = importlib.import_module("dumas_h3_latent_upscale")
|
|
smaller = latent._shrink_model_tile_param({
|
|
"tile_size_mode": "rows_cols",
|
|
"grid_rows": 4,
|
|
"grid_cols": 4,
|
|
"spatial_w_overlap": 128,
|
|
"spatial_h_overlap": 128,
|
|
"fade_width": 64,
|
|
"fade_height": 64,
|
|
"min_tile_size": 256,
|
|
})
|
|
self.assertIsNotNone(smaller)
|
|
self.assertEqual(smaller["grid_rows"], 8)
|
|
self.assertEqual(smaller["grid_cols"], 8)
|
|
self.assertEqual(smaller["spatial_w_overlap"], 0)
|
|
self.assertEqual(smaller["spatial_h_overlap"], 0)
|
|
self.assertEqual(smaller["fade_width"], 0)
|
|
self.assertEqual(smaller["fade_height"], 0)
|
|
self.assertEqual(smaller["min_tile_size"], 32)
|
|
|
|
def test_shrink_model_tile_param_rows_cols_can_reach_thirty_two(self):
|
|
latent = importlib.import_module("dumas_h3_latent_upscale")
|
|
smaller = latent._shrink_model_tile_param({
|
|
"tile_size_mode": "rows_cols",
|
|
"grid_rows": 16,
|
|
"grid_cols": 16,
|
|
"spatial_w_overlap": 0,
|
|
"spatial_h_overlap": 0,
|
|
"fade_width": 0,
|
|
"fade_height": 0,
|
|
"min_tile_size": 32,
|
|
})
|
|
self.assertIsNotNone(smaller)
|
|
self.assertEqual(smaller["grid_rows"], 32)
|
|
self.assertEqual(smaller["grid_cols"], 32)
|
|
|
|
def test_temporal_segments_split_long_sequences(self):
|
|
latent = importlib.import_module("dumas_h3_latent_upscale")
|
|
bounds = latent._temporal_segments(36, 85, 17)
|
|
self.assertGreater(len(bounds), 1)
|
|
self.assertEqual(bounds[0][0], 0)
|
|
self.assertEqual(bounds[-1][2], 36)
|
|
|
|
def test_shrink_temporal_param_reduces_chunk_length(self):
|
|
latent = importlib.import_module("dumas_h3_latent_upscale")
|
|
smaller = latent._shrink_temporal_param({
|
|
"chunk_length": 85,
|
|
"temporal_overlap": 17,
|
|
})
|
|
self.assertIsNotNone(smaller)
|
|
self.assertEqual(smaller["chunk_length"], 17)
|
|
self.assertEqual(smaller["temporal_overlap"], 0)
|
|
|
|
def test_cuda_model_temporal_params_cap_saved_workflows(self):
|
|
latent = importlib.import_module("dumas_h3_latent_upscale")
|
|
chunk_length, temporal_overlap = latent._effective_temporal_params({
|
|
"mode": "model",
|
|
"device": "cuda",
|
|
"chunk_length": 85,
|
|
"temporal_overlap": 17,
|
|
})
|
|
self.assertEqual(chunk_length, 17)
|
|
self.assertEqual(temporal_overlap, 0)
|
|
|
|
def test_interp_temporal_params_preserve_upstream_defaults(self):
|
|
latent = importlib.import_module("dumas_h3_latent_upscale")
|
|
chunk_length, temporal_overlap = latent._effective_temporal_params({
|
|
"mode": "interp",
|
|
"device": "cuda",
|
|
"chunk_length": 85,
|
|
"temporal_overlap": 17,
|
|
})
|
|
self.assertEqual(chunk_length, 85)
|
|
self.assertEqual(temporal_overlap, 17)
|
|
|
|
def test_unload_upscale_model_defers_while_held(self):
|
|
latent = importlib.import_module("dumas_h3_latent_upscale")
|
|
|
|
class FakeParam:
|
|
device = "cuda"
|
|
|
|
class FakeModel:
|
|
def __init__(self):
|
|
self.moves = []
|
|
|
|
def parameters(self):
|
|
return iter((FakeParam(),))
|
|
|
|
def to(self, device):
|
|
self.moves.append(device)
|
|
return self
|
|
|
|
cache_key = "upscale.safetensors::cuda::fp16"
|
|
original_cache_value = latent._MODEL_CACHE.get(cache_key)
|
|
original_hold_depth = latent._MODEL_HOLD_DEPTH
|
|
fake_model = FakeModel()
|
|
try:
|
|
latent._MODEL_CACHE[cache_key] = fake_model
|
|
latent._MODEL_HOLD_DEPTH = 0
|
|
with latent._hold_upscale_model_loaded():
|
|
latent.unload_upscale_model("upscale.safetensors", "cuda", "fp16")
|
|
self.assertEqual(fake_model.moves, [])
|
|
|
|
latent.unload_upscale_model("upscale.safetensors", "cuda", "fp16")
|
|
self.assertEqual(fake_model.moves, ["cpu"])
|
|
finally:
|
|
latent._MODEL_HOLD_DEPTH = original_hold_depth
|
|
if original_cache_value is None:
|
|
latent._MODEL_CACHE.pop(cache_key, None)
|
|
else:
|
|
latent._MODEL_CACHE[cache_key] = original_cache_value
|
|
|
|
def test_model_upscale_releases_cached_model_after_pass(self):
|
|
latent = importlib.import_module("dumas_h3_latent_upscale")
|
|
calls = []
|
|
|
|
original_temporal = latent._upscale_video_temporal_chunks
|
|
original_unload_now = latent._unload_upscale_model_now
|
|
original_cuda = latent.torch.cuda
|
|
original_device = getattr(latent.torch, "device", None)
|
|
try:
|
|
latent.torch.cuda = types.SimpleNamespace(is_available=lambda: True)
|
|
latent.torch.device = lambda value: value
|
|
|
|
def temporal(video, param, upscaler):
|
|
calls.append(("temporal", latent._MODEL_HOLD_DEPTH))
|
|
return "video", 8, 16
|
|
|
|
def unload_now(name, device, precision):
|
|
calls.append(("unload", name, device, precision, latent._MODEL_HOLD_DEPTH))
|
|
|
|
latent._upscale_video_temporal_chunks = temporal
|
|
latent._unload_upscale_model_now = unload_now
|
|
|
|
result = latent.upscale_latent_video("source", {
|
|
"mode": "model",
|
|
"model_name": "upscale.safetensors",
|
|
"device": "cuda",
|
|
"precision": "fp16",
|
|
})
|
|
|
|
self.assertEqual(result, ("video", 8, 16))
|
|
self.assertEqual(calls[0], ("temporal", 1))
|
|
self.assertEqual(calls[1], ("unload", "upscale.safetensors", "cuda", "fp16", 1))
|
|
self.assertEqual(latent._MODEL_HOLD_DEPTH, 0)
|
|
finally:
|
|
latent._upscale_video_temporal_chunks = original_temporal
|
|
latent._unload_upscale_model_now = original_unload_now
|
|
latent.torch.cuda = original_cuda
|
|
if original_device is None:
|
|
delattr(latent.torch, "device")
|
|
else:
|
|
latent.torch.device = original_device
|
|
|
|
def test_upscale_video_model_raises_when_gpu_cannot_shrink(self):
|
|
latent = importlib.import_module("dumas_h3_latent_upscale")
|
|
|
|
original_tiled = latent._upscale_video_model_tiled
|
|
original_shrink = latent._shrink_model_tile_param
|
|
try:
|
|
latent._shrink_model_tile_param = lambda _param: None
|
|
latent._upscale_video_model_tiled = lambda *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError("out of memory"))
|
|
|
|
with self.assertRaisesRegex(RuntimeError, "H3 latent upscale exhausted its GPU spatial fallbacks"):
|
|
latent.upscale_video_model("video", {"device": "cuda", "precision": "fp16"})
|
|
finally:
|
|
latent._upscale_video_model_tiled = original_tiled
|
|
latent._shrink_model_tile_param = original_shrink
|
|
|
|
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: ([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("Prompt 1", result[3])
|
|
self.assertIn("Beat 1 info\nAnatomy guard: not injected for this prompt\nReferences used: none", result[3])
|
|
self.assertIn("Beat 2 info\nAnatomy guard: injected into this prompt\nReferences used: none", 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()
|