[Nodes UI] Invocation Component Updates

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

@@ -1,7 +1,7 @@
 import { v4 as uuidv4 } from 'uuid';
 
 import 'reactflow/dist/style.css';
-import { useCallback } from 'react';
+import { memo, useCallback } from 'react';
 import {
   Tooltip,
   Menu,
@@ -10,16 +10,18 @@ import {
   MenuItem,
   IconButton,
 } from '@chakra-ui/react';
-import { FaPlus } from 'react-icons/fa';
+import { FaEllipsisV, FaPlus } from 'react-icons/fa';
 import { useAppDispatch, useAppSelector } from 'app/storeHooks';
 import { nodeAdded } from '../store/nodesSlice';
 import { cloneDeep, map } from 'lodash';
 import { RootState } from 'app/store';
 import { useBuildInvocation } from '../hooks/useBuildInvocation';
 import { addToast } from 'features/system/store/systemSlice';
 import { makeToast } from 'features/system/hooks/useToastWatcher';
+import { IAIIconButton } from 'exports';
+import { AnyInvocationType } from 'services/events/types';
 
-export const AddNodeMenu = () => {
+const AddNodeMenu = () => {
   const dispatch = useAppDispatch();
 
   const invocationTemplates = useAppSelector(
@@ -29,7 +31,7 @@ export const AddNodeMenu = () => {
   const buildInvocation = useBuildInvocation();
 
   const addNode = useCallback(
-    (nodeType: string) => {
+    (nodeType: AnyInvocationType) => {
       const invocation = buildInvocation(nodeType);
 
       if (!invocation) {
@@ -47,9 +49,13 @@ export const AddNodeMenu = () => {
   );
 
   return (
-    <Menu>
-      <MenuButton as={IconButton} aria-label="Add Node" icon={<FaPlus />} />
-      <MenuList>
+    <Menu isLazy>
+      <MenuButton
+        as={IAIIconButton}
+        aria-label="Add Node"
+        icon={<FaEllipsisV />}
+      />
+      <MenuList overflowY="scroll" height={400}>
         {map(invocationTemplates, ({ title, description, type }, key) => {
           return (
             <Tooltip key={key} label={description} placement="end" hasArrow>
@@ -61,3 +67,5 @@ export const AddNodeMenu = () => {
     </Menu>
   );
 };
+
+export default memo(AddNodeMenu);

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

@@ -1,5 +1,5 @@
 import { Tooltip } from '@chakra-ui/react';
-import { CSSProperties, useMemo } from 'react';
+import { CSSProperties, memo, useMemo } from 'react';
 import {
   Handle,
   Position,
@@ -19,11 +19,11 @@ const handleBaseStyles: CSSProperties = {
 };
 
 const inputHandleStyles: CSSProperties = {
-  left: '-1.7rem',
+  left: '-1rem',
 };
 
 const outputHandleStyles: CSSProperties = {
-  right: '-1.7rem',
+  right: '-0.5rem',
 };
 
 const requiredConnectionStyles: CSSProperties = {
@@ -38,13 +38,14 @@ type FieldHandleProps = {
   styles?: CSSProperties;
 };
 
-export const FieldHandle = (props: FieldHandleProps) => {
+const FieldHandle = (props: FieldHandleProps) => {
   const { nodeId, field, isValidConnection, handleType, styles } = props;
   const { name, title, type, description } = field;
 
+  console.log(props);
+
   return (
     <Tooltip
-      key={name}
       label={type}
       placement={handleType === 'target' ? 'start' : 'end'}
       hasArrow
@@ -67,3 +68,5 @@ export const FieldHandle = (props: FieldHandleProps) => {
     </Tooltip>
   );
 };
+
+export default memo(FieldHandle);

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

@@ -2,8 +2,9 @@ import 'reactflow/dist/style.css';
 import { Tooltip, Badge, HStack } from '@chakra-ui/react';
 import { map } from 'lodash';
 import { FIELDS } from '../types/constants';
+import { memo } from 'react';
 
-export const FieldTypeLegend = () => {
+const FieldTypeLegend = () => {
   return (
     <HStack>
       {map(FIELDS, ({ title, description, color }, key) => (
@@ -16,3 +17,5 @@ export const FieldTypeLegend = () => {
     </HStack>
   );
 };
+
+export default memo(FieldTypeLegend);

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

@@ -1,15 +1,12 @@
 import {
   Background,
-  Controls,
   MiniMap,
   OnConnect,
   OnEdgesChange,
   OnNodesChange,
   ReactFlow,
-  ConnectionLineType,
   OnConnectStart,
   OnConnectEnd,
-  Panel,
 } from 'reactflow';
 import { useAppDispatch, useAppSelector } from 'app/storeHooks';
 import { RootState } from 'app/store';
@@ -22,10 +19,10 @@ import {
 } from '../store/nodesSlice';
 import { useCallback } from 'react';
 import { InvocationComponent } from './InvocationComponent';
-import { AddNodeMenu } from './AddNodeMenu';
-import { FieldTypeLegend } from './FieldTypeLegend';
-import { Button } from '@chakra-ui/react';
-import { nodesGraphBuilt } from 'services/thunks/session';
+import TopLeftPanel from './panels/TopLeftPanel';
+import TopRightPanel from './panels/TopRightPanel';
+import TopCenterPanel from './panels/TopCenterPanel';
+import BottomLeftPanel from './panels/BottomLeftPanel.tsx';
 
 const nodeTypes = { invocation: InvocationComponent };
 
@@ -69,10 +66,6 @@ export const Flow = () => {
     [dispatch]
   );
 
-  const handleInvoke = useCallback(() => {
-    dispatch(nodesGraphBuilt());
-  }, [dispatch]);
-
   return (
     <ReactFlow
       nodeTypes={nodeTypes}
@@ -87,17 +80,11 @@ export const Flow = () => {
         style: { strokeWidth: 2 },
       }}
     >
-      <Panel position="top-left">
-        <AddNodeMenu />
-      </Panel>
-      <Panel position="top-center">
-        <Button onClick={handleInvoke}>Will it blend?</Button>
-      </Panel>
-      <Panel position="top-right">
-        <FieldTypeLegend />
-      </Panel>
+      <TopLeftPanel />
+      <TopCenterPanel />
+      <TopRightPanel />
+      <BottomLeftPanel />
       <Background />
-      <Controls />
       <MiniMap nodeStrokeWidth={3} zoomable pannable />
     </ReactFlow>
   );

invokeai/frontend/web/src/features/nodes/components/IAINode/IAINodeHeader.tsx

@@ -0,0 +1,39 @@
+import { Flex, Heading, Tooltip, Icon } from '@chakra-ui/react';
+import { InvocationTemplate } from 'features/nodes/types/types';
+import { memo, MutableRefObject } from 'react';
+import { FaInfoCircle } from 'react-icons/fa';
+
+interface IAINodeHeaderProps {
+  nodeId: string;
+  template: InvocationTemplate;
+}
+
+const IAINodeHeader = (props: IAINodeHeaderProps) => {
+  const { nodeId, template } = props;
+  return (
+    <Flex
+      borderTopRadius="md"
+      justifyContent="space-between"
+      background="base.700"
+      px={2}
+      py={1}
+      alignItems="center"
+    >
+      <Tooltip label={nodeId}>
+        <Heading size="xs" fontWeight={600} color="base.100">
+          {template.title}
+        </Heading>
+      </Tooltip>
+      <Tooltip
+        label={template.description}
+        placement="top"
+        hasArrow
+        shouldWrapChildren
+      >
+        <Icon color="base.300" as={FaInfoCircle} h="min-content" />
+      </Tooltip>
+    </Flex>
+  );
+};
+
+export default memo(IAINodeHeader);

invokeai/frontend/web/src/features/nodes/components/IAINode/IAINodeInputs.tsx

@@ -0,0 +1,146 @@
+import {
+  InputFieldTemplate,
+  InputFieldValue,
+  InvocationTemplate,
+} from 'features/nodes/types/types';
+import { memo, ReactNode, useCallback } from 'react';
+import { map } from 'lodash';
+import { useAppSelector } from 'app/storeHooks';
+import { RootState } from 'app/store';
+import {
+  Box,
+  Flex,
+  FormControl,
+  FormLabel,
+  HStack,
+  Tooltip,
+  Divider,
+} from '@chakra-ui/react';
+import FieldHandle from '../FieldHandle';
+import { useIsValidConnection } from 'features/nodes/hooks/useIsValidConnection';
+import InputFieldComponent from '../InputFieldComponent';
+import { HANDLE_TOOLTIP_OPEN_DELAY } from 'features/nodes/types/constants';
+
+interface IAINodeInputProps {
+  nodeId: string;
+
+  input: InputFieldValue;
+  template?: InputFieldTemplate | undefined;
+  connected: boolean;
+}
+
+function IAINodeInput(props: IAINodeInputProps) {
+  const { nodeId, input, template, connected } = props;
+  const isValidConnection = useIsValidConnection();
+
+  return (
+    <Box
+      position="relative"
+      borderColor={
+        !template
+          ? 'error.400'
+          : !connected &&
+            ['always', 'connectionOnly'].includes(
+              String(template?.inputRequirement)
+            ) &&
+            input.value === undefined
+          ? 'warning.400'
+          : undefined
+      }
+    >
+      <FormControl isDisabled={!template ? true : connected} pl={2}>
+        {!template ? (
+          <HStack justifyContent="space-between" alignItems="center">
+            <FormLabel>Unknown input: {input.name}</FormLabel>
+          </HStack>
+        ) : (
+          <>
+            <HStack justifyContent="space-between" alignItems="center">
+              <HStack>
+                <Tooltip
+                  label={template?.description}
+                  placement="top"
+                  hasArrow
+                  shouldWrapChildren
+                  openDelay={HANDLE_TOOLTIP_OPEN_DELAY}
+                >
+                  <FormLabel>{template?.title}</FormLabel>
+                </Tooltip>
+              </HStack>
+              <InputFieldComponent
+                nodeId={nodeId}
+                field={input}
+                template={template}
+              />
+            </HStack>
+
+            {!['never', 'directOnly'].includes(
+              template?.inputRequirement ?? ''
+            ) && (
+              <FieldHandle
+                nodeId={nodeId}
+                field={template}
+                isValidConnection={isValidConnection}
+                handleType="target"
+              />
+            )}
+          </>
+        )}
+      </FormControl>
+    </Box>
+  );
+}
+
+interface IAINodeInputsProps {
+  nodeId: string;
+  template: InvocationTemplate;
+  inputs: Record<string, InputFieldValue>;
+}
+
+const IAINodeInputs = (props: IAINodeInputsProps) => {
+  const { nodeId, template, inputs } = props;
+
+  const edges = useAppSelector((state: RootState) => state.nodes.edges);
+
+  const renderIAINodeInputs = useCallback(() => {
+    const IAINodeInputsToRender: ReactNode[] = [];
+    const inputSockets = map(inputs);
+
+    inputSockets.forEach((inputSocket, index) => {
+      const inputTemplate = template.inputs[inputSocket.name];
+
+      const isConnected = Boolean(
+        edges.filter((connectedInput) => {
+          return (
+            connectedInput.target === nodeId &&
+            connectedInput.targetHandle === inputSocket.name
+          );
+        }).length
+      );
+
+      if (index < inputSockets.length) {
+        IAINodeInputsToRender.push(<Divider />);
+      }
+
+      IAINodeInputsToRender.push(
+        <IAINodeInput
+          key={inputSocket.id}
+          nodeId={nodeId}
+          input={inputSocket}
+          template={inputTemplate}
+          connected={isConnected}
+        />
+      );
+    });
+
+    return (
+      <Flex flexDir="column" gap={2} p={2}>
+        {IAINodeInputsToRender}
+      </Flex>
+    );
+  }, [edges, inputs, nodeId, template.inputs]);
+
+  return renderIAINodeInputs();
+};
+
+export default memo(IAINodeInputs);