staging area multi-canvas support

Author: Attila Cseh

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelSelected.ts

@@ -1,7 +1,10 @@
 import { logger } from 'app/logging/logger';
 import type { AppStartListening } from 'app/store/store';
 import { bboxSyncedToOptimalDimension, rgRefImageModelChanged } from 'features/controlLayers/store/canvasSlice';
-import { buildSelectIsStaging, selectCanvasSessionId } from 'features/controlLayers/store/canvasStagingAreaSlice';
+import {
+  buildSelectIsStagingBySessionId,
+  selectSelectedCanvasSessionId,
+} from 'features/controlLayers/store/canvasStagingAreaSlice';
 import { loraIsEnabledChanged } from 'features/controlLayers/store/lorasSlice';
 import { modelChanged, syncedToOptimalDimension, vaeSelected } from 'features/controlLayers/store/paramsSlice';
 import { refImageModelChanged, selectReferenceImageEntities } from 'features/controlLayers/store/refImagesSlice';
@@ -159,7 +162,9 @@ export const addModelSelectedListener = (startAppListening: AppStartListening) =
       if (modelBase !== state.params.model?.base) {
         // Sync generate tab settings whenever the model base changes
         dispatch(syncedToOptimalDimension());
-        const isStaging = buildSelectIsStaging(selectCanvasSessionId(state))(state);
+        const sessionId = selectSelectedCanvasSessionId(state);
+        const selectIsStaging = buildSelectIsStagingBySessionId(sessionId);
+        const isStaging = selectIsStaging(state);
         if (!isStaging) {
           // Canvas tab only syncs if not staging
           dispatch(bboxSyncedToOptimalDimension());

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/setDefaultSettings.ts

@@ -1,7 +1,10 @@
 import type { AppStartListening } from 'app/store/store';
 import { isNil } from 'es-toolkit';
 import { bboxHeightChanged, bboxWidthChanged } from 'features/controlLayers/store/canvasSlice';
-import { buildSelectIsStaging, selectCanvasSessionId } from 'features/controlLayers/store/canvasStagingAreaSlice';
+import {
+  buildSelectIsStagingBySessionId,
+  selectSelectedCanvasSessionId,
+} from 'features/controlLayers/store/canvasStagingAreaSlice';
 import {
   heightChanged,
   setCfgRescaleMultiplier,
@@ -115,7 +118,9 @@ export const addSetDefaultSettingsListener = (startAppListening: AppStartListeni
         }
         const setSizeOptions = { updateAspectRatio: true, clamp: true };
 
-        const isStaging = buildSelectIsStaging(selectCanvasSessionId(state))(state);
+        const sessionId = selectSelectedCanvasSessionId(state);
+        const selectIsStaging = buildSelectIsStagingBySessionId(sessionId);
+        const isStaging = selectIsStaging(state);
 
         const activeTab = selectActiveTab(getState());
         if (activeTab === 'generate') {

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

@@ -22,7 +22,7 @@ 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, undoableCanvasesReducer } from 'features/controlLayers/store/canvasSlice';
+import { canvasSliceConfig, migrateCanvas, 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';
@@ -219,8 +219,9 @@ export const createStore = (options?: { persist?: boolean; persistDebounce?: num
   // Once-off listener to support waiting for rehydration before rendering the app
   startAppListening({
     actionCreator: createAction(REMEMBER_REHYDRATED),
-    effect: (action, { unsubscribe }) => {
+    effect: (action, { dispatch, unsubscribe }) => {
       unsubscribe();
+      dispatch(migrateCanvas());
       options?.onRehydrated?.();
     },
   });

invokeai/frontend/web/src/common/components/SessionMenuItems.tsx

@@ -1,6 +1,6 @@
 import { MenuItem } from '@invoke-ai/ui-library';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
-import { useSelectedCanvasManagerSafe } from 'features/controlLayers/contexts/CanvasManagerProviderGate';
+import { useCanvasManagerSafe } from 'features/controlLayers/contexts/CanvasManagerProviderGate';
 import { allEntitiesDeleted, inpaintMaskAdded } from 'features/controlLayers/store/canvasSlice';
 import { paramsReset } from 'features/controlLayers/store/paramsSlice';
 import { selectActiveTab } from 'features/ui/store/uiSelectors';
@@ -12,7 +12,7 @@ export const SessionMenuItems = memo(() => {
   const { t } = useTranslation();
   const dispatch = useAppDispatch();
   const tab = useAppSelector(selectActiveTab);
-  const canvasManager = useSelectedCanvasManagerSafe();
+  const canvasManager = useCanvasManagerSafe();
 
   const resetCanvasLayers = useCallback(() => {
     dispatch(allEntitiesDeleted());

invokeai/frontend/web/src/features/controlLayers/components/RefImage/RefImageImage.tsx

@@ -4,8 +4,8 @@ import { objectEquals } from '@observ33r/object-equals';
 import { skipToken } from '@reduxjs/toolkit/query';
 import { useAppSelector, useAppStore } from 'app/store/storeHooks';
 import { UploadImageIconButton } from 'common/hooks/useImageUploadButton';
+import { useCanvasIsStaging } from 'features/controlLayers/hooks/useCanvasIsStaging';
 import { bboxSizeOptimized, bboxSizeRecalled } from 'features/controlLayers/store/canvasSlice';
-import { useCanvasIsStaging } from 'features/controlLayers/store/canvasStagingAreaSlice';
 import { sizeOptimized, sizeRecalled } from 'features/controlLayers/store/paramsSlice';
 import type { CroppableImageWithDims } from 'features/controlLayers/store/types';
 import { imageDTOToCroppableImage, imageDTOToImageWithDims } from 'features/controlLayers/store/util';

invokeai/frontend/web/src/features/controlLayers/components/RefImage/RefImageSettings.tsx

@@ -10,7 +10,10 @@ import { IPAdapterMethod } from 'features/controlLayers/components/RefImage/IPAd
 import { RefImageModel } from 'features/controlLayers/components/RefImage/RefImageModel';
 import { RefImageNoImageState } from 'features/controlLayers/components/RefImage/RefImageNoImageState';
 import { RefImageNoImageStateWithCanvasOptions } from 'features/controlLayers/components/RefImage/RefImageNoImageStateWithCanvasOptions';
-import { CanvasManagerProviderGate, useCanvasManagerSafe } from 'features/controlLayers/contexts/CanvasManagerProviderGate';
+import {
+  CanvasManagerProviderGate,
+  useCanvasManagerSafe,
+} from 'features/controlLayers/contexts/CanvasManagerProviderGate';
 import { useRefImageIdContext } from 'features/controlLayers/contexts/RefImageIdContext';
 import { selectIsFLUX } from 'features/controlLayers/store/paramsSlice';
 import {

invokeai/frontend/web/src/features/controlLayers/components/RegionalGuidance/RegionalGuidanceRefImageImage.tsx

@@ -3,8 +3,8 @@ import { useStore } from '@nanostores/react';
 import { skipToken } from '@reduxjs/toolkit/query';
 import { useAppSelector, useAppStore } from 'app/store/storeHooks';
 import { UploadImageIconButton } from 'common/hooks/useImageUploadButton';
+import { useCanvasIsStaging } from 'features/controlLayers/hooks/useCanvasIsStaging';
 import { bboxSizeOptimized, bboxSizeRecalled } from 'features/controlLayers/store/canvasSlice';
-import { useCanvasIsStaging } from 'features/controlLayers/store/canvasStagingAreaSlice';
 import { sizeOptimized, sizeRecalled } from 'features/controlLayers/store/paramsSlice';
 import type { ImageWithDims } from 'features/controlLayers/store/types';
 import type { setRegionalGuidanceReferenceImageDndTarget } from 'features/dnd/dnd';

invokeai/frontend/web/src/features/controlLayers/components/StagingArea/context.tsx

@@ -1,12 +1,13 @@
 import { useStore } from '@nanostores/react';
 import { useAppStore } from 'app/store/storeHooks';
+import { useScopedCanvasSessionId } from 'features/controlLayers/hooks/useCanvasSessionId';
 import {
   selectStagingAreaAutoSwitch,
   settingsStagingAreaAutoSwitchChanged,
 } from 'features/controlLayers/store/canvasSettingsSlice';
 import { rasterLayerAdded } from 'features/controlLayers/store/canvasSlice';
 import {
-  buildSelectCanvasQueueItems,
+  buildSelectCanvasQueueItemsBySessionId,
   canvasQueueItemDiscarded,
   canvasSessionReset,
 } from 'features/controlLayers/store/canvasStagingAreaSlice';
@@ -26,11 +27,12 @@ import { getInitialProgressData, StagingAreaApi } from './state';
 
 const StagingAreaContext = createContext<StagingAreaApi | null>(null);
 
-export const StagingAreaContextProvider = memo(({ children, sessionId }: PropsWithChildren<{ sessionId: string }>) => {
+export const StagingAreaContextProvider = memo(({ canvasId, children }: PropsWithChildren<{ canvasId: string }>) => {
   const store = useAppStore();
   const socket = useStore($socket);
+  const sessionId = useScopedCanvasSessionId(canvasId);
   const stagingAreaAppApi = useMemo<StagingAreaAppApi>(() => {
-    const selectQueueItems = buildSelectCanvasQueueItems(sessionId);
+    const selectQueueItems = buildSelectCanvasQueueItemsBySessionId(sessionId);
 
     const _stagingAreaAppApi: StagingAreaAppApi = {
       getAutoSwitch: () => selectStagingAreaAutoSwitch(store.getState()),
@@ -58,16 +60,18 @@ export const StagingAreaContextProvider = memo(({ children, sessionId }: PropsWi
         });
       },
       onDiscard: ({ item_id, status }) => {
-        store.dispatch(canvasQueueItemDiscarded({ itemId: item_id }));
+        store.dispatch(canvasQueueItemDiscarded({ canvasId, itemId: item_id }));
         if (status === 'in_progress' || status === 'pending') {
           store.dispatch(queueApi.endpoints.cancelQueueItem.initiate({ item_id }, { track: false }));
         }
       },
       onDiscardAll: () => {
-        store.dispatch(canvasSessionReset());
-        store.dispatch(
-          queueApi.endpoints.cancelQueueItemsByDestination.initiate({ destination: sessionId }, { track: false })
-        );
+        store.dispatch(canvasSessionReset({ canvasId }));
+        if (sessionId) {
+          store.dispatch(
+            queueApi.endpoints.cancelQueueItemsByDestination.initiate({ destination: sessionId }, { track: false })
+          );
+        }
       },
       onAccept: (item, imageDTO) => {
         const bboxRect = selectBboxRect(store.getState());
@@ -80,22 +84,30 @@ export const StagingAreaContextProvider = memo(({ children, sessionId }: PropsWi
         };
 
         store.dispatch(rasterLayerAdded({ overrides, isSelected: selectedEntityIdentifier?.type === 'raster_layer' }));
-        store.dispatch(canvasSessionReset());
-        store.dispatch(
-          queueApi.endpoints.cancelQueueItemsByDestination.initiate({ destination: sessionId }, { track: false })
-        );
+        store.dispatch(canvasSessionReset({ canvasId }));
+        if (sessionId) {
+          store.dispatch(
+            queueApi.endpoints.cancelQueueItemsByDestination.initiate({ destination: sessionId }, { track: false })
+          );
+        }
       },
       onAutoSwitchChange: (mode) => {
         store.dispatch(settingsStagingAreaAutoSwitchChanged(mode));
       },
     };
 
     return _stagingAreaAppApi;
-  }, [sessionId, socket, store]);
+  }, [canvasId, sessionId, socket, store]);
 
   const [stagingAreaApi] = useState(() => new StagingAreaApi());
 
   useEffect(() => {
+    if (!sessionId) {
+      return () => {
+        stagingAreaApi.cleanup();
+      };
+    }
+
     stagingAreaApi.connectToApp(sessionId, stagingAreaAppApi);
 
     // We need to subscribe to the queue items query manually to ensure the staging area actually gets the items

invokeai/frontend/web/src/features/controlLayers/contexts/CanvasInstanceContextProvider.tsx

@@ -1,6 +1,5 @@
 import type { PropsWithChildren } from 'react';
 import { createContext, memo, useContext } from 'react';
-import { assert } from 'tsafe';
 
 const CanvasInstanceContext = createContext<string | null>(null);
 
@@ -9,18 +8,6 @@ export const CanvasInstanceContextProvider = memo(({ canvasId, children }: Props
 });
 CanvasInstanceContextProvider.displayName = 'CanvasInstanceContextProvider';
 
-export const useScopedCanvas = () => {
-  const canvasId = useContext(CanvasInstanceContext);
-  assert(canvasId, 'useCanvasInstanceContext must be used within a CanvasInstanceContext');
-  return canvasId;
-};
-
-export const useScopedCanvasSafe = () => {
+export const useScopedCanvasIdSafe = () => {
   return useContext(CanvasInstanceContext);
 };
-
-export const useHasScopedCanvas = () => {
-  const canvasId = useContext(CanvasInstanceContext);
-
-  return !!canvasId;
-};

invokeai/frontend/web/src/features/controlLayers/contexts/CanvasManagerProviderGate.tsx

@@ -1,21 +1,20 @@
 import { useStore } from '@nanostores/react';
 import { useAppSelector } from 'app/store/storeHooks';
+import { useCanvasId } from 'features/controlLayers/hooks/useCanvasId';
 import type { CanvasManager } from 'features/controlLayers/konva/CanvasManager';
 import { $canvasManagers } from 'features/controlLayers/store/ephemeral';
-import { selectSelectedCanvas } from 'features/controlLayers/store/selectors';
+import { selectSelectedCanvasId } from 'features/controlLayers/store/selectors';
 import type { PropsWithChildren } from 'react';
-import { createContext, memo, useContext, useMemo } from 'react';
+import { createContext, memo } from 'react';
 import { assert } from 'tsafe';
 
-import { useScopedCanvas, useScopedCanvasSafe } from './CanvasInstanceContextProvider';
-
 const CanvasManagerContext = createContext<{ [canvasId: string]: CanvasManager } | null>(null);
 
 export const CanvasManagerProviderGate = memo(({ children }: PropsWithChildren) => {
   const canvasManagers = useStore($canvasManagers);
-  const selectedCanvas = useAppSelector(selectSelectedCanvas);
+  const selectedCanvasId = useAppSelector(selectSelectedCanvasId);
 
-  if((Object.keys(canvasManagers).length === 0) || !canvasManagers[selectedCanvas.id]) {
+  if (Object.keys(canvasManagers).length === 0 || !canvasManagers[selectedCanvasId]) {
     return null;
   }
 
@@ -24,68 +23,6 @@ export const CanvasManagerProviderGate = memo(({ children }: PropsWithChildren)
 
 CanvasManagerProviderGate.displayName = 'CanvasManagerProviderGate';
 
-/**
- * Consumes the scoped CanvasManager from the context. This hook must be used within the CanvasManagerProviderGate, otherwise
- * it will throw an error.
- */
-export const useScopedCanvasManager = (): CanvasManager => {
-  const canvasManagers = useContext(CanvasManagerContext);
-  assert(canvasManagers, 'useScopedCanvasManager must be used within a CanvasManagerProviderGate');
-
-  const scopedCanvasId = useScopedCanvas();
-  const canvasManager = useMemo(() => {
-    return canvasManagers[scopedCanvasId];
-  }, [canvasManagers, scopedCanvasId]);
-  assert(canvasManager, 'Scoped canvas manager not initialised');
-
-  return canvasManager;
-};
-
-/**
- * Consumes the selected CanvasManager from the context. This hook must be used within the CanvasManagerProviderGate, otherwise
- * it will throw an error.
- */
-export const useSelectedCanvasManager = (): CanvasManager => {
-  const canvasManagers = useContext(CanvasManagerContext);
-  assert(canvasManagers, 'useScopedCanvasManager must be used within a CanvasManagerProviderGate');
-
-  const selectedCanvas = useAppSelector(selectSelectedCanvas);
-  const canvasManager = useMemo(() => {
-    return canvasManagers[selectedCanvas.id];
-  }, [canvasManagers, selectedCanvas]);
-  assert(canvasManager, 'Selected canvas manager not initialised');
-
-  return canvasManager;
-};
-
-/**
- * Consumes the scoped CanvasManager from the context. If the CanvasManager is not available, it will return null.
- */
-export const useScopedCanvasManagerSafe = (): CanvasManager | null => {
-  const canvasManagers = useStore($canvasManagers);
-  const scopedCanvasId = useScopedCanvasSafe();
-
-  const canvasManager = useMemo(() => {
-    return scopedCanvasId ? canvasManagers[scopedCanvasId] : undefined;
-  }, [canvasManagers, scopedCanvasId]);
-
-  return canvasManager ?? null;
-};
-
-/**
- * Consumes the selected CanvasManager from the context. If the CanvasManager is not available, it will return null.
- */
-export const useSelectedCanvasManagerSafe = (): CanvasManager | null => {
-  const canvasManagers = useStore($canvasManagers);
-  const selectedCanvas = useAppSelector(selectSelectedCanvas);
-
-  const canvasManager = useMemo(() => {
-    return canvasManagers[selectedCanvas.id];
-  }, [canvasManagers, selectedCanvas]);
-
-  return canvasManager ?? null;
-};
-
 /**
  * Consumes the CanvasManager from the context. If the CanvasManager is not available, it will throw an error.
  */
@@ -101,12 +38,7 @@ export const useCanvasManager = (): CanvasManager => {
  */
 export const useCanvasManagerSafe = (): CanvasManager | null => {
   const canvasManagers = useStore($canvasManagers);
-  const scopedCanvasId = useScopedCanvasSafe();
-  const selectedCanvas = useAppSelector(selectSelectedCanvas);
-
-  const canvasManager = useMemo(() => {
-    return scopedCanvasId ? canvasManagers[scopedCanvasId] : canvasManagers[selectedCanvas.id];
-  }, [canvasManagers, selectedCanvas, scopedCanvasId]);
+  const canvasId = useCanvasId();
 
-  return canvasManager ?? null;
+  return canvasManagers[canvasId] ?? null;
 };