diff --git a/comfy_execution/utils.py b/comfy_execution/utils.py index 3bca3e995..60e707881 100644 --- a/comfy_execution/utils.py +++ b/comfy_execution/utils.py @@ -23,6 +23,19 @@ current_executing_context: contextvars.ContextVar[ExecutionContext | None] = con def get_executing_context() -> ExecutionContext | None: return current_executing_context.get(None) + +def is_output_needed(output_index: int) -> bool: + """Check if an output at the given index is connected downstream. + + Returns True if the output might be used (should be computed). + Returns False if the output is definitely not connected (safe to skip). + """ + ctx = get_executing_context() + if ctx is None or ctx.expected_outputs is None: + return True + return output_index in ctx.expected_outputs + + class CurrentNodeContext: """ Context manager for setting the current executing node context. diff --git a/tests-unit/execution_test/expected_outputs_test.py b/tests-unit/execution_test/expected_outputs_test.py index e3b610674..b96902826 100644 --- a/tests-unit/execution_test/expected_outputs_test.py +++ b/tests-unit/execution_test/expected_outputs_test.py @@ -10,6 +10,7 @@ from comfy_execution.utils import ( CurrentNodeContext, ExecutionContext, get_executing_context, + is_output_needed, ) @@ -267,3 +268,31 @@ class TestSchemaLazyOutputs: return IO.NodeOutput(1.0) assert TestNodeWithoutLazyOutputs.LAZY_OUTPUTS is False + + +class TestIsOutputNeeded: + """Tests for is_output_needed() helper function.""" + + def test_output_needed_when_in_expected(self): + """Test that output is needed when in expected_outputs.""" + with CurrentNodeContext("prompt-1", "node-1", 0, frozenset({0, 2})): + assert is_output_needed(0) is True + assert is_output_needed(2) is True + + def test_output_not_needed_when_not_in_expected(self): + """Test that output is not needed when not in expected_outputs.""" + with CurrentNodeContext("prompt-1", "node-1", 0, frozenset({0, 2})): + assert is_output_needed(1) is False + assert is_output_needed(3) is False + + def test_output_needed_when_no_context(self): + """Test that output is needed when no context.""" + assert get_executing_context() is None + assert is_output_needed(0) is True + assert is_output_needed(1) is True + + def test_output_needed_when_expected_outputs_is_none(self): + """Test that output is needed when expected_outputs is None.""" + with CurrentNodeContext("prompt-1", "node-1", 0, None): + assert is_output_needed(0) is True + assert is_output_needed(1) is True