112 lines
5.6 KiB
Python
112 lines
5.6 KiB
Python
import json
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
|
|
COMFY_ROOT = Path(__file__).resolve().parents[3]
|
|
PACKAGE_ROOT = Path(__file__).resolve().parents[1]
|
|
WORKFLOW_DIR = PACKAGE_ROOT / "examples" / "wan22"
|
|
|
|
|
|
class WorkflowJson:
|
|
def __init__(self, filename):
|
|
self.path = WORKFLOW_DIR / filename
|
|
self.workflow = json.loads(self.path.read_text(encoding="utf-8"))
|
|
self.nodes = {int(node["id"]): node for node in self.workflow["nodes"]}
|
|
self.links = {int(link[0]): link for link in self.workflow["links"]}
|
|
|
|
def nodes_of_type(self, node_type):
|
|
return [node for node in self.nodes.values() if node["type"] == node_type]
|
|
|
|
def link(self, source_id, source_slot, target_id, target_slot):
|
|
return next(
|
|
(
|
|
link
|
|
for link in self.links.values()
|
|
if (int(link[1]), int(link[2]), int(link[3]), int(link[4]))
|
|
== (source_id, source_slot, target_id, target_slot)
|
|
),
|
|
None,
|
|
)
|
|
|
|
def assert_links_valid(self, test_case):
|
|
for node_id, node in self.nodes.items():
|
|
for slot, input_data in enumerate(node.get("inputs", [])):
|
|
link_id = input_data.get("link")
|
|
if link_id is None:
|
|
continue
|
|
test_case.assertIn(link_id, self.links, f"{self.path.name} input link {link_id}")
|
|
test_case.assertEqual((int(self.links[link_id][3]), int(self.links[link_id][4])), (node_id, slot))
|
|
for slot, output_data in enumerate(node.get("outputs", [])):
|
|
for link_id in output_data.get("links") or []:
|
|
test_case.assertIn(link_id, self.links, f"{self.path.name} output link {link_id}")
|
|
test_case.assertEqual((int(self.links[link_id][1]), int(self.links[link_id][2])), (node_id, slot))
|
|
|
|
|
|
class Wan22SegmentWorkflowJsonTests(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.start = WorkflowJson("anima_to_wan2.2_segment_start.json")
|
|
cls.extend = WorkflowJson("wan2.2_segment_extend.json")
|
|
cls.assemble = WorkflowJson("wan2.2_segment_assemble.json")
|
|
|
|
def test_all_segment_workflow_links_are_bidirectional(self):
|
|
self.start.assert_links_valid(self)
|
|
self.extend.assert_links_valid(self)
|
|
self.assemble.assert_links_valid(self)
|
|
|
|
def test_new_stable_path_does_not_use_easy_loop_nodes(self):
|
|
for workflow in (self.start, self.extend, self.assemble):
|
|
self.assertFalse(
|
|
[node for node in workflow.nodes.values() if node["type"] in {"easy forLoopStart", "easy forLoopEnd"}],
|
|
workflow.path.name,
|
|
)
|
|
|
|
def test_start_workflow_stores_decode_frames_as_segment_zero(self):
|
|
decode = self.start.nodes_of_type("VAEDecodeTiled")[0]
|
|
store = self.start.nodes_of_type("NukunWan22SegmentStore")[0]
|
|
settings = self.start.nodes_of_type("NukunWan22VideoSettings")[0]
|
|
self.assertEqual(store["widgets_values"][1], 0)
|
|
self.assertEqual(store["widgets_values"][3:5], [False, True])
|
|
self.assertEqual(settings["widgets_values"][:5], ["image_to_video", "balanced", "portrait", 5, 16])
|
|
self.assertIsNotNone(self.start.link(decode["id"], 0, store["id"], 0))
|
|
|
|
def test_start_workflow_has_no_direct_video_save(self):
|
|
self.assertFalse(self.start.nodes_of_type("CreateVideo"))
|
|
self.assertFalse(self.start.nodes_of_type("SaveVideo"))
|
|
self.assertFalse(self.start.nodes_of_type("SaveAnimatedWEBP"))
|
|
|
|
def test_extend_workflow_loads_last_frame_caption_renders_and_trims(self):
|
|
loader = self.extend.nodes_of_type("NukunWan22SegmentLoader")[0]
|
|
vision = self.extend.nodes_of_type("NukunOllamaVisionCaptioner")[0]
|
|
latent = self.extend.nodes_of_type("NukunWan22TI2VLatent")[0]
|
|
decode = self.extend.nodes_of_type("VAEDecodeTiled")[0]
|
|
store = self.extend.nodes_of_type("NukunWan22SegmentStore")[0]
|
|
self.assertEqual(loader["widgets_values"], ["wan_run_001", -1])
|
|
self.assertEqual(store["widgets_values"][3:5], [True, True])
|
|
self.assertIsNotNone(self.extend.link(loader["id"], 0, vision["id"], 0))
|
|
self.assertIsNotNone(self.extend.link(loader["id"], 0, latent["id"], 2))
|
|
self.assertIsNotNone(self.extend.link(decode["id"], 0, store["id"], 0))
|
|
self.assertIsNotNone(self.extend.link(loader["id"], 2, store["id"], 2))
|
|
|
|
def test_extend_workflow_keeps_prompt_bridge_and_seed_connections_visible(self):
|
|
loader = self.extend.nodes_of_type("NukunWan22SegmentLoader")[0]
|
|
refiner = self.extend.nodes_of_type("NukunOllamaPromptRefiner")[0]
|
|
sampler = self.extend.nodes_of_type("KSampler")[0]
|
|
seed_nodes = self.extend.nodes_of_type("Seed (rgthree)")
|
|
self.assertEqual(len(seed_nodes), 3)
|
|
self.assertEqual({node["widgets_values"][0] for node in seed_nodes}, {33004, 44005, 55006})
|
|
self.assertIsNotNone(self.extend.link(loader["id"], 6, self.extend.nodes_of_type("NukunWan22SegmentStore")[0]["id"], 13))
|
|
self.assertIn("wan2_2_video", refiner["widgets_values"])
|
|
self.assertEqual(sampler["widgets_values"][:7], [55006, "fixed", 25, 5, "dpmpp_2m", "bong_tangent", 1])
|
|
|
|
def test_assemble_workflow_uses_frame_sequence_assembler_only(self):
|
|
assembler = self.assemble.nodes_of_type("NukunWan22FrameSequenceAssembler")[0]
|
|
self.assertEqual(assembler["widgets_values"], ["wan_run_001", 16, "final", "libx264", 8, True])
|
|
self.assertFalse(self.assemble.nodes_of_type("CreateVideo"))
|
|
self.assertFalse(self.assemble.nodes_of_type("SaveVideo"))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|