diff --git a/nodes/keep_face_upscale_image.py b/nodes/keep_face_upscale_image.py index 9aea415..2c17967 100644 --- a/nodes/keep_face_upscale_image.py +++ b/nodes/keep_face_upscale_image.py @@ -19,12 +19,12 @@ from comfy_api.latest import io try: # the package import (ComfyUI loads the node package) from ..src.keep_image_utils import comfy_image_to_cv2, cv2_to_comfy_image - from ..src.keep_model_loader import KEEPModelPack + from ..src.keep_model_loader import KEEPModelPack, keep_model_pack_error from ..src.keep_processor import KEEPFaceProcessor from ..src.keep_progress import report_progress except ImportError: # the flat import (the repository root on sys.path) from src.keep_image_utils import comfy_image_to_cv2, cv2_to_comfy_image - from src.keep_model_loader import KEEPModelPack + from src.keep_model_loader import KEEPModelPack, keep_model_pack_error from src.keep_processor import KEEPFaceProcessor from src.keep_progress import report_progress @@ -96,15 +96,22 @@ class KEEP_FaceUpscaleImageNode(io.ComfyNode): def validate_inputs(cls, **kwargs) -> bool | str: """Refuse a model pack that is not a ``KEEPModelPack``. - The returned string is the error message ComfyUI shows to the user. + ComfyUI runs ``validate_inputs`` during the **prompt validation**, + where the links are not resolved yet: a connected pack arrives as + ``None``. An unresolved pack is therefore accepted here and checked + in ``execute``, which raises the actionable error. """ keep_model = kwargs.get("keep_model") + if keep_model is None or ( + isinstance(keep_model, tuple) and len(keep_model) == 1 and keep_model[0] is None + ): + # A linked pack is not resolved yet during the prompt validation: + # ComfyUI hands it over as ``None`` (after its list unwrapping) or + # as the ``(None,)`` marker of a missing link. The real type check + # happens in ``execute``. + return True if not isinstance(keep_model, KEEPModelPack): - return ( - "Invalid KEEP Model Pack provided: expected a KEEPModelPack, " - f"got {type(keep_model).__name__}. Connect the output of the " - "'Load KEEP Models' node." - ) + return keep_model_pack_error(keep_model) return True @classmethod @@ -122,6 +129,10 @@ class KEEP_FaceUpscaleImageNode(io.ComfyNode): The pack is always offloaded - also when the processing fails (the exception is **not** swallowed: a failed node must fail loudly). """ + error = keep_model_pack_error(keep_model) + if error is not None: + raise RuntimeError(error) + try: keep_model.load_device() report_progress(0.0, 1.0) diff --git a/nodes/keep_process_image_sequence.py b/nodes/keep_process_image_sequence.py index c1fa6f4..7d6dfda 100644 --- a/nodes/keep_process_image_sequence.py +++ b/nodes/keep_process_image_sequence.py @@ -16,10 +16,10 @@ checkpoint behind it. from comfy_api.latest import io try: # the package import (ComfyUI loads the node package) - from ..src.keep_model_loader import KEEPModelPack + from ..src.keep_model_loader import KEEPModelPack, keep_model_pack_error from ..src.keep_processor import KEEPFaceProcessor except ImportError: # the flat import (the repository root on sys.path) - from src.keep_model_loader import KEEPModelPack + from src.keep_model_loader import KEEPModelPack, keep_model_pack_error from src.keep_processor import KEEPFaceProcessor @@ -100,14 +100,24 @@ class KEEP_ProcessImageSequenceNode(io.ComfyNode): @classmethod def validate_inputs(cls, **kwargs) -> bool | str: - """Refuse a model pack that is not a ``KEEPModelPack``.""" + """Refuse a model pack that is not a ``KEEPModelPack``. + + ComfyUI runs ``validate_inputs`` during the **prompt validation**, + where the links are not resolved yet: a connected pack arrives as + ``None``. An unresolved pack is therefore accepted here and checked + in ``execute``, which raises the actionable error. + """ keep_model = kwargs.get("keep_model") + if keep_model is None or ( + isinstance(keep_model, tuple) and len(keep_model) == 1 and keep_model[0] is None + ): + # A linked pack is not resolved yet during the prompt validation: + # ComfyUI hands it over as ``None`` (after its list unwrapping) or + # as the ``(None,)`` marker of a missing link. The real type check + # happens in ``execute``. + return True if not isinstance(keep_model, KEEPModelPack): - return ( - "Invalid KEEP Model Pack provided: expected a KEEPModelPack, " - f"got {type(keep_model).__name__}. Connect the output of the " - "'Load KEEP Models' node." - ) + return keep_model_pack_error(keep_model) return True @classmethod @@ -126,6 +136,10 @@ class KEEP_ProcessImageSequenceNode(io.ComfyNode): The pack is always offloaded - also when the processing fails (the exception is **not** swallowed: a failed node must fail loudly). """ + error = keep_model_pack_error(keep_model) + if error is not None: + raise RuntimeError(error) + try: num_frames = int(images.shape[0]) if hasattr(images, "shape") else len(images) diff --git a/src/keep_model_loader.py b/src/keep_model_loader.py index b00372b..6158017 100644 --- a/src/keep_model_loader.py +++ b/src/keep_model_loader.py @@ -118,6 +118,23 @@ def _load_flownet_weights(flownet, path: str) -> None: target.load_state_dict(weights, strict=True) +def keep_model_pack_error(pack) -> str | None: + """The actionable message when ``pack`` is not a ``KEEPModelPack``. + + ``None`` means "no error". A ``None`` pack is **not** an error here: + ComfyUI runs the prompt validation before the links are resolved, so a + connected pack arrives as ``None`` in that phase - it is checked in + ``execute`` instead, which raises the message. + """ + if isinstance(pack, KEEPModelPack): + return None + return ( + "Invalid KEEP Model Pack provided: expected a KEEPModelPack, " + f"got {type(pack).__name__}. Connect the output of the " + "'Load KEEP Models' node." + ) + + class KEEPModelPack: """The models of the pack and their residency (P4.5). diff --git a/tests/test_nodes_execute.py b/tests/test_nodes_execute.py index 7f4b65b..3ab3fcf 100644 --- a/tests/test_nodes_execute.py +++ b/tests/test_nodes_execute.py @@ -226,6 +226,56 @@ def test_validate_inputs_accepts_a_real_pack(node_module, stub_pack): ) +@pytest.mark.unit +def test_validate_inputs_accepts_an_unresolved_pack(node_module): + """A linked pack has no value during the prompt validation. + + ComfyUI resolves the links only during the execution, so the pack arrives + as ``None`` in the validation phase. Rejecting it there broke every + prompt - the error was reported once per input of the node. + """ + assert ( + node_module.KEEP_FaceUpscaleImageNode.validate_inputs( + image=torch.rand((1, 4, 4, 3), dtype=torch.float32), + keep_model=None, + final_upscale_factor=1.0, + has_aligned_face=True, + only_center_face=True, + draw_bounding_box=False, + ) + is True + ) + # ComfyUI marks an unresolved link with a one element tuple + # (``execution.get_input_data`` -> ``mark_missing``). + assert ( + node_module.KEEP_FaceUpscaleImageNode.validate_inputs( + image=torch.rand((1, 4, 4, 3), dtype=torch.float32), + keep_model=(None,), + final_upscale_factor=1.0, + has_aligned_face=True, + only_center_face=True, + draw_bounding_box=False, + ) + is True + ) + + +@pytest.mark.unit +def test_execute_raises_for_a_foreign_pack(node_module): + """The pack is type checked in ``execute`` - which raises, never returns None.""" + with pytest.raises(RuntimeError) as error: + node_module.KEEP_FaceUpscaleImageNode.execute( + image=torch.rand((1, 4, 4, 3), dtype=torch.float32), + keep_model=object(), + final_upscale_factor=1.0, + has_aligned_face=True, + only_center_face=True, + draw_bounding_box=False, + ) + assert "KEEPModelPack" in str(error.value), str(error.value) + assert "Load KEEP Models" in str(error.value), str(error.value) + + @pytest.mark.mock_integration def test_upscale_face_image_returns_the_processed_image(node_module, stub_pack, monkeypatch): """The node returns the processed image and uses the pack once.""" @@ -377,6 +427,39 @@ def test_the_sequence_node_refuses_a_foreign_pack(sequence_node_module): assert "KEEPModelPack" in message, message +@pytest.mark.unit +def test_the_sequence_node_accepts_an_unresolved_pack(sequence_node_module): + """A linked pack has no value during the prompt validation - it is fine.""" + assert ( + sequence_node_module.KEEP_ProcessImageSequenceNode.validate_inputs( + images=torch.rand((2, 4, 4, 3), dtype=torch.float32), + keep_model=None, + final_upscale_factor=1.0, + has_aligned_frames=False, + only_center_face=True, + draw_bounding_box=False, + max_clip_length=20, + ) + is True + ) + + +@pytest.mark.unit +def test_the_sequence_node_raises_for_a_foreign_pack(sequence_node_module): + """The pack is type checked in ``execute`` - which raises, never returns None.""" + with pytest.raises(RuntimeError) as error: + sequence_node_module.KEEP_ProcessImageSequenceNode.execute( + images=torch.rand((2, 4, 4, 3), dtype=torch.float32), + keep_model=object(), + final_upscale_factor=1.0, + has_aligned_frames=False, + only_center_face=True, + draw_bounding_box=False, + max_clip_length=20, + ) + assert "KEEPModelPack" in str(error.value), str(error.value) + + @pytest.mark.mock_integration def test_the_sequence_node_calls_the_processor_once( sequence_node_module, stub_pack, monkeypatch