feat(ui): tidy node-related types

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

@@ -9,7 +9,7 @@ import {
   buildNotesNode,
 } from '../store/util/buildNodeData';
 import { DRAG_HANDLE_CLASSNAME, NODE_WIDTH } from '../types/constants';
-import { AnyNodeData, InvocationTemplate } from '../types/invocation';
+import { AnyNode, InvocationTemplate } from '../types/invocation';
 const templatesSelector = createSelector(
   [(state: RootState) => state.nodes],
   (nodes) => nodes.nodeTemplates
@@ -26,7 +26,7 @@ export const useBuildNodeData = () => {
 
   return useCallback(
     // string here is "any invocation type"
-    (type: string | 'current_image' | 'notes'): Node<AnyNodeData> => {
+    (type: string | 'current_image' | 'notes'): AnyNode => {
       let _x = window.innerWidth / 2;
       let _y = window.innerHeight / 2;
 

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

@@ -28,7 +28,7 @@ import {
   appSocketQueueItemStatusChanged,
 } from 'services/events/actions';
 import { v4 as uuidv4 } from 'uuid';
-import { DRAG_HANDLE_CLASSNAME } from '../types/constants';
+import { SHARED_NODE_PROPERTIES } from '../types/constants';
 import {
   BoardFieldValue,
   BooleanFieldValue,
@@ -50,7 +50,7 @@ import {
   VAEModelFieldValue,
 } from '../types/field';
 import {
-  AnyNodeData,
+  AnyNode,
   InvocationTemplate,
   isInvocationNode,
   isNotesNode,
@@ -157,7 +157,7 @@ const nodesSlice = createSlice({
       }
       state.nodes[nodeIndex] = action.payload.node;
     },
-    nodeAdded: (state, action: PayloadAction<Node<AnyNodeData>>) => {
+    nodeAdded: (state, action: PayloadAction<AnyNode>) => {
       const node = action.payload;
       const position = findUnoccupiedPosition(
         state.nodes,
@@ -520,7 +520,7 @@ const nodesSlice = createSlice({
         state.edges = applyEdgeChanges(edgeChanges, state.edges);
       }
     },
-    nodesDeleted: (state, action: PayloadAction<Node<AnyNodeData>[]>) => {
+    nodesDeleted: (state, action: PayloadAction<AnyNode[]>) => {
       action.payload.forEach((node) => {
         state.workflow.exposedFields = state.workflow.exposedFields.filter(
           (f) => f.nodeId !== node.id
@@ -731,7 +731,7 @@ const nodesSlice = createSlice({
 
       state.nodes = applyNodeChanges(
         nodes.map((node) => ({
-          item: { ...node, dragHandle: `.${DRAG_HANDLE_CLASSNAME}` },
+          item: { ...node, ...SHARED_NODE_PROPERTIES },
           type: 'add',
         })),
         []

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

@@ -1,6 +1,4 @@
 import {
-  Edge,
-  Node,
   OnConnectStartParams,
   SelectionMode,
   Viewport,
@@ -8,16 +6,16 @@ import {
 } from 'reactflow';
 import { FieldIdentifier, FieldType } from '../types/field';
 import {
-  AnyNodeData,
-  InvocationEdgeExtra,
+  AnyNode,
+  InvocationNodeEdge,
   InvocationTemplate,
   NodeExecutionState,
 } from '../types/invocation';
 import { WorkflowV2 } from '../types/workflow';
 
 export type NodesState = {
-  nodes: Node<AnyNodeData>[];
-  edges: Edge<InvocationEdgeExtra>[];
+  nodes: AnyNode[];
+  edges: InvocationNodeEdge[];
   nodeTemplates: Record<string, InvocationTemplate>;
   connectionStartParams: OnConnectStartParams | null;
   connectionStartFieldType: FieldType | null;
@@ -37,8 +35,8 @@ export type NodesState = {
   isReady: boolean;
   mouseOverField: FieldIdentifier | null;
   mouseOverNode: string | null;
-  nodesToCopy: Node<AnyNodeData>[];
-  edgesToCopy: Edge<InvocationEdgeExtra>[];
+  nodesToCopy: AnyNode[];
+  edgesToCopy: InvocationNodeEdge[];
   isAddNodePopoverOpen: boolean;
   addNewNodePosition: XYPosition | null;
   selectionMode: SelectionMode;

invokeai/frontend/web/src/features/nodes/store/util/buildNodeData.ts

@@ -1,4 +1,4 @@
-import { DRAG_HANDLE_CLASSNAME } from 'features/nodes/types/constants';
+import { SHARED_NODE_PROPERTIES } from 'features/nodes/types/constants';
 import {
   FieldInputInstance,
   FieldOutputInstance,
@@ -14,10 +14,6 @@ import { reduce } from 'lodash-es';
 import { Node, XYPosition } from 'reactflow';
 import { v4 as uuidv4 } from 'uuid';
 
-export const SHARED_NODE_PROPERTIES: Partial<Node> = {
-  dragHandle: `.${DRAG_HANDLE_CLASSNAME}`,
-};
-
 export const buildNotesNode = (position: XYPosition): Node<NotesNodeData> => {
   const nodeId = uuidv4();
   const node: Node<NotesNodeData> = {

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

@@ -1,3 +1,5 @@
+import { Node } from 'reactflow';
+
 /**
  * How long to wait before showing a tooltip when hovering a field handle.
  */
@@ -14,6 +16,13 @@ export const NODE_WIDTH = 320;
  */
 export const DRAG_HANDLE_CLASSNAME = 'node-drag-handle';
 
+/**
+ * reactflow-specifc properties shared between all node types.
+ */
+export const SHARED_NODE_PROPERTIES: Partial<Node> = {
+  dragHandle: `.${DRAG_HANDLE_CLASSNAME}`,
+};
+
 /**
  * Helper for getting the kind of a field.
  */

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

@@ -1,4 +1,4 @@
-import { Node } from 'reactflow';
+import { Edge, Node } from 'reactflow';
 import { z } from 'zod';
 import { zProgressImage } from './common';
 import {
@@ -64,16 +64,16 @@ export type InvocationNodeData = z.infer<typeof zInvocationNodeData>;
 export type CurrentImageNodeData = z.infer<typeof zCurrentImageNodeData>;
 export type AnyNodeData = z.infer<typeof zAnyNodeData>;
 
-export const isInvocationNode = (
-  node?: Node<AnyNodeData>
-): node is Node<InvocationNodeData> =>
+export type InvocationNode = Node<InvocationNodeData, 'invocation'>;
+export type NotesNode = Node<NotesNodeData, 'notes'>;
+export type CurrentImageNode = Node<CurrentImageNodeData, 'current_image'>;
+export type AnyNode = Node<AnyNodeData>;
+
+export const isInvocationNode = (node?: AnyNode): node is InvocationNode =>
   Boolean(node && node.type === 'invocation');
-export const isNotesNode = (
-  node?: Node<AnyNodeData>
-): node is Node<NotesNodeData> => Boolean(node && node.type === 'notes');
-export const isCurrentImageNode = (
-  node?: Node<AnyNodeData>
-): node is Node<CurrentImageNodeData> =>
+export const isNotesNode = (node?: AnyNode): node is NotesNode =>
+  Boolean(node && node.type === 'notes');
+export const isCurrentImageNode = (node?: AnyNode): node is CurrentImageNode =>
   Boolean(node && node.type === 'current_image');
 export const isInvocationNodeData = (
   node?: AnyNodeData
@@ -101,8 +101,9 @@ export type NodeStatus = z.infer<typeof zNodeStatus>;
 // #endregion
 
 // #region Edges
-export const zInvocationEdgeExtra = z.object({
+export const zInvocationNodeEdgeExtra = z.object({
   type: z.union([z.literal('default'), z.literal('collapsed')]),
 });
-export type InvocationEdgeExtra = z.infer<typeof zInvocationEdgeExtra>;
+export type InvocationNodeEdgeExtra = z.infer<typeof zInvocationNodeEdgeExtra>;
+export type InvocationNodeEdge = Edge<InvocationNodeEdgeExtra>;
 // #endregion