canvas dependencies initialization fixed

Author: Attila Cseh

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 { canvasSettingsReducer, canvasSettingsSliceConfig } from 'features/controlLayers/store/canvasSettingsSlice';
-import { canvasSliceConfig, migrateCanvas, undoableCanvasesReducer } from 'features/controlLayers/store/canvasSlice';
+import { canvasSliceConfig, initializeCanvasDependencies, migrateCanvas, undoableCanvasesReducer } from 'features/controlLayers/store/canvasSlice';
 import { canvasSessionSliceConfig } from 'features/controlLayers/store/canvasStagingAreaSlice';
 import { lorasSliceConfig } from 'features/controlLayers/store/lorasSlice';
 import { paramsSliceConfig, paramsSliceReducer } from 'features/controlLayers/store/paramsSlice';
@@ -222,6 +222,7 @@ export const createStore = (options?: { persist?: boolean; persistDebounce?: num
     effect: (action, { dispatch, unsubscribe }) => {
       unsubscribe();
       dispatch(migrateCanvas());
+      dispatch(initializeCanvasDependencies());
       options?.onRehydrated?.();
     },
   });

invokeai/frontend/web/src/features/controlLayers/store/canvasSettingsSlice.ts

@@ -11,9 +11,11 @@ import { z } from 'zod';
 import {
   canvasAdding,
   canvasDeleted,
+  canvasInitialized,
   canvasMultiCanvasMigrated,
   MIGRATION_MULTI_CANVAS_ID_PLACEHOLDER,
 } from './canvasSlice';
+import { selectActiveCanvasId } from './selectors';
 
 const zAutoSwitchMode = z.enum(['off', 'switch_on_start', 'switch_on_finish']);
 export type AutoSwitchMode = z.infer<typeof zAutoSwitchMode>;
@@ -154,7 +156,7 @@ type CanvasPayloadAction<P> = PayloadAction<PayloadWithCanvasId<P>>;
 
 const canvasSettingsSlice = createSlice({
   name: 'canvasSettings',
-  initialState: getInitialCanvasSettingsState(),
+  initialState: getInitialCanvasSettingsState,
   reducers: {
     settingsClipToBboxChanged: (state, action: PayloadAction<{ clipToBbox: boolean }>) => {
       const { clipToBbox } = action.payload;
@@ -231,10 +233,16 @@ const canvasSettingsSlice = createSlice({
       state.canvases[settings.canvasId] = settings;
       delete state.canvases[MIGRATION_MULTI_CANVAS_ID_PLACEHOLDER];
     });
+    builder.addCase(canvasInitialized, (state, action) => {
+      const canvasId = action.payload.canvasId;
+      if(!state.canvases[canvasId]) {
+        state.canvases[canvasId] = getInitialCanvasInstanceSettings(canvasId);
+      }
+    });
   },
 });
 
-const canvasInstanceSettingsSlice = createSlice({
+const canvasInstanceSettingsFragment = createSlice({
   name: 'canvasSettings',
   initialState: {} as CanvasInstanceSettings,
   reducers: {
@@ -294,11 +302,11 @@ export const {
   settingsBgColorChanged,
   settingsFgColorChanged,
   settingsColorsSetToDefault,
-} = canvasInstanceSettingsSlice.actions;
+} = canvasInstanceSettingsFragment.actions;
 
-const isCanvasInstanceSettingsAction = isAnyOf(...Object.values(canvasInstanceSettingsSlice.actions));
+const isCanvasInstanceSettingsAction = isAnyOf(...Object.values(canvasInstanceSettingsFragment.actions));
 
-export const canvasSettingsReducer = (state: CanvasSettingsState, action: UnknownAction): CanvasSettingsState => {
+export const canvasSettingsReducer = (state: CanvasSettingsState | undefined, action: UnknownAction): CanvasSettingsState => {
   state = canvasSettingsSlice.reducer(state, action);
 
   if (!isCanvasInstanceSettingsAction(action)) {
@@ -311,7 +319,7 @@ export const canvasSettingsReducer = (state: CanvasSettingsState, action: Unknow
     ...state,
     canvases: {
       ...state.canvases,
-      [canvasId]: canvasInstanceSettingsSlice.reducer(state.canvases[canvasId], action),
+      [canvasId]: canvasInstanceSettingsFragment.reducer(state.canvases[canvasId], action),
     },
   };
 };

invokeai/frontend/web/src/features/controlLayers/store/canvasSlice.ts

@@ -193,10 +193,6 @@ const canvasesSlice = createSlice({
       },
     },
     canvasAdding: (_state, _action: CanvasPayloadAction<unknown>) => {},
-    canvasMigrated: (state) => {
-      delete state.migration;
-    },
-    canvasMultiCanvasMigrated: (_state, _action: CanvasPayloadAction<unknown>) => {},
     canvasActivated: (state, action: CanvasPayloadAction<unknown>) => {
       const { canvasId } = action.payload;
 
@@ -227,10 +223,15 @@ const canvasesSlice = createSlice({
       delete state.canvases[canvas.id];
     },
     canvasDeleted: (_state, _action: CanvasPayloadAction<unknown>) => {},
+    canvasMigrated: (state) => {
+      delete state.migration;
+    },
+    canvasMultiCanvasMigrated: (_state, _action: CanvasPayloadAction<unknown>) => {},
+    canvasInitialized: (_state, _action: CanvasPayloadAction<unknown>) => {},
   },
 });
 
-const canvasSlice = createSlice({
+const canvasFragment = createSlice({
   name: 'canvas',
   initialState: {} as CanvasState,
   reducers: {
@@ -1890,6 +1891,11 @@ export const addCanvas = (payload: { isSelected?: boolean }) => (dispatch: AppDi
 
 export const MIGRATION_MULTI_CANVAS_ID_PLACEHOLDER = 'multi-canvas-id-placeholder';
 
+export const deleteCanvas = (payload: { canvasId: string }) => (dispatch: AppDispatch) => {
+  dispatch(canvasesSlice.actions.canvasDeleting(payload));
+  dispatch(canvasesSlice.actions.canvasDeleted(payload));
+};
+
 export const migrateCanvas = () => (dispatch: AppDispatch, getState: () => RootState) => {
   const state = getState();
 
@@ -1900,17 +1906,19 @@ export const migrateCanvas = () => (dispatch: AppDispatch, getState: () => RootS
   dispatch(canvasesSlice.actions.canvasMigrated());
 };
 
-export const deleteCanvas = (payload: { canvasId: string }) => (dispatch: AppDispatch) => {
-  dispatch(canvasesSlice.actions.canvasDeleting(payload));
-  dispatch(canvasesSlice.actions.canvasDeleted(payload));
-};
+export const initializeCanvasDependencies = () => (dispatch: AppDispatch, getState: () => RootState) => {
+  const state = getState();
+
+  dispatch(canvasInitialized({ canvasId: state.canvas.activeCanvasId }));
+}
 
 export const {
   // Canvas
   canvasAdding,
-  canvasMultiCanvasMigrated,
   canvasDeleted,
   canvasActivated,
+  canvasMultiCanvasMigrated,
+  canvasInitialized,
 } = canvasesSlice.actions;
 
 export const {
@@ -2010,9 +2018,9 @@ export const {
   inpaintMaskDenoiseLimitChanged,
   inpaintMaskDenoiseLimitDeleted,
   // inpaintMaskRecalled,
-} = canvasSlice.actions;
+} = canvasFragment.actions;
 
-const isCanvasSliceAction = isAnyOf(...Object.values(canvasSlice.actions));
+const isCanvasAction = isAnyOf(...Object.values(canvasFragment.actions));
 
 let filter = true;
 const isActionFileterd = isAnyOf(canvasNameChanged, entitySelected);
@@ -2024,7 +2032,7 @@ const canvasUndoableConfig: UndoableOptions<CanvasState, UnknownAction> = {
   clearHistoryType: canvasClearHistory.type,
   filter: (action, _state, _history) => {
     // Ignore both all actions from other slices and canvas management actions
-    if (!action.type.startsWith(canvasSlice.name) || isActionFileterd(action)) {
+    if (!action.type.startsWith(canvasFragment.name) || isActionFileterd(action)) {
       return false;
     }
     // Throttle rapid actions of the same type
@@ -2035,15 +2043,15 @@ const canvasUndoableConfig: UndoableOptions<CanvasState, UnknownAction> = {
   // debug: import.meta.env.MODE === 'development',
 };
 
-const undoableCanvasReducer = undoable(canvasSlice.reducer, canvasUndoableConfig);
+const undoableCanvasReducer = undoable(canvasFragment.reducer, canvasUndoableConfig);
 
 export const undoableCanvasesReducer = (
-  state: CanvasesStateWithHistory,
+  state: CanvasesStateWithHistory | undefined,
   action: UnknownAction
 ): CanvasesStateWithHistory => {
   state = canvasesSlice.reducer(state, action);
 
-  if (!isCanvasSliceAction(action)) {
+  if (!isCanvasAction(action)) {
     return state;
   }
 

invokeai/frontend/web/src/features/controlLayers/store/canvasStagingAreaSlice.ts

@@ -1,4 +1,5 @@
-import { createSelector, createSlice, type PayloadAction } from '@reduxjs/toolkit';
+import type { PayloadAction, UnknownAction } from '@reduxjs/toolkit';
+import { createSelector, createSlice, isAnyOf } from '@reduxjs/toolkit';
 import { EMPTY_ARRAY } from 'app/store/constants';
 import type { RootState } from 'app/store/store';
 import type { SliceConfig } from 'app/store/types';
@@ -11,64 +12,85 @@ import z from 'zod';
 import {
   canvasAdding,
   canvasDeleted,
+  canvasInitialized,
   canvasMultiCanvasMigrated,
   MIGRATION_MULTI_CANVAS_ID_PLACEHOLDER,
 } from './canvasSlice';
 import { selectActiveCanvasId } from './selectors';
 
-const zCanvasSessionState = z.object({
+const zCanvasSession = z.object({
   canvasId: z.string().min(1),
   canvasSessionId: z.string(),
   canvasDiscardedQueueItems: z.array(z.number().int()),
 });
-type CanvasSessionState = z.infer<typeof zCanvasSessionState>;
+type CanvasSession = z.infer<typeof zCanvasSession>;
 const zCanvasStagingAreaState = z.object({
   _version: z.literal(2),
-  sessions: z.record(z.string(), zCanvasSessionState),
+  sessions: z.record(z.string(), zCanvasSession),
 });
 type CanvasStagingAreaState = z.infer<typeof zCanvasStagingAreaState>;
 
 type CanvasPayload<T> = { canvasId: string } & T;
 type CanvasPayloadAction<T> = PayloadAction<CanvasPayload<T>>;
 
-const getInitialCanvasSessionState = (canvasId: string): CanvasSessionState => ({
+const getInitialCanvasSessionState = (canvasId: string): CanvasSession => ({
   canvasId,
   canvasSessionId: getPrefixedId('canvas'),
   canvasDiscardedQueueItems: [],
 });
 
-const getInitialState = (): CanvasStagingAreaState => ({
+const getInitialCanvasStagingAreaState = (): CanvasStagingAreaState => ({
   _version: 2,
   sessions: {},
 });
 
 const canvasStagingAreaSlice = createSlice({
   name: 'canvasSession',
-  initialState: getInitialState(),
-  reducers: {
-    canvasQueueItemDiscarded: (state, action: CanvasPayloadAction<{ itemId: number }>) => {
-      const { canvasId, itemId } = action.payload;
-
-      const session = state.sessions[canvasId];
+  initialState: getInitialCanvasStagingAreaState,
+  reducers: {},
+  extraReducers(builder) {
+    builder.addCase(canvasAdding, (state, action) => {
+      const session = getInitialCanvasSessionState(action.payload.canvasId);
+      state.sessions[session.canvasId] = session;
+    });
+    builder.addCase(canvasDeleted, (state, action) => {
+      delete state.sessions[action.payload.canvasId];
+    });
+    builder.addCase(canvasMultiCanvasMigrated, (state, action) => {
+      const session = state.sessions[MIGRATION_MULTI_CANVAS_ID_PLACEHOLDER];
       if (!session) {
         return;
       }
+      session.canvasId = action.payload.canvasId;
+      state.sessions[session.canvasId] = session;
+      delete state.sessions[MIGRATION_MULTI_CANVAS_ID_PLACEHOLDER];
+    });
+    builder.addCase(canvasInitialized, (state, action) => {
+      const canvasId = action.payload.canvasId;
+      if(!state.sessions[canvasId]) {
+        state.sessions[canvasId] = getInitialCanvasSessionState(canvasId);
+      }
+    });
+  },
+});
+
+const canvasSessionFragment = createSlice({
+  name: 'canvasSession',
+  initialState: {} as CanvasSession,
+  reducers: {
+    canvasQueueItemDiscarded: (state, action: CanvasPayloadAction<{ itemId: number }>) => {
+      const { itemId } = action.payload;
 
-      if (!session.canvasDiscardedQueueItems.includes(itemId)) {
-        session.canvasDiscardedQueueItems.push(itemId);
+      if (!state.canvasDiscardedQueueItems.includes(itemId)) {
+        state.canvasDiscardedQueueItems.push(itemId);
       }
     },
     canvasSessionReset: {
       reducer: (state, action: CanvasPayloadAction<{ canvasSessionId: string }>) => {
-        const { canvasId, canvasSessionId } = action.payload;
-
-        const session = state.sessions[canvasId];
-        if (!session) {
-          return;
-        }
+        const { canvasSessionId } = action.payload;
 
-        session.canvasSessionId = canvasSessionId;
-        session.canvasDiscardedQueueItems = [];
+        state.canvasSessionId = canvasSessionId;
+        state.canvasDiscardedQueueItems = [];
       },
       prepare: (payload: CanvasPayload<object>) => {
         return {
@@ -80,32 +102,34 @@ const canvasStagingAreaSlice = createSlice({
       },
     },
   },
-  extraReducers(builder) {
-    builder.addCase(canvasAdding, (state, action) => {
-      const session = getInitialCanvasSessionState(action.payload.canvasId);
-      state.sessions[session.canvasId] = session;
-    });
-    builder.addCase(canvasDeleted, (state, action) => {
-      delete state.sessions[action.payload.canvasId];
-    });
-    builder.addCase(canvasMultiCanvasMigrated, (state, action) => {
-      const session = state.sessions[MIGRATION_MULTI_CANVAS_ID_PLACEHOLDER];
-      if (!session) {
-        return;
-      }
-      session.canvasId = action.payload.canvasId;
-      state.sessions[session.canvasId] = session;
-      delete state.sessions[MIGRATION_MULTI_CANVAS_ID_PLACEHOLDER];
-    });
-  },
 });
 
-export const { canvasSessionReset, canvasQueueItemDiscarded } = canvasStagingAreaSlice.actions;
+export const { canvasSessionReset, canvasQueueItemDiscarded } = canvasSessionFragment.actions;
+
+const isCanvasSessionAction = isAnyOf(...Object.values(canvasSessionFragment.actions));
+
+export const canvasSessionReducer = (state: CanvasStagingAreaState | undefined, action: UnknownAction): CanvasStagingAreaState => {
+    state = canvasStagingAreaSlice.reducer(state, action);
+
+  if (!isCanvasSessionAction(action)) {
+    return state;
+  }
+
+  const canvasId = action.payload.canvasId;
+
+  return {
+    ...state,
+    sessions: {
+      ...state.sessions,
+      [canvasId]: canvasSessionFragment.reducer(state.sessions[canvasId], action),
+    },
+  };
+};
 
 export const canvasSessionSliceConfig: SliceConfig<typeof canvasStagingAreaSlice> = {
   slice: canvasStagingAreaSlice,
   schema: zCanvasStagingAreaState,
-  getInitialState,
+  getInitialState: getInitialCanvasStagingAreaState,
   persistConfig: {
     migrate: (state) => {
       assert(isPlainObject(state));
@@ -117,7 +141,7 @@ export const canvasSessionSliceConfig: SliceConfig<typeof canvasStagingAreaSlice
         const session = {
           canvasId: MIGRATION_MULTI_CANVAS_ID_PLACEHOLDER,
           ...state,
-        } as CanvasSessionState;
+        } as CanvasSession;
 
         state = {
           _version: 2,
@@ -130,7 +154,7 @@ export const canvasSessionSliceConfig: SliceConfig<typeof canvasStagingAreaSlice
   },
 };
 
-const findSessionByCanvasId = (sessions: Record<string, CanvasSessionState>, canvasId: string) => {
+const findSessionByCanvasId = (sessions: Record<string, CanvasSession>, canvasId: string) => {
   const session = sessions[canvasId];
   assert(session, 'Session must exist for a canvas once the canvas has been created');
   return session;