feat(ui): tabbed canvases

[csehatt741]

Summary

A new feature was implemented to enable multiple canvases in the canvas editor.

• SliceConfig was updated in types.js to handle partially undoable slices:
  • new generic parametes, TInternalState and TSerializedState, were added to strongly type states used in Redux and the persisted into storage
  • new functions, wrapState and unwrapState, were added to the migrate field for wrapping/unwrapping state during serialization/deserialization
  • undoableConfig was deleted, as reduxUndoOptions is used only in slices, so this field became redundant
• store.ts was refactored to use the updated SliceConfig
• new schemas, zStateWithHistory, zCanvasStateWithHistory, zCanvasesState, zCanvasesStateWithHistory and zCanvasesStateWithoutHistory, were created in types.ts to represent types in the partially undoable canvases slice
• new selectors, selectCanvases selectSelectedCanvasWithHistory and selectSelectedCanvas, were added to selectors.ts to isolate changes due to refactoring in the canvas slice from components
• canvasSlice was split into two parts representing the canvases slice without history and the canvas with undoable history
• undoableCanvasesReducer higher order reducer was created to combine canvases and canvas reducers
• RTK was updated to version 2.9.0

Related Issues / Discussions

Closes #8380

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable)
• ❗Changes to a redux slice have a corresponding migration
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)

[csehatt741]

@psychedelicious, here is the current status of multi-canvas development and refactoring:

• to limit the scope of refactoring I still kept the notion of the active canvas
  • it always provides scope for any components even not in CanvasInstanceContext, without refactoring the component hierarchy
  • functionality can still be improved further in the future
• canvasSlice:
  • state is split into per canvas instance state
  • undo/redo per canvas implemented
• canvasStagingAreaSlice state is split into per canvas instance state
• canvasSettingsSlice state is split into shared and per canvas instance state
• paramsSlice state is split into shared and per canvas instance state

Outstanding tasks:

• revise code where selectSelectedCanvas selector is used
  • there might be context that overrides the default selected canvas one
  • functionality might use all the canvases now, eg.: state.getImageUsage

[csehatt741]

Since there is always only one canvas rendered, the active one, so the CanvasManager is created and initialized when the canvas is displayed and destroyed when the canvas becomes inactive.
Though, I haven't seen any issues caused by this lifecycle of the CanvasManager yet, but I expect this might lead to problems:

• what is the effect of destroying a CanvasManager when the canvas becomes inactive and then reusing it when the canvas goes active again: https://github.com/invoke-ai/InvokeAI/pull/8553/files#diff-59388fa7d596ff6926d86545074a35049822341021aa0fb6df10e5a6f512ce12R48
• how about queued and pending tasks when the UI is reopened? Do they expect the CanvasManager is available when they complete?

[psychedelicious]

I'm getting a UI crash when starting from a fresh install (presumably the same would happen from fresh UI state).

Uncaught Error: Wrong assertion encountered: "Session must exist for a canvas once the canvas has been created"

I logged the canvasId in useCanvasSessionId(), and see that a CanvasState with this ID is created in redux when I check devtools.

---

There are a few items that jump out at me as I start a high-level review.

1. Each canvas instance's state is split across 4 slices (canvas, canvasSettings, canvasSession and params). This feels fragile and hard to reason about. We need to coordinate between these slices. We'll need to be very defensive to prevent race conditions (I suspect the crash I'm running into is one such race condition).

   I'm guessing the motivation for this design is to reduce the amount of code change. But if we are introducing a new major core concept like multiple workspace instances, it seems appropriate to change the architecture to support it as simply as possible. The coordination seems simpler on the surface but I think it makes the design more complicated.

   I wonder if consolidating the slices could work better. We could use nested reducers to scope undo functionality to only the parts of state that need it (e.g. layers).

2. There are now two dispatch hooks (useAppDispatch() and useParamsDispatch()). The developer now needs to understand the app state intimately to know which to use. Also, it's possible that some user actions might need access to both the dispatchers.

   Can we instead update useAppDispatch() to include the logic from useParamsDispatch()? That way, contributors don't need to worry about which to use.

3. I'm not seeing where canvas actions, like adding a layer, are being handled. I suppose this is WIP?

[psychedelicious]

After more contemplation and some experimentation, I think consolidating the slices into canvas instances is the right move. That said, I'm certainly open to discussion and being wrong about this!

Here's a draft proposal for consolidated canvas state. I think it covers most of the required touch points, except interacting with the CanvasManager class.

Consolidated root canvas state and instance state structure:

// this gets added to root store
type RootCanvasState = {
  _version: number;
  activeCanvasId: string;
  canvases: Record<string, CanvasInstance>;
};

// these instance-scoped slices do not get added to root store
type CanvasInstance = {
  id: string;
  name: string;
  canvas: StateWithHistory<CanvasEntityState>; // previously `CanvasState`
  params: ParamsState;
  settings: SettingsState;
  staging: StagingState;
};

Utilities for working with it:

const isCanvasInstanceAction = isAnyOf(...Object.values(canvasSlice.actions).map((a) => a.match));

// use a Symbol to prevent key collisions
const CANVAS_ID_META_KEY = Symbol('canvasId');

const injectCanvasId = (canvasId: string, action: CanvasInstanceAction) => {
  // mutates the action - could inject into payload instead
  Object.assign(action, { meta: { [CANVAS_ID_META_KEY]: canvasId } });
}

const extractCanvasId = (action: CanvasInstanceAction): string | undefined => {
  return action?.meta?.[CANVAS_ID_META_KEY];
}

Use react context to provide current canvasId:

const CanvasIdContext = createContext<string | null>(null);
const useCanvasIdContextSafe = () => useContext(CanvasIdContext);
const useCanvasIdContextRequired = () => {
  const canvasId = useCanvasIdContext();
  if (canvasId === null) {
    throw new Error("useCanvasIdContextRequired must be used within a CanvasIdProvider");
  }
  return canvasId;
}

Wrapper around useDispatch to auto-inject canvasId (could also be middleware I suppose):

const useAppDispatch = () => {
  const _dispatch = useDispatch();
  const canvasId = useCanvasIdContextSafe() // string | null
  return useCallback((action: UnknownAction) => {
    if (canvasId !== null && isCanvasInstanceAction(action)) {
      injectCanvasId(canvasId, action)
    }

    dispatch(action)
  })
}

Wrap the undoable slice reducers:

const canvasEntitySlice = createSlice(...)
const undoableConfig = { /* as before */ }
export const undoableCanvasEntityReducer = undoable(canvasEntitySlice.reducer, undoableConfig)

Use an "action router" pattern in the root canvas reducer to route actions to the correct canvas instance:

extraReducers: (builder) => {

  builder.addMatcher(isCanvasInstanceAction, (state, action) => {
    const canvasId = extractCanvasId(action);
    if (!canvasId) {
      return;
    }

    const instance = state.canvases[canvasId];
    if (!instance) {
      return;
    }

    // delegate to sub-reducers (each handles own undo/redo)
    if (isCanvasEntityAction(action)) {
      instance.canvas = undoableCanvasEntityReducer(instance.canvas, action);
    }
    if (isParamsAction(action)) {
      instance.params = paramsReducer(instance.params, action);
    }
    // etc
  });
}

To handle key-specific persist denylists, we could update PersistConfig to be something like this, instead of having a denylist and methods to wrap history:

type PersistConfig<TInternal, TSerialized> = {
  serialize?: (state: TInternal) => TSerialized;
  deserialize?: (state: TSerialized) => TInternal;
}

const wrapHistory = <T,>(state: T): StateWithHistory<T> => ({ /* ... */ })
const unwrapHistory = <T,>(state: StateWithHistory<T>): T => ({ /* ... */ })

const canvasEntityPersistConfig: PersistConfig<StateWithHistory<CanvasEntityState>, CanvasEntityState> = {
  migrate: (state) => zConsolidatedCanvasState.parse(state),

  serialize: (state) => ({
    _version: state._version,
    activeCanvasId: state.activeCanvasId,
    canvases: Object.fromEntries(
      Object.entries(state.canvases).map(([id, instance]) => [
        id,
        {
          ...instance,
          canvas: unwrapHistory(instance.canvas)
          params: omit(instance.params, ['sensitiveKey']) // key-specific denylist
          staging: undefined, // don't persist staging (can't remember if this is needed in the actual codebase but you get the idea)
        }
      ])
    ),
  }),

  deserialize: (state, initialState) => ({
    _version: state._version,
    activeCanvasId: state.activeCanvasId,
    canvases: Object.fromEntries(
      Object.entries(state.canvases).map(([id, instance]) => [
        id,
        {
          ...instance,
          canvas: wrapHistory(instance.canvas),
          params: merge(instance.params, defaults), // inject defaults
          staging: getInitialStagingState(id), // re-initialize staging
        }
      ])
    ),
  }),
}

In components, to select state, so long as we use createSelector, we'll get weakmap-memoized selector caches for free:

export const selectCanvasInstance = createSelector(
  [
    (state: RootState) => state.canvas.canvases,
    (state: RootState, canvasId: string) => canvasId
  ],
  (canvases, canvasId) => canvases[canvasId]
);

export const selectCanvasSettings = createSelector(
  [selectCanvasInstance],
  (instance) => instance.settings
);

// because this is composed w/ the other selectors, it will have a second arg for canvasId
export const selectBrushWidth = createSelector(
  [selectCanvasSettings],
  (settings) => settings.brushWidth
);

// in the component
const canvasId = useCanvasIdContextRequired();
const brushWidth = useAppSelector((state) => selectBrushWidth(state, canvasId));

[csehatt741]

@psychedelicious
Actually I've explicitly invested into migration, but even did not try it out with a clean install before.
The root cause was missing canvas states in auxiliary canvas slices when the canvas is initialized the first time.
Issue has been fixed and PR updated.

1. I just also had the same idea, as I was making progress with this feature and had to cope with multiple related issues, so wanted to discuss merging related slices into canvasSlice.
   The issue above was definitely related to this problem too.

   The reason why I have not combined these slices yet are, because I want to reduce the amount of code change and I just realized the complexity of issues caused by keeping the current design yesterday. I did not have time to think it over and discuss earlier.

   I'll investigate this issue and will get back soon!

2. I chose not to implicitly encapsulate context information in useAppDispatch to give developers the option to explicitly deal with other contexts different from the active one, eg.: to dispatch actions to any canvases. Maybe it's unnecessary, as there's always only one canvas rendered.

3. I've tried out invoking the canvas instance reducer in extraReducers as suggested, but undo/redo did not work there. It might be related to proxies within reducers. So I explicitly created the undoableCanvasesReducer invoking canvasesSlice and undoableCanvasReducer there.

[psychedelicious]

Actually I've explicitly invested into migration, but even did not try it out with a clean install before.

No worries, I understand. Thinking about migrations... Currently, they are exclusively intra-slice, which means we have to jump through hoops to do inter-slice migration (as you found out). I'm not sure if those hoops are worth it 🤔. A redux action that we dispatch exactly once just doesn't feel quite right.

I think for now, we should set aside data migration. Once we have everything working correctly, we can figure out a pattern for inter-slice migration.

I chose not to implicitly encapsulate context information in useAppDispatch to give developers the option to explicitly deal with other contexts different from the active one, eg.: to dispatch actions to any canvases. Maybe it's unnecessary, as there's always only one canvas rendered.

Maybe I am misunderstanding, but it looks like the useParamsDispatch() and paramsDispatch() utilities have the same behaviour where they implicitly choose where to dispatch actions.

"Dispatch to the active canvas" will be the typical case, while "dispatch to some other canvas" will be the special case. Suggest to prioritize the DX for the typical case over the special case; docstrings can clarify the behaviour.

For the special case, the proposed useAppDispatch() wrapper could check for the existence of the canvas ID key before injecting the canvas ID. Then callers can inject the canvas ID before dispatching to trivially override the implicit behaviour.

I've tried out invoking the canvas instance reducer in extraReducers as suggested, but undo/redo did not work there. It might be related to proxies within reducers. So I explicitly created the undoableCanvasesReducer invoking canvasesSlice and undoableCanvasReducer there.

That's unexpected. It worked for me when I experimented with this pattern a while ago. My understanding of RTK is that extraReducers is the "official" way to do this kinda thing, preferred over manual reducer composition. That said, I don't have a problem with how you set up the undoable reducer.

PS - Great work on this. Its gonna be a huge improvement to the app!