[nodes] Add Edge data type (invoke-ai#2958)

Adds an `Edge` data type, replacing the current tuple used for edges.

Author: psychedelicious

invokeai/app/api/routers/sessions.py

@@ -10,6 +10,7 @@
 from ...invocations import *
 from ...invocations.baseinvocation import BaseInvocation
 from ...services.graph import (
+    Edge,
     EdgeConnection,
     Graph,
     GraphExecutionState,
@@ -92,7 +93,7 @@ async def get_session(
 async def add_node(
     session_id: str = Path(description="The id of the session"),
     node: Annotated[
-        Union[BaseInvocation.get_invocations()], Field(discriminator="type")
+        Union[BaseInvocation.get_invocations()], Field(discriminator="type") # type: ignore
     ] = Body(description="The node to add"),
 ) -> str:
     """Adds a node to the graph"""
@@ -125,7 +126,7 @@ async def update_node(
     session_id: str = Path(description="The id of the session"),
     node_path: str = Path(description="The path to the node in the graph"),
     node: Annotated[
-        Union[BaseInvocation.get_invocations()], Field(discriminator="type")
+        Union[BaseInvocation.get_invocations()], Field(discriminator="type") # type: ignore
     ] = Body(description="The new node"),
 ) -> GraphExecutionState:
     """Updates a node in the graph and removes all linked edges"""
@@ -186,7 +187,7 @@ async def delete_node(
 )
 async def add_edge(
     session_id: str = Path(description="The id of the session"),
-    edge: tuple[EdgeConnection, EdgeConnection] = Body(description="The edge to add"),
+    edge: Edge = Body(description="The edge to add"),
 ) -> GraphExecutionState:
     """Adds an edge to the graph"""
     session = ApiDependencies.invoker.services.graph_execution_manager.get(session_id)
@@ -228,9 +229,9 @@ async def delete_edge(
         return Response(status_code=404)
 
     try:
-        edge = (
-            EdgeConnection(node_id=from_node_id, field=from_field),
-            EdgeConnection(node_id=to_node_id, field=to_field),
+        edge = Edge(
+            source=EdgeConnection(node_id=from_node_id, field=from_field),
+            destination=EdgeConnection(node_id=to_node_id, field=to_field)
         )
         session.delete_edge(edge)
         ApiDependencies.invoker.services.graph_execution_manager.set(

invokeai/app/cli_app.py

@@ -19,7 +19,7 @@
 from .services.events import EventServiceBase
 from .services.model_manager_initializer import get_model_manager
 from .services.restoration_services import RestorationServices
-from .services.graph import EdgeConnection, GraphExecutionState
+from .services.graph import Edge, EdgeConnection, GraphExecutionState
 from .services.image_storage import DiskImageStorage
 from .services.invocation_queue import MemoryInvocationQueue
 from .services.invocation_services import InvocationServices
@@ -77,7 +77,7 @@ def exit(*args, **kwargs):
 
 def generate_matching_edges(
     a: BaseInvocation, b: BaseInvocation
-) -> list[tuple[EdgeConnection, EdgeConnection]]:
+) -> list[Edge]:
     """Generates all possible edges between two invocations"""
     atype = type(a)
     btype = type(b)
@@ -94,9 +94,9 @@ def generate_matching_edges(
     matching_fields = matching_fields.difference(invalid_fields)
 
     edges = [
-        (
-            EdgeConnection(node_id=a.id, field=field),
-            EdgeConnection(node_id=b.id, field=field),
+        Edge(
+            source=EdgeConnection(node_id=a.id, field=field),
+            destination=EdgeConnection(node_id=b.id, field=field)
         )
         for field in matching_fields
     ]
@@ -111,16 +111,15 @@ class SessionError(Exception):
 def invoke_all(context: CliContext):
     """Runs all invocations in the specified session"""
     context.invoker.invoke(context.session, invoke_all=True)
-    while not context.session.is_complete():
+    while not context.get_session().is_complete():
         # Wait some time
-        session = context.get_session()
         time.sleep(0.1)
 
     # Print any errors
     if context.session.has_error():
         for n in context.session.errors:
             print(
-                f"Error in node {n} (source node {context.session.prepared_source_mapping[n]}): {session.errors[n]}"
+                f"Error in node {n} (source node {context.session.prepared_source_mapping[n]}): {context.session.errors[n]}"
             )
 
         raise SessionError()
@@ -203,7 +202,7 @@ def invoke_cli():
                     continue
 
                 # Pipe previous command output (if there was a previous command)
-                edges = []
+                edges: list[Edge] = list()
                 if len(history) > 0 or current_id != start_id:
                     from_id = (
                         history[0] if current_id == start_id else str(current_id - 1)
@@ -225,19 +224,19 @@ def invoke_cli():
                         matching_edges = generate_matching_edges(
                             link_node, command.command
                         )
-                        matching_destinations = [e[1] for e in matching_edges]
-                        edges = [e for e in edges if e[1] not in matching_destinations]
+                        matching_destinations = [e.destination for e in matching_edges]
+                        edges = [e for e in edges if e.destination not in matching_destinations]
                         edges.extend(matching_edges)
 
                 if "link" in args and args["link"]:
                     for link in args["link"]:
-                        edges = [e for e in edges if e[1].node_id != command.command.id and e[1].field != link[2]]
+                        edges = [e for e in edges if e.destination.node_id != command.command.id and e.destination.field != link[2]]
                         edges.append(
-                            (
-                                EdgeConnection(node_id=link[1], field=link[0]),
-                                EdgeConnection(
+                            Edge(
+                                source=EdgeConnection(node_id=link[1], field=link[0]),
+                                destination=EdgeConnection(
                                     node_id=command.command.id, field=link[2]
-                                ),
+                                )
                             )
                         )