feat(nodes): support bottleneck flag for nodes

Author: psychedelicious

invokeai/app/invocations/baseinvocation.py

@@ -72,6 +72,17 @@ class Classification(str, Enum, metaclass=MetaEnum):
     Special = "special"
 
 
+class Bottleneck(str, Enum, metaclass=MetaEnum):
+    """
+    The bottleneck of an invocation.
+    - `Network`: The invocation's execution is network-bound.
+    - `GPU`: The invocation's execution is GPU-bound.
+    """
+
+    Network = "network"
+    GPU = "gpu"
+
+
 class UIConfigBase(BaseModel):
     """
     Provides additional node configuration to the UI.
@@ -241,6 +252,8 @@ def invoke_internal(self, context: InvocationContext, services: "InvocationServi
         json_schema_extra={"field_kind": FieldKind.NodeAttribute},
     )
 
+    bottleneck: ClassVar[Bottleneck]
+
     UIConfig: ClassVar[UIConfigBase]
 
     model_config = ConfigDict(
@@ -399,6 +412,7 @@ def get_output_for_type(cls, output_type: str) -> type[BaseInvocationOutput] | N
     "use_cache",
     "type",
     "workflow",
+    "bottleneck",
 }
 
 RESERVED_INPUT_FIELD_NAMES = {"metadata", "board"}
@@ -483,6 +497,7 @@ def invocation(
     version: Optional[str] = None,
     use_cache: Optional[bool] = True,
     classification: Classification = Classification.Stable,
+    bottleneck: Bottleneck = Bottleneck.GPU,
 ) -> Callable[[Type[TBaseInvocation]], Type[TBaseInvocation]]:
     """
     Registers an invocation.
@@ -494,6 +509,7 @@ def invocation(
     :param Optional[str] version: Adds a version to the invocation. Must be a valid semver string. Defaults to None.
     :param Optional[bool] use_cache: Whether or not to use the invocation cache. Defaults to True. The user may override this in the workflow editor.
     :param Classification classification: The classification of the invocation. Defaults to FeatureClassification.Stable. Use Beta or Prototype if the invocation is unstable.
+    :param Bottleneck bottleneck: The bottleneck of the invocation. Defaults to Bottleneck.GPU. Use Network if the invocation is network-bound.
     """
 
     def wrapper(cls: Type[TBaseInvocation]) -> Type[TBaseInvocation]:
@@ -530,6 +546,8 @@ def wrapper(cls: Type[TBaseInvocation]) -> Type[TBaseInvocation]:
         if use_cache is not None:
             cls.model_fields["use_cache"].default = use_cache
 
+        cls.bottleneck = bottleneck
+
         # Add the invocation type to the model.
 
         # You'd be tempted to just add the type field and rebuild the model, like this:

invokeai/invocation_api/__init__.py

@@ -7,6 +7,7 @@
 from invokeai.app.invocations.baseinvocation import (
     BaseInvocation,
     BaseInvocationOutput,
+    Bottleneck,
     Classification,
     invocation,
     invocation_output,
@@ -86,6 +87,7 @@
     # invokeai.app.invocations.baseinvocation
     "BaseInvocation",
     "BaseInvocationOutput",
+    "Bottleneck",
     "Classification",
     "invocation",
     "invocation_output",