438 lines
16 KiB
Python
438 lines
16 KiB
Python
"""P4.1 - the V3 schema of the nodes keeps the compatibility contract.
|
|
|
|
The node ids, the display names, the category, the socket types and every
|
|
widget (name, default, min/max/step and order) are compatibility guarantees
|
|
of the pack: existing user workflows must keep working unchanged.
|
|
|
|
The contract below is the **recorded V1 ``INPUT_TYPES``** (P0.1). After the
|
|
V3 migration the widgets live in the schema, so the test converts the V3
|
|
schema into the very same shape and compares it with the record. A V3
|
|
schema that drops a widget, renames one, changes a default or reorders the
|
|
inputs fails here.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parent.parent
|
|
NODES_DIR = REPO_ROOT / "nodes"
|
|
EXAMPLES_DIR = REPO_ROOT / "example_workflows"
|
|
|
|
#: The node ids of the pack - they must never change.
|
|
PACK_NODE_IDS = (
|
|
"KEEP_ModelLoader",
|
|
"KEEP_FaceUpscaleImage",
|
|
"KEEP_ProcessImageSequence",
|
|
)
|
|
|
|
#: The display names of the pack - they must never change.
|
|
PACK_DISPLAY_NAMES = {
|
|
"KEEP_ModelLoader": "Load KEEP Models",
|
|
"KEEP_FaceUpscaleImage": "KEEP Single Image",
|
|
"KEEP_ProcessImageSequence": "KEEP Image Sequence",
|
|
}
|
|
|
|
#: The category of the pack - it must never change.
|
|
PACK_CATEGORY = "ComfyUI-KEEP"
|
|
|
|
#: The expected widget contract of every node (the recorded V1
|
|
#: ``INPUT_TYPES`` - the source of truth of this file).
|
|
EXPECTED_INPUT_TYPES = {
|
|
"KEEP_ModelLoader": {
|
|
"required": {
|
|
"model": (["KEEP", "Asian"], {"default": "KEEP"}),
|
|
"detection_model": (
|
|
[
|
|
"retinaface_resnet50",
|
|
"retinaface_mobile0.25",
|
|
"YOLOv5l",
|
|
"YOLOv5n",
|
|
],
|
|
{"default": "retinaface_resnet50"},
|
|
),
|
|
},
|
|
"optional": {
|
|
"bg_upscale_model": ("UPSCALE_MODEL",),
|
|
"face_upscale_model": ("UPSCALE_MODEL",),
|
|
},
|
|
},
|
|
"KEEP_FaceUpscaleImage": {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"keep_model": ("KEEP_MODEL_PACK",),
|
|
"final_upscale_factor": (
|
|
"FLOAT",
|
|
{
|
|
"default": 1.0,
|
|
"min": 0.5,
|
|
"max": 4.0,
|
|
"step": 0.1,
|
|
"tooltip": "The final upscaling factor for the output image. The image will be resized to this scale after processing.",
|
|
},
|
|
),
|
|
"has_aligned_face": (
|
|
"BOOLEAN",
|
|
{
|
|
"default": False,
|
|
"tooltip": "Check if the input image is an already aligned 512x512 face.",
|
|
},
|
|
),
|
|
"only_center_face": (
|
|
"BOOLEAN",
|
|
{
|
|
"default": True,
|
|
"tooltip": "If the image has multiple faces, only process the one closest to the center.",
|
|
},
|
|
),
|
|
"draw_bounding_box": (
|
|
"BOOLEAN",
|
|
{
|
|
"default": False,
|
|
"tooltip": "Draw a bounding box around the detected face on the output image.",
|
|
},
|
|
),
|
|
}
|
|
},
|
|
"KEEP_ProcessImageSequence": {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"keep_model": ("KEEP_MODEL_PACK",),
|
|
"final_upscale_factor": (
|
|
"FLOAT",
|
|
{
|
|
"default": 1.0,
|
|
"min": 0.5,
|
|
"max": 4.0,
|
|
"step": 0.1,
|
|
"tooltip": "The final upscaling factor for the output frames. They will be resized to this scale after processing.",
|
|
},
|
|
),
|
|
"has_aligned_frames": (
|
|
"BOOLEAN",
|
|
{
|
|
"default": False,
|
|
"tooltip": "Check if the input frames are already aligned 512x512 faces.",
|
|
},
|
|
),
|
|
"only_center_face": (
|
|
"BOOLEAN",
|
|
{
|
|
"default": True,
|
|
"tooltip": "If frames have multiple faces, only process the one closest to the center.",
|
|
},
|
|
),
|
|
"draw_bounding_box": (
|
|
"BOOLEAN",
|
|
{
|
|
"default": False,
|
|
"tooltip": "Draw a bounding box around the detected face on the output frames.",
|
|
},
|
|
),
|
|
"max_clip_length": (
|
|
"INT",
|
|
{
|
|
"default": 20,
|
|
"min": 1,
|
|
"max": 100,
|
|
"step": 1,
|
|
"tooltip": "Maximum number of frames to process in a single batch to manage VRAM.",
|
|
},
|
|
),
|
|
}
|
|
},
|
|
}
|
|
|
|
#: The expected output sockets of every node (the recorded V1
|
|
#: ``RETURN_TYPES`` / ``RETURN_NAMES``).
|
|
EXPECTED_OUTPUTS = {
|
|
"KEEP_ModelLoader": {
|
|
"types": ("KEEP_MODEL_PACK",),
|
|
"names": ("keep_model_pack",),
|
|
},
|
|
"KEEP_FaceUpscaleImage": {
|
|
"types": ("IMAGE",),
|
|
# the V1 node had no RETURN_NAMES - the socket type is the name
|
|
"names": ("IMAGE",),
|
|
},
|
|
"KEEP_ProcessImageSequence": {
|
|
"types": ("IMAGE",),
|
|
"names": ("processed_images",),
|
|
},
|
|
}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# helpers
|
|
# ---------------------------------------------------------------------------
|
|
#: The class of every node id (the names are stable - they are the class
|
|
#: names of the pack).
|
|
NODE_CLASS_NAMES = {
|
|
"KEEP_ModelLoader": "KEEP_ModelLoaderNode",
|
|
"KEEP_FaceUpscaleImage": "KEEP_FaceUpscaleImageNode",
|
|
"KEEP_ProcessImageSequence": "KEEP_ProcessImageSequenceNode",
|
|
}
|
|
|
|
#: The nodes that are **already** migrated to the V3 schema. The list grows
|
|
#: with the migration steps (P4.1 the loader, P4.2 the single image node,
|
|
#: P4.3 the sequence node); the contract test below works with both
|
|
#: representations, so it stays meaningful in every step.
|
|
V3_NODE_IDS = (
|
|
"KEEP_ModelLoader",
|
|
"KEEP_FaceUpscaleImage",
|
|
"KEEP_ProcessImageSequence",
|
|
)
|
|
|
|
|
|
def _pack_nodes(nodes_package) -> dict:
|
|
"""``{node_id: node class}`` - read from the exports of ``nodes/``."""
|
|
return {
|
|
node_id: getattr(nodes_package, class_name)
|
|
for node_id, class_name in NODE_CLASS_NAMES.items()
|
|
}
|
|
|
|
|
|
def _is_v3(node_class) -> bool:
|
|
from comfy_api.latest import io
|
|
|
|
return isinstance(node_class, type) and issubclass(node_class, io.ComfyNode)
|
|
|
|
|
|
def _v1_contract(input_types: dict) -> dict:
|
|
"""The recorded V1 ``INPUT_TYPES`` in the shape of the node contract."""
|
|
contract: dict[str, list[dict]] = {"required": [], "optional": []}
|
|
for section, widgets in input_types.items():
|
|
for name, spec in widgets.items():
|
|
io_type = spec[0]
|
|
options = spec[1] if len(spec) > 1 else {}
|
|
entry = {
|
|
"id": name,
|
|
"type": io_type if isinstance(io_type, str) else "COMBO",
|
|
"tooltip": options.get("tooltip"),
|
|
}
|
|
if not isinstance(io_type, str):
|
|
entry["options"] = list(io_type)
|
|
for field in ("default", "min", "max", "step"):
|
|
if field in options:
|
|
entry[field] = options[field]
|
|
contract[section].append(entry)
|
|
return contract
|
|
|
|
|
|
def _schema_contract(schema) -> dict:
|
|
"""The V3 schema in the shape of the node contract."""
|
|
contract: dict[str, list[dict]] = {"required": [], "optional": []}
|
|
for item in schema.inputs:
|
|
section = "optional" if item.optional else "required"
|
|
entry = {
|
|
"id": item.id,
|
|
"type": item.get_io_type(),
|
|
"tooltip": item.tooltip,
|
|
}
|
|
for field in ("default", "min", "max", "step", "options"):
|
|
value = getattr(item, field, None)
|
|
if value is not None:
|
|
entry[field] = value
|
|
contract[section].append(entry)
|
|
return contract
|
|
|
|
|
|
def _node_contract(node_class) -> dict:
|
|
"""The contract of a node - from its V3 schema or from its V1 class.
|
|
|
|
The V3 migration moves the widgets from ``INPUT_TYPES`` into the
|
|
schema; the contract of the pack is the same either way, so the test
|
|
compares the *contract*, not the representation.
|
|
"""
|
|
if _is_v3(node_class):
|
|
schema = node_class.GET_SCHEMA()
|
|
return {
|
|
"v3": True,
|
|
"node_id": schema.node_id,
|
|
"display_name": schema.display_name,
|
|
"category": schema.category,
|
|
"description": schema.description,
|
|
"inputs": _schema_contract(schema),
|
|
"output_types": tuple(output.get_io_type() for output in schema.outputs),
|
|
"output_names": tuple(
|
|
output.display_name if output.display_name else output.get_io_type()
|
|
for output in schema.outputs
|
|
),
|
|
"tooltips": {item.id: item.tooltip for item in schema.inputs},
|
|
}
|
|
input_types = node_class.INPUT_TYPES()
|
|
return {
|
|
"v3": False,
|
|
"node_id": None,
|
|
"display_name": None,
|
|
"category": node_class.CATEGORY,
|
|
"description": getattr(node_class, "DESCRIPTION", ""),
|
|
"inputs": _v1_contract(input_types),
|
|
"output_types": tuple(node_class.RETURN_TYPES),
|
|
"output_names": tuple(
|
|
getattr(node_class, "RETURN_NAMES", node_class.RETURN_TYPES)
|
|
),
|
|
"tooltips": {
|
|
entry["id"]: entry["tooltip"]
|
|
for entries in _v1_contract(input_types).values()
|
|
for entry in entries
|
|
},
|
|
}
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def nodes_package():
|
|
"""The ``nodes`` package of the node pack (imported under the alias)."""
|
|
return importlib.import_module("ComfyUI_KEEP.nodes")
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def pack_nodes(nodes_package):
|
|
"""``{node_id: node class}`` of the three pack nodes."""
|
|
return _pack_nodes(nodes_package)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# the layout
|
|
# ---------------------------------------------------------------------------
|
|
@pytest.mark.unit
|
|
def test_every_node_has_its_own_file():
|
|
"""One file per node, and the old root ``nodes.py`` is gone."""
|
|
for filename in (
|
|
"keep_model_loader.py",
|
|
"keep_face_upscale_image.py",
|
|
"keep_process_image_sequence.py",
|
|
"__init__.py",
|
|
):
|
|
assert (NODES_DIR / filename).is_file(), f"missing {NODES_DIR / filename}"
|
|
assert not (REPO_ROOT / "nodes.py").exists(), (
|
|
"the root nodes.py must be deleted - the classes live in nodes/ now"
|
|
)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_the_pack_registers_exactly_the_three_node_ids(pack_nodes):
|
|
"""The three node ids of the pack - and no other node."""
|
|
assert sorted(pack_nodes) == sorted(PACK_NODE_IDS), sorted(pack_nodes)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_every_node_defines_a_v3_schema_and_execute(pack_nodes):
|
|
"""Every migrated node is a V3 node: ``define_schema`` + ``execute``."""
|
|
for node_id in V3_NODE_IDS:
|
|
node_class = pack_nodes[node_id]
|
|
assert _is_v3(node_class), f"{node_id} is not a V3 node yet"
|
|
schema = node_class.GET_SCHEMA()
|
|
assert schema.node_id == node_id, node_id
|
|
assert callable(node_class.execute), node_id
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# the schema contract
|
|
# ---------------------------------------------------------------------------
|
|
@pytest.mark.unit
|
|
def test_the_display_names_and_the_category_are_unchanged(pack_nodes, nodes_package):
|
|
"""The display names of the three nodes and the category are unchanged."""
|
|
v1_display_names = getattr(nodes_package, "NODE_DISPLAY_NAME_MAPPINGS", {})
|
|
for node_id, node_class in pack_nodes.items():
|
|
contract = _node_contract(node_class)
|
|
display_name = contract["display_name"] or v1_display_names.get(node_id)
|
|
assert display_name == PACK_DISPLAY_NAMES[node_id], node_id
|
|
assert contract["category"] == PACK_CATEGORY, node_id
|
|
if contract["v3"]:
|
|
assert contract["description"], f"{node_id} needs a schema description"
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_the_widget_contract_is_unchanged(pack_nodes):
|
|
"""Every widget keeps its name, type, default, min/max/step and order."""
|
|
for node_id, expected in EXPECTED_INPUT_TYPES.items():
|
|
contract = _node_contract(pack_nodes[node_id])
|
|
actual = contract["inputs"]
|
|
expected_contract = _v1_contract(expected)
|
|
|
|
assert list(actual) == list(expected_contract), (
|
|
f"{node_id}: the input sections changed ({list(actual)} != "
|
|
f"{list(expected_contract)})"
|
|
)
|
|
for section in expected_contract:
|
|
actual_entries = actual[section]
|
|
expected_entries = expected_contract[section]
|
|
assert [entry["id"] for entry in actual_entries] == [
|
|
entry["id"] for entry in expected_entries
|
|
], f"{node_id}/{section}: the widget order changed"
|
|
for actual_entry, expected_entry in zip(actual_entries, expected_entries):
|
|
for key, value in expected_entry.items():
|
|
if key == "tooltip" and value is None:
|
|
# the V1 node had no tooltip here - the V3 schema
|
|
# requires one (the plan's P4.1: every input gets a
|
|
# tooltip), but the text is new
|
|
if contract["v3"]:
|
|
assert actual_entry["tooltip"], (
|
|
f"{node_id}/{actual_entry['id']}: "
|
|
"the input needs a tooltip"
|
|
)
|
|
continue
|
|
assert actual_entry.get(key) == value, (
|
|
f"{node_id}/{actual_entry['id']}/{key}: "
|
|
f"{actual_entry.get(key)!r} != {value!r}"
|
|
)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_every_input_has_a_tooltip(pack_nodes):
|
|
"""Every input of every V3 node carries a tooltip."""
|
|
for node_id, node_class in pack_nodes.items():
|
|
if not _is_v3(node_class):
|
|
continue
|
|
schema = node_class.GET_SCHEMA()
|
|
assert schema.inputs, node_id
|
|
for item in schema.inputs:
|
|
assert isinstance(item.tooltip, str) and item.tooltip.strip(), (
|
|
f"{node_id}/{item.id}: the input has no tooltip"
|
|
)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_the_output_sockets_are_unchanged(pack_nodes):
|
|
"""The output socket types and their names are unchanged."""
|
|
for node_id, expected in EXPECTED_OUTPUTS.items():
|
|
contract = _node_contract(pack_nodes[node_id])
|
|
assert contract["output_types"] == expected["types"], (
|
|
f"{node_id}: {contract['output_types']} != {expected['types']}"
|
|
)
|
|
assert contract["output_names"] == expected["names"], (
|
|
f"{node_id}: {contract['output_names']} != {expected['names']}"
|
|
)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_the_example_workflow_ids_exist_in_the_pack(pack_nodes):
|
|
"""The ids of the shipped example are registered by the node."""
|
|
workflow = json.loads(
|
|
(EXAMPLES_DIR / "keep-workflow.json").read_text(encoding="utf-8")
|
|
)
|
|
used = {
|
|
node.get("type")
|
|
for node in workflow["nodes"]
|
|
if isinstance(node.get("type"), str) and node["type"].startswith("KEEP_")
|
|
}
|
|
assert used, "the example must use the KEEP nodes"
|
|
assert used <= set(pack_nodes), sorted(used - set(pack_nodes))
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_the_node_modules_are_importable_under_their_own_names():
|
|
"""Every node file is importable as a module of the package."""
|
|
for module_name, class_name in (
|
|
("keep_model_loader", "KEEP_ModelLoaderNode"),
|
|
("keep_face_upscale_image", "KEEP_FaceUpscaleImageNode"),
|
|
("keep_process_image_sequence", "KEEP_ProcessImageSequenceNode"),
|
|
):
|
|
module = importlib.import_module(f"ComfyUI_KEEP.nodes.{module_name}")
|
|
assert hasattr(module, class_name), f"{module_name} has no {class_name}"
|