feat(ui): improved node parsing

invokeai/frontend/web/config/common.ts

@@ -0,0 +1,14 @@
+import react from '@vitejs/plugin-react-swc';
+import { visualizer } from 'rollup-plugin-visualizer';
+import { PluginOption, UserConfig } from 'vite';
+import eslint from 'vite-plugin-eslint';
+import tsconfigPaths from 'vite-tsconfig-paths';
+import { nodePolyfills } from 'vite-plugin-node-polyfills';
+
+export const commonPlugins: UserConfig['plugins'] = [
+  react(),
+  eslint(),
+  tsconfigPaths(),
+  visualizer() as unknown as PluginOption,
+  nodePolyfills(),
+];

invokeai/frontend/web/config/vite.app.config.ts

@@ -1,17 +1,9 @@
-import react from '@vitejs/plugin-react-swc';
-import { visualizer } from 'rollup-plugin-visualizer';
-import { PluginOption, UserConfig } from 'vite';
-import eslint from 'vite-plugin-eslint';
-import tsconfigPaths from 'vite-tsconfig-paths';
+import { UserConfig } from 'vite';
+import { commonPlugins } from './common';
 
 export const appConfig: UserConfig = {
   base: './',
-  plugins: [
-    react(),
-    eslint(),
-    tsconfigPaths(),
-    visualizer() as unknown as PluginOption,
-  ],
+  plugins: [...commonPlugins],
   build: {
     chunkSizeWarningLimit: 1500,
   },

invokeai/frontend/web/config/vite.package.config.ts

@@ -1,19 +1,13 @@
-import react from '@vitejs/plugin-react-swc';
 import path from 'path';
-import { visualizer } from 'rollup-plugin-visualizer';
-import { PluginOption, UserConfig } from 'vite';
+import { UserConfig } from 'vite';
 import dts from 'vite-plugin-dts';
-import eslint from 'vite-plugin-eslint';
-import tsconfigPaths from 'vite-tsconfig-paths';
 import cssInjectedByJsPlugin from 'vite-plugin-css-injected-by-js';
+import { commonPlugins } from './common';
 
 export const packageConfig: UserConfig = {
   base: './',
   plugins: [
-    react(),
-    eslint(),
-    tsconfigPaths(),
-    visualizer() as unknown as PluginOption,
+    ...commonPlugins,
     dts({
       insertTypesEntry: true,
     }),

invokeai/frontend/web/package.json

@@ -53,6 +53,7 @@
     ]
   },
   "dependencies": {
+    "@apidevtools/swagger-parser": "^10.1.0",
     "@chakra-ui/anatomy": "^2.1.1",
     "@chakra-ui/icons": "^2.0.19",
     "@chakra-ui/react": "^2.7.1",
@@ -154,6 +155,7 @@
     "vite-plugin-css-injected-by-js": "^3.1.1",
     "vite-plugin-dts": "^2.3.0",
     "vite-plugin-eslint": "^1.8.1",
+    "vite-plugin-node-polyfills": "^0.9.0",
     "vite-tsconfig-paths": "^4.2.0",
     "yarn": "^1.22.19"
   }

invokeai/frontend/web/src/app/contexts/DeleteImageContext.tsx

@@ -15,7 +15,7 @@ import { ImageDTO } from 'services/api/types';
 import { RootState } from 'app/store/store';
 import { canvasSelector } from 'features/canvas/store/canvasSelectors';
 import { controlNetSelector } from 'features/controlNet/store/controlNetSlice';
-import { nodesSelecter } from 'features/nodes/store/nodesSlice';
+import { nodesSelector } from 'features/nodes/store/nodesSlice';
 import { generationSelector } from 'features/parameters/store/generationSelectors';
 import { some } from 'lodash-es';
 
@@ -30,7 +30,7 @@ export const selectImageUsage = createSelector(
   [
     generationSelector,
     canvasSelector,
-    nodesSelecter,
+    nodesSelector,
     controlNetSelector,
     (state: RootState, image_name?: string) => image_name,
   ],

invokeai/frontend/web/src/app/store/middleware/devtools/actionSanitizer.ts

@@ -1,30 +1,34 @@
 import { AnyAction } from '@reduxjs/toolkit';
 import { isAnyGraphBuilt } from 'features/nodes/store/actions';
-import { forEach } from 'lodash-es';
+import { nodeTemplatesBuilt } from 'features/nodes/store/nodesSlice';
+import { receivedOpenAPISchema } from 'services/api/thunks/schema';
 import { Graph } from 'services/api/types';
 
 export const actionSanitizer = <A extends AnyAction>(action: A): A => {
   if (isAnyGraphBuilt(action)) {
     if (action.payload.nodes) {
       const sanitizedNodes: Graph['nodes'] = {};
 
-      // Sanitize nodes as needed
-      forEach(action.payload.nodes, (node, key) => {
-        // Don't log the whole freaking dataURL
-        if (node.type === 'dataURL_image') {
-          const { dataURL, ...rest } = node;
-          sanitizedNodes[key] = { ...rest, dataURL: '<dataURL>' };
-        } else {
-          sanitizedNodes[key] = { ...node };
-        }
-      });
-
       return {
         ...action,
         payload: { ...action.payload, nodes: sanitizedNodes },
       };
     }
   }
 
+  if (receivedOpenAPISchema.fulfilled.match(action)) {
+    return {
+      ...action,
+      payload: '<OpenAPI schema omitted>',
+    };
+  }
+
+  if (nodeTemplatesBuilt.match(action)) {
+    return {
+      ...action,
+      payload: '<Node templates omitted>',
+    };
+  }
+
   return action;
 };

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/index.ts

@@ -82,6 +82,7 @@ import {
   addImageRemovedFromBoardFulfilledListener,
   addImageRemovedFromBoardRejectedListener,
 } from './listeners/imageRemovedFromBoard';
+import { addReceivedOpenAPISchemaListener } from './listeners/receivedOpenAPISchema';
 
 export const listenerMiddleware = createListenerMiddleware();
 
@@ -205,3 +206,6 @@ addImageAddedToBoardRejectedListener();
 addImageRemovedFromBoardFulfilledListener();
 addImageRemovedFromBoardRejectedListener();
 addBoardIdSelectedListener();
+
+// Node schemas
+addReceivedOpenAPISchemaListener();

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

@@ -0,0 +1,35 @@
+import { receivedOpenAPISchema } from 'services/api/thunks/schema';
+import { startAppListening } from '..';
+import { log } from 'app/logging/useLogger';
+import { parseSchema } from 'features/nodes/util/parseSchema';
+import { nodeTemplatesBuilt } from 'features/nodes/store/nodesSlice';
+import { size } from 'lodash-es';
+
+const schemaLog = log.child({ namespace: 'schema' });
+
+export const addReceivedOpenAPISchemaListener = () => {
+  startAppListening({
+    actionCreator: receivedOpenAPISchema.fulfilled,
+    effect: (action, { dispatch, getState }) => {
+      const schemaJSON = action.payload;
+
+      schemaLog.info({ data: { schemaJSON } }, 'Dereferenced OpenAPI schema');
+
+      const nodeTemplates = parseSchema(schemaJSON);
+
+      schemaLog.info(
+        { data: { nodeTemplates } },
+        `Built ${size(nodeTemplates)} node templates`
+      );
+
+      dispatch(nodeTemplatesBuilt(nodeTemplates));
+    },
+  });
+
+  startAppListening({
+    actionCreator: receivedOpenAPISchema.rejected,
+    effect: (action, { dispatch, getState }) => {
+      schemaLog.error('Problem dereferencing OpenAPI Schema');
+    },
+  });
+};

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

@@ -3,7 +3,7 @@ import { startAppListening } from '..';
 import { createSelector } from '@reduxjs/toolkit';
 import { generationSelector } from 'features/parameters/store/generationSelectors';
 import { canvasSelector } from 'features/canvas/store/canvasSelectors';
-import { nodesSelecter } from 'features/nodes/store/nodesSlice';
+import { nodesSelector } from 'features/nodes/store/nodesSlice';
 import { controlNetSelector } from 'features/controlNet/store/controlNetSlice';
 import { forEach, uniqBy } from 'lodash-es';
 import { imageUrlsReceived } from 'services/api/thunks/image';
@@ -16,7 +16,7 @@ const selectAllUsedImages = createSelector(
   [
     generationSelector,
     canvasSelector,
-    nodesSelecter,
+    nodesSelector,
     controlNetSelector,
     selectImagesEntities,
   ],

invokeai/frontend/web/src/common/components/IAIMantineMultiSelect.tsx

@@ -27,7 +27,6 @@ const IAIMantineMultiSelect = (props: IAIMultiSelectProps) => {
             borderWidth: '2px',
             borderColor: 'var(--invokeai-colors-base-800)',
             color: 'var(--invokeai-colors-base-100)',
-            padding: 10,
             paddingRight: 24,
             fontWeight: 600,
             '&:hover': { borderColor: 'var(--invokeai-colors-base-700)' },

invokeai/frontend/web/src/common/components/IAIMantineSelect.tsx

@@ -34,6 +34,10 @@ const IAIMantineSelect = (props: IAISelectProps) => {
             '&:focus': {
               borderColor: 'var(--invokeai-colors-accent-600)',
             },
+            '&:disabled': {
+              backgroundColor: 'var(--invokeai-colors-base-700)',
+              color: 'var(--invokeai-colors-base-400)',
+            },
           },
           dropdown: {
             backgroundColor: 'var(--invokeai-colors-base-800)',
@@ -64,7 +68,7 @@ const IAIMantineSelect = (props: IAISelectProps) => {
             },
           },
           rightSection: {
-            width: 24,
+            width: 32,
           },
         })}
         {...rest}

invokeai/frontend/web/src/features/nodes/components/AddNodeMenu.tsx

@@ -1,28 +1,41 @@
 import 'reactflow/dist/style.css';
-import { memo, useCallback } from 'react';
-import {
-  Tooltip,
-  Menu,
-  MenuButton,
-  MenuList,
-  MenuItem,
-} from '@chakra-ui/react';
-import { FaEllipsisV } from 'react-icons/fa';
+import { useCallback, forwardRef } from 'react';
+import { Flex, Text } from '@chakra-ui/react';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
-import { nodeAdded } from '../store/nodesSlice';
+import { nodeAdded, nodesSelector } from '../store/nodesSlice';
 import { map } from 'lodash-es';
-import { RootState } from 'app/store/store';
 import { useBuildInvocation } from '../hooks/useBuildInvocation';
 import { AnyInvocationType } from 'services/events/types';
-import IAIIconButton from 'common/components/IAIIconButton';
 import { useAppToaster } from 'app/components/Toaster';
+import { createSelector } from '@reduxjs/toolkit';
+import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
+import IAIMantineMultiSelect from 'common/components/IAIMantineMultiSelect';
+
+type NodeTemplate = {
+  label: string;
+  value: string;
+  description: string;
+};
+
+const selector = createSelector(
+  nodesSelector,
+  (nodes) => {
+    const data: NodeTemplate[] = map(nodes.invocationTemplates, (template) => {
+      return {
+        label: template.title,
+        value: template.type,
+        description: template.description,
+      };
+    });
+
+    return { data };
+  },
+  defaultSelectorOptions
+);
 
 const AddNodeMenu = () => {
   const dispatch = useAppDispatch();
-
-  const invocationTemplates = useAppSelector(
-    (state: RootState) => state.nodes.invocationTemplates
-  );
+  const { data } = useAppSelector(selector);
 
   const buildInvocation = useBuildInvocation();
 
@@ -46,23 +59,52 @@ const AddNodeMenu = () => {
   );
 
   return (
-    <Menu isLazy>
-      <MenuButton
-        as={IAIIconButton}
-        aria-label="Add Node"
-        icon={<FaEllipsisV />}
+    <Flex sx={{ gap: 2, alignItems: 'center' }}>
+      <IAIMantineMultiSelect
+        selectOnBlur={false}
+        placeholder="Add Node"
+        value={[]}
+        data={data}
+        maxDropdownHeight={400}
+        nothingFound="No matching nodes"
+        itemComponent={SelectItem}
+        filter={(value, selected, item: NodeTemplate) =>
+          item.label.toLowerCase().includes(value.toLowerCase().trim()) ||
+          item.value.toLowerCase().includes(value.toLowerCase().trim()) ||
+          item.description.toLowerCase().includes(value.toLowerCase().trim())
+        }
+        onChange={(v) => {
+          v[0] && addNode(v[0] as AnyInvocationType);
+        }}
+        sx={{
+          width: '18rem',
+        }}
       />
-      <MenuList overflowY="scroll" height={400}>
-        {map(invocationTemplates, ({ title, description, type }, key) => {
-          return (
-            <Tooltip key={key} label={description} placement="end" hasArrow>
-              <MenuItem onClick={() => addNode(type)}>{title}</MenuItem>
-            </Tooltip>
-          );
-        })}
-      </MenuList>
-    </Menu>
+    </Flex>
   );
 };
 
-export default memo(AddNodeMenu);
+interface ItemProps extends React.ComponentPropsWithoutRef<'div'> {
+  value: string;
+  label: string;
+  description: string;
+}
+
+const SelectItem = forwardRef<HTMLDivElement, ItemProps>(
+  ({ label, description, ...others }: ItemProps, ref) => {
+    return (
+      <div ref={ref} {...others}>
+        <div>
+          <Text>{label}</Text>
+          <Text size="xs" color="base.600">
+            {description}
+          </Text>
+        </div>
+      </div>
+    );
+  }
+);
+
+SelectItem.displayName = 'SelectItem';
+
+export default AddNodeMenu;

invokeai/frontend/web/src/features/nodes/components/panels/TopLeftPanel.tsx

@@ -1,10 +1,10 @@
 import { memo } from 'react';
 import { Panel } from 'reactflow';
-import NodeSearch from '../search/NodeSearch';
+import AddNodeMenu from '../AddNodeMenu';
 
 const TopLeftPanel = () => (
   <Panel position="top-left">
-    <NodeSearch />
+    <AddNodeMenu />
   </Panel>
 );