diff --git a/invokeai/app/services/shared/graph_validation.py b/invokeai/app/services/shared/graph_validation.py index 1430c1b0722..da48151e761 100644 --- a/invokeai/app/services/shared/graph_validation.py +++ b/invokeai/app/services/shared/graph_validation.py @@ -12,6 +12,7 @@ Callable, Concatenate, Iterable, + Literal, Optional, ParamSpec, TypeVar, @@ -2244,14 +2245,20 @@ def _get_collector_port_types( return item_input_field_types, collection_input_field_types, output_field_types def _resolve_item_input_types(self, item_input_field_types: list[Any]) -> set[Any]: - return { - resolved_type - for input_field_type in item_input_field_types - for resolved_type in ( - [input_field_type] if get_origin(input_field_type) is None else get_args(input_field_type) + resolved: set[Any] = set() + for input_field_type in item_input_field_types: + branches = ( + [input_field_type] if get_origin(input_field_type) in (None, Literal) else get_args(input_field_type) ) - if resolved_type != NoneType - } + for branch in branches: + if branch == NoneType: + continue + # A Literal's arguments are values, not types: `Literal[256, 512]` sends ints. + if get_origin(branch) is Literal: + resolved.update(type(value) for value in get_args(branch)) + else: + resolved.add(branch) + return resolved def _resolve_collection_input_types( self, diff --git a/tests/test_node_graph.py b/tests/test_node_graph.py index 73b8dfafc20..e33727f03da 100644 --- a/tests/test_node_graph.py +++ b/tests/test_node_graph.py @@ -33,6 +33,7 @@ StringInvocation, ) from invokeai.app.invocations.upscale import ESRGANInvocation +from invokeai.app.services.session_queue.session_queue_common import Batch from invokeai.app.services.shared.graph import ( CollectInvocation, CollectInvocationOutput, @@ -52,6 +53,7 @@ DerivedFieldTestInvocation, ImageToImageTestInvocation, ListPassThroughInvocation, + LiteralTestInvocation, OptionalCollectionTestInvocation, Pep604UnionItemCollectionTestInvocation, PolymorphicStringTestInvocation, @@ -1179,6 +1181,62 @@ def test_graph_rejects_if_branch_that_makes_collector_items_incompatible(): assert graph.is_valid() +def test_graph_collector_accepts_literal_items_as_the_type_of_their_values(): + # A Literal output (FLUX loaders' max_seq_len: Literal[256, 512]) fed to a collector used to raise + # TypeError from issubclass() on its values, which enqueue_batch returned as a 500. + graph = Graph() + graph.add_node(LiteralTestInvocation(id="literal")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_node(IntegerCollectionInvocation(id="ints")) + graph.add_edge(create_edge("literal", "value", "collect", "item")) + graph.add_edge(create_edge("collect", "collection", "ints", "collection")) + assert graph.is_valid() + + +def test_graph_collector_rejects_literal_items_for_a_collection_of_another_type(): + graph = Graph() + graph.add_node(LiteralTestInvocation(id="literal")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_node(StringCollectionInvocation(id="strings")) + graph.add_edge(create_edge("literal", "value", "collect", "item")) + with pytest.raises(InvalidEdgeError, match="matching type"): + graph.add_edge(create_edge("collect", "collection", "strings", "collection")) + + +def test_graph_collector_accepts_optional_literal_items(): + graph = Graph() + graph.add_node(LiteralTestInvocation(id="literal")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_node(StringCollectionInvocation(id="strings")) + graph.add_edge(create_edge("literal", "optional_value", "collect", "item")) + graph.add_edge(create_edge("collect", "collection", "strings", "collection")) + assert graph.is_valid() + + +def test_batch_with_literal_collector_items_is_validated_without_type_error(): + # The enqueue path: a graph posted in a Batch is validated while the request body is parsed. + graph = { + "id": "g", + "nodes": { + "literal": {"id": "literal", "type": "test_literal"}, + "collect": {"id": "collect", "type": "collect"}, + "strings": {"id": "strings", "type": "string_collection"}, + }, + "edges": [ + { + "source": {"node_id": "literal", "field": "value"}, + "destination": {"node_id": "collect", "field": "item"}, + }, + { + "source": {"node_id": "collect", "field": "collection"}, + "destination": {"node_id": "strings", "field": "collection"}, + }, + ], + } + with pytest.raises(ValidationError, match="matching type"): + Batch.model_validate({"graph": graph}) + + def test_graph_rejects_mixed_if_branch_types_for_collector_items(): graph = Graph() string_value = StringInvocation(id="string", value="text") diff --git a/tests/test_nodes.py b/tests/test_nodes.py index 71d3ffc1216..d826041ed1a 100644 --- a/tests/test_nodes.py +++ b/tests/test_nodes.py @@ -1,4 +1,4 @@ -from typing import Any, Callable, Union +from typing import Any, Callable, Literal, Optional, Union from unittest.mock import MagicMock from pydantic import BaseModel @@ -144,6 +144,20 @@ def invoke(self, context: InvocationContext) -> UnionCollectionTestInvocationOut return UnionCollectionTestInvocationOutput(value=self.value) +@invocation_output("test_literal_output") +class LiteralTestInvocationOutput(BaseInvocationOutput): + value: Literal[256, 512] = OutputField(default=256) + optional_value: Optional[Literal["a", "b"]] = OutputField(default=None) + + +@invocation("test_literal", version="1.0.0") +class LiteralTestInvocation(BaseInvocation): + """Outputs Literal-typed values, as the FLUX model loaders' max_seq_len does.""" + + def invoke(self, context: InvocationContext) -> LiteralTestInvocationOutput: + return LiteralTestInvocationOutput() + + @invocation("test_optional_collection", version="1.0.0") class OptionalCollectionTestInvocation(BaseInvocation): value: Union[list[str], None] = InputField(default=None)