undo/redo handles multiple canvases

Author: Attila Cseh

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

@@ -22,15 +22,15 @@ import { merge } from 'es-toolkit';
 import { omit, pick } from 'es-toolkit/compat';
 import { changeBoardModalSliceConfig } from 'features/changeBoardModal/store/slice';
 import { canvasSettingsSliceConfig } from 'features/controlLayers/store/canvasSettingsSlice';
-import { canvasSliceConfig } from 'features/controlLayers/store/canvasSlice';
+import { canvasSliceConfig, undoableCanvasesReducer } from 'features/controlLayers/store/canvasSlice';
 import { canvasSessionSliceConfig } from 'features/controlLayers/store/canvasStagingAreaSlice';
 import { lorasSliceConfig } from 'features/controlLayers/store/lorasSlice';
 import { paramsSliceConfig } from 'features/controlLayers/store/paramsSlice';
 import { refImagesSliceConfig } from 'features/controlLayers/store/refImagesSlice';
 import { dynamicPromptsSliceConfig } from 'features/dynamicPrompts/store/dynamicPromptsSlice';
 import { gallerySliceConfig } from 'features/gallery/store/gallerySlice';
 import { modelManagerSliceConfig } from 'features/modelManagerV2/store/modelManagerV2Slice';
-import { nodesSliceConfig } from 'features/nodes/store/nodesSlice';
+import { nodesSliceConfig, undoableNodesSliceReducer } from 'features/nodes/store/nodesSlice';
 import { workflowLibrarySliceConfig } from 'features/nodes/store/workflowLibrarySlice';
 import { workflowSettingsSliceConfig } from 'features/nodes/store/workflowSettingsSlice';
 import { upscaleSliceConfig } from 'features/parameters/store/upscaleSlice';
@@ -44,7 +44,6 @@ import { diff } from 'jsondiffpatch';
 import dynamicMiddlewares from 'redux-dynamic-middlewares';
 import type { SerializeFunction, UnserializeFunction } from 'redux-remember';
 import { REMEMBER_REHYDRATED, rememberEnhancer, rememberReducer } from 'redux-remember';
-import undoable, { newHistory } from 'redux-undo';
 import { serializeError } from 'serialize-error';
 import { api } from 'services/api';
 import { authToastMiddleware } from 'services/api/authToastMiddleware';
@@ -91,22 +90,14 @@ const ALL_REDUCERS = {
   [api.reducerPath]: api.reducer,
   [canvasSessionSliceConfig.slice.reducerPath]: canvasSessionSliceConfig.slice.reducer,
   [canvasSettingsSliceConfig.slice.reducerPath]: canvasSettingsSliceConfig.slice.reducer,
-  // Undoable!
-  [canvasSliceConfig.slice.reducerPath]: undoable(
-    canvasSliceConfig.slice.reducer,
-    canvasSliceConfig.undoableConfig?.reduxUndoOptions
-  ),
+  [canvasSliceConfig.slice.reducerPath]: undoableCanvasesReducer,
   [changeBoardModalSliceConfig.slice.reducerPath]: changeBoardModalSliceConfig.slice.reducer,
   [configSliceConfig.slice.reducerPath]: configSliceConfig.slice.reducer,
   [dynamicPromptsSliceConfig.slice.reducerPath]: dynamicPromptsSliceConfig.slice.reducer,
   [gallerySliceConfig.slice.reducerPath]: gallerySliceConfig.slice.reducer,
   [lorasSliceConfig.slice.reducerPath]: lorasSliceConfig.slice.reducer,
   [modelManagerSliceConfig.slice.reducerPath]: modelManagerSliceConfig.slice.reducer,
-  // Undoable!
-  [nodesSliceConfig.slice.reducerPath]: undoable(
-    nodesSliceConfig.slice.reducer,
-    nodesSliceConfig.undoableConfig?.reduxUndoOptions
-  ),
+  [nodesSliceConfig.slice.reducerPath]: undoableNodesSliceReducer,
   [paramsSliceConfig.slice.reducerPath]: paramsSliceConfig.slice.reducer,
   [queueSliceConfig.slice.reducerPath]: queueSliceConfig.slice.reducer,
   [refImagesSliceConfig.slice.reducerPath]: refImagesSliceConfig.slice.reducer,
@@ -128,7 +119,7 @@ const unserialize: UnserializeFunction = (data, key) => {
   if (!sliceConfig?.persistConfig) {
     throw new Error(`No persist config for slice "${key}"`);
   }
-  const { getInitialState, persistConfig, undoableConfig } = sliceConfig;
+  const { getInitialState, persistConfig } = sliceConfig;
   let state;
   try {
     const initialState = getInitialState();
@@ -160,12 +151,7 @@ const unserialize: UnserializeFunction = (data, key) => {
     state = getInitialState();
   }
 
-  // Undoable slices must be wrapped in a history!
-  if (undoableConfig) {
-    return newHistory([], state, []);
-  } else {
-    return state;
-  }
+  return persistConfig.wrapState ? persistConfig.wrapState(state) : state;
 };
 
 const serialize: SerializeFunction = (data, key) => {
@@ -175,7 +161,7 @@ const serialize: SerializeFunction = (data, key) => {
   }
 
   const result = omit(
-    sliceConfig.undoableConfig ? data.present : data,
+    sliceConfig.persistConfig.unwrapState ? sliceConfig.persistConfig.unwrapState(data) : data,
     sliceConfig.persistConfig.persistDenylist ?? []
   );
 

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

@@ -1,10 +1,9 @@
 import type { Slice } from '@reduxjs/toolkit';
-import type { UndoableOptions } from 'redux-undo';
 import type { ZodType } from 'zod';
 
 type StateFromSlice<T extends Slice> = T extends Slice<infer U> ? U : never;
 
-export type SliceConfig<T extends Slice> = {
+export type SliceConfig<T extends Slice, TInternalState = StateFromSlice<T>, TSerializedState = StateFromSlice<T>> = {
   /**
    * The redux slice (return of createSlice).
    */
@@ -16,7 +15,7 @@ export type SliceConfig<T extends Slice> = {
   /**
    * A function that returns the initial state of the slice.
    */
-  getInitialState: () => StateFromSlice<T>;
+  getInitialState: () => TSerializedState;
   /**
    * The optional persist configuration for this slice. If omitted, the slice will not be persisted.
    */
@@ -28,19 +27,24 @@ export type SliceConfig<T extends Slice> = {
      * @param state The rehydrated state.
      * @returns A correctly-shaped state.
      */
-    migrate: (state: unknown) => StateFromSlice<T>;
+    migrate: (state: unknown) => TSerializedState;
     /**
      * Keys to omit from the persisted state.
      */
     persistDenylist?: (keyof StateFromSlice<T>)[];
-  };
-  /**
-   * The optional undoable configuration for this slice. If omitted, the slice will not be undoable.
-   */
-  undoableConfig?: {
     /**
-     * The options to be passed into redux-undo.
+     * Wraps state into state with history
+     *
+     * @param state The state without history
+     * @returns The state with history
+     */
+    wrapState?: (state: unknown) => TInternalState;
+    /**
+     * Unwraps state with history
+     *
+     * @param state The state with history
+     * @returns The state without history
      */
-    reduxUndoOptions: UndoableOptions<StateFromSlice<T>>;
+    unwrapState?: (state: TInternalState) => TSerializedState;
   };
 };

invokeai/frontend/web/src/features/controlLayers/components/RasterLayer/RasterLayerAdjustmentsPanel.tsx

@@ -13,7 +13,7 @@ import {
   rasterLayerAdjustmentsReset,
   rasterLayerAdjustmentsSet,
 } from 'features/controlLayers/store/canvasSlice';
-import { selectCanvasSlice, selectEntity } from 'features/controlLayers/store/selectors';
+import { selectEntity, selectSelectedCanvas } from 'features/controlLayers/store/selectors';
 import React, { memo, useCallback, useMemo } from 'react';
 import { useTranslation } from 'react-i18next';
 import { PiArrowCounterClockwiseBold, PiCaretDownBold, PiCheckBold, PiTrashBold } from 'react-icons/pi';
@@ -25,30 +25,32 @@ export const RasterLayerAdjustmentsPanel = memo(() => {
   const canvasManager = useCanvasManager();
 
   const selectHasAdjustments = useMemo(() => {
-    return createSelector(selectCanvasSlice, (canvas) => Boolean(selectEntity(canvas, entityIdentifier)?.adjustments));
+    return createSelector(selectSelectedCanvas, (canvas) =>
+      Boolean(selectEntity(canvas, entityIdentifier)?.adjustments)
+    );
   }, [entityIdentifier]);
 
   const hasAdjustments = useAppSelector(selectHasAdjustments);
 
   const selectMode = useMemo(() => {
     return createSelector(
-      selectCanvasSlice,
+      selectSelectedCanvas,
       (canvas) => selectEntity(canvas, entityIdentifier)?.adjustments?.mode ?? 'simple'
     );
   }, [entityIdentifier]);
   const mode = useAppSelector(selectMode);
 
   const selectEnabled = useMemo(() => {
     return createSelector(
-      selectCanvasSlice,
+      selectSelectedCanvas,
       (canvas) => selectEntity(canvas, entityIdentifier)?.adjustments?.enabled ?? false
     );
   }, [entityIdentifier]);
   const enabled = useAppSelector(selectEnabled);
 
   const selectCollapsed = useMemo(() => {
     return createSelector(
-      selectCanvasSlice,
+      selectSelectedCanvas,
       (canvas) => selectEntity(canvas, entityIdentifier)?.adjustments?.collapsed ?? false
     );
   }, [entityIdentifier]);

invokeai/frontend/web/src/features/controlLayers/components/RasterLayer/RasterLayerCurvesAdjustmentsEditor.tsx

@@ -5,7 +5,7 @@ import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import { useEntityAdapterContext } from 'features/controlLayers/contexts/EntityAdapterContext';
 import { useEntityIdentifierContext } from 'features/controlLayers/contexts/EntityIdentifierContext';
 import { rasterLayerAdjustmentsCurvesUpdated } from 'features/controlLayers/store/canvasSlice';
-import { selectCanvasSlice, selectEntity } from 'features/controlLayers/store/selectors';
+import { selectEntity, selectSelectedCanvas } from 'features/controlLayers/store/selectors';
 import type { ChannelName, ChannelPoints, CurvesAdjustmentsConfig } from 'features/controlLayers/store/types';
 import { memo, useCallback, useEffect, useMemo, useState } from 'react';
 import { useTranslation } from 'react-i18next';
@@ -72,15 +72,15 @@ export const RasterLayerCurvesAdjustmentsEditor = memo(() => {
   const { t } = useTranslation();
   const selectCurves = useMemo(() => {
     return createSelector(
-      selectCanvasSlice,
+      selectSelectedCanvas,
       (canvas) => selectEntity(canvas, entityIdentifier)?.adjustments?.curves ?? DEFAULT_CURVES
     );
   }, [entityIdentifier]);
   const curves = useAppSelector(selectCurves);
 
   const selectIsDisabled = useMemo(() => {
     return createSelector(
-      selectCanvasSlice,
+      selectSelectedCanvas,
       (canvas) => selectEntity(canvas, entityIdentifier)?.adjustments?.enabled !== true
     );
   }, [entityIdentifier]);

invokeai/frontend/web/src/features/controlLayers/components/RasterLayer/RasterLayerMenuItemsAdjustments.tsx

@@ -1,20 +1,25 @@
 import { MenuItem } from '@invoke-ai/ui-library';
+import { createSelector } from '@reduxjs/toolkit';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import { useEntityIdentifierContext } from 'features/controlLayers/contexts/EntityIdentifierContext';
 import { rasterLayerAdjustmentsCancel, rasterLayerAdjustmentsSet } from 'features/controlLayers/store/canvasSlice';
+import { selectSelectedCanvas } from 'features/controlLayers/store/selectors';
 import type { CanvasRasterLayerState } from 'features/controlLayers/store/types';
 import { makeDefaultRasterLayerAdjustments } from 'features/controlLayers/store/util';
-import { memo, useCallback } from 'react';
+import { memo, useCallback, useMemo } from 'react';
 import { useTranslation } from 'react-i18next';
 import { PiSlidersHorizontalBold } from 'react-icons/pi';
 
 export const RasterLayerMenuItemsAdjustments = memo(() => {
   const dispatch = useAppDispatch();
   const entityIdentifier = useEntityIdentifierContext<'raster_layer'>();
   const { t } = useTranslation();
-  const layer = useAppSelector((s) =>
-    s.canvas.present.rasterLayers.entities.find((e: CanvasRasterLayerState) => e.id === entityIdentifier.id)
-  );
+  const selectRasterLayer = useMemo(() => {
+    return createSelector(selectSelectedCanvas, (canvas) =>
+      canvas.rasterLayers.entities.find((e: CanvasRasterLayerState) => e.id === entityIdentifier.id)
+    );
+  }, [entityIdentifier]);
+  const layer = useAppSelector(selectRasterLayer);
   const hasAdjustments = Boolean(layer?.adjustments);
   const onToggleAdjustmentsPresence = useCallback(() => {
     if (hasAdjustments) {

invokeai/frontend/web/src/features/controlLayers/components/RasterLayer/RasterLayerSimpleAdjustmentsEditor.tsx

@@ -3,7 +3,7 @@ import { createSelector } from '@reduxjs/toolkit';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import { useEntityIdentifierContext } from 'features/controlLayers/contexts/EntityIdentifierContext';
 import { rasterLayerAdjustmentsSimpleUpdated } from 'features/controlLayers/store/canvasSlice';
-import { selectCanvasSlice, selectEntity } from 'features/controlLayers/store/selectors';
+import { selectEntity, selectSelectedCanvas } from 'features/controlLayers/store/selectors';
 import type { SimpleAdjustmentsConfig } from 'features/controlLayers/store/types';
 import React, { memo, useCallback, useMemo } from 'react';
 import { useTranslation } from 'react-i18next';
@@ -21,7 +21,7 @@ const AdjustmentSliderRow = ({ label, name, onChange, min = -1, max = 1, step =
   const entityIdentifier = useEntityIdentifierContext<'raster_layer'>();
   const selectValue = useMemo(() => {
     return createSelector(
-      selectCanvasSlice,
+      selectSelectedCanvas,
       (canvas) =>
         selectEntity(canvas, entityIdentifier)?.adjustments?.simple?.[name] ?? DEFAULT_SIMPLE_ADJUSTMENTS[name]
     );
@@ -54,7 +54,7 @@ export const RasterLayerSimpleAdjustmentsEditor = memo(() => {
   const { t } = useTranslation();
   const selectIsDisabled = useMemo(() => {
     return createSelector(
-      selectCanvasSlice,
+      selectSelectedCanvas,
       (canvas) => selectEntity(canvas, entityIdentifier)?.adjustments?.enabled !== true
     );
   }, [entityIdentifier]);