nodes phase 5: workflow saving and loading

invokeai/app/invocations/baseinvocation.py

@@ -2,15 +2,18 @@
 
 from __future__ import annotations
 
+import json
 from abc import ABC, abstractmethod
 from enum import Enum
 from inspect import signature
+import re
 from typing import (
     TYPE_CHECKING,
     AbstractSet,
     Any,
     Callable,
     ClassVar,
+    Literal,
     Mapping,
     Optional,
     Type,
@@ -20,8 +23,8 @@
     get_type_hints,
 )
 
-from pydantic import BaseModel, Field
-from pydantic.fields import Undefined
+from pydantic import BaseModel, Field, validator
+from pydantic.fields import Undefined, ModelField
 from pydantic.typing import NoArgAnyCallable
 
 if TYPE_CHECKING:
@@ -141,9 +144,11 @@ class UIType(str, Enum):
     # endregion
 
     # region Misc
-    FilePath = "FilePath"
     Enum = "enum"
     Scheduler = "Scheduler"
+    WorkflowField = "WorkflowField"
+    IsIntermediate = "IsIntermediate"
+    MetadataField = "MetadataField"
     # endregion
 
 
@@ -365,12 +370,12 @@ def OutputField(
 class UIConfigBase(BaseModel):
     """
     Provides additional node configuration to the UI.
-    This is used internally by the @tags and @title decorator logic. You probably want to use those
-    decorators, though you may add this class to a node definition to specify the title and tags.
+    This is used internally by the @invocation decorator logic. Do not use this directly.
     """
 
-    tags: Optional[list[str]] = Field(default_factory=None, description="The tags to display in the UI")
-    title: Optional[str] = Field(default=None, description="The display name of the node")
+    tags: Optional[list[str]] = Field(default_factory=None, description="The node's tags")
+    title: Optional[str] = Field(default=None, description="The node's display name")
+    category: Optional[str] = Field(default=None, description="The node's category")
 
 
 class InvocationContext:
@@ -383,10 +388,11 @@ def __init__(self, services: InvocationServices, graph_execution_state_id: str):
 
 
 class BaseInvocationOutput(BaseModel):
-    """Base class for all invocation outputs"""
+    """
+    Base class for all invocation outputs.
 
-    # All outputs must include a type name like this:
-    # type: Literal['your_output_name'] # noqa f821
+    All invocation outputs must use the `@invocation_output` decorator to provide their unique type.
+    """
 
     @classmethod
     def get_all_subclasses_tuple(cls):
@@ -422,12 +428,12 @@ def __init__(self, node_id: str, field_name: str):
 
 
 class BaseInvocation(ABC, BaseModel):
-    """A node to process inputs and produce outputs.
-    May use dependency injection in __init__ to receive providers.
     """
+    A node to process inputs and produce outputs.
+    May use dependency injection in __init__ to receive providers.
 
-    # All invocations must include a type name like this:
-    # type: Literal['your_output_name'] # noqa f821
+    All invocations must use the `@invocation` decorator to provide their unique type.
+    """
 
     @classmethod
     def get_all_subclasses(cls):
@@ -466,6 +472,8 @@ def schema_extra(schema: dict[str, Any], model_class: Type[BaseModel]) -> None:
                 schema["title"] = uiconfig.title
             if uiconfig and hasattr(uiconfig, "tags"):
                 schema["tags"] = uiconfig.tags
+            if uiconfig and hasattr(uiconfig, "category"):
+                schema["category"] = uiconfig.category
             if "required" not in schema or not isinstance(schema["required"], list):
                 schema["required"] = list()
             schema["required"].extend(["type", "id"])
@@ -505,37 +513,110 @@ def invoke_internal(self, context: InvocationContext) -> BaseInvocationOutput:
                     raise MissingInputException(self.__fields__["type"].default, field_name)
         return self.invoke(context)
 
-    id: str = Field(description="The id of this node. Must be unique among all nodes.")
+    id: str = Field(
+        description="The id of this instance of an invocation. Must be unique among all instances of invocations."
+    )
     is_intermediate: bool = InputField(
-        default=False, description="Whether or not this node is an intermediate node.", input=Input.Direct
+        default=False, description="Whether or not this is an intermediate invocation.", ui_type=UIType.IsIntermediate
+    )
+    workflow: Optional[str] = InputField(
+        default=None,
+        description="The workflow to save with the image",
+        ui_type=UIType.WorkflowField,
     )
+
+    @validator("workflow", pre=True)
+    def validate_workflow_is_json(cls, v):
+        if v is None:
+            return None
+        try:
+            json.loads(v)
+        except json.decoder.JSONDecodeError:
+            raise ValueError("Workflow must be valid JSON")
+        return v
+
     UIConfig: ClassVar[Type[UIConfigBase]]
 
 
-T = TypeVar("T", bound=BaseInvocation)
+GenericBaseInvocation = TypeVar("GenericBaseInvocation", bound=BaseInvocation)
+
+
+def invocation(
+    invocation_type: str, title: Optional[str] = None, tags: Optional[list[str]] = None, category: Optional[str] = None
+) -> Callable[[Type[GenericBaseInvocation]], Type[GenericBaseInvocation]]:
+    """
+    Adds metadata to an invocation.
 
+    :param str invocation_type: The type of the invocation. Must be unique among all invocations.
+    :param Optional[str] title: Adds a title to the invocation. Use if the auto-generated title isn't quite right. Defaults to None.
+    :param Optional[list[str]] tags: Adds tags to the invocation. Invocations may be searched for by their tags. Defaults to None.
+    :param Optional[str] category: Adds a category to the invocation. Used to group the invocations in the UI. Defaults to None.
+    """
 
-def title(title: str) -> Callable[[Type[T]], Type[T]]:
-    """Adds a title to the invocation. Use this to override the default title generation, which is based on the class name."""
+    def wrapper(cls: Type[GenericBaseInvocation]) -> Type[GenericBaseInvocation]:
+        # Validate invocation types on creation of invocation classes
+        # TODO: ensure unique?
+        if re.compile(r"^\S+$").match(invocation_type) is None:
+            raise ValueError(f'"invocation_type" must consist of non-whitespace characters, got "{invocation_type}"')
 
-    def wrapper(cls: Type[T]) -> Type[T]:
+        # Add OpenAPI schema extras
         uiconf_name = cls.__qualname__ + ".UIConfig"
         if not hasattr(cls, "UIConfig") or cls.UIConfig.__qualname__ != uiconf_name:
             cls.UIConfig = type(uiconf_name, (UIConfigBase,), dict())
-        cls.UIConfig.title = title
+        if title is not None:
+            cls.UIConfig.title = title
+        if tags is not None:
+            cls.UIConfig.tags = tags
+        if category is not None:
+            cls.UIConfig.category = category
+
+        # Add the invocation type to the pydantic model of the invocation
+        invocation_type_annotation = Literal[invocation_type]  # type: ignore
+        invocation_type_field = ModelField.infer(
+            name="type",
+            value=invocation_type,
+            annotation=invocation_type_annotation,
+            class_validators=None,
+            config=cls.__config__,
+        )
+        cls.__fields__.update({"type": invocation_type_field})
+        cls.__annotations__.update({"type": invocation_type_annotation})
+
         return cls
 
     return wrapper
 
 
-def tags(*tags: str) -> Callable[[Type[T]], Type[T]]:
-    """Adds tags to the invocation. Use this to improve the streamline finding the invocation in the UI."""
+GenericBaseInvocationOutput = TypeVar("GenericBaseInvocationOutput", bound=BaseInvocationOutput)
+
+
+def invocation_output(
+    output_type: str,
+) -> Callable[[Type[GenericBaseInvocationOutput]], Type[GenericBaseInvocationOutput]]:
+    """
+    Adds metadata to an invocation output.
+
+    :param str output_type: The type of the invocation output. Must be unique among all invocation outputs.
+    """
+
+    def wrapper(cls: Type[GenericBaseInvocationOutput]) -> Type[GenericBaseInvocationOutput]:
+        # Validate output types on creation of invocation output classes
+        # TODO: ensure unique?
+        if re.compile(r"^\S+$").match(output_type) is None:
+            raise ValueError(f'"output_type" must consist of non-whitespace characters, got "{output_type}"')
+
+        # Add the output type to the pydantic model of the invocation output
+        output_type_annotation = Literal[output_type]  # type: ignore
+        output_type_field = ModelField.infer(
+            name="type",
+            value=output_type,
+            annotation=output_type_annotation,
+            class_validators=None,
+            config=cls.__config__,
+        )
+        cls.__fields__.update({"type": output_type_field})
+        cls.__annotations__.update({"type": output_type_annotation})
 
-    def wrapper(cls: Type[T]) -> Type[T]:
-        uiconf_name = cls.__qualname__ + ".UIConfig"
-        if not hasattr(cls, "UIConfig") or cls.UIConfig.__qualname__ != uiconf_name:
-            cls.UIConfig = type(uiconf_name, (UIConfigBase,), dict())
-        cls.UIConfig.tags = list(tags)
         return cls
 
     return wrapper

invokeai/app/invocations/collections.py

@@ -1,24 +1,19 @@
 # Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654) and the InvokeAI Team
 
-from typing import Literal
 
 import numpy as np
 from pydantic import validator
 
 from invokeai.app.invocations.primitives import IntegerCollectionOutput
 from invokeai.app.util.misc import SEED_MAX, get_random_seed
 
-from .baseinvocation import BaseInvocation, InputField, InvocationContext, tags, title
+from .baseinvocation import BaseInvocation, InputField, InvocationContext, invocation
 
 
-@title("Integer Range")
-@tags("collection", "integer", "range")
+@invocation("range", title="Integer Range", tags=["collection", "integer", "range"], category="collections")
 class RangeInvocation(BaseInvocation):
     """Creates a range of numbers from start to stop with step"""
 
-    type: Literal["range"] = "range"
-
-    # Inputs
     start: int = InputField(default=0, description="The start of the range")
     stop: int = InputField(default=10, description="The stop of the range")
     step: int = InputField(default=1, description="The step of the range")
@@ -33,14 +28,15 @@ def invoke(self, context: InvocationContext) -> IntegerCollectionOutput:
         return IntegerCollectionOutput(collection=list(range(self.start, self.stop, self.step)))
 
 
-@title("Integer Range of Size")
-@tags("range", "integer", "size", "collection")
+@invocation(
+    "range_of_size",
+    title="Integer Range of Size",
+    tags=["collection", "integer", "size", "range"],
+    category="collections",
+)
 class RangeOfSizeInvocation(BaseInvocation):
     """Creates a range from start to start + size with step"""
 
-    type: Literal["range_of_size"] = "range_of_size"
-
-    # Inputs
     start: int = InputField(default=0, description="The start of the range")
     size: int = InputField(default=1, description="The number of values")
     step: int = InputField(default=1, description="The step of the range")
@@ -49,14 +45,15 @@ def invoke(self, context: InvocationContext) -> IntegerCollectionOutput:
         return IntegerCollectionOutput(collection=list(range(self.start, self.start + self.size, self.step)))
 
 
-@title("Random Range")
-@tags("range", "integer", "random", "collection")
+@invocation(
+    "random_range",
+    title="Random Range",
+    tags=["range", "integer", "random", "collection"],
+    category="collections",
+)
 class RandomRangeInvocation(BaseInvocation):
     """Creates a collection of random numbers"""
 
-    type: Literal["random_range"] = "random_range"
-
-    # Inputs
     low: int = InputField(default=0, description="The inclusive low value")
     high: int = InputField(default=np.iinfo(np.int32).max, description="The exclusive high value")
     size: int = InputField(default=1, description="The number of values to generate")

invokeai/app/invocations/compel.py

@@ -1,6 +1,6 @@
 import re
 from dataclasses import dataclass
-from typing import List, Literal, Union
+from typing import List, Union
 
 import torch
 from compel import Compel, ReturnedEmbeddingsType
@@ -26,8 +26,8 @@
     InvocationContext,
     OutputField,
     UIComponent,
-    tags,
-    title,
+    invocation,
+    invocation_output,
 )
 from .model import ClipField
 
@@ -44,13 +44,10 @@ class ConditioningFieldData:
 #    PerpNeg = "perp_neg"
 
 
-@title("Compel Prompt")
-@tags("prompt", "compel")
+@invocation("compel", title="Prompt", tags=["prompt", "compel"], category="conditioning")
 class CompelInvocation(BaseInvocation):
     """Parse prompt using compel package to conditioning."""
 
-    type: Literal["compel"] = "compel"
-
     prompt: str = InputField(
         default="",
         description=FieldDescriptions.compel_prompt,
@@ -265,13 +262,15 @@ def _lora_loader():
         return c, c_pooled, ec
 
 
-@title("SDXL Compel Prompt")
-@tags("sdxl", "compel", "prompt")
+@invocation(
+    "sdxl_compel_prompt",
+    title="SDXL Prompt",
+    tags=["sdxl", "compel", "prompt"],
+    category="conditioning",
+)
 class SDXLCompelPromptInvocation(BaseInvocation, SDXLPromptInvocationBase):
     """Parse prompt using compel package to conditioning."""
 
-    type: Literal["sdxl_compel_prompt"] = "sdxl_compel_prompt"
-
     prompt: str = InputField(default="", description=FieldDescriptions.compel_prompt, ui_component=UIComponent.Textarea)
     style: str = InputField(default="", description=FieldDescriptions.compel_prompt, ui_component=UIComponent.Textarea)
     original_width: int = InputField(default=1024, description="")
@@ -324,13 +323,15 @@ def invoke(self, context: InvocationContext) -> ConditioningOutput:
         )
 
 
-@title("SDXL Refiner Compel Prompt")
-@tags("sdxl", "compel", "prompt")
+@invocation(
+    "sdxl_refiner_compel_prompt",
+    title="SDXL Refiner Prompt",
+    tags=["sdxl", "compel", "prompt"],
+    category="conditioning",
+)
 class SDXLRefinerCompelPromptInvocation(BaseInvocation, SDXLPromptInvocationBase):
     """Parse prompt using compel package to conditioning."""
 
-    type: Literal["sdxl_refiner_compel_prompt"] = "sdxl_refiner_compel_prompt"
-
     style: str = InputField(
         default="", description=FieldDescriptions.compel_prompt, ui_component=UIComponent.Textarea
     )  # TODO: ?
@@ -372,20 +373,17 @@ def invoke(self, context: InvocationContext) -> ConditioningOutput:
         )
 
 
+@invocation_output("clip_skip_output")
 class ClipSkipInvocationOutput(BaseInvocationOutput):
     """Clip skip node output"""
 
-    type: Literal["clip_skip_output"] = "clip_skip_output"
     clip: ClipField = OutputField(default=None, description=FieldDescriptions.clip, title="CLIP")
 
 
-@title("CLIP Skip")
-@tags("clipskip", "clip", "skip")
+@invocation("clip_skip", title="CLIP Skip", tags=["clipskip", "clip", "skip"], category="conditioning")
 class ClipSkipInvocation(BaseInvocation):
     """Skip layers in clip text_encoder model."""
 
-    type: Literal["clip_skip"] = "clip_skip"
-
     clip: ClipField = InputField(description=FieldDescriptions.clip, input=Input.Connection, title="CLIP")
     skipped_layers: int = InputField(default=0, description=FieldDescriptions.skipped_layers)