feat(nodes): allow node clobbering

Author: psychedelicious

invokeai/app/invocations/baseinvocation.py

@@ -262,6 +262,26 @@ class InvocationRegistry:
     @classmethod
     def register_invocation(cls, invocation: type[BaseInvocation]) -> None:
         """Registers an invocation."""
+
+        invocation_type = invocation.get_type()
+        node_pack = invocation.UIConfig.node_pack
+
+        # Log a warning when an existing invocation is being clobbered by the one we are registering
+        clobbered_invocation = InvocationRegistry.get_invocation_for_type(invocation_type)
+        if clobbered_invocation is not None:
+            # This should always be true - we just checked if the invocation type was in the set
+            clobbered_node_pack = clobbered_invocation.UIConfig.node_pack
+
+            if clobbered_node_pack == "invokeai":
+                # The invocation being clobbered is a core invocation
+                logger.warning(f'Overriding core node "{invocation_type}" with node from "{node_pack}"')
+            else:
+                # The invocation being clobbered is a custom invocation
+                logger.warning(
+                    f'Overriding node "{invocation_type}" from "{node_pack}" with node from "{clobbered_node_pack}"'
+                )
+            cls._invocation_classes.remove(clobbered_invocation)
+
         cls._invocation_classes.add(invocation)
         cls.invalidate_invocation_typeadapter()
 
@@ -320,6 +340,15 @@ def get_invocation_for_type(cls, invocation_type: str) -> type[BaseInvocation] |
     @classmethod
     def register_output(cls, output: "type[TBaseInvocationOutput]") -> None:
         """Registers an invocation output."""
+        output_type = output.get_type()
+
+        # Log a warning when an existing invocation is being clobbered by the one we are registering
+        clobbered_output = InvocationRegistry.get_output_for_type(output_type)
+        if clobbered_output is not None:
+            # TODO(psyche): We do not record the node pack of the output, so we cannot log it here
+            logger.warning(f'Overriding invocation output "{output_type}"')
+            cls._output_classes.remove(clobbered_output)
+
         cls._output_classes.add(output)
         cls.invalidate_output_typeadapter()
 
@@ -328,6 +357,11 @@ def get_output_classes(cls) -> Iterable[type[BaseInvocationOutput]]:
         """Gets all invocation outputs."""
         return cls._output_classes
 
+    @classmethod
+    def get_outputs_map(cls) -> dict[str, type[BaseInvocationOutput]]:
+        """Gets a map of all output types to their output classes."""
+        return {i.get_type(): i for i in cls.get_output_classes()}
+
     @classmethod
     @lru_cache(maxsize=1)
     def get_output_typeadapter(cls) -> TypeAdapter[Any]:
@@ -353,6 +387,11 @@ def get_output_types(cls) -> Iterable[str]:
         """Gets all invocation output types."""
         return (i.get_type() for i in cls.get_output_classes())
 
+    @classmethod
+    def get_output_for_type(cls, output_type: str) -> type[BaseInvocationOutput] | None:
+        """Gets the output class for a given output type."""
+        return cls.get_outputs_map().get(output_type)
+
 
 RESERVED_NODE_ATTRIBUTE_FIELD_NAMES = {
     "id",
@@ -466,25 +505,6 @@ def wrapper(cls: Type[TBaseInvocation]) -> Type[TBaseInvocation]:
         # The node pack is the module name - will be "invokeai" for built-in nodes
         node_pack = cls.__module__.split(".")[0]
 
-        # Handle the case where an existing node is being clobbered by the one we are registering
-        if invocation_type in InvocationRegistry.get_invocation_types():
-            clobbered_invocation = InvocationRegistry.get_invocation_for_type(invocation_type)
-            # This should always be true - we just checked if the invocation type was in the set
-            assert clobbered_invocation is not None
-
-            clobbered_node_pack = clobbered_invocation.UIConfig.node_pack
-
-            if clobbered_node_pack == "invokeai":
-                # The node being clobbered is a core node
-                raise ValueError(
-                    f'Cannot load node "{invocation_type}" from node pack "{node_pack}" - a core node with the same type already exists'
-                )
-            else:
-                # The node being clobbered is a custom node
-                raise ValueError(
-                    f'Cannot load node "{invocation_type}" from node pack "{node_pack}" - a node with the same type already exists in node pack "{clobbered_node_pack}"'
-                )
-
         validate_fields(cls.model_fields, invocation_type)
 
         # Add OpenAPI schema extras
@@ -578,13 +598,9 @@ def wrapper(cls: Type[TBaseInvocationOutput]) -> Type[TBaseInvocationOutput]:
         if re.compile(r"^\S+$").match(output_type) is None:
             raise ValueError(f'"output_type" must consist of non-whitespace characters, got "{output_type}"')
 
-        if output_type in InvocationRegistry.get_output_types():
-            raise ValueError(f'Invocation type "{output_type}" already exists')
-
         validate_fields(cls.model_fields, output_type)
 
         # Add the output type to the model.
-
         output_type_annotation = Literal[output_type]  # type: ignore
         output_type_field = Field(
             title="type", default=output_type, json_schema_extra={"field_kind": FieldKind.NodeAttribute}