Add auto layout controls to node editor

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

@@ -1125,7 +1125,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/ViewportControls.tsx

@@ -1,24 +1,71 @@
-import { ButtonGroup, IconButton } from '@invoke-ai/ui-library';
+import {
+  Button,
+  ButtonGroup,
+  CompositeSlider,
+  Divider,
+  Flex,
+  FormControl,
+  FormLabel,
+  Grid,
+  IconButton,
+  NumberDecrementStepper,
+  NumberIncrementStepper,
+  NumberInput,
+  NumberInputField,
+  NumberInputStepper,
+  Popover,
+  PopoverArrow,
+  PopoverBody,
+  PopoverContent,
+  PopoverFooter,
+  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 {
+  type LayeringStrategy,
+  layeringStrategyChanged,
+  layerSpacingChanged,
+  type LayoutDirection,
+  layoutDirectionChanged,
+  type NodeAlignment,
+  nodeAlignmentChanged,
+  nodeSpacingChanged,
+  selectLayeringStrategy,
+  selectLayerSpacing,
+  selectLayoutDirection,
+  selectNodeAlignment,
+  selectNodeSpacing,
   selectShouldShowMinimapPanel,
   shouldShowMinimapPanelChanged,
 } from 'features/nodes/store/workflowSettingsSlice';
-import { memo, useCallback } from 'react';
+import { type ChangeEvent, memo, useCallback } from 'react';
 import { useTranslation } from 'react-i18next';
 import {
   PiFrameCornersBold,
+  PiGitDiffBold,
   PiMagnifyingGlassMinusBold,
   PiMagnifyingGlassPlusBold,
   PiMapPinBold,
 } from 'react-icons/pi';
 
+const [useLayoutSettingsPopover] = buildUseBoolean(false);
+
 const ViewportControls = () => {
   const { t } = useTranslation();
   const { zoomIn, zoomOut, fitView } = useReactFlow();
+  const autoLayout = useAutoLayout();
   const dispatch = useAppDispatch();
+  const popover = useLayoutSettingsPopover();
   const shouldShowMinimapPanel = useAppSelector(selectShouldShowMinimapPanel);
+  const layeringStrategy = useAppSelector(selectLayeringStrategy);
+  const nodeSpacing = useAppSelector(selectNodeSpacing);
+  const layerSpacing = useAppSelector(selectLayerSpacing);
+  const layoutDirection = useAppSelector(selectLayoutDirection);
+  const nodeAlignment = useAppSelector(selectNodeAlignment);
 
   const handleClickedZoomIn = useCallback(() => {
     zoomIn({ duration: 300 });
@@ -36,6 +83,62 @@ const ViewportControls = () => {
     dispatch(shouldShowMinimapPanelChanged(!shouldShowMinimapPanel));
   }, [shouldShowMinimapPanel, dispatch]);
 
+  const handleLayeringStrategyChanged = useCallback(
+    (e: ChangeEvent<HTMLSelectElement>) => {
+      dispatch(layeringStrategyChanged(e.target.value as LayeringStrategy));
+    },
+    [dispatch]
+  );
+
+  const handleNodeSpacingSliderChange = useCallback(
+    (v: number) => {
+      dispatch(nodeSpacingChanged(v));
+    },
+    [dispatch]
+  );
+
+  const handleNodeSpacingInputChange = useCallback(
+    (_: string, v: number) => {
+      dispatch(nodeSpacingChanged(v));
+    },
+    [dispatch]
+  );
+
+  const handleLayerSpacingSliderChange = useCallback(
+    (v: number) => {
+      dispatch(layerSpacingChanged(v));
+    },
+    [dispatch]
+  );
+
+  const handleLayerSpacingInputChange = useCallback(
+    (_: string, v: number) => {
+      dispatch(layerSpacingChanged(v));
+    },
+    [dispatch]
+  );
+
+  const handleLayoutDirectionChanged = useCallback(
+    (e: ChangeEvent<HTMLSelectElement>) => {
+      dispatch(layoutDirectionChanged(e.target.value as LayoutDirection));
+    },
+    [dispatch]
+  );
+
+  const handleNodeAlignmentChanged = useCallback(
+    (e: ChangeEvent<HTMLSelectElement>) => {
+      const value = e.target.value as NodeAlignment;
+      dispatch(nodeAlignmentChanged(value));
+    },
+    [dispatch]
+  );
+
+  const handleApplyAutoLayout = useCallback(async () => {
+    await autoLayout();
+    fitView({ duration: 300 });
+    popover.setFalse();
+  }, [autoLayout, fitView, popover]);
+
   return (
     <ButtonGroup orientation="vertical">
       <IconButton
@@ -56,6 +159,105 @@ const ViewportControls = () => {
         onClick={handleClickedFitView}
         icon={<PiFrameCornersBold />}
       />
+      <Popover isOpen={popover.isTrue} onClose={popover.setFalse} placement="top">
+        <PopoverTrigger>
+          <IconButton
+            tooltip={t('nodes.layout.autoLayout')}
+            aria-label={t('nodes.layout.autoLayout')}
+            icon={<PiGitDiffBold />}
+            onClick={popover.toggle}
+          />
+        </PopoverTrigger>
+        <PopoverContent>
+          <PopoverArrow />
+
+          <PopoverBody>
+            <Flex direction="column" gap={2}>
+              <FormControl>
+                <FormLabel>{t('nodes.layout.layoutDirection')}</FormLabel>
+                <Select 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 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 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
+                  />
+                  <NumberInput
+                    size="sm"
+                    value={nodeSpacing}
+                    min={0}
+                    max={200}
+                    onChange={handleNodeSpacingInputChange}
+                    w={24}
+                  >
+                    <NumberInputField />
+                    <NumberInputStepper>
+                      <NumberIncrementStepper />
+                      <NumberDecrementStepper />
+                    </NumberInputStepper>
+                  </NumberInput>
+                </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
+                  />
+                  <NumberInput
+                    size="sm"
+                    value={layerSpacing}
+                    min={0}
+                    max={200}
+                    onChange={handleLayerSpacingInputChange}
+                    w={24}
+                  >
+                    <NumberInputField />
+                    <NumberInputStepper>
+                      <NumberIncrementStepper />
+                      <NumberDecrementStepper />
+                    </NumberInputStepper>
+                  </NumberInput>
+                </Grid>
+              </FormControl>
+            </Flex>
+          </PopoverBody>
+          <PopoverFooter>
+            <Button w="full" onClick={handleApplyAutoLayout}>
+              {t('common.apply')}
+            </Button>
+          </PopoverFooter>
+        </PopoverContent>
+      </Popover>
       {/* <Tooltip
         label={
           shouldShowFieldTypeLegend

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

@@ -0,0 +1,128 @@
+import { graphlib, layout } from '@dagrejs/dagre';
+import type { Edge, NodeChange } 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;
+
+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(() => {
+    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);
+      }
+      // update the Height based on the node's measured height or use a default value
+      const measuredHeight = node.measured?.height;
+      const height =
+        typeof measuredHeight === 'number'
+          ? measuredHeight
+          : isNotesNode(node)
+            ? ESTIMATED_NOTES_NODE_HEIGHT
+            : DEFAULT_NODE_HEIGHT;
+
+      g.setNode(node.id, {
+        width: node.width ?? NODE_WIDTH,
+        height: height,
+      });
+    });
+
+    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 position changes for each node based on the new layout
+    const positionChanges: NodeChange<AnyNode>[] = 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;
+};