137 lines
5.1 KiB
Python
137 lines
5.1 KiB
Python
import importlib
|
|
import sys
|
|
import types
|
|
import unittest
|
|
|
|
|
|
class DumasH3LongVideosUpstreamWrapperTests(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls._saved_modules = {
|
|
name: sys.modules.get(name)
|
|
for name in (
|
|
"torch",
|
|
"nodes",
|
|
"comfy",
|
|
"comfy.utils",
|
|
"comfy.sample",
|
|
"comfy.samplers",
|
|
"comfy.nested_tensor",
|
|
"comfy.model_management",
|
|
"latent_preview",
|
|
"node_helpers",
|
|
"folder_paths",
|
|
"dumas_h3_longvideos",
|
|
"dumas_h3_longvideos_upstream",
|
|
)
|
|
}
|
|
|
|
fake_torch = types.SimpleNamespace(
|
|
cuda=types.SimpleNamespace(OutOfMemoryError=RuntimeError),
|
|
float32="float32",
|
|
float16="float16",
|
|
bfloat16="bfloat16",
|
|
zeros=lambda *args, **kwargs: None,
|
|
empty=lambda *args, **kwargs: None,
|
|
cat=lambda *args, **kwargs: None,
|
|
stack=lambda *args, **kwargs: None,
|
|
tensor=lambda *args, **kwargs: None,
|
|
no_grad=lambda: _NullContext(),
|
|
inference_mode=lambda: _NullContext(),
|
|
)
|
|
fake_nodes = types.SimpleNamespace(
|
|
NODE_CLASS_MAPPINGS={},
|
|
common_ksampler=lambda *args, **kwargs: ({},),
|
|
)
|
|
fake_comfy_samplers = types.SimpleNamespace(
|
|
KSampler=types.SimpleNamespace(
|
|
SAMPLERS=("res_multistep", "euler"),
|
|
SCHEDULERS=("simple", "normal"),
|
|
)
|
|
)
|
|
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,
|
|
sample=types.SimpleNamespace(),
|
|
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.sample"] = fake_comfy.sample
|
|
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["latent_preview"] = types.SimpleNamespace()
|
|
sys.modules["node_helpers"] = types.SimpleNamespace()
|
|
sys.modules["folder_paths"] = types.SimpleNamespace(
|
|
get_folder_paths=lambda name: [],
|
|
get_filename_list=lambda name: [],
|
|
get_full_path=lambda name, filename: None,
|
|
get_temp_directory=lambda: "/tmp",
|
|
get_output_directory=lambda: "/tmp",
|
|
models_dir="/tmp",
|
|
)
|
|
|
|
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_dumas_key_wraps_upstream_node(self):
|
|
mappings = self.module.NODE_CLASS_MAPPINGS
|
|
|
|
self.assertIs(mappings["DumasH3LongVideos"], self.module.H3LongVideos)
|
|
self.assertIs(mappings["H3LongVideos"], self.module.H3LongVideos)
|
|
self.assertIs(mappings["H3LongVideosREF2VA"], self.module.H3LongVideos)
|
|
|
|
def test_upstream_schema_is_exposed_under_dumas_key(self):
|
|
node_cls = self.module.NODE_CLASS_MAPPINGS["DumasH3LongVideos"]
|
|
schema = node_cls.INPUT_TYPES()
|
|
|
|
self.assertIn("prompt", schema["required"])
|
|
self.assertTrue(schema["required"]["prompt"][1]["forceInput"])
|
|
self.assertIn("first_frame", schema["optional"])
|
|
self.assertIn("ref_image_1", schema["optional"])
|
|
self.assertIn("latent_upscale", schema["optional"])
|
|
self.assertIn("handoff_frames", schema["optional"])
|
|
self.assertEqual(schema["optional"]["handoff_frames"][1]["default"], 1)
|
|
self.assertEqual(node_cls.RETURN_NAMES[0:4], ("images", "audio", "info", "script"))
|
|
|
|
def test_handoff_context_claim_names_reference_range(self):
|
|
upstream = importlib.import_module("dumas_h3_longvideos_upstream")
|
|
|
|
self.assertIn("<Picture 2> through <Picture 22>", upstream.handoff_context_claim(2, 22))
|
|
self.assertIn("no new subjects", upstream.handoff_context_claim(2, 22))
|
|
self.assertIn("<Picture 5>", upstream.handoff_context_claim(5, 5))
|
|
|
|
|
|
class _NullContext:
|
|
def __enter__(self):
|
|
return None
|
|
|
|
def __exit__(self, *_exc):
|
|
return False
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|