perf(ui): optimize redux selectors for workflow editor

- Build selectors for each node in a react context so components can
re-use the same selectors
- Cache the selectors in the context

Author: psychedelicious

invokeai/frontend/web/src/app/components/GlobalHookIsolator.tsx

@@ -15,6 +15,8 @@ import { useDndMonitor } from 'features/dnd/useDndMonitor';
 import { useDynamicPromptsWatcher } from 'features/dynamicPrompts/hooks/useDynamicPromptsWatcher';
 import { useStarterModelsToast } from 'features/modelManagerV2/hooks/useStarterModelsToast';
 import { useWorkflowBuilderWatcher } from 'features/nodes/components/sidePanel/workflow/IsolatedWorkflowBuilderWatcher';
+import { useSyncExecutionState } from 'features/nodes/hooks/useNodeExecutionState';
+import { useSyncNodeErrors } from 'features/nodes/store/util/fieldValidators';
 import { useReadinessWatcher } from 'features/queue/store/readiness';
 import { configChanged } from 'features/system/store/configSlice';
 import { selectLanguage } from 'features/system/store/systemSelectors';
@@ -47,10 +49,12 @@ export const GlobalHookIsolator = memo(
     useCloseChakraTooltipsOnDragFix();
     useNavigationApi();
     useDndMonitor();
+    useSyncNodeErrors();
 
     // Persistent subscription to the queue counts query - canvas relies on this to know if there are pending
     // and/or in progress canvas sessions.
     useGetQueueCountsByDestinationQuery(queueCountArg);
+    useSyncExecutionState();
 
     useEffect(() => {
       i18n.changeLanguage(language);

invokeai/frontend/web/src/app/store/createMemoizedSelector.ts

@@ -18,3 +18,5 @@ export const getSelectorsOptions = {
     argsMemoize: lruMemoize,
   }),
 };
+
+export const createLruSelector = createSelectorCreator(lruMemoize);

invokeai/frontend/web/src/features/changeBoardModal/components/ChangeBoardModal.tsx

@@ -1,7 +1,6 @@
 import type { ComboboxOnChange, ComboboxOption } from '@invoke-ai/ui-library';
 import { Combobox, ConfirmationAlertDialog, Flex, FormControl, Text } from '@invoke-ai/ui-library';
 import { createSelector } from '@reduxjs/toolkit';
-import { createMemoizedSelector } from 'app/store/createMemoizedSelector';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import { useAssertSingleton } from 'common/hooks/useAssertSingleton';
 import {
@@ -14,7 +13,7 @@ import { useTranslation } from 'react-i18next';
 import { useListAllBoardsQuery } from 'services/api/endpoints/boards';
 import { useAddImagesToBoardMutation, useRemoveImagesFromBoardMutation } from 'services/api/endpoints/images';
 
-const selectImagesToChange = createMemoizedSelector(
+const selectImagesToChange = createSelector(
   selectChangeBoardModalSlice,
   (changeBoardModal) => changeBoardModal.image_names
 );

invokeai/frontend/web/src/features/nodes/components/flow/Flow.tsx

@@ -22,7 +22,6 @@ import { useConnection } from 'features/nodes/hooks/useConnection';
 import { useIsValidConnection } from 'features/nodes/hooks/useIsValidConnection';
 import { useIsWorkflowEditorLocked } from 'features/nodes/hooks/useIsWorkflowEditorLocked';
 import { useNodeCopyPaste } from 'features/nodes/hooks/useNodeCopyPaste';
-import { useSyncExecutionState } from 'features/nodes/hooks/useNodeExecutionState';
 import {
   $addNodeCmdk,
   $cursorPos,
@@ -83,23 +82,16 @@ export const Flow = memo(() => {
   const nodes = useAppSelector(selectNodes);
   const edges = useAppSelector(selectEdges);
   const viewport = useStore($viewport);
-  const needsFit = useStore($needsFit);
-  const mayUndo = useAppSelector(selectMayUndo);
-  const mayRedo = useAppSelector(selectMayRedo);
   const shouldSnapToGrid = useAppSelector(selectShouldSnapToGrid);
   const selectionMode = useAppSelector(selectSelectionMode);
   const { onConnectStart, onConnect, onConnectEnd } = useConnection();
   const flowWrapper = useRef<HTMLDivElement>(null);
   const isValidConnection = useIsValidConnection();
-  const cancelConnection = useReactFlowStore(selectCancelConnection);
   const updateNodeInternals = useUpdateNodeInternals();
-  const store = useAppStore();
-  const isWorkflowsFocused = useIsRegionFocused('workflows');
   const isLocked = useIsWorkflowEditorLocked();
 
   useFocusRegion('workflows', flowWrapper);
 
-  useSyncExecutionState();
   const [borderRadius] = useToken('radii', ['base']);
   const flowStyles = useMemo<CSSProperties>(() => ({ borderRadius }), [borderRadius]);
 
@@ -110,12 +102,12 @@ export const Flow = memo(() => {
       if (!flow) {
         return;
       }
-      if (needsFit) {
+      if ($needsFit.get()) {
         $needsFit.set(false);
         flow.fitView();
       }
     },
-    [dispatch, needsFit]
+    [dispatch]
   );
 
   const onEdgesChange: OnEdgesChange<AnyEdge> = useCallback(
@@ -214,6 +206,83 @@ export const Flow = memo(() => {
 
   // #endregion
 
+  const onNodeClick = useCallback<NodeMouseHandler<AnyNode>>((e, node) => {
+    if (!$isSelectingOutputNode.get()) {
+      return;
+    }
+    if (!isInvocationNode(node)) {
+      return;
+    }
+    const { id } = node.data;
+    $outputNodeId.set(id);
+    $isSelectingOutputNode.set(false);
+  }, []);
+
+  return (
+    <>
+      <ReactFlow<AnyNode, AnyEdge>
+        id="workflow-editor"
+        ref={flowWrapper}
+        defaultViewport={viewport}
+        nodeTypes={nodeTypes}
+        edgeTypes={edgeTypes}
+        nodes={nodes}
+        edges={edges}
+        onInit={onInit}
+        onNodeClick={onNodeClick}
+        onMouseMove={onMouseMove}
+        onNodesChange={onNodesChange}
+        onEdgesChange={onEdgesChange}
+        onReconnect={onReconnect}
+        onReconnectStart={onReconnectStart}
+        onReconnectEnd={onReconnectEnd}
+        onConnectStart={onConnectStart}
+        onConnect={onConnect}
+        onConnectEnd={onConnectEnd}
+        onMoveEnd={handleMoveEnd}
+        connectionLineComponent={CustomConnectionLine}
+        isValidConnection={isValidConnection}
+        edgesFocusable={!isLocked}
+        edgesReconnectable={!isLocked}
+        nodesDraggable={!isLocked}
+        nodesConnectable={!isLocked}
+        nodesFocusable={!isLocked}
+        elementsSelectable={!isLocked}
+        minZoom={0.1}
+        snapToGrid={shouldSnapToGrid}
+        snapGrid={snapGrid}
+        connectionRadius={30}
+        proOptions={proOptions}
+        style={flowStyles}
+        onPaneClick={handlePaneClick}
+        deleteKeyCode={null}
+        selectionMode={selectionMode}
+        elevateEdgesOnSelect
+        nodeDragThreshold={1}
+        noDragClassName={NO_DRAG_CLASS}
+        noWheelClassName={NO_WHEEL_CLASS}
+        noPanClassName={NO_PAN_CLASS}
+      >
+        <Background />
+      </ReactFlow>
+      <HotkeyIsolator />
+    </>
+  );
+});
+
+Flow.displayName = 'Flow';
+
+const HotkeyIsolator = memo(() => {
+  const isLocked = useIsWorkflowEditorLocked();
+
+  const mayUndo = useAppSelector(selectMayUndo);
+  const mayRedo = useAppSelector(selectMayRedo);
+
+  const cancelConnection = useReactFlowStore(selectCancelConnection);
+
+  const store = useAppStore();
+  const isWorkflowsFocused = useIsRegionFocused('workflows');
+
   const { copySelection, pasteSelection, pasteSelectionWithEdges } = useNodeCopyPaste();
 
   useRegisteredHotkeys({
@@ -239,12 +308,12 @@ export const Flow = memo(() => {
       }
     });
     if (nodeChanges.length > 0) {
-      dispatch(nodesChanged(nodeChanges));
+      store.dispatch(nodesChanged(nodeChanges));
     }
     if (edgeChanges.length > 0) {
-      dispatch(edgesChanged(edgeChanges));
+      store.dispatch(edgesChanged(edgeChanges));
     }
-  }, [dispatch, store]);
+  }, [store]);
   useRegisteredHotkeys({
     id: 'selectAll',
     category: 'workflows',
@@ -273,20 +342,20 @@ export const Flow = memo(() => {
     id: 'undo',
     category: 'workflows',
     callback: () => {
-      dispatch(undo());
+      store.dispatch(undo());
     },
     options: { enabled: isWorkflowsFocused && !isLocked && mayUndo, preventDefault: true },
-    dependencies: [mayUndo, isLocked, isWorkflowsFocused],
+    dependencies: [store, mayUndo, isLocked, isWorkflowsFocused],
   });
 
   useRegisteredHotkeys({
     id: 'redo',
     category: 'workflows',
     callback: () => {
-      dispatch(redo());
+      store.dispatch(redo());
     },
     options: { enabled: isWorkflowsFocused && !isLocked && mayRedo, preventDefault: true },
-    dependencies: [mayRedo, isLocked, isWorkflowsFocused],
+    dependencies: [store, mayRedo, isLocked, isWorkflowsFocused],
   });
 
   const onEscapeHotkey = useCallback(() => {
@@ -313,12 +382,12 @@ export const Flow = memo(() => {
         edgeChanges.push({ type: 'remove', id });
       });
     if (nodeChanges.length > 0) {
-      dispatch(nodesChanged(nodeChanges));
+      store.dispatch(nodesChanged(nodeChanges));
     }
     if (edgeChanges.length > 0) {
-      dispatch(edgesChanged(edgeChanges));
+      store.dispatch(edgesChanged(edgeChanges));
     }
-  }, [dispatch, store]);
+  }, [store]);
   useRegisteredHotkeys({
     id: 'deleteSelection',
     category: 'workflows',
@@ -327,65 +396,6 @@ export const Flow = memo(() => {
     dependencies: [deleteSelection, isWorkflowsFocused, isLocked],
   });
 
-  const onNodeClick = useCallback<NodeMouseHandler<AnyNode>>((e, node) => {
-    if (!$isSelectingOutputNode.get()) {
-      return;
-    }
-    if (!isInvocationNode(node)) {
-      return;
-    }
-    const { id } = node.data;
-    $outputNodeId.set(id);
-    $isSelectingOutputNode.set(false);
-  }, []);
-
-  return (
-    <ReactFlow<AnyNode, AnyEdge>
-      id="workflow-editor"
-      ref={flowWrapper}
-      defaultViewport={viewport}
-      nodeTypes={nodeTypes}
-      edgeTypes={edgeTypes}
-      nodes={nodes}
-      edges={edges}
-      onInit={onInit}
-      onNodeClick={onNodeClick}
-      onMouseMove={onMouseMove}
-      onNodesChange={onNodesChange}
-      onEdgesChange={onEdgesChange}
-      onReconnect={onReconnect}
-      onReconnectStart={onReconnectStart}
-      onReconnectEnd={onReconnectEnd}
-      onConnectStart={onConnectStart}
-      onConnect={onConnect}
-      onConnectEnd={onConnectEnd}
-      onMoveEnd={handleMoveEnd}
-      connectionLineComponent={CustomConnectionLine}
-      isValidConnection={isValidConnection}
-      edgesFocusable={!isLocked}
-      edgesReconnectable={!isLocked}
-      nodesDraggable={!isLocked}
-      nodesConnectable={!isLocked}
-      nodesFocusable={!isLocked}
-      elementsSelectable={!isLocked}
-      minZoom={0.1}
-      snapToGrid={shouldSnapToGrid}
-      snapGrid={snapGrid}
-      connectionRadius={30}
-      proOptions={proOptions}
-      style={flowStyles}
-      onPaneClick={handlePaneClick}
-      deleteKeyCode={null}
-      selectionMode={selectionMode}
-      elevateEdgesOnSelect
-      nodeDragThreshold={1}
-      noDragClassName={NO_DRAG_CLASS}
-      noWheelClassName={NO_WHEEL_CLASS}
-      noPanClassName={NO_PAN_CLASS}
-    >
-      <Background />
-    </ReactFlow>
-  );
+  return null;
 });
-
-Flow.displayName = 'Flow';
+HotkeyIsolator.displayName = 'HotkeyIsolator';

invokeai/frontend/web/src/features/nodes/components/flow/edges/util/buildEdgeSelectors.ts

@@ -1,16 +1,16 @@
 import { createSelector } from '@reduxjs/toolkit';
 import { colorTokenToCssVar } from 'common/util/colorTokenToCssVar';
-import { selectNodesSlice } from 'features/nodes/store/selectors';
+import { selectNodes } from 'features/nodes/store/selectors';
 import type { Templates } from 'features/nodes/store/types';
 import { selectWorkflowSettingsSlice } from 'features/nodes/store/workflowSettingsSlice';
 import { isInvocationNode } from 'features/nodes/types/invocation';
 
 import { getFieldColor } from './getEdgeColor';
 
 export const buildSelectAreConnectedNodesSelected = (source: string, target: string) =>
-  createSelector(selectNodesSlice, (nodes): boolean => {
-    const sourceNode = nodes.nodes.find((node) => node.id === source);
-    const targetNode = nodes.nodes.find((node) => node.id === target);
+  createSelector(selectNodes, (nodes): boolean => {
+    const sourceNode = nodes.find((node) => node.id === source);
+    const targetNode = nodes.find((node) => node.id === target);
 
     return Boolean(sourceNode?.selected || targetNode?.selected);
   });
@@ -22,10 +22,13 @@ export const buildSelectEdgeColor = (
   target: string,
   targetHandleId: string | null | undefined
 ) =>
-  createSelector(selectNodesSlice, selectWorkflowSettingsSlice, (nodes, workflowSettings): string => {
+  createSelector(selectNodes, selectWorkflowSettingsSlice, (nodes, workflowSettings): string => {
     const { shouldColorEdges } = workflowSettings;
-    const sourceNode = nodes.nodes.find((node) => node.id === source);
-    const targetNode = nodes.nodes.find((node) => node.id === target);
+    if (!shouldColorEdges) {
+      return colorTokenToCssVar('base.500');
+    }
+    const sourceNode = nodes.find((node) => node.id === source);
+    const targetNode = nodes.find((node) => node.id === target);
 
     if (!sourceNode || !sourceHandleId || !targetNode || !targetHandleId) {
       return colorTokenToCssVar('base.500');
@@ -37,7 +40,7 @@ export const buildSelectEdgeColor = (
     const outputFieldTemplate = sourceNodeTemplate?.outputs[sourceHandleId];
     const sourceType = isInvocationToInvocationEdge ? outputFieldTemplate?.type : undefined;
 
-    return sourceType && shouldColorEdges ? getFieldColor(sourceType) : colorTokenToCssVar('base.500');
+    return sourceType ? getFieldColor(sourceType) : colorTokenToCssVar('base.500');
   });
 
 export const buildSelectEdgeLabel = (
@@ -47,9 +50,9 @@ export const buildSelectEdgeLabel = (
   target: string,
   targetHandleId: string | null | undefined
 ) =>
-  createSelector(selectNodesSlice, (nodes): string | null => {
-    const sourceNode = nodes.nodes.find((node) => node.id === source);
-    const targetNode = nodes.nodes.find((node) => node.id === target);
+  createSelector(selectNodes, (nodes): string | null => {
+    const sourceNode = nodes.find((node) => node.id === source);
+    const targetNode = nodes.find((node) => node.id === target);
 
     if (!sourceNode || !sourceHandleId || !targetNode || !targetHandleId) {
       return null;