nodes/wildminder-ComfyUI-KEEP
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(nodes): do not type check the linked model pack during prompt validation

Browse Source 
ComfyUI runs a V3 node's validate_inputs during the prompt validation, where
the links are not resolved yet: a connected pack arrives as None (after the
list unwrapping) or as the (None,) marker of a missing link.  Both nodes
rejected that unresolved pack, so every prompt failed with the same error
once per input.

- validate_inputs now accepts an unresolved pack and only refuses a value
  that is present but of the wrong type.
- execute raises the actionable error for a pack that is really wrong
  (the plan already requires execute to raise instead of returning None).
- the shared message lives in src/keep_model_loader.keep_model_pack_error.
- tests cover the unresolved pack (None and (None,)) and the raising execute.

Verified against the real ComfyUI validation path: execution.get_input_data
with a linked pack returns (None,) and validate_inputs now returns True.
This commit is contained in:
WildAi
2026-10-24 20:55:56 +03:00
parent 1af15430a0
commit 726da45f65
4 changed files with 141 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes/keep_face_upscale_image.py
+19 -8
View File
@@ -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)
nodes/keep_process_image_sequence.py
+22 -8
View File
@@ -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)
src/keep_model_loader.py
+17
View File
@@ -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).
tests/test_nodes_execute.py
+83
View File
@@ -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