fix(ui): canvas staging area & batch handling fixes

Handful of intertwined fixes.

- Create and use helper function to reset staging area.
- Clear staging area when queue items are canceled, failed, cleared, etc. Fixes a bug where the bbox ends up offset and images are put into the wrong spot.
- Fix a number of similar bugs where canvas would "forget" it had pending generations, but they continued to generate. Canvas needs to track batches that should be displayed in it using `state.canvas.batchIds`, and this was getting cleared without actually canceling those batches.
- Disable the `discard current image` button on canvas if there is only one image. Prevents accidentally canceling all canvas batches if you spam the button.

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

@@ -1,12 +1,18 @@
 import { isAnyOf } from '@reduxjs/toolkit';
 import { logger } from 'app/logging/logger';
 import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
-import { canvasBatchIdsReset, commitStagingAreaImage, discardStagedImages } from 'features/canvas/store/canvasSlice';
+import {
+  canvasBatchIdsReset,
+  commitStagingAreaImage,
+  discardStagedImages,
+  resetCanvas,
+  setInitialCanvasImage,
+} from 'features/canvas/store/canvasSlice';
 import { addToast } from 'features/system/store/systemSlice';
 import { t } from 'i18next';
 import { queueApi } from 'services/api/endpoints/queue';
 
-const matcher = isAnyOf(commitStagingAreaImage, discardStagedImages);
+const matcher = isAnyOf(commitStagingAreaImage, discardStagedImages, resetCanvas, setInitialCanvasImage);
 
 export const addCommitStagingAreaImageListener = (startAppListening: AppStartListening) => {
   startAppListening({

invokeai/frontend/web/src/features/canvas/components/IAICanvasStagingAreaToolbar.tsx

@@ -49,14 +49,20 @@ const selector = createMemoizedSelector(selectCanvasSlice, (canvas) => {
 const ClearStagingIntermediatesIconButton = () => {
   const dispatch = useAppDispatch();
   const { t } = useTranslation();
+  const totalStagedImages = useAppSelector((s) => s.canvas.layerState.stagingArea.images.length);
 
   const handleDiscardStagingArea = useCallback(() => {
     dispatch(discardStagedImages());
   }, [dispatch]);
 
   const handleDiscardStagingImage = useCallback(() => {
-    dispatch(discardStagedImage());
-  }, [dispatch]);
+    // Discarding all staged images triggers cancelation of all canvas batches. It's too easy to accidentally
+    // click the discard button, so to prevent accidental cancelation of all batches, we only discard the current
+    // image if there are more than one staged images.
+    if (totalStagedImages > 1) {
+      dispatch(discardStagedImage());
+    }
+  }, [dispatch, totalStagedImages]);
 
   return (
     <>
@@ -67,6 +73,7 @@ const ClearStagingIntermediatesIconButton = () => {
         onClick={handleDiscardStagingImage}
         colorScheme="invokeBlue"
         fontSize={16}
+        isDisabled={totalStagedImages <= 1}
       />
       <IconButton
         tooltip={`${t('unifiedCanvas.discardAll')} (Esc)`}

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

@@ -190,7 +190,6 @@ export const canvasSlice = createSlice({
           ],
         };
         state.futureLayerStates = [];
-        state.batchIds = [];
 
         const newScale = calculateScale(
           stageDimensions.width,
@@ -286,40 +285,14 @@ export const canvasSlice = createSlice({
     },
     discardStagedImages: (state) => {
       pushToPrevLayerStates(state);
-
-      state.layerState.stagingArea = deepClone(initialLayerState.stagingArea);
-
+      resetStagingArea(state);
       state.futureLayerStates = [];
-      state.shouldShowStagingOutline = true;
-      state.shouldShowStagingImage = true;
-      state.batchIds = [];
     },
     discardStagedImage: (state) => {
       const { images, selectedImageIndex } = state.layerState.stagingArea;
       pushToPrevLayerStates(state);
-
       images.splice(selectedImageIndex, 1);
-
-      if (images.length === 0) {
-        pushToPrevLayerStates(state);
-
-        state.layerState.stagingArea = deepClone(initialLayerState.stagingArea);
-
-        state.futureLayerStates = [];
-        state.shouldShowStagingOutline = true;
-        state.shouldShowStagingImage = true;
-        state.batchIds = [];
-      }
-
-      if (selectedImageIndex >= images.length) {
-        state.layerState.stagingArea.selectedImageIndex = images.length - 1;
-      }
-
-      if (!images.length) {
-        state.shouldShowStagingImage = false;
-        state.shouldShowStagingOutline = false;
-      }
-
+      state.layerState.stagingArea.selectedImageIndex = Math.max(0, images.length - 1);
       state.futureLayerStates = [];
     },
     addFillRect: (state) => {
@@ -433,7 +406,6 @@ export const canvasSlice = createSlice({
       pushToPrevLayerStates(state);
       state.layerState = deepClone(initialLayerState);
       state.futureLayerStates = [];
-      state.batchIds = [];
       state.boundingBoxCoordinates = {
         ...initialCanvasState.boundingBoxCoordinates,
       };
@@ -534,12 +506,9 @@ export const canvasSlice = createSlice({
           ...imageToCommit,
         });
       }
-      state.layerState.stagingArea = deepClone(initialLayerState.stagingArea);
 
+      resetStagingArea(state);
       state.futureLayerStates = [];
-      state.shouldShowStagingOutline = true;
-      state.shouldShowStagingImage = true;
-      state.batchIds = [];
     },
     setBoundingBoxScaleMethod: {
       reducer: (state, action: PayloadActionWithOptimalDimension<BoundingBoxScaleMethod>) => {
@@ -647,12 +616,19 @@ export const canvasSlice = createSlice({
       if (batch_status.in_progress === 0 && batch_status.pending === 0) {
         state.batchIds = state.batchIds.filter((id) => id !== batch_status.batch_id);
       }
+
+      const queueItemStatus = action.payload.data.queue_item.status;
+      if (queueItemStatus === 'canceled' || queueItemStatus === 'failed') {
+        resetStagingAreaIfEmpty(state);
+      }
     });
     builder.addMatcher(queueApi.endpoints.clearQueue.matchFulfilled, (state) => {
       state.batchIds = [];
+      resetStagingAreaIfEmpty(state);
     });
     builder.addMatcher(queueApi.endpoints.cancelByBatchIds.matchFulfilled, (state, action) => {
       state.batchIds = state.batchIds.filter((id) => !action.meta.arg.originalArgs.batch_ids.includes(id));
+      resetStagingAreaIfEmpty(state);
     });
   },
 });
@@ -726,7 +702,7 @@ export const canvasPersistConfig: PersistConfig<CanvasState> = {
   name: canvasSlice.name,
   initialState: initialCanvasState,
   migrate: migrateCanvasState,
-  persistDenylist: [],
+  persistDenylist: ['shouldShowStagingImage', 'shouldShowStagingOutline'],
 };
 
 const pushToPrevLayerStates = (state: CanvasState) => {
@@ -742,3 +718,15 @@ const pushToFutureLayerStates = (state: CanvasState) => {
     state.futureLayerStates = state.futureLayerStates.slice(0, MAX_HISTORY);
   }
 };
+
+const resetStagingAreaIfEmpty = (state: CanvasState) => {
+  if (state.batchIds.length === 0 && state.layerState.stagingArea.images.length === 0) {
+    resetStagingArea(state);
+  }
+};
+
+const resetStagingArea = (state: CanvasState) => {
+  state.layerState.stagingArea = { ...initialCanvasState.layerState.stagingArea };
+  state.shouldShowStagingImage = initialCanvasState.shouldShowStagingImage;
+  state.shouldShowStagingOutline = initialCanvasState.shouldShowStagingOutline;
+};