feat(ui): undo/redo in regional prompts

using the `redux-undo` library

invokeai/frontend/web/package.json

@@ -95,6 +95,7 @@
     "reactflow": "^11.10.4",
     "redux-dynamic-middlewares": "^2.2.0",
     "redux-remember": "^5.1.0",
+    "redux-undo": "^1.1.0",
     "rfdc": "^1.3.1",
     "roarr": "^7.21.1",
     "serialize-error": "^11.0.3",

invokeai/frontend/web/pnpm-lock.yaml

@@ -140,6 +140,9 @@ dependencies:
   redux-remember:
     specifier: ^5.1.0
     version: 5.1.0(redux@5.0.1)
+  redux-undo:
+    specifier: ^1.1.0
+    version: 1.1.0
   rfdc:
     specifier: ^1.3.1
     version: 1.3.1
@@ -11962,6 +11965,10 @@ packages:
       redux: 5.0.1
     dev: false
 
+  /redux-undo@1.1.0:
+    resolution: {integrity: sha512-zzLFh2qeF0MTIlzDhDLm9NtkfBqCllQJ3OCuIl5RKlG/ayHw6GUdIFdMhzMS9NnrnWdBX5u//ExMOHpfudGGOg==}
+    dev: false
+
   /redux@5.0.1:
     resolution: {integrity: sha512-M9/ELqF6fy8FwmkpnF0S3YKOqMyoWJ4+CS5Efg2ct3oY9daQvd/Pc71FpGZsVsbl3Cpb+IIcjBDUnnyBdQbq4w==}
     dev: false

invokeai/frontend/web/src/app/store/store.ts

@@ -24,6 +24,7 @@ import { queueSlice } from 'features/queue/store/queueSlice';
 import {
   regionalPromptsPersistConfig,
   regionalPromptsSlice,
+  regionalPromptsUndoableConfig,
 } from 'features/regionalPrompts/store/regionalPromptsSlice';
 import { sdxlPersistConfig, sdxlSlice } from 'features/sdxl/store/sdxlSlice';
 import { configSlice } from 'features/system/store/configSlice';
@@ -34,6 +35,7 @@ import { defaultsDeep, keys, omit, pick } from 'lodash-es';
 import dynamicMiddlewares from 'redux-dynamic-middlewares';
 import type { SerializeFunction, UnserializeFunction } from 'redux-remember';
 import { rememberEnhancer, rememberReducer } from 'redux-remember';
+import undoable from 'redux-undo';
 import { serializeError } from 'serialize-error';
 import { api } from 'services/api';
 import { authToastMiddleware } from 'services/api/authToastMiddleware';
@@ -63,7 +65,7 @@ const allReducers = {
   [queueSlice.name]: queueSlice.reducer,
   [workflowSlice.name]: workflowSlice.reducer,
   [hrfSlice.name]: hrfSlice.reducer,
-  [regionalPromptsSlice.name]: regionalPromptsSlice.reducer,
+  [regionalPromptsSlice.name]: undoable(regionalPromptsSlice.reducer, regionalPromptsUndoableConfig),
   [api.reducerPath]: api.reducer,
 };
 
@@ -120,6 +122,7 @@ const unserialize: UnserializeFunction = (data, key) => {
   try {
     const { initialState, migrate } = persistConfig;
     const parsed = JSON.parse(data);
+
     // strip out old keys
     const stripped = pick(parsed, keys(initialState));
     // run (additive) migrations
@@ -147,7 +150,8 @@ const serialize: SerializeFunction = (data, key) => {
   if (!persistConfig) {
     throw new Error(`No persist config for slice "${key}"`);
   }
-  const result = omit(data, persistConfig.persistDenylist);
+  const isUndoable = 'present' in data && 'past' in data && 'future' in data && '_latestUnfiltered' in data;
+  const result = omit(isUndoable ? data.present : data, persistConfig.persistDenylist);
   return JSON.stringify(result);
 };
 

invokeai/frontend/web/src/features/nodes/util/graph/addRegionalPromptsToGraph.ts

@@ -22,8 +22,8 @@ export const addRegionalPromptsToGraph = async (state: RootState, graph: NonNull
   const { dispatch } = getStore();
   // TODO: Handle non-SDXL
   // const isSDXL = state.generation.model?.base === 'sdxl';
-  const { autoNegative } = state.regionalPrompts;
+  const { autoNegative } = state.regionalPrompts.present;
-  const layers = state.regionalPrompts.layers
+  const layers = state.regionalPrompts.present.layers
     .filter((l) => l.kind === 'promptRegionLayer') // We only want the prompt region layers
     .filter((l) => l.isVisible); // Only visible layers are rendered on the canvas
 

invokeai/frontend/web/src/features/regionalPrompts/components/BrushSize.tsx

@@ -7,7 +7,7 @@ import { useTranslation } from 'react-i18next';
 export const BrushSize = memo(() => {
   const dispatch = useAppDispatch();
   const { t } = useTranslation();
-  const brushSize = useAppSelector((s) => s.regionalPrompts.brushSize);
+  const brushSize = useAppSelector((s) => s.regionalPrompts.present.brushSize);
   const onChange = useCallback(
     (v: number) => {
       dispatch(brushSizeChanged(v));

invokeai/frontend/web/src/features/regionalPrompts/components/LayerColorPicker.tsx

@@ -19,7 +19,7 @@ export const LayerColorPicker = memo(({ id }: Props) => {
   const selectColor = useMemo(
     () =>
       createMemoizedSelector(selectRegionalPromptsSlice, (regionalPrompts) => {
-        const layer = regionalPrompts.layers.find((l) => l.id === id);
+        const layer = regionalPrompts.present.layers.find((l) => l.id === id);
         assert(layer);
         return layer.color;
       }),

invokeai/frontend/web/src/features/regionalPrompts/components/LayerListItem.tsx

@@ -17,9 +17,9 @@ type Props = {
 export const LayerListItem = memo(({ id }: Props) => {
   const dispatch = useAppDispatch();
   const { t } = useTranslation();
-  const selectedLayer = useAppSelector((s) => s.regionalPrompts.selectedLayer);
+  const selectedLayer = useAppSelector((s) => s.regionalPrompts.present.selectedLayer);
   const color = useAppSelector((s) => {
-    const color = s.regionalPrompts.layers.find((l) => l.id === id)?.color;
+    const color = s.regionalPrompts.present.layers.find((l) => l.id === id)?.color;
     if (color) {
       return rgbaColorToString({ ...color, a: selectedLayer === id ? 1 : 0.35 });
     }

invokeai/frontend/web/src/features/regionalPrompts/components/LayerMenu.tsx

@@ -30,8 +30,8 @@ export const LayerMenu = memo(({ id }: Props) => {
   const selectValidActions = useMemo(
     () =>
       createMemoizedSelector(selectRegionalPromptsSlice, (regionalPrompts) => {
-        const layerIndex = regionalPrompts.layers.findIndex((l) => l.id === id);
+        const layerIndex = regionalPrompts.present.layers.findIndex((l) => l.id === id);
-        const layerCount = regionalPrompts.layers.length;
+        const layerCount = regionalPrompts.present.layers.length;
         return {
           canMoveForward: layerIndex < layerCount - 1,
           canMoveBackward: layerIndex > 0,

invokeai/frontend/web/src/features/regionalPrompts/components/NegativeModeCombobox.tsx

@@ -14,7 +14,7 @@ const options: ComboboxOption[] = [
 const AutoNegativeCombobox = () => {
   const dispatch = useAppDispatch();
   const { t } = useTranslation();
-  const autoNegative = useAppSelector((s) => s.regionalPrompts.autoNegative);
+  const autoNegative = useAppSelector((s) => s.regionalPrompts.present.autoNegative);
 
   const onChange = useCallback<ComboboxOnChange>(
     (v) => {

invokeai/frontend/web/src/features/regionalPrompts/components/PromptLayerOpacity.tsx

@@ -10,7 +10,7 @@ import { useTranslation } from 'react-i18next';
 export const PromptLayerOpacity = memo(() => {
   const dispatch = useAppDispatch();
   const { t } = useTranslation();
-  const promptLayerOpacity = useAppSelector((s) => s.regionalPrompts.promptLayerOpacity);
+  const promptLayerOpacity = useAppSelector((s) => s.regionalPrompts.present.promptLayerOpacity);
   const onChange = useCallback(
     (v: number) => {
       dispatch(promptLayerOpacityChanged(v));

invokeai/frontend/web/src/features/regionalPrompts/components/RegionalPromptsEditor.tsx

@@ -10,12 +10,13 @@ import AutoNegativeCombobox from 'features/regionalPrompts/components/NegativeMo
 import { PromptLayerOpacity } from 'features/regionalPrompts/components/PromptLayerOpacity';
 import { StageComponent } from 'features/regionalPrompts/components/StageComponent';
 import { ToolChooser } from 'features/regionalPrompts/components/ToolChooser';
+import { UndoRedoButtonGroup } from 'features/regionalPrompts/components/UndoRedoButtonGroup';
 import { selectRegionalPromptsSlice } from 'features/regionalPrompts/store/regionalPromptsSlice';
 import { getRegionalPromptLayerBlobs } from 'features/regionalPrompts/util/getLayerBlobs';
 import { memo } from 'react';
 
 const selectLayerIdsReversed = createMemoizedSelector(selectRegionalPromptsSlice, (regionalPrompts) =>
-  regionalPrompts.layers.map((l) => l.id).reverse()
+  regionalPrompts.present.layers.map((l) => l.id).reverse()
 );
 
 const debugBlobs = () => {
@@ -29,10 +30,11 @@ export const RegionalPromptsEditor = memo(() => {
       <Flex flexDir="column" gap={4} flexShrink={0}>
         <Flex gap={3}>
           <ButtonGroup isAttached={false}>
-            <Button onClick={debugBlobs}>DEBUG</Button>
+            <Button onClick={debugBlobs}>🐛</Button>
             <AddLayerButton />
             <DeleteAllLayersButton />
           </ButtonGroup>
+          <UndoRedoButtonGroup />
           <ToolChooser />
         </Flex>
         <BrushSize />

invokeai/frontend/web/src/features/regionalPrompts/components/StageComponent.tsx

@@ -19,14 +19,14 @@ import { useCallback, useLayoutEffect } from 'react';
 const log = logger('regionalPrompts');
 const $stage = atom<Konva.Stage | null>(null);
 const selectSelectedLayerColor = createMemoizedSelector(selectRegionalPromptsSlice, (regionalPrompts) => {
-  return regionalPrompts.layers.find((l) => l.id === regionalPrompts.selectedLayer)?.color ?? null;
+  return regionalPrompts.present.layers.find((l) => l.id === regionalPrompts.present.selectedLayer)?.color ?? null;
 });
 
 const useStageRenderer = (container: HTMLDivElement | null, wrapper: HTMLDivElement | null) => {
   const dispatch = useAppDispatch();
   const width = useAppSelector((s) => s.generation.width);
   const height = useAppSelector((s) => s.generation.height);
-  const state = useAppSelector((s) => s.regionalPrompts);
+  const state = useAppSelector((s) => s.regionalPrompts.present);
   const stage = useStore($stage);
   const { onMouseDown, onMouseUp, onMouseMove, onMouseEnter, onMouseLeave } = useMouseEvents();
   const cursorPosition = useStore($cursorPosition);

invokeai/frontend/web/src/features/regionalPrompts/components/ToolChooser.tsx

@@ -5,7 +5,7 @@ import { useCallback } from 'react';
 import { PiArrowsOutCardinalBold, PiEraserBold, PiPaintBrushBold } from 'react-icons/pi';
 
 export const ToolChooser: React.FC = () => {
-  const tool = useAppSelector((s) => s.regionalPrompts.tool);
+  const tool = useAppSelector((s) => s.regionalPrompts.present.tool);
   const dispatch = useAppDispatch();
   const setToolToBrush = useCallback(() => {
     dispatch(toolChanged('brush'));

invokeai/frontend/web/src/features/regionalPrompts/components/UndoRedoButtonGroup.tsx

@@ -0,0 +1,32 @@
+/* eslint-disable i18next/no-literal-string */
+import { ButtonGroup, IconButton } from '@invoke-ai/ui-library';
+import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
+import { redoRegionalPrompts, undoRegionalPrompts } from 'features/regionalPrompts/store/regionalPromptsSlice';
+import { memo, useCallback } from 'react';
+import { useHotkeys } from 'react-hotkeys-hook';
+import { PiArrowClockwiseBold, PiArrowCounterClockwiseBold } from 'react-icons/pi';
+
+export const UndoRedoButtonGroup = memo(() => {
+  const dispatch = useAppDispatch();
+
+  const mayUndo = useAppSelector((s) => s.regionalPrompts.past.length > 0);
+  const undo = useCallback(() => {
+    dispatch(undoRegionalPrompts());
+  }, [dispatch]);
+  useHotkeys(['meta+z', 'ctrl+z'], undo, { enabled: mayUndo, preventDefault: true }, [mayUndo, undo]);
+
+  const mayRedo = useAppSelector((s) => s.regionalPrompts.future.length > 0);
+  const redo = useCallback(() => {
+    dispatch(redoRegionalPrompts());
+  }, [dispatch]);
+  useHotkeys(['meta+shift+z', 'ctrl+shift+z'], redo, { enabled: mayRedo, preventDefault: true }, [mayRedo, redo]);
+
+  return (
+    <ButtonGroup>
+      <IconButton aria-label="undo" onClick={undo} icon={<PiArrowCounterClockwiseBold />} isDisabled={!mayUndo} />
+      <IconButton aria-label="redo" onClick={redo} icon={<PiArrowClockwiseBold />} isDisabled={!mayRedo} />
+    </ButtonGroup>
+  );
+});
+
+UndoRedoButtonGroup.displayName = 'UndoRedoButtonGroup';

invokeai/frontend/web/src/features/regionalPrompts/hooks/layerStateHooks.ts

@@ -9,7 +9,7 @@ export const useLayerPositivePrompt = (layerId: string) => {
     () =>
       createSelector(
         selectRegionalPromptsSlice,
-        (regionalPrompts) => regionalPrompts.layers.find((l) => l.id === layerId)?.positivePrompt
+        (regionalPrompts) => regionalPrompts.present.layers.find((l) => l.id === layerId)?.positivePrompt
       ),
     [layerId]
   );
@@ -23,7 +23,7 @@ export const useLayerNegativePrompt = (layerId: string) => {
     () =>
       createSelector(
         selectRegionalPromptsSlice,
-        (regionalPrompts) => regionalPrompts.layers.find((l) => l.id === layerId)?.negativePrompt
+        (regionalPrompts) => regionalPrompts.present.layers.find((l) => l.id === layerId)?.negativePrompt
       ),
     [layerId]
   );
@@ -37,7 +37,7 @@ export const useLayerIsVisible = (layerId: string) => {
     () =>
       createSelector(
         selectRegionalPromptsSlice,
-        (regionalPrompts) => regionalPrompts.layers.find((l) => l.id === layerId)?.isVisible
+        (regionalPrompts) => regionalPrompts.present.layers.find((l) => l.id === layerId)?.isVisible
       ),
     [layerId]
   );

invokeai/frontend/web/src/features/regionalPrompts/hooks/mouseEventHooks.ts

@@ -12,7 +12,7 @@ import type Konva from 'konva';
 import type { KonvaEventObject } from 'konva/lib/Node';
 import { useCallback } from 'react';
 
-const getTool = () => getStore().getState().regionalPrompts.tool;
+const getTool = () => getStore().getState().regionalPrompts.present.tool;
 
 const getIsFocused = (stage: Konva.Stage) => {
   return stage.container().contains(document.activeElement);
@@ -49,17 +49,19 @@ export const useMouseEvents = () => {
     [dispatch]
   );
 
-  const onMouseUp = useCallback((e: KonvaEventObject<MouseEvent | TouchEvent>) => {
+  const onMouseUp = useCallback(
+    (e: KonvaEventObject<MouseEvent | TouchEvent>) => {
       const stage = e.target.getStage();
       if (!stage) {
         return;
       }
       const tool = getTool();
       if ((tool === 'brush' || tool === 'eraser') && $isMouseDown.get()) {
-      // Add another point to the last line.
         $isMouseDown.set(false);
       }
-  }, []);
+    },
+    []
+  );
 
   const onMouseMove = useCallback(
     (e: KonvaEventObject<MouseEvent | TouchEvent>) => {

invokeai/frontend/web/src/features/regionalPrompts/store/regionalPromptsSlice.ts

@@ -1,11 +1,12 @@
 import type { PayloadAction } from '@reduxjs/toolkit';
-import { createSlice } from '@reduxjs/toolkit';
+import { createAction, createSlice, isAnyOf } from '@reduxjs/toolkit';
 import type { PersistConfig, RootState } from 'app/store/store';
 import { moveBackward, moveForward, moveToBack, moveToFront } from 'common/util/arrayUtils';
 import type { ParameterAutoNegative } from 'features/parameters/types/parameterSchemas';
 import type { IRect, Vector2d } from 'konva/lib/types';
 import { atom } from 'nanostores';
 import type { RgbColor } from 'react-colorful';
+import type { StateWithHistory, UndoableOptions } from 'redux-undo';
 import { assert } from 'tsafe';
 import { v4 as uuidv4 } from 'uuid';
 
@@ -67,6 +68,7 @@ type RegionalPromptsState = {
   brushSize: number;
   promptLayerOpacity: number;
   autoNegative: ParameterAutoNegative;
+  lastActionType: string | null;
 };
 
 export const initialRegionalPromptsState: RegionalPromptsState = {
@@ -77,6 +79,7 @@ export const initialRegionalPromptsState: RegionalPromptsState = {
   layers: [],
   promptLayerOpacity: 0.5, // This currently doesn't work
   autoNegative: 'off',
+  lastActionType: null,
 };
 
 const isLine = (obj: LayerObject): obj is LineObject => obj.kind === 'line';
@@ -235,6 +238,10 @@ export const regionalPromptsSlice = createSlice({
     autoNegativeChanged: (state, action: PayloadAction<ParameterAutoNegative>) => {
       state.autoNegative = action.payload;
     },
+    lineFinished: (state) => {
+      console.log('lineFinished');
+      return state;
+    },
   },
 });
 
@@ -293,13 +300,6 @@ const migrateRegionalPromptsState = (state: any): any => {
   return state;
 };
 
-export const regionalPromptsPersistConfig: PersistConfig<RegionalPromptsState> = {
-  name: regionalPromptsSlice.name,
-  initialState: initialRegionalPromptsState,
-  migrate: migrateRegionalPromptsState,
-  persistDenylist: [],
-};
-
 export const $isMouseDown = atom(false);
 export const $isMouseOver = atom(false);
 export const $cursorPosition = atom<Vector2d | null>(null);
@@ -322,3 +322,55 @@ const getLayerLineId = (layerId: string, lineId: string) => `${layerId}.line_${l
 export const getLayerObjectGroupId = (layerId: string, groupId: string) => `${layerId}.objectGroup_${groupId}`;
 export const getLayerBboxId = (layerId: string) => `${layerId}.bbox`;
 export const getLayerTransparencyRectId = (layerId: string) => `${layerId}.transparency_rect`;
+
+export const regionalPromptsPersistConfig: PersistConfig<RegionalPromptsState> = {
+  name: regionalPromptsSlice.name,
+  initialState: initialRegionalPromptsState,
+  migrate: migrateRegionalPromptsState,
+  persistDenylist: ['tool', 'lastActionType'],
+};
+
+// Payload-less actions for `redux-undo`
+export const undoRegionalPrompts = createAction(`${regionalPromptsSlice.name}/undo`);
+export const redoRegionalPrompts = createAction(`${regionalPromptsSlice.name}/redo`);
+export const clearHistoryRegionalPrompts = createAction(`${regionalPromptsSlice.name}/clearHistory`);
+
+// These actions are grouped together as single undoable actions
+const undoableGroupByMatcher = isAnyOf(
+  pointsAdded,
+  positivePromptChanged,
+  negativePromptChanged,
+  brushSizeChanged,
+  layerTranslated,
+  promptLayerOpacityChanged
+);
+
+export const regionalPromptsUndoableConfig: UndoableOptions = {
+  limit: 64,
+  undoType: undoRegionalPrompts.type,
+  redoType: redoRegionalPrompts.type,
+  clearHistoryType: clearHistoryRegionalPrompts.type,
+  groupBy: (action, _currentState: RegionalPromptsState, _previousHistory: StateWithHistory<RegionalPromptsState>) => {
+    if (undoableGroupByMatcher(action)) {
+      return action.type;
+    }
+    return null;
+  },
+  filter: (action, _currentState: RegionalPromptsState, _previousHistory: StateWithHistory<RegionalPromptsState>) => {
+    // Return `true` if we should record state in history, `false` if not
+    if (layerBboxChanged.match(action)) {
+      // This action is triggered on state changes, including when we undo. If we do not ignore this action, when we
+      // undo, this action triggers and empties the future states array. Therefore, we must ignore this action.
+      return false;
+    }
+    if (toolChanged.match(action)) {
+      // We don't want to record tool changes in the undo history
+      return false;
+    }
+    if (!action.type.startsWith('regionalPrompts/')) {
+      // Ignore all actions from other slices
+      return false;
+    }
+    return true;
+  },
+};

invokeai/frontend/web/src/features/regionalPrompts/util/getLayerBlobs.ts

@@ -17,24 +17,26 @@ export const getRegionalPromptLayerBlobs = async (
   preview: boolean = false
 ): Promise<Record<string, Blob>> => {
   const state = getStore().getState();
+  const reduxLayers = state.regionalPrompts.present.layers;
+  const selectedLayerId = state.regionalPrompts.present.selectedLayer;
   const container = document.createElement('div');
   const stage = new Konva.Stage({ container, width: state.generation.width, height: state.generation.height });
-  renderLayers(stage, state.regionalPrompts.layers, state.regionalPrompts.selectedLayer, 1, 'brush');
+  renderLayers(stage, reduxLayers, selectedLayerId, 1, 'brush');
 
-  const layers = stage.find<Konva.Layer>(`.${REGIONAL_PROMPT_LAYER_NAME}`);
+  const konvaLayers = stage.find<Konva.Layer>(`.${REGIONAL_PROMPT_LAYER_NAME}`);
   const blobs: Record<string, Blob> = {};
 
   // First remove all layers
-  for (const layer of layers) {
+  for (const layer of konvaLayers) {
     layer.remove();
   }
 
   // Next render each layer to a blob
-  for (const layer of layers) {
+  for (const layer of konvaLayers) {
     if (layerIds && !layerIds.includes(layer.id())) {
       continue;
     }
-    const reduxLayer = state.regionalPrompts.layers.find((l) => l.id === layer.id());
+    const reduxLayer = reduxLayers.find((l) => l.id === layer.id());
     assert(reduxLayer, `Redux layer ${layer.id()} not found`);
     stage.add(layer);
     const blob = await new Promise<Blob>((resolve) => {

invokeai/frontend/web/src/features/regionalPrompts/util/renderers.ts

@@ -240,7 +240,7 @@ export const renderLayers = (
     for (const reduxObject of reduxLayer.objects) {
       // TODO: Handle rects, images, etc
       if (reduxObject.kind !== 'line') {
-        return;
+        continue;
       }
 
       let konvaObject = stage.findOne<Konva.Line>(`#${reduxObject.id}`);