fix(ui): fix issue with node editor state not restoring correctly on mount

If `reactflow` initializes before the node templates are parsed, edges may not be rendered and the viewport may get reset.

- Add `isReady` state to `NodesState`. This is false when we are loading or parsing node templates and true when that is finished.
- Conditionally render `reactflow` based on `isReady`.
- Add `viewport` to `NodesState` & handlers to keep it synced. This allows `reactflow` to mount and unmount freely and not lose viewport.

invokeai/frontend/web/src/features/dnd/hooks/useScaledCenteredModifer.ts

@@ -8,7 +8,8 @@ import { useCallback } from 'react';
 
 const selectZoom = createSelector(
   [stateSelector, activeTabNameSelector],
-  ({ nodes }, activeTabName) => (activeTabName === 'nodes' ? nodes.zoom : 1)
+  ({ nodes }, activeTabName) =>
+    activeTabName === 'nodes' ? nodes.viewport.zoom : 1
 );
 
 /**

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

@@ -1,3 +1,4 @@
+import { useToken } from '@chakra-ui/react';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import { useCallback } from 'react';
 import {
@@ -7,8 +8,7 @@ import {
   OnConnectStart,
   OnEdgesChange,
   OnEdgesDelete,
-  OnInit,
-  OnMove,
+  OnMoveEnd,
   OnNodesChange,
   OnNodesDelete,
   OnSelectionChangeFunc,
@@ -26,7 +26,7 @@ import {
   nodesDeleted,
   selectedEdgesChanged,
   selectedNodesChanged,
-  zoomChanged,
+  viewportChanged,
 } from '../store/nodesSlice';
 import { CustomConnectionLine } from './CustomConnectionLine';
 import { edgeTypes } from './CustomEdges';
@@ -44,12 +44,15 @@ export const Flow = () => {
   const dispatch = useAppDispatch();
   const nodes = useAppSelector((state) => state.nodes.nodes);
   const edges = useAppSelector((state) => state.nodes.edges);
+  const viewport = useAppSelector((state) => state.nodes.viewport);
   const shouldSnapToGrid = useAppSelector(
     (state) => state.nodes.shouldSnapToGrid
   );
 
   const isValidConnection = useIsValidConnection();
 
+  const [borderRadius] = useToken('radii', ['base']);
+
   const onNodesChange: OnNodesChange = useCallback(
     (changes) => {
       dispatch(nodesChanged(changes));
@@ -82,10 +85,6 @@ export const Flow = () => {
     dispatch(connectionEnded());
   }, [dispatch]);
 
-  const onInit: OnInit = useCallback((v) => {
-    v.fitView();
-  }, []);
-
   const onEdgesDelete: OnEdgesDelete = useCallback(
     (edges) => {
       dispatch(edgesDeleted(edges));
@@ -108,16 +107,16 @@ export const Flow = () => {
     [dispatch]
   );
 
-  const handleMove: OnMove = useCallback(
+  const handleMoveEnd: OnMoveEnd = useCallback(
     (e, viewport) => {
-      const { zoom } = viewport;
-      dispatch(zoomChanged(zoom));
+      dispatch(viewportChanged(viewport));
     },
     [dispatch]
   );
 
   return (
     <ReactFlow
+      defaultViewport={viewport}
       nodeTypes={nodeTypes}
       edgeTypes={edgeTypes}
       nodes={nodes}
@@ -129,16 +128,16 @@ export const Flow = () => {
       onConnectStart={onConnectStart}
       onConnect={onConnect}
       onConnectEnd={onConnectEnd}
-      onMove={handleMove}
+      onMoveEnd={handleMoveEnd}
       connectionLineComponent={CustomConnectionLine}
       onSelectionChange={handleSelectionChange}
-      onInit={onInit}
       isValidConnection={isValidConnection}
       minZoom={0.2}
       snapToGrid={shouldSnapToGrid}
       snapGrid={[25, 25]}
       connectionRadius={30}
       proOptions={proOptions}
+      style={{ borderRadius }}
     >
       <TopLeftPanel />
       <TopCenterPanel />

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

@@ -1,13 +1,18 @@
-import { Box } from '@chakra-ui/react';
+import { Flex } from '@chakra-ui/react';
+import { useAppSelector } from 'app/store/storeHooks';
+import { IAINoContentFallback } from 'common/components/IAIImageFallback';
 import ResizeHandle from 'features/ui/components/tabs/ResizeHandle';
 import { memo, useState } from 'react';
+import { MdDeviceHub } from 'react-icons/md';
 import { Panel, PanelGroup } from 'react-resizable-panels';
 import 'reactflow/dist/style.css';
-import { Flow } from './Flow';
 import NodeEditorPanelGroup from './panel/NodeEditorPanelGroup';
+import { Flow } from './Flow';
+import { AnimatePresence, motion } from 'framer-motion';
 
 const NodeEditor = () => {
   const [isPanelCollapsed, setIsPanelCollapsed] = useState(false);
+  const isReady = useAppSelector((state) => state.nodes.isReady);
   return (
     <PanelGroup
       id="node-editor"
@@ -27,17 +32,76 @@ const NodeEditor = () => {
         collapsedDirection={isPanelCollapsed ? 'left' : undefined}
       />
       <Panel id="node-editor-content">
-        <Box
+        <Flex
           layerStyle={'first'}
           sx={{
             position: 'relative',
             width: 'full',
             height: 'full',
             borderRadius: 'base',
+            alignItems: 'center',
+            justifyContent: 'center',
           }}
         >
-          <Flow />
-        </Box>
+          <AnimatePresence>
+            {isReady && (
+              <motion.div
+                layoutId="node-editor-flow"
+                initial={{
+                  opacity: 0,
+                }}
+                animate={{
+                  opacity: 1,
+                  transition: { duration: 0.2 },
+                }}
+                exit={{
+                  opacity: 0,
+                  transition: { duration: 0.2 },
+                }}
+                style={{ width: '100%', height: '100%' }}
+              >
+                <Flow />
+              </motion.div>
+            )}
+          </AnimatePresence>
+          <AnimatePresence>
+            {!isReady && (
+              <motion.div
+                layoutId="node-editor-loading"
+                initial={{
+                  opacity: 0,
+                }}
+                animate={{
+                  opacity: 1,
+                  transition: { duration: 0.2 },
+                }}
+                exit={{
+                  opacity: 0,
+                  transition: { duration: 0.2 },
+                }}
+                style={{ position: 'absolute', width: '100%', height: '100%' }}
+              >
+                <Flex
+                  layerStyle={'first'}
+                  sx={{
+                    position: 'relative',
+                    width: 'full',
+                    height: 'full',
+                    borderRadius: 'base',
+                    alignItems: 'center',
+                    justifyContent: 'center',
+                    pointerEvents: 'none',
+                  }}
+                >
+                  <IAINoContentFallback
+                    label="Loading Nodes..."
+                    icon={MdDeviceHub}
+                  />
+                </Flex>
+              </motion.div>
+            )}
+          </AnimatePresence>
+        </Flex>
       </Panel>
     </PanelGroup>
   );

invokeai/frontend/web/src/features/nodes/store/nodesPersistDenylist.ts

@@ -10,4 +10,5 @@ export const nodesPersistDenylist: (keyof NodesState)[] = [
   'currentConnectionFieldType',
   'selectedNodes',
   'selectedEdges',
+  'isReady',
 ];

invokeai/frontend/web/src/features/nodes/store/nodesSlice.ts

@@ -14,6 +14,7 @@ import {
   Node,
   NodeChange,
   OnConnectStartParams,
+  Viewport,
 } from 'reactflow';
 import { receivedOpenAPISchema } from 'services/api/thunks/schema';
 import { sessionInvoked } from 'services/api/thunks/session';
@@ -56,6 +57,7 @@ export const initialNodesState: NodesState = {
   edges: [],
   schema: null,
   nodeTemplates: {},
+  isReady: false,
   connectionStartParams: null,
   currentConnectionFieldType: null,
   shouldShowFieldTypeLegend: false,
@@ -78,7 +80,7 @@ export const initialNodesState: NodesState = {
     exposedFields: [],
   },
   nodeExecutionStates: {},
-  zoom: 1,
+  viewport: { x: 0, y: 0, zoom: 1 },
 };
 
 type FieldValueAction<T extends InputFieldValue> = PayloadAction<{
@@ -521,6 +523,7 @@ const nodesSlice = createSlice({
       action: PayloadAction<Record<string, InvocationTemplate>>
     ) => {
       state.nodeTemplates = action.payload;
+      state.isReady = true;
     },
     nodeEditorReset: (state) => {
       state.nodes = [];
@@ -587,11 +590,14 @@ const nodesSlice = createSlice({
         []
       );
     },
-    zoomChanged: (state, action: PayloadAction<number>) => {
-      state.zoom = action.payload;
+    viewportChanged: (state, action: PayloadAction<Viewport>) => {
+      state.viewport = action.payload;
     },
   },
   extraReducers: (builder) => {
+    builder.addCase(receivedOpenAPISchema.pending, (state) => {
+      state.isReady = false;
+    });
     builder.addCase(receivedOpenAPISchema.fulfilled, (state, action) => {
       state.schema = action.payload;
     });
@@ -693,7 +699,7 @@ export const {
   workflowExposedFieldAdded,
   workflowExposedFieldRemoved,
   fieldLabelChanged,
-  zoomChanged,
+  viewportChanged,
 } = nodesSlice.actions;
 
 export default nodesSlice.reducer;

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

@@ -1,5 +1,5 @@
 import { OpenAPIV3 } from 'openapi-types';
-import { Edge, Node, OnConnectStartParams } from 'reactflow';
+import { Edge, Node, OnConnectStartParams, Viewport } from 'reactflow';
 import {
   FieldType,
   InvocationEdgeExtra,
@@ -27,5 +27,6 @@ export type NodesState = {
   selectedEdges: string[];
   workflow: Omit<Workflow, 'nodes' | 'edges'>;
   nodeExecutionStates: Record<string, NodeExecutionState>;
-  zoom: number;
+  viewport: Viewport;
+  isReady: boolean;
 };