feat(nodes,ui): add detection of custom nodes

Custom nodes have a new attribute `node_pack` indicating the node pack they came from.

- This is displayed in the UI in the icon icon tooltip.
- If a workflow is loaded and a node is unavailable, its node pack will be displayed (if it is known).
- If a workflow is migrated from v1 to v2, and the node is unknown, it falls back to "Unknown". If the missing node pack is installed and the node is updated, the node pack will be updated as expected.

invokeai/frontend/web/public/locales/en.json

@@ -160,6 +160,7 @@
         "trainingDesc2": "InvokeAI already supports training custom embeddourings using Textual Inversion using the main script.",
         "txt2img": "Text To Image",
         "unifiedCanvas": "Unified Canvas",
+        "unknown": "Unknown",
         "upload": "Upload"
     },
     "controlnet": {
@@ -802,6 +803,7 @@
         "cannotConnectOutputToOutput": "Cannot connect output to output",
         "cannotConnectToSelf": "Cannot connect to self",
         "cannotDuplicateConnection": "Cannot create duplicate connections",
+        "nodePack": "Node pack",
         "clipField": "Clip",
         "clipFieldDescription": "Tokenizer and text_encoder submodels.",
         "collection": "Collection",
@@ -966,6 +968,7 @@
         "unableToParseNode": "Unable to parse node",
         "unableToUpdateNode": "Unable to update node",
         "unableToValidateWorkflow": "Unable to Validate Workflow",
+        "unableToMigrateWorkflow": "Unable to Migrate Workflow",
         "unknownErrorValidatingWorkflow": "Unknown error validating workflow",
         "inputFieldTypeParseError": "Unable to parse type of input field {{node}}.{{field}} ({{message}})",
         "outputFieldTypeParseError": "Unable to parse type of output field {{node}}.{{field}} ({{message}})",
@@ -979,9 +982,9 @@
         "unhandledInputProperty": "Unhandled input property",
         "unhandledOutputProperty": "Unhandled output property",
         "unknownField": "Unknown field",
-        "unknownFieldType": "$t(nodes.unknownField) type",
+        "unknownFieldType": "$t(nodes.unknownField) type: {{type}}",
         "unknownNode": "Unknown Node",
-        "unknownNodeType": "$t(nodes.unknownNode) type",
+        "unknownNodeType": "Unknown node type",
         "unknownTemplate": "Unknown Template",
         "unknownInput": "Unknown input: {{name}}",
         "unkownInvocation": "Unknown Invocation type",

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/workflowLoadRequested.ts

@@ -3,7 +3,10 @@ import { parseify } from 'common/util/serialize';
 import { workflowLoadRequested } from 'features/nodes/store/actions';
 import { workflowLoaded } from 'features/nodes/store/nodesSlice';
 import { $flow } from 'features/nodes/store/reactFlowInstance';
-import { WorkflowVersionError } from 'features/nodes/types/error';
+import {
+  WorkflowMigrationError,
+  WorkflowVersionError,
+} from 'features/nodes/types/error';
 import { validateWorkflow } from 'features/nodes/util/workflow/validateWorkflow';
 import { addToast } from 'features/system/store/systemSlice';
 import { makeToast } from 'features/system/util/makeToast';
@@ -67,6 +70,18 @@ export const addWorkflowLoadRequestedListener = () => {
               })
             )
           );
+        } else if (e instanceof WorkflowMigrationError) {
+          // There was a problem migrating the workflow to the latest version
+          log.error({ error: parseify(e) }, e.message);
+          dispatch(
+            addToast(
+              makeToast({
+                title: t('nodes.unableToValidateWorkflow'),
+                status: 'error',
+                description: e.message,
+              })
+            )
+          );
         } else if (e instanceof z.ZodError) {
           // There was a problem validating the workflow itself
           const { message } = fromZodError(e, {

invokeai/frontend/web/src/features/nodes/components/flow/nodes/Invocation/InvocationNodeInfoIcon.tsx

@@ -24,6 +24,7 @@ const InvocationNodeInfoIcon = ({ nodeId }: Props) => {
       <Icon
         as={FaInfoCircle}
         sx={{
+          display: 'block',
           boxSize: 4,
           w: 8,
           color: needsUpdate ? 'error.400' : 'base.400',
@@ -109,6 +110,11 @@ const TooltipContent = memo(({ nodeId }: { nodeId: string }) => {
       <Text as="span" sx={{ fontWeight: 600 }}>
         {title}
       </Text>
+      {nodeTemplate?.nodePack && (
+        <Text opacity={0.7}>
+          {t('nodes.nodePack')}: {nodeTemplate.nodePack}
+        </Text>
+      )}
       <Text sx={{ opacity: 0.7, fontStyle: 'oblique 5deg' }}>
         {nodeTemplate?.description}
       </Text>

invokeai/frontend/web/src/features/nodes/components/flow/nodes/Invocation/InvocationNodeUnknownFallback.tsx

@@ -1,9 +1,10 @@
-import { Box, Flex, Text } from '@chakra-ui/react';
+import { Flex, Text } from '@chakra-ui/react';
+import { useNodePack } from 'features/nodes/hooks/useNodePack';
 import { DRAG_HANDLE_CLASSNAME } from 'features/nodes/types/constants';
 import { memo } from 'react';
+import { useTranslation } from 'react-i18next';
 import NodeCollapseButton from '../common/NodeCollapseButton';
 import NodeWrapper from '../common/NodeWrapper';
-import { useTranslation } from 'react-i18next';
 
 type Props = {
   nodeId: string;
@@ -21,6 +22,7 @@ const InvocationNodeUnknownFallback = ({
   selected,
 }: Props) => {
   const { t } = useTranslation();
+  const nodePack = useNodePack(nodeId);
   return (
     <NodeWrapper nodeId={nodeId} selected={selected}>
       <Flex
@@ -62,12 +64,22 @@ const InvocationNodeUnknownFallback = ({
             fontSize: 'sm',
           }}
         >
-          <Box>
-            <Text as="span">{t('nodes.unknownNodeType')}: </Text>
-            <Text as="span" fontWeight={600}>
-              {type}
+          <Flex gap={2} flexDir="column">
+            <Text as="span">
+              {t('nodes.unknownNodeType')}:{' '}
+              <Text as="span" fontWeight={600}>
+                {type}
+              </Text>
             </Text>
-          </Box>
+            {nodePack && (
+              <Text as="span">
+                {t('nodes.nodePack')}:{' '}
+                <Text as="span" fontWeight={600}>
+                  {nodePack}
+                </Text>
+              </Text>
+            )}
+          </Flex>
         </Flex>
       )}
     </NodeWrapper>

invokeai/frontend/web/src/features/nodes/components/flow/nodes/Invocation/fields/InputFieldRenderer.tsx

@@ -298,7 +298,7 @@ const InputFieldRenderer = ({ nodeId, fieldName }: InputFieldProps) => {
           _dark: { color: 'error.300' },
         }}
       >
-        {t('nodes.unknownFieldType')}: {fieldInstance?.type.name}
+        {t('nodes.unknownFieldType', { type: fieldInstance?.type.name })}
       </Text>
     </Box>
   );

invokeai/frontend/web/src/features/nodes/hooks/useNodePack.ts

@@ -0,0 +1,27 @@
+import { createSelector } from '@reduxjs/toolkit';
+import { stateSelector } from 'app/store/store';
+import { useAppSelector } from 'app/store/storeHooks';
+import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
+import { useMemo } from 'react';
+import { isInvocationNode } from '../types/invocation';
+
+export const useNodePack = (nodeId: string) => {
+  const selector = useMemo(
+    () =>
+      createSelector(
+        stateSelector,
+        ({ nodes }) => {
+          const node = nodes.nodes.find((node) => node.id === nodeId);
+          if (!isInvocationNode(node)) {
+            return false;
+          }
+          return node.data.nodePack;
+        },
+        defaultSelectorOptions
+      ),
+    [nodeId]
+  );
+
+  const title = useAppSelector(selector);
+  return title;
+};

invokeai/frontend/web/src/features/nodes/types/error.ts

@@ -12,6 +12,20 @@ export class WorkflowVersionError extends Error {
     this.name = this.constructor.name;
   }
 }
+/**
+ * Workflow Migration Error
+ * Raised when a workflow migration fails.
+ */
+export class WorkflowMigrationError extends Error {
+  /**
+   * Create WorkflowMigrationError
+   * @param {String} message
+   */
+  constructor(message: string) {
+    super(message);
+    this.name = this.constructor.name;
+  }
+}
 
 /**
  * Unable to Update Node Error

invokeai/frontend/web/src/features/nodes/types/invocation.ts

@@ -21,6 +21,7 @@ export const zInvocationTemplate = z.object({
   withWorkflow: z.boolean(),
   version: zSemVer,
   useCache: z.boolean(),
+  nodePack: z.string().min(1).nullish(),
 });
 export type InvocationTemplate = z.infer<typeof zInvocationTemplate>;
 // #endregion
@@ -36,6 +37,7 @@ export const zInvocationNodeData = z.object({
   isIntermediate: z.boolean(),
   useCache: z.boolean(),
   version: zSemVer,
+  nodePack: z.string().min(1).nullish(),
   inputs: z.record(zFieldInputInstance),
   outputs: z.record(zFieldOutputInstance),
 });

invokeai/frontend/web/src/features/nodes/util/schema/parseSchema.ts

@@ -82,6 +82,7 @@ export const parseSchema = (
     const tags = schema.tags ?? [];
     const description = schema.description ?? '';
     const version = schema.version;
+    const nodePack = schema.node_pack;
     let withWorkflow = false;
 
     const inputs = reduce(
@@ -257,6 +258,7 @@ export const parseSchema = (
       outputs,
       useCache,
       withWorkflow,
+      nodePack,
     };
 
     Object.assign(invocationsAccumulator, { [type]: invocation });

invokeai/frontend/web/src/features/nodes/util/workflow/migrations.ts

@@ -1,7 +1,14 @@
+import { $store } from 'app/store/nanostores/store';
+import { RootState } from 'app/store/store';
+import { FieldType } from 'features/nodes/types/field';
+import { InvocationNodeData } from 'features/nodes/types/invocation';
 import { t } from 'i18next';
-import { forEach, isString } from 'lodash-es';
+import { forEach } from 'lodash-es';
 import { z } from 'zod';
-import { WorkflowVersionError } from '../../types/error';
+import {
+  WorkflowMigrationError,
+  WorkflowVersionError,
+} from '../../types/error';
 import { zSemVer } from '../../types/semver';
 import { FIELD_TYPE_V1_TO_FIELD_TYPE_V2_MAPPING } from '../../types/v1/fieldTypeMap';
 import { WorkflowV1, zWorkflowV1 } from '../../types/v1/workflowV1';
@@ -20,22 +27,39 @@ const zWorkflowMetaVersion = z.object({
  * Migrates a workflow from V1 to V2.
  */
 const migrateV1toV2 = (workflowToMigrate: WorkflowV1): WorkflowV2 => {
+  const invocationTemplates = ($store.get()?.getState() as RootState).nodes
+    .nodeTemplates;
   workflowToMigrate.nodes.forEach((node) => {
     if (node.type === 'invocation') {
+      // Migrate field types
       forEach(node.data.inputs, (input) => {
-        if (!isString(input.type)) {
-          return;
+        const newFieldType = FIELD_TYPE_V1_TO_FIELD_TYPE_V2_MAPPING[input.type];
+        if (!newFieldType) {
+          throw new WorkflowMigrationError(
+            t('nodes.unknownFieldType', { type: input.type })
+          );
         }
-        (input.type as unknown) =
-          FIELD_TYPE_V1_TO_FIELD_TYPE_V2_MAPPING[input.type];
+        // Cast as the V2 type
+        (input.type as unknown as FieldType) = newFieldType;
       });
       forEach(node.data.outputs, (output) => {
-        if (!isString(output.type)) {
-          return;
-        }
-        (output.type as unknown) =
+        const newFieldType =
           FIELD_TYPE_V1_TO_FIELD_TYPE_V2_MAPPING[output.type];
+        if (!newFieldType) {
+          throw new WorkflowMigrationError(
+            t('nodes.unknownFieldType', { type: output.type })
+          );
+        }
+        // Cast as the V2 type
+        (output.type as unknown as FieldType) = newFieldType;
       });
+      // Migrate nodePack
+      const invocationTemplate = invocationTemplates[node.data.type];
+      const nodePack = invocationTemplate
+        ? invocationTemplate.nodePack
+        : t('common.unknown');
+      // Cast as the V2 type
+      (node.data as unknown as InvocationNodeData).nodePack = nodePack;
     }
   });
   (workflowToMigrate.meta.version as WorkflowV2['meta']['version']) = '2.0.0';
@@ -49,6 +73,7 @@ export const parseAndMigrateWorkflow = (data: unknown): WorkflowV2 => {
   const workflowVersionResult = zWorkflowMetaVersion.safeParse(data);
 
   if (!workflowVersionResult.success) {
+    console.log(data);
     throw new WorkflowVersionError(t('nodes.unableToGetWorkflowVersion'));
   }