fix(graph): resolve Literal collector items to the types of their values

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,

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")

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)