Add auto layout controls to node editor

invokeai/frontend/web/public/locales/en.json

@@ -1133,7 +1133,23 @@
         "addItem": "Add Item",
         "generateValues": "Generate Values",
         "floatRangeGenerator": "Float Range Generator",
-        "integerRangeGenerator": "Integer Range Generator"
+        "integerRangeGenerator": "Integer Range Generator",
+        "layout": {
+            "autoLayout": "Auto Layout",
+            "layeringStrategy": "Layering Strategy",
+            "networkSimplex": "Network Simplex",
+            "longestPath": "Longest Path",
+            "nodeSpacing": "Node Spacing",
+            "layerSpacing": "Layer Spacing",
+            "layoutDirection": "Layout Direction",
+            "layoutDirectionRight": "Right",
+            "layoutDirectionDown": "Down",
+            "alignment": "Node Alignment",
+            "alignmentUL": "Top Left",
+            "alignmentDL": "Bottom Left",
+            "alignmentUR": "Top Right",
+            "alignmentDR": "Bottom Right"
+        }
     },
     "parameters": {
         "aspect": "Aspect",

invokeai/frontend/web/src/features/nodes/components/flow/panels/BottomLeftPanel/AutoLayoutPopover.tsx

@@ -0,0 +1,194 @@
+import {
+  Button,
+  CompositeNumberInput,
+  CompositeSlider,
+  Divider,
+  Flex,
+  FormControl,
+  FormLabel,
+  Grid,
+  IconButton,
+  Popover,
+  PopoverArrow,
+  PopoverBody,
+  PopoverContent,
+  PopoverTrigger,
+  Select,
+} from '@invoke-ai/ui-library';
+import { useReactFlow } from '@xyflow/react';
+import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
+import { buildUseBoolean } from 'common/hooks/useBoolean';
+import { useAutoLayout } from 'features/nodes/hooks/useAutoLayout';
+import {
+  layeringStrategyChanged,
+  layerSpacingChanged,
+  layoutDirectionChanged,
+  nodeAlignmentChanged,
+  nodeSpacingChanged,
+  selectLayeringStrategy,
+  selectLayerSpacing,
+  selectLayoutDirection,
+  selectNodeAlignment,
+  selectNodeSpacing,
+  zLayeringStrategy,
+  zLayoutDirection,
+  zNodeAlignment,
+} from 'features/nodes/store/workflowSettingsSlice';
+import { type ChangeEvent, memo, useCallback } from 'react';
+import { useTranslation } from 'react-i18next';
+import { PiMagicWandBold } from 'react-icons/pi';
+
+const [useLayoutSettingsPopover] = buildUseBoolean(false);
+
+export const AutoLayoutPopover = memo(() => {
+  const { t } = useTranslation();
+  const { fitView } = useReactFlow();
+  const autoLayout = useAutoLayout();
+  const dispatch = useAppDispatch();
+  const popover = useLayoutSettingsPopover();
+  const layeringStrategy = useAppSelector(selectLayeringStrategy);
+  const nodeSpacing = useAppSelector(selectNodeSpacing);
+  const layerSpacing = useAppSelector(selectLayerSpacing);
+  const layoutDirection = useAppSelector(selectLayoutDirection);
+  const nodeAlignment = useAppSelector(selectNodeAlignment);
+
+  const handleLayeringStrategyChanged = useCallback(
+    (e: ChangeEvent<HTMLSelectElement>) => {
+      const val = zLayeringStrategy.parse(e.target.value);
+      dispatch(layeringStrategyChanged(val));
+    },
+    [dispatch]
+  );
+
+  const handleNodeSpacingSliderChange = useCallback(
+    (v: number) => {
+      dispatch(nodeSpacingChanged(v));
+    },
+    [dispatch]
+  );
+
+  const handleNodeSpacingInputChange = useCallback(
+    (v: number) => {
+      dispatch(nodeSpacingChanged(v));
+    },
+    [dispatch]
+  );
+
+  const handleLayerSpacingSliderChange = useCallback(
+    (v: number) => {
+      dispatch(layerSpacingChanged(v));
+    },
+    [dispatch]
+  );
+
+  const handleLayerSpacingInputChange = useCallback(
+    (v: number) => {
+      dispatch(layerSpacingChanged(v));
+    },
+    [dispatch]
+  );
+
+  const handleLayoutDirectionChanged = useCallback(
+    (e: ChangeEvent<HTMLSelectElement>) => {
+      const val = zLayoutDirection.parse(e.target.value);
+      dispatch(layoutDirectionChanged(val));
+    },
+    [dispatch]
+  );
+
+  const handleNodeAlignmentChanged = useCallback(
+    (e: ChangeEvent<HTMLSelectElement>) => {
+      const val = zNodeAlignment.parse(e.target.value);
+      dispatch(nodeAlignmentChanged(val));
+    },
+    [dispatch]
+  );
+
+  const handleApplyAutoLayout = useCallback(() => {
+    autoLayout();
+    fitView({ duration: 300 });
+    popover.setFalse();
+  }, [autoLayout, fitView, popover]);
+
+  return (
+    <Popover isOpen={popover.isTrue} onClose={popover.setFalse} placement="top">
+      <PopoverTrigger>
+        <IconButton
+          tooltip={t('nodes.layout.autoLayout')}
+          aria-label={t('nodes.layout.autoLayout')}
+          icon={<PiMagicWandBold />}
+          onClick={popover.toggle}
+        />
+      </PopoverTrigger>
+      <PopoverContent>
+        <PopoverArrow />
+
+        <PopoverBody>
+          <Flex direction="column" gap={2}>
+            <FormControl>
+              <FormLabel>{t('nodes.layout.layoutDirection')}</FormLabel>
+              <Select size="sm" value={layoutDirection} onChange={handleLayoutDirectionChanged}>
+                <option value="LR">{t('nodes.layout.layoutDirectionRight')}</option>
+                <option value="TB">{t('nodes.layout.layoutDirectionDown')}</option>
+              </Select>
+            </FormControl>
+            <FormControl>
+              <FormLabel>{t('nodes.layout.layeringStrategy')}</FormLabel>
+              <Select size="sm" value={layeringStrategy} onChange={handleLayeringStrategyChanged}>
+                <option value="network-simplex">{t('nodes.layout.networkSimplex')}</option>
+                <option value="longest-path">{t('nodes.layout.longestPath')}</option>
+              </Select>
+            </FormControl>
+            <FormControl>
+              <FormLabel>{t('nodes.layout.alignment')}</FormLabel>
+              <Select size="sm" value={nodeAlignment} onChange={handleNodeAlignmentChanged}>
+                <option value="UL">{t('nodes.layout.alignmentUL')}</option>
+                <option value="DL">{t('nodes.layout.alignmentDL')}</option>
+                <option value="UR">{t('nodes.layout.alignmentUR')}</option>
+                <option value="DR">{t('nodes.layout.alignmentDR')}</option>
+              </Select>
+            </FormControl>
+            <Divider />
+            <FormControl>
+              <FormLabel>{t('nodes.layout.nodeSpacing')}</FormLabel>
+              <Grid w="full" gap={2} templateColumns="1fr auto">
+                <CompositeSlider min={0} max={200} value={nodeSpacing} onChange={handleNodeSpacingSliderChange} marks />
+                <CompositeNumberInput
+                  value={nodeSpacing}
+                  min={0}
+                  max={200}
+                  onChange={handleNodeSpacingInputChange}
+                  w={24}
+                />
+              </Grid>
+            </FormControl>
+            <FormControl>
+              <FormLabel>{t('nodes.layout.layerSpacing')}</FormLabel>
+              <Grid w="full" gap={2} templateColumns="1fr auto">
+                <CompositeSlider
+                  min={0}
+                  max={200}
+                  value={layerSpacing}
+                  onChange={handleLayerSpacingSliderChange}
+                  marks
+                />
+                <CompositeNumberInput
+                  value={layerSpacing}
+                  min={0}
+                  max={200}
+                  onChange={handleLayerSpacingInputChange}
+                  w={24}
+                />
+              </Grid>
+            </FormControl>
+            <Divider />
+            <Button w="full" onClick={handleApplyAutoLayout}>
+              {t('common.apply')}
+            </Button>
+          </Flex>
+        </PopoverBody>
+      </PopoverContent>
+    </Popover>
+  );
+});
+AutoLayoutPopover.displayName = 'AutoLayoutPopover';

invokeai/frontend/web/src/features/nodes/components/flow/panels/BottomLeftPanel/ViewportControls.tsx

@@ -14,6 +14,8 @@ import {
   PiMapPinBold,
 } from 'react-icons/pi';
 
+import { AutoLayoutPopover } from './AutoLayoutPopover';
+
 const ViewportControls = () => {
   const { t } = useTranslation();
   const { zoomIn, zoomOut, fitView } = useReactFlow();
@@ -56,20 +58,7 @@ const ViewportControls = () => {
         onClick={handleClickedFitView}
         icon={<PiFrameCornersBold />}
       />
-      {/* <Tooltip
-        label={
-          shouldShowFieldTypeLegend
-            ? t('nodes.hideLegendNodes')
-            : t('nodes.showLegendNodes')
-        }
-      >
-        <IconButton
-          aria-label="Toggle field type legend"
-          isChecked={shouldShowFieldTypeLegend}
-          onClick={handleClickedToggleFieldTypeLegend}
-          icon={<FaInfo />}
-        />
-      </Tooltip> */}
+      <AutoLayoutPopover />
       <IconButton
         tooltip={shouldShowMinimapPanel ? t('nodes.hideMinimapnodes') : t('nodes.showMinimapnodes')}
         aria-label={shouldShowMinimapPanel ? t('nodes.hideMinimapnodes') : t('nodes.showMinimapnodes')}

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

@@ -0,0 +1,138 @@
+import { graphlib, layout } from '@dagrejs/dagre';
+import type { Edge, NodePositionChange } from '@xyflow/react';
+import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
+import { nodesChanged } from 'features/nodes/store/nodesSlice';
+import { selectEdges, selectNodes } from 'features/nodes/store/selectors';
+import {
+  selectLayeringStrategy,
+  selectLayerSpacing,
+  selectLayoutDirection,
+  selectNodeAlignment,
+  selectNodeSpacing,
+} from 'features/nodes/store/workflowSettingsSlice';
+import { NODE_WIDTH } from 'features/nodes/types/constants';
+import type { AnyNode } from 'features/nodes/types/invocation';
+import { isNotesNode } from 'features/nodes/types/invocation';
+import { useCallback } from 'react';
+
+const ESTIMATED_NOTES_NODE_HEIGHT = 200;
+const DEFAULT_NODE_HEIGHT = NODE_WIDTH;
+
+const getNodeHeight = (node: AnyNode): number => {
+  if (node.measured?.height) {
+    return node.measured.height;
+  }
+  if (isNotesNode(node)) {
+    return ESTIMATED_NOTES_NODE_HEIGHT;
+  }
+  return DEFAULT_NODE_HEIGHT;
+};
+
+const getNodeWidth = (node: AnyNode): number => {
+  if (node.measured?.width) {
+    return node.measured.width;
+  }
+  return NODE_WIDTH;
+};
+
+export const useAutoLayout = (): (() => void) => {
+  const dispatch = useAppDispatch();
+  const nodes = useAppSelector(selectNodes);
+  const edges = useAppSelector(selectEdges);
+  const nodeSpacing = useAppSelector(selectNodeSpacing);
+  const layerSpacing = useAppSelector(selectLayerSpacing);
+  const layeringStrategy = useAppSelector(selectLayeringStrategy);
+  const layoutDirection = useAppSelector(selectLayoutDirection);
+  const nodeAlignment = useAppSelector(selectNodeAlignment);
+
+  const autoLayout = useCallback(() => {
+    // We'll do graph layout using dagre, then convert the results to reactflow position changes
+    const g = new graphlib.Graph();
+
+    g.setGraph({
+      rankdir: layoutDirection,
+      nodesep: nodeSpacing,
+      ranksep: layerSpacing,
+      ranker: layeringStrategy,
+      align: nodeAlignment,
+    });
+
+    g.setDefaultEdgeLabel(() => ({}));
+
+    const selectedNodes = nodes.filter((n) => n.selected);
+    const isLayoutSelection = selectedNodes.length > 1 && nodes.length > selectedNodes.length;
+    const nodesToLayout = isLayoutSelection ? selectedNodes : nodes;
+
+    //Anchor of the selected nodes
+    const selectionAnchor = {
+      minX: Infinity,
+      minY: Infinity,
+    };
+
+    nodesToLayout.forEach((node) => {
+      // If we're laying out a selection, adjust the anchor to the top-left of the selection
+      if (isLayoutSelection) {
+        selectionAnchor.minX = Math.min(selectionAnchor.minX, node.position.x);
+        selectionAnchor.minY = Math.min(selectionAnchor.minY, node.position.y);
+      }
+
+      g.setNode(node.id, {
+        width: getNodeWidth(node),
+        height: getNodeHeight(node),
+      });
+    });
+
+    let edgesToLayout: Edge[] = edges;
+    if (isLayoutSelection) {
+      const nodesToLayoutIds = new Set(nodesToLayout.map((n) => n.id));
+      edgesToLayout = edges.filter((edge) => nodesToLayoutIds.has(edge.source) && nodesToLayoutIds.has(edge.target));
+    }
+
+    edgesToLayout.forEach((edge) => {
+      g.setEdge(edge.source, edge.target);
+    });
+
+    layout(g);
+
+    // anchor for the new layout
+    const layoutAnchor = {
+      minX: Infinity,
+      minY: Infinity,
+    };
+    let offsetX = 0;
+    let offsetY = 0;
+
+    if (isLayoutSelection) {
+      // Get the top-left position of the new layout
+      nodesToLayout.forEach((node) => {
+        const nodeInfo = g.node(node.id);
+        // Convert from center to top-left
+        const topLeftX = nodeInfo.x - nodeInfo.width / 2;
+        const topLeftY = nodeInfo.y - nodeInfo.height / 2;
+        // Use the top-left coordinates to find the bounding box
+        layoutAnchor.minX = Math.min(layoutAnchor.minX, topLeftX);
+        layoutAnchor.minY = Math.min(layoutAnchor.minY, topLeftY);
+      });
+      // Calculate the offset needed to move the new layout to the original position
+      offsetX = selectionAnchor.minX - layoutAnchor.minX;
+      offsetY = selectionAnchor.minY - layoutAnchor.minY;
+    }
+
+    // Create reactflow position changes for each node based on the new layout
+    const positionChanges: NodePositionChange[] = nodesToLayout.map((node) => {
+      const nodeInfo = g.node(node.id);
+      // Convert from center-based position to top-left-based position
+      const x = nodeInfo.x - nodeInfo.width / 2;
+      const y = nodeInfo.y - nodeInfo.height / 2;
+      const newPosition = {
+        x: isLayoutSelection ? x + offsetX : x,
+        y: isLayoutSelection ? y + offsetY : y,
+      };
+      return { id: node.id, type: 'position', position: newPosition };
+    });
+
+    dispatch(nodesChanged(positionChanges));
+  }, [dispatch, edges, nodes, nodeSpacing, layerSpacing, layeringStrategy, layoutDirection, nodeAlignment]);
+
+  return autoLayout;
+};