feat(ui): iterate on builder (WIP)

invokeai/frontend/web/src/features/nodes/components/sidePanel/builder/ContainerContext.ts

@@ -0,0 +1,8 @@
+import { createContext, useContext } from 'react';
+
+export const ContainerDirectionContext = createContext<'row' | 'column' | null>(null);
+
+export const useContainerDirectionContext = () => {
+  const containerDirection = useContext(ContainerDirectionContext);
+  return containerDirection;
+};

invokeai/frontend/web/src/features/nodes/components/sidePanel/builder/ContainerElementComponent.tsx

@@ -1,55 +1,95 @@
-import { Flex, Grid, GridItem } from '@invoke-ai/ui-library';
+import { ContainerDirectionContext } from 'features/nodes/components/sidePanel/builder/ContainerContext';
 import { DividerElementComponent } from 'features/nodes/components/sidePanel/builder/DividerElementComponent';
+import { ElementWrapper } from 'features/nodes/components/sidePanel/builder/ElementWrapper';
 import { HeadingElementComponent } from 'features/nodes/components/sidePanel/builder/HeadingElementComponent';
 import { NodeFieldElementComponent } from 'features/nodes/components/sidePanel/builder/NodeFieldElementComponent';
 import { TextElementComponent } from 'features/nodes/components/sidePanel/builder/TextElementComponent';
-import type { ContainerElement, FormElement } from 'features/nodes/types/workflow';
-import { Fragment, memo } from 'react';
+import { useElement } from 'features/nodes/types/workflow';
+import { memo } from 'react';
 import type { Equals } from 'tsafe';
 import { assert } from 'tsafe';
 
 const getGridTemplateColumns = (count: number) => {
-  return Array.from({ length: count }, () => '1fr').join(' auto ');
+  return Array.from({ length: count }, () => '1fr').join(' ');
+};
+const fill = (count: number, val: string, last?: string) => {
+  return Array.from({ length: count }, (_, i) => {
+    if (last && i === count - 1) {
+      return last;
+    }
+    return val;
+  }).join(' ');
 };
 
-export const ContainerElementComponent = memo(({ element }: { element: ContainerElement }) => {
-  const { id, data } = element;
-  const { columns } = data;
+export const ContainerElementComponent = memo(({ id }: { id: string }) => {
+  const element = useElement(id);
+
+  if (!element || element.type !== 'container') {
+    return null;
+  }
+
+  const { children, direction } = element.data;
 
   return (
-    <Grid id={id} gap={4} gridTemplateColumns={getGridTemplateColumns(columns.length)} gridAutoFlow="column">
-      {columns.map((elements, columnIndex) => {
-        const key = `${element.id}_${columnIndex}`;
-        const withDivider = columnIndex < columns.length - 1;
-        return (
-          <Fragment key={key}>
-            <GridItem as={Grid} id={key} gap={4} gridAutoRows="min-content" gridAutoFlow="row">
-              {elements.map((element) => (
-                <FormElementComponent key={element.id} element={element} />
-              ))}
-            </GridItem>
-            {withDivider && <Flex w="1px" bg="base.800" flexShrink={0} />}
-          </Fragment>
-        );
-      })}
-    </Grid>
+    <ContainerDirectionContext.Provider value={direction}>
+      <ElementWrapper id={id} gap={4} flexDir={direction}>
+        {children.map((childId) => (
+          <FormElementComponent key={childId} id={childId} />
+        ))}
+      </ElementWrapper>
+    </ContainerDirectionContext.Provider>
   );
 });
 ContainerElementComponent.displayName = 'ContainerElementComponent';
+// export const ContainerElementComponent = memo(({ id }: { id: string }) => {
+//   const element = useElement(id);
+
+//   if (!element || element.type !== 'container') {
+//     return null;
+//   }
+
+//   const { children, direction } = element.data;
+
+//   return (
+//     <GridItem
+//       as={Grid}
+//       id={id}
+//       gap={4}
+//       w="full"
+//       h="full"
+//       gridTemplateColumns={direction === 'row' ? fill(children.length, '1fr') : undefined}
+//       gridTemplateRows={direction === 'column' ? fill(children.length, 'min-content', '1fr') : undefined}
+//       gridAutoFlow={direction === 'column' ? 'row' : 'column'}
+//       alignItems="baseline"
+//     >
+//       {children.map((childId) => (
+//         <FormElementComponent key={childId} id={childId} />
+//       ))}
+//     </GridItem>
+//   );
+// });
+// ContainerElementComponent.displayName = 'ContainerElementComponent';
+
+export const FormElementComponent = memo(({ id }: { id: string }) => {
+  const element = useElement(id);
+
+  if (!element) {
+    return null;
+  }
+
+  const { type } = element;
 
-export const FormElementComponent = memo(({ element }: { element: FormElement }) => {
-  const { type, id } = element;
   switch (type) {
     case 'container':
-      return <ContainerElementComponent key={id} element={element} />;
+      return <ContainerElementComponent key={id} id={id} />;
     case 'node-field':
-      return <NodeFieldElementComponent key={id} element={element} />;
+      return <NodeFieldElementComponent key={id} id={id} />;
     case 'divider':
-      return <DividerElementComponent key={id} element={element} />;
+      return <DividerElementComponent key={id} id={id} />;
     case 'heading':
-      return <HeadingElementComponent key={id} element={element} />;
+      return <HeadingElementComponent key={id} id={id} />;
     case 'text':
-      return <TextElementComponent key={id} element={element} />;
+      return <TextElementComponent key={id} id={id} />;
     default:
       assert<Equals<typeof type, never>>(false, `Unhandled type ${type}`);
   }

invokeai/frontend/web/src/features/nodes/components/sidePanel/builder/DividerElementComponent.tsx

@@ -1,11 +1,20 @@
 import { Flex } from '@invoke-ai/ui-library';
-import type { DividerElement } from 'features/nodes/types/workflow';
+import { useContainerDirectionContext } from 'features/nodes/components/sidePanel/builder/ContainerContext';
 import { memo } from 'react';
 
-export const DividerElementComponent = memo(({ element }: { element: DividerElement }) => {
-  const { id } = element;
+export const DividerElementComponent = memo(({ id }: { id: string }) => {
+  const containerDirection = useContainerDirectionContext();
 
-  return <Flex id={id} h="1px" bg="base.800" flexShrink={0} />;
+  return (
+    <Flex
+      flex="0 0 auto"
+      id={id}
+      h={containerDirection === 'column' ? '1px' : undefined}
+      w={containerDirection === 'column' ? undefined : '1px'}
+      bg="base.700"
+      flexShrink={0}
+    />
+  );
 });
 
 DividerElementComponent.displayName = 'DividerElementComponent';

invokeai/frontend/web/src/features/nodes/components/sidePanel/builder/ElementWrapper.tsx

@@ -0,0 +1,12 @@
+import type { FlexProps } from '@invoke-ai/ui-library';
+import { Flex } from '@invoke-ai/ui-library';
+import { useContainerDirectionContext } from 'features/nodes/components/sidePanel/builder/ContainerContext';
+import type { PropsWithChildren } from 'react';
+import { memo } from 'react';
+
+export const ElementWrapper = memo((props: PropsWithChildren<FlexProps>) => {
+  const containerDirection = useContainerDirectionContext();
+  return <Flex flex={containerDirection === 'column' ? '1 1 0' : undefined} {...props} />;
+});
+
+ElementWrapper.displayName = 'ElementWrapper';

invokeai/frontend/web/src/features/nodes/components/sidePanel/builder/HeadingElementComponent.tsx

@@ -1,5 +1,6 @@
 import { Heading } from '@invoke-ai/ui-library';
-import type { HeadingElement } from 'features/nodes/types/workflow';
+import { ElementWrapper } from 'features/nodes/components/sidePanel/builder/ElementWrapper';
+import { useElement } from 'features/nodes/types/workflow';
 import { memo } from 'react';
 
 const LEVEL_TO_SIZE = {
@@ -10,14 +11,19 @@ const LEVEL_TO_SIZE = {
   5: 'xs',
 } as const;
 
-export const HeadingElementComponent = memo(({ element }: { element: HeadingElement }) => {
-  const { id, data } = element;
+export const HeadingElementComponent = memo(({ id }: { id: string }) => {
+  const element = useElement(id);
+
+  if (!element || element.type !== 'heading') {
+    return null;
+  }
+  const { data } = element;
   const { content, level } = data;
 
   return (
-    <Heading id={id} size={LEVEL_TO_SIZE[level]}>
-      {content}
-    </Heading>
+    <ElementWrapper id={id}>
+      <Heading size={LEVEL_TO_SIZE[level]}>{content}</Heading>
+    </ElementWrapper>
   );
 });
 

invokeai/frontend/web/src/features/nodes/components/sidePanel/builder/NodeFieldElementComponent.tsx

@@ -1,19 +1,24 @@
-import { Flex } from '@invoke-ai/ui-library';
 import { InputFieldGate } from 'features/nodes/components/flow/nodes/Invocation/fields/InputFieldGate';
 import { InputFieldViewMode } from 'features/nodes/components/flow/nodes/Invocation/fields/InputFieldViewMode';
-import type { NodeFieldElement } from 'features/nodes/types/workflow';
+import { ElementWrapper } from 'features/nodes/components/sidePanel/builder/ElementWrapper';
+import { useElement } from 'features/nodes/types/workflow';
 import { memo } from 'react';
 
-export const NodeFieldElementComponent = memo(({ element }: { element: NodeFieldElement }) => {
-  const { id, data } = element;
+export const NodeFieldElementComponent = memo(({ id }: { id: string }) => {
+  const element = useElement(id);
+
+  if (!element || element.type !== 'node-field') {
+    return null;
+  }
+  const { data } = element;
   const { fieldIdentifier } = data;
 
   return (
-    <Flex id={id} flexBasis="100%">
+    <ElementWrapper id={id}>
       <InputFieldGate nodeId={fieldIdentifier.nodeId} fieldName={fieldIdentifier.fieldName}>
         <InputFieldViewMode nodeId={fieldIdentifier.nodeId} fieldName={fieldIdentifier.fieldName} />
       </InputFieldGate>
-    </Flex>
+    </ElementWrapper>
   );
 });
 

invokeai/frontend/web/src/features/nodes/components/sidePanel/builder/TextElementComponent.tsx

@@ -1,15 +1,21 @@
 import { Text } from '@invoke-ai/ui-library';
-import type { TextElement } from 'features/nodes/types/workflow';
+import { ElementWrapper } from 'features/nodes/components/sidePanel/builder/ElementWrapper';
+import { useElement } from 'features/nodes/types/workflow';
 import { memo } from 'react';
 
-export const TextElementComponent = memo(({ element }: { element: TextElement }) => {
-  const { id, data } = element;
+export const TextElementComponent = memo(({ id }: { id: string }) => {
+  const element = useElement(id);
+
+  if (!element || element.type !== 'text') {
+    return null;
+  }
+  const { data } = element;
   const { content, fontSize } = data;
 
   return (
-    <Text id={id} fontSize={fontSize}>
-      {content}
-    </Text>
+    <ElementWrapper id={id}>
+      <Text fontSize={fontSize}>{content}</Text>
+    </ElementWrapper>
   );
 });
 

invokeai/frontend/web/src/features/nodes/components/sidePanel/builder/WorkflowBuilder.tsx

@@ -1,7 +1,7 @@
 import { Flex } from '@invoke-ai/ui-library';
 import ScrollableContent from 'common/components/OverlayScrollbars/ScrollableContent';
 import { FormElementComponent } from 'features/nodes/components/sidePanel/builder/ContainerElementComponent';
-import { data } from 'features/nodes/types/workflow';
+import { rootId } from 'features/nodes/types/workflow';
 import { memo } from 'react';
 
 export const WorkflowBuilder = memo(() => {
@@ -9,7 +9,7 @@ export const WorkflowBuilder = memo(() => {
     <ScrollableContent>
       <Flex w="full" h="full" justifyContent="center">
         <Flex w="full" h="full" maxW={512}>
-          <FormElementComponent element={data} />
+          <FormElementComponent id={rootId} />
         </Flex>
       </Flex>
     </ScrollableContent>

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

@@ -28,6 +28,17 @@ const blankWorkflow: Omit<WorkflowV3, 'nodes' | 'edges'> = {
   exposedFields: [],
   meta: { version: '3.0.0', category: 'user' },
   id: undefined,
+  form: {
+    elements: {},
+    structure: {
+      id: 'root',
+      type: 'container',
+      data: {
+        direction: 'column',
+        children: [],
+      },
+    },
+  },
 };
 
 const initialWorkflowState: WorkflowState = {

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

@@ -1,4 +1,6 @@
-import { nanoid } from 'nanoid';
+import { SyncableMap } from 'common/util/SyncableMap/SyncableMap';
+import { getPrefixedId } from 'features/controlLayers/konva/util';
+import { useMemo, useSyncExternalStore } from 'react';
 import { z } from 'zod';
 
 import { zFieldIdentifier } from './field';
@@ -73,18 +75,46 @@ export const zWorkflowV3 = z.object({
     category: zWorkflowCategory.default('user'),
     version: z.literal('3.0.0'),
   }),
+  form: z.object({
+    elements: z.record(z.lazy(() => zFormElement)),
+    structure: z.lazy(() => zContainerElement),
+  }),
 });
 export type WorkflowV3 = z.infer<typeof zWorkflowV3>;
 // #endregion
 
 // #region Workflow Builder
 
+export const elements = new SyncableMap<string, FormElement>();
+
+export const addElement = (element: FormElement) => {
+  elements.set(element.id, element);
+};
+
+export const removeElement = (id: ElementId) => {
+  return elements.delete(id);
+};
+
+export const getElement = (id: ElementId) => {
+  return elements.get(id);
+};
+
+export const useElement = <T extends FormElement>(id: string) => {
+  const map = useSyncExternalStore(elements.subscribe, elements.getSnapshot);
+  const element = useMemo(() => map.get(id), [id, map]);
+  return element as T | undefined;
+};
+
+const zElementId = z.string().trim().min(1);
+type ElementId = z.infer<typeof zElementId>;
+
 const zElementBase = z.object({
-  id: z.string().trim().min(1),
+  id: zElementId,
 });
 
+const NODE_FIELD_TYPE = 'node-field';
 const zNodeFieldElement = zElementBase.extend({
-  type: z.literal('node-field'),
+  type: z.literal(NODE_FIELD_TYPE),
   data: z.object({
     fieldIdentifier: zFieldIdentifier,
   }),
@@ -93,16 +123,21 @@ export type NodeFieldElement = z.infer<typeof zNodeFieldElement>;
 const nodeField = (
   nodeId: NodeFieldElement['data']['fieldIdentifier']['nodeId'],
   fieldName: NodeFieldElement['data']['fieldIdentifier']['fieldName']
-): NodeFieldElement => ({
-  id: nanoid(),
-  type: 'node-field',
-  data: {
-    fieldIdentifier: { nodeId, fieldName },
-  },
-});
+): NodeFieldElement => {
+  const element: NodeFieldElement = {
+    id: getPrefixedId(NODE_FIELD_TYPE),
+    type: NODE_FIELD_TYPE,
+    data: {
+      fieldIdentifier: { nodeId, fieldName },
+    },
+  };
+  addElement(element);
+  return element;
+};
 
+const HEADING_TYPE = 'heading';
 const zHeadingElement = zElementBase.extend({
-  type: z.literal('heading'),
+  type: z.literal(HEADING_TYPE),
   data: z.object({
     content: z.string(),
     level: z.union([z.literal(1), z.literal(2), z.literal(3), z.literal(4), z.literal(5)]),
@@ -112,117 +147,140 @@ export type HeadingElement = z.infer<typeof zHeadingElement>;
 const heading = (
   content: HeadingElement['data']['content'],
   level: HeadingElement['data']['level']
-): HeadingElement => ({
-  id: nanoid(),
-  type: 'heading',
-  data: {
-    content,
-    level,
-  },
-});
+): HeadingElement => {
+  const element: HeadingElement = {
+    id: getPrefixedId(HEADING_TYPE),
+    type: HEADING_TYPE,
+    data: {
+      content,
+      level,
+    },
+  };
+  addElement(element);
+  return element;
+};
 
+const TEXT_TYPE = 'text';
 const zTextElement = zElementBase.extend({
-  type: z.literal('text'),
+  type: z.literal(TEXT_TYPE),
   data: z.object({
     content: z.string(),
     fontSize: z.enum(['sm', 'md', 'lg']),
   }),
 });
 export type TextElement = z.infer<typeof zTextElement>;
-const text = (content: TextElement['data']['content'], fontSize: TextElement['data']['fontSize']): TextElement => ({
-  id: nanoid(),
-  type: 'text',
-  data: {
-    content,
-    fontSize,
-  },
-});
+const text = (content: TextElement['data']['content'], fontSize: TextElement['data']['fontSize']): TextElement => {
+  const element: TextElement = {
+    id: getPrefixedId(TEXT_TYPE),
+    type: TEXT_TYPE,
+    data: {
+      content,
+      fontSize,
+    },
+  };
+  addElement(element);
+  return element;
+};
 
+const DIVIDER_TYPE = 'divider';
 const zDividerElement = zElementBase.extend({
-  type: z.literal('divider'),
+  type: z.literal(DIVIDER_TYPE),
 });
 export type DividerElement = z.infer<typeof zDividerElement>;
-const divider = (): DividerElement => ({
-  id: nanoid(),
-  type: 'divider',
-});
+const divider = (): DividerElement => {
+  const element: DividerElement = {
+    id: getPrefixedId(DIVIDER_TYPE),
+    type: DIVIDER_TYPE,
+  };
+  addElement(element);
+  return element;
+};
 
 export type ContainerElement = {
   id: string;
-  type: 'container';
+  type: typeof CONTAINER_TYPE;
   data: {
-    columns: FormElement[][];
+    direction: 'row' | 'column';
+    children: ElementId[];
   };
 };
 
+const CONTAINER_TYPE = 'container';
 const zContainerElement: z.ZodType<ContainerElement> = zElementBase.extend({
-  type: z.literal('container'),
+  type: z.literal(CONTAINER_TYPE),
   data: z.object({
-    columns: z.lazy(() => z.array(z.array(zFormElement))),
+    direction: z.enum(['row', 'column']),
+    children: z.array(zElementId),
   }),
 });
-const container = (columns: ContainerElement['data']['columns']): ContainerElement => ({
-  id: nanoid(),
-  type: 'container',
-  data: {
-    columns,
-  },
-});
+const container = (
+  direction: ContainerElement['data']['direction'],
+  children: ContainerElement['data']['children']
+): ContainerElement => {
+  const element: ContainerElement = {
+    id: getPrefixedId(CONTAINER_TYPE),
+    type: CONTAINER_TYPE,
+    data: {
+      direction,
+      children,
+    },
+  };
+  addElement(element);
+  return element;
+};
 
 const zFormElement = z.union([zContainerElement, zNodeFieldElement, zHeadingElement, zTextElement, zDividerElement]);
 
 export type FormElement = z.infer<typeof zFormElement>;
 
-export const data: ContainerElement = container([
-  [
-    heading('My Cool Workflow', 1),
-    text('This is a description of what my workflow does. It does things.', 'md'),
-    divider(),
-    heading('First Section', 2),
-    text('The first section includes fields relevant to the first section. This note describes that fact.', 'sm'),
-    container([
-      [nodeField('7aed1a5f-7fd7-4184-abe8-ddea0ea5e706', 'image')],
-      [nodeField('7aed1a5f-7fd7-4184-abe8-ddea0ea5e706', 'image')],
-      [nodeField('7aed1a5f-7fd7-4184-abe8-ddea0ea5e706', 'image')],
-    ]),
-    nodeField('9c058600-8d73-4702-912b-0ccf37403bfd', 'value'),
-    nodeField('7a8bbab2-6919-4cfc-bd7c-bcfda3c79ecf', 'value'),
-    nodeField('4e16cbf6-457c-46fb-9ab7-9cb262fa1e03', 'value'),
-    nodeField('39cb5272-a9d7-4da9-9c35-32e02b46bb34', 'color'),
-    container([
-      [
-        nodeField('4f609a81-0e25-47d1-ba0d-f24fedd5273f', 'value'),
-        nodeField('4f609a81-0e25-47d1-ba0d-f24fedd5273f', 'value'),
-      ],
-      [
-        nodeField('4f609a81-0e25-47d1-ba0d-f24fedd5273f', 'value'),
-        nodeField('4f609a81-0e25-47d1-ba0d-f24fedd5273f', 'value'),
-        nodeField('4f609a81-0e25-47d1-ba0d-f24fedd5273f', 'value'),
-        nodeField('4f609a81-0e25-47d1-ba0d-f24fedd5273f', 'value'),
-      ],
-      [
-        container([
-          [
-            nodeField('4f609a81-0e25-47d1-ba0d-f24fedd5273f', 'value'),
-            nodeField('4f609a81-0e25-47d1-ba0d-f24fedd5273f', 'value'),
-          ],
-          [
-            nodeField('4f609a81-0e25-47d1-ba0d-f24fedd5273f', 'value'),
-            nodeField('4f609a81-0e25-47d1-ba0d-f24fedd5273f', 'value'),
-          ],
-        ]),
-      ],
-    ]),
-    nodeField('14744f68-9000-4694-b4d6-cbe83ee231ee', 'model'),
-    divider(),
-    text('These are some text that are definitely super helpful.', 'sm'),
-    divider(),
-    container([
-      [
-        nodeField('7aed1a5f-7fd7-4184-abe8-ddea0ea5e706', 'image'),
-        nodeField('7aed1a5f-7fd7-4184-abe8-ddea0ea5e706', 'image'),
-      ],
-      [nodeField('7a8bbab2-6919-4cfc-bd7c-bcfda3c79ecf', 'value')],
-    ]),
-  ],
-]);
+export const rootId: string = container('column', [
+  heading('My Cool Workflow', 1).id,
+  text('This is a description of what my workflow does. It does things.', 'md').id,
+  divider().id,
+  heading('First Section', 2).id,
+  text('The first section includes fields relevant to the first section. This note describes that fact.', 'sm').id,
+  container('row', [
+    nodeField('7aed1a5f-7fd7-4184-abe8-ddea0ea5e706', 'image').id,
+    divider().id,
+    nodeField('7aed1a5f-7fd7-4184-abe8-ddea0ea5e706', 'image').id,
+    divider().id,
+    nodeField('7aed1a5f-7fd7-4184-abe8-ddea0ea5e706', 'image').id,
+  ]).id,
+  nodeField('9c058600-8d73-4702-912b-0ccf37403bfd', 'value').id,
+  nodeField('7a8bbab2-6919-4cfc-bd7c-bcfda3c79ecf', 'value').id,
+  nodeField('4e16cbf6-457c-46fb-9ab7-9cb262fa1e03', 'value').id,
+  nodeField('39cb5272-a9d7-4da9-9c35-32e02b46bb34', 'color').id,
+  container('row', [
+    container('column', [
+      nodeField('4f609a81-0e25-47d1-ba0d-f24fedd5273f', 'value').id,
+      nodeField('4f609a81-0e25-47d1-ba0d-f24fedd5273f', 'value').id,
+    ]).id,
+    container('column', [
+      nodeField('4f609a81-0e25-47d1-ba0d-f24fedd5273f', 'value').id,
+      nodeField('4f609a81-0e25-47d1-ba0d-f24fedd5273f', 'value').id,
+      nodeField('4f609a81-0e25-47d1-ba0d-f24fedd5273f', 'value').id,
+      nodeField('4f609a81-0e25-47d1-ba0d-f24fedd5273f', 'value').id,
+    ]).id,
+    container('column', [
+      container('row', [
+        nodeField('4f609a81-0e25-47d1-ba0d-f24fedd5273f', 'value').id,
+        nodeField('4f609a81-0e25-47d1-ba0d-f24fedd5273f', 'value').id,
+      ]).id,
+      container('row', [
+        nodeField('4f609a81-0e25-47d1-ba0d-f24fedd5273f', 'value').id,
+        nodeField('4f609a81-0e25-47d1-ba0d-f24fedd5273f', 'value').id,
+      ]).id,
+    ]).id,
+  ]).id,
+  nodeField('14744f68-9000-4694-b4d6-cbe83ee231ee', 'model').id,
+  divider().id,
+  text('These are some text that are definitely super helpful.', 'sm').id,
+  divider().id,
+  container('row', [
+    container('column', [
+      nodeField('7aed1a5f-7fd7-4184-abe8-ddea0ea5e706', 'image').id,
+      nodeField('7aed1a5f-7fd7-4184-abe8-ddea0ea5e706', 'image').id,
+    ]).id,
+    container('column', [nodeField('7a8bbab2-6919-4cfc-bd7c-bcfda3c79ecf', 'value').id]).id,
+  ]).id,
+]).id;