feat(ui): revised canvas progress & staging image handling

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/index.ts

@@ -35,7 +35,7 @@ import { addInvocationErrorEventListener } from 'app/store/middleware/listenerMi
 import { addInvocationStartedEventListener } from 'app/store/middleware/listenerMiddleware/listeners/socketio/socketInvocationStarted';
 import { addModelInstallEventListener } from 'app/store/middleware/listenerMiddleware/listeners/socketio/socketModelInstall';
 import { addModelLoadEventListener } from 'app/store/middleware/listenerMiddleware/listeners/socketio/socketModelLoad';
-import { addSocketQueueItemStatusChangedEventListener } from 'app/store/middleware/listenerMiddleware/listeners/socketio/socketQueueItemStatusChanged';
+import { addSocketQueueEventsListeners } from 'app/store/middleware/listenerMiddleware/listeners/socketio/socketQueueEvents';
 import { addUpdateAllNodesRequestedListener } from 'app/store/middleware/listenerMiddleware/listeners/updateAllNodesRequested';
 import { addWorkflowLoadRequestedListener } from 'app/store/middleware/listenerMiddleware/listeners/workflowLoadRequested';
 import type { AppDispatch, RootState } from 'app/store/store';
@@ -99,7 +99,7 @@ addSocketConnectedEventListener(startAppListening);
 addSocketDisconnectedEventListener(startAppListening);
 addModelLoadEventListener(startAppListening);
 addModelInstallEventListener(startAppListening);
-addSocketQueueItemStatusChangedEventListener(startAppListening);
+addSocketQueueEventsListeners(startAppListening);
 addBulkDownloadListeners(startAppListening);
 
 // Boards

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

@@ -1,10 +1,11 @@
+import { isAnyOf } from '@reduxjs/toolkit';
 import { logger } from 'app/logging/logger';
 import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
 import {
   layerAdded,
   layerImageAdded,
+  stagingAreaCanceledStaging,
   stagingAreaImageAccepted,
-  stagingAreaReset,
 } from 'features/controlLayers/store/canvasV2Slice';
 import { toast } from 'features/toast/toast';
 import { t } from 'i18next';
@@ -13,25 +14,15 @@ import { assert } from 'tsafe';
 
 export const addStagingListeners = (startAppListening: AppStartListening) => {
   startAppListening({
-    actionCreator: stagingAreaReset,
-    effect: async (_, { dispatch, getState }) => {
+    matcher: isAnyOf(stagingAreaCanceledStaging, stagingAreaImageAccepted),
+    effect: async (_, { dispatch }) => {
       const log = logger('canvas');
-      const stagingArea = getState().canvasV2.stagingArea;
-
-      if (!stagingArea) {
-        // Should not happen
-        return;
-      }
-
-      if (stagingArea.batchIds.length === 0) {
-        return;
-      }
 
       try {
         const req = dispatch(
-          queueApi.endpoints.cancelByBatchIds.initiate(
-            { batch_ids: stagingArea.batchIds },
-            { fixedCacheKey: 'cancelByBatchIds' }
+          queueApi.endpoints.cancelByBatchOrigin.initiate(
+            { origin: 'canvas' },
+            { fixedCacheKey: 'cancelByBatchOrigin' }
           )
         );
         const { canceled } = await req.unwrap();
@@ -59,7 +50,7 @@ export const addStagingListeners = (startAppListening: AppStartListening) => {
     actionCreator: stagingAreaImageAccepted,
     effect: async (action, api) => {
       const { imageDTO } = action.payload;
-      const { layers, stagingArea, selectedEntityIdentifier } = api.getState().canvasV2;
+      const { layers, selectedEntityIdentifier, bbox } = api.getState().canvasV2;
       let layer = layers.entities.find((layer) => layer.id === selectedEntityIdentifier?.id);
 
       if (!layer) {
@@ -73,13 +64,11 @@ export const addStagingListeners = (startAppListening: AppStartListening) => {
       }
 
       assert(layer, 'No layer found to stage image');
-      assert(stagingArea, 'Staging should be defined');
 
-      const { x, y } = stagingArea.bbox;
+      const { x, y } = bbox;
       const { id } = layer;
 
       api.dispatch(layerImageAdded({ id, imageDTO, pos: { x, y } }));
-      api.dispatch(stagingAreaReset());
     },
   });
 };

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

@@ -1,12 +1,7 @@
 import { enqueueRequested } from 'app/store/actions';
 import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
 import { getNodeManager } from 'features/controlLayers/konva/nodeManager';
-import {
-  stagingAreaBatchIdAdded,
-  stagingAreaInitialized,
-  stagingAreaReset,
-} from 'features/controlLayers/store/canvasV2Slice';
-import { isImageViewerOpenChanged } from 'features/gallery/store/gallerySlice';
+import { stagingAreaCanceledStaging, stagingAreaStartedStaging } from 'features/controlLayers/store/canvasV2Slice';
 import { prepareLinearUIBatch } from 'features/nodes/util/graph/buildLinearBatchConfig';
 import { buildSD1Graph } from 'features/nodes/util/graph/generation/buildSD1Graph';
 import { buildSDXLGraph } from 'features/nodes/util/graph/generation/buildSDXLGraph';
@@ -19,20 +14,13 @@ export const addEnqueueRequestedLinear = (startAppListening: AppStartListening)
       enqueueRequested.match(action) && action.payload.tabName === 'generation',
     effect: async (action, { getState, dispatch }) => {
       const state = getState();
-      const { shouldShowProgressInViewer } = state.ui;
       const model = state.canvasV2.params.model;
       const { prepend } = action.payload;
 
-      let didInitializeStagingArea = false;
-
-      if (state.canvasV2.stagingArea === null) {
-        dispatch(
-          stagingAreaInitialized({
-            batchIds: [],
-            bbox: state.canvasV2.bbox,
-          })
-        );
-        didInitializeStagingArea = true;
+      let didStartStaging = false;
+      if (!state.canvasV2.stagingArea.isStaging) {
+        dispatch(stagingAreaStartedStaging());
+        didStartStaging = true;
       }
 
       try {
@@ -57,23 +45,11 @@ export const addEnqueueRequestedLinear = (startAppListening: AppStartListening)
             fixedCacheKey: 'enqueueBatch',
           })
         );
-
-        const enqueueResult = await req.unwrap();
         req.reset();
-
-        if (shouldShowProgressInViewer) {
-          dispatch(isImageViewerOpenChanged(true));
-        }
-        // TODO(psyche): update the backend schema, this is always provided
-        const batchId = enqueueResult.batch.batch_id;
-        assert(batchId, 'No batch ID found in enqueue result');
-        dispatch(stagingAreaBatchIdAdded({ batchId }));
+        await req.unwrap();
       } catch {
-        if (didInitializeStagingArea) {
-          // We initialized the staging area in this listener, and there was a problem at some point. This means
-          // there only possible canvas batch id is the one we just added, so we can reset the staging area without
-          // losing any data.
-          dispatch(stagingAreaReset());
+        if (didStartStaging && getState().canvasV2.stagingArea.isStaging) {
+          dispatch(stagingAreaCanceledStaging());
         }
       }
     },

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

@@ -30,6 +30,7 @@ export const addEnqueueRequestedNodes = (startAppListening: AppStartListening) =
           graph,
           workflow: builtWorkflow,
           runs: state.canvasV2.params.iterations,
+          origin: 'workflows',
         },
         prepend: action.payload.prepend,
       };

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

@@ -12,19 +12,21 @@ const log = logger('socketio');
 export const addGeneratorProgressEventListener = (startAppListening: AppStartListening) => {
   startAppListening({
     actionCreator: socketGeneratorProgress,
-    effect: (action, { getState }) => {
+    effect: (action) => {
       log.trace(parseify(action.payload), `Generator progress`);
-      const { invocation_source_id, step, total_steps, progress_image, batch_id } = action.payload.data;
-      const nes = deepClone($nodeExecutionStates.get()[invocation_source_id]);
-      if (nes) {
-        nes.status = zNodeStatus.enum.IN_PROGRESS;
-        nes.progress = (step + 1) / total_steps;
-        nes.progressImage = progress_image ?? null;
-        upsertExecutionState(nes.nodeId, nes);
+      const { invocation_source_id, step, total_steps, progress_image, origin } = action.payload.data;
+
+      if (origin === 'workflows') {
+        const nes = deepClone($nodeExecutionStates.get()[invocation_source_id]);
+        if (nes) {
+          nes.status = zNodeStatus.enum.IN_PROGRESS;
+          nes.progress = (step + 1) / total_steps;
+          nes.progressImage = progress_image ?? null;
+          upsertExecutionState(nes.nodeId, nes);
+        }
       }
 
-      const isCanvasQueueItem = getState().canvasV2.stagingArea?.batchIds.includes(batch_id);
-      if (isCanvasQueueItem) {
+      if (origin === 'canvas') {
         $lastProgressEvent.set(action.payload.data);
       }
     },

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

@@ -3,13 +3,7 @@ import type { AppStartListening } from 'app/store/middleware/listenerMiddleware'
 import { deepClone } from 'common/util/deepClone';
 import { parseify } from 'common/util/serialize';
 import { stagingAreaImageAdded } from 'features/controlLayers/store/canvasV2Slice';
-import {
-  boardIdSelected,
-  galleryViewChanged,
-  imageSelected,
-  isImageViewerOpenChanged,
-  offsetChanged,
-} from 'features/gallery/store/gallerySlice';
+import { boardIdSelected, galleryViewChanged, imageSelected, offsetChanged } from 'features/gallery/store/gallerySlice';
 import { $nodeExecutionStates, upsertExecutionState } from 'features/nodes/hooks/useExecutionState';
 import { zNodeStatus } from 'features/nodes/types/invocation';
 import { CANVAS_OUTPUT } from 'features/nodes/util/graph/constants';
@@ -17,7 +11,6 @@ import { boardsApi } from 'services/api/endpoints/boards';
 import { imagesApi } from 'services/api/endpoints/images';
 import { getCategories, getListImagesUrl } from 'services/api/util';
 import { socketInvocationComplete } from 'services/events/actions';
-import { assert } from 'tsafe';
 
 // These nodes output an image, but do not actually *save* an image, so we don't want to handle the gallery logic on them
 const nodeTypeDenylist = ['load_image', 'image'];
@@ -35,7 +28,7 @@ export const addInvocationCompleteEventListener = (startAppListening: AppStartLi
       // This complete event has an associated image output
       if (data.result.type === 'image_output' && !nodeTypeDenylist.includes(data.invocation.type)) {
         const { image_name } = data.result.image;
-        const { canvasV2, gallery } = getState();
+        const { gallery, canvasV2 } = getState();
 
         // This populates the `getImageDTO` cache
         const imageDTORequest = dispatch(
@@ -47,11 +40,21 @@ export const addInvocationCompleteEventListener = (startAppListening: AppStartLi
         const imageDTO = await imageDTORequest.unwrap();
         imageDTORequest.unsubscribe();
 
-        // Add canvas images to the staging area
-        if (canvasV2.stagingArea?.batchIds.includes(data.batch_id) && data.invocation_source_id === CANVAS_OUTPUT) {
-          const stagingArea = getState().canvasV2.stagingArea;
-          assert(stagingArea, 'Staging should be defined');
-          dispatch(stagingAreaImageAdded({ imageDTO }));
+        // handle tab-specific logic
+        if (data.origin === 'canvas') {
+          if (data.invocation_source_id === CANVAS_OUTPUT && canvasV2.stagingArea.isStaging) {
+            dispatch(stagingAreaImageAdded({ imageDTO }));
+          }
+        } else if (data.origin === 'workflows') {
+          const nes = deepClone($nodeExecutionStates.get()[invocation_source_id]);
+          if (nes) {
+            nes.status = zNodeStatus.enum.COMPLETED;
+            if (nes.progress !== null) {
+              nes.progress = 1;
+            }
+            nes.outputs.push(result);
+            upsertExecutionState(nes.nodeId, nes);
+          }
         }
 
         if (!imageDTO.is_intermediate) {
@@ -106,20 +109,9 @@ export const addInvocationCompleteEventListener = (startAppListening: AppStartLi
             }
 
             dispatch(imageSelected(imageDTO));
-            dispatch(isImageViewerOpenChanged(true));
           }
         }
       }
-
-      const nes = deepClone($nodeExecutionStates.get()[invocation_source_id]);
-      if (nes) {
-        nes.status = zNodeStatus.enum.COMPLETED;
-        if (nes.progress !== null) {
-          nes.progress = 1;
-        }
-        nes.outputs.push(result);
-        upsertExecutionState(nes.nodeId, nes);
-      }
     },
   });
 };

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketQueueEvents.tsx

@@ -12,7 +12,15 @@ import { socketQueueItemStatusChanged } from 'services/events/actions';
 
 const log = logger('socketio');
 
-export const addSocketQueueItemStatusChangedEventListener = (startAppListening: AppStartListening) => {
+export const addSocketQueueEventsListeners = (startAppListening: AppStartListening) => {
+  // When the queue is cleared or canvas batch is canceled, we should clear the last canvas progress event
+  startAppListening({
+    matcher: queueApi.endpoints.clearQueue.matchFulfilled,
+    effect: () => {
+      $lastProgressEvent.set(null);
+    },
+  });
+
   startAppListening({
     actionCreator: socketQueueItemStatusChanged,
     effect: async (action, { dispatch, getState }) => {
@@ -29,13 +37,11 @@ export const addSocketQueueItemStatusChangedEventListener = (startAppListening:
         error_type,
         error_message,
         error_traceback,
-        batch_id,
+        origin,
       } = action.payload.data;
 
       log.debug(action.payload, `Queue item ${item_id} status updated: ${status}`);
 
-      const isCanvasQueueItem = getState().canvasV2.stagingArea?.batchIds.includes(batch_id);
-
       // Update this specific queue item in the list of queue items (this is the queue item DTO, without the session)
       dispatch(
         queueApi.util.updateQueryData('listQueueItems', undefined, (draft) => {
@@ -96,7 +102,7 @@ export const addSocketQueueItemStatusChangedEventListener = (startAppListening:
       } else if (status === 'failed' && error_type) {
         const isLocal = getState().config.isLocal ?? true;
         const sessionId = session_id;
-        if (isCanvasQueueItem) {
+        if (origin === 'canvas') {
           $lastProgressEvent.set(null);
         }
 
@@ -115,9 +121,7 @@ export const addSocketQueueItemStatusChangedEventListener = (startAppListening:
             />
           ),
         });
-      } else if (status === 'completed' && isCanvasQueueItem) {
-        $lastProgressEvent.set(null);
-      } else if (status === 'canceled' && isCanvasQueueItem) {
+      } else if (status === 'canceled' && origin === 'canvas') {
         $lastProgressEvent.set(null);
       }
     },

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

@@ -3,13 +3,12 @@ import { useStore } from '@nanostores/react';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import {
   $shouldShowStagedImage,
+  stagingAreaCanceledStaging,
   stagingAreaImageAccepted,
   stagingAreaImageDiscarded,
   stagingAreaNextImageSelected,
   stagingAreaPreviousImageSelected,
-  stagingAreaReset,
 } from 'features/controlLayers/store/canvasV2Slice';
-import type { CanvasV2State } from 'features/controlLayers/store/types';
 import { memo, useCallback, useMemo } from 'react';
 import { useHotkeys } from 'react-hotkeys-hook';
 import { useTranslation } from 'react-i18next';
@@ -25,29 +24,23 @@ import {
 } from 'react-icons/pi';
 
 export const StagingAreaToolbar = memo(() => {
-  const stagingArea = useAppSelector((s) => s.canvasV2.stagingArea);
+  const isStaging = useAppSelector((s) => s.canvasV2.stagingArea.isStaging);
 
-  if (!stagingArea || stagingArea.images.length === 0) {
+  if (!isStaging) {
     return null;
   }
 
-  return <StagingAreaToolbarContent stagingArea={stagingArea} />;
+  return <StagingAreaToolbarContent />;
 });
 
 StagingAreaToolbar.displayName = 'StagingAreaToolbar';
 
-type Props = {
-  stagingArea: NonNullable<CanvasV2State['stagingArea']>;
-};
-
-export const StagingAreaToolbarContent = memo(({ stagingArea }: Props) => {
+export const StagingAreaToolbarContent = memo(() => {
   const dispatch = useAppDispatch();
+  const stagingArea = useAppSelector((s) => s.canvasV2.stagingArea);
   const shouldShowStagedImage = useStore($shouldShowStagedImage);
   const images = useMemo(() => stagingArea.images, [stagingArea]);
-  const imageDTO = useMemo(() => {
-    if (stagingArea.selectedImageIndex === null) {
-      return null;
-    }
+  const selectedImageDTO = useMemo(() => {
     return images[stagingArea.selectedImageIndex] ?? null;
   }, [images, stagingArea.selectedImageIndex]);
 
@@ -64,29 +57,26 @@ export const StagingAreaToolbarContent = memo(({ stagingArea }: Props) => {
   }, [dispatch]);
 
   const onAccept = useCallback(() => {
-    if (!imageDTO || !stagingArea) {
+    if (!selectedImageDTO) {
       return;
     }
-    dispatch(stagingAreaImageAccepted({ imageDTO }));
-  }, [dispatch, imageDTO, stagingArea]);
+    dispatch(stagingAreaImageAccepted({ imageDTO: selectedImageDTO }));
+  }, [dispatch, selectedImageDTO]);
 
   const onDiscardOne = useCallback(() => {
-    if (!imageDTO || !stagingArea) {
+    if (!selectedImageDTO) {
       return;
     }
     if (images.length === 1) {
-      dispatch(stagingAreaReset());
+      dispatch(stagingAreaCanceledStaging());
     } else {
-      dispatch(stagingAreaImageDiscarded({ imageDTO }));
+      dispatch(stagingAreaImageDiscarded({ imageDTO: selectedImageDTO }));
     }
-  }, [dispatch, imageDTO, images.length, stagingArea]);
+  }, [dispatch, selectedImageDTO, images.length]);
 
   const onDiscardAll = useCallback(() => {
-    if (!stagingArea) {
-      return;
-    }
-    dispatch(stagingAreaReset());
-  }, [dispatch, stagingArea]);
+    dispatch(stagingAreaCanceledStaging());
+  }, [dispatch]);
 
   const onToggleShouldShowStagedImage = useCallback(() => {
     $shouldShowStagedImage.set(!shouldShowStagedImage);
@@ -117,6 +107,14 @@ export const StagingAreaToolbarContent = memo(({ stagingArea }: Props) => {
     preventDefault: true,
   });
 
+  const counterText = useMemo(() => {
+    if (images.length > 0) {
+      return `${(stagingArea.selectedImageIndex ?? 0) + 1} of ${images.length}`;
+    } else {
+      return `0 of 0`;
+    }
+  }, [images.length, stagingArea.selectedImageIndex]);
+
   return (
     <>
       <ButtonGroup borderRadius="base" shadow="dark-lg">
@@ -128,11 +126,9 @@ export const StagingAreaToolbarContent = memo(({ stagingArea }: Props) => {
           colorScheme="invokeBlue"
           isDisabled={images.length <= 1 || !shouldShowStagedImage}
         />
-        <Button
-          colorScheme="base"
-          pointerEvents="none"
-          minW={20}
-        >{`${(stagingArea.selectedImageIndex ?? 0) + 1}/${images.length}`}</Button>
+        <Button colorScheme="base" pointerEvents="none" minW={28}>
+          {counterText}
+        </Button>
         <IconButton
           tooltip={`${t('unifiedCanvas.next')} (Right)`}
           aria-label={`${t('unifiedCanvas.next')} (Right)`}
@@ -149,6 +145,7 @@ export const StagingAreaToolbarContent = memo(({ stagingArea }: Props) => {
           icon={<PiCheckBold />}
           onClick={onAccept}
           colorScheme="invokeBlue"
+          isDisabled={!selectedImageDTO}
         />
         <IconButton
           tooltip={shouldShowStagedImage ? t('unifiedCanvas.showResultsOn') : t('unifiedCanvas.showResultsOff')}
@@ -161,10 +158,10 @@ export const StagingAreaToolbarContent = memo(({ stagingArea }: Props) => {
         <IconButton
           tooltip={`${t('unifiedCanvas.saveToGallery')} (Shift+S)`}
           aria-label={t('unifiedCanvas.saveToGallery')}
-          isDisabled={!imageDTO || !imageDTO.is_intermediate}
           icon={<PiFloppyDiskBold />}
           onClick={onSaveStagingImage}
           colorScheme="invokeBlue"
+          isDisabled={!selectedImageDTO || !selectedImageDTO.is_intermediate}
         />
         <IconButton
           tooltip={`${t('unifiedCanvas.discardCurrent')}`}
@@ -173,7 +170,7 @@ export const StagingAreaToolbarContent = memo(({ stagingArea }: Props) => {
           onClick={onDiscardOne}
           colorScheme="invokeBlue"
           fontSize={16}
-          isDisabled={images.length <= 1}
+          isDisabled={!selectedImageDTO}
         />
         <IconButton
           tooltip={`${t('unifiedCanvas.discardAll')} (Esc)`}
@@ -182,7 +179,6 @@ export const StagingAreaToolbarContent = memo(({ stagingArea }: Props) => {
           onClick={onDiscardAll}
           colorScheme="error"
           fontSize={16}
-          isDisabled={images.length === 0}
         />
       </ButtonGroup>
     </>

invokeai/frontend/web/src/features/controlLayers/components/ToolChooser.tsx

@@ -43,7 +43,7 @@ export const ToolChooser: React.FC = () => {
   const { t } = useTranslation();
   const dispatch = useAppDispatch();
   const selectedEntityIdentifier = useAppSelector(selectSelectedEntityIdentifier);
-  const isStaging = useAppSelector((s) => s.canvasV2.stagingArea !== null);
+  const isStaging = useAppSelector((s) => s.canvasV2.stagingArea.isStaging);
   const isDrawingToolDisabled = useMemo(
     () => !getIsDrawingToolEnabled(selectedEntityIdentifier),
     [selectedEntityIdentifier]

invokeai/frontend/web/src/features/controlLayers/konva/nodeManager.ts

@@ -76,6 +76,7 @@ export type StateApi = {
   getInpaintMaskState: () => CanvasV2State['inpaintMask'];
   getStagingAreaState: () => CanvasV2State['stagingArea'];
   getLastProgressEvent: () => InvocationDenoiseProgressEvent | null;
+  resetLastProgressEvent: () => void;
   onInpaintMaskImageCached: (imageDTO: ImageDTO) => void;
   onRegionMaskImageCached: (id: string, imageDTO: ImageDTO) => void;
   onLayerImageCached: (imageDTO: ImageDTO) => void;
@@ -280,8 +281,10 @@ export class KonvaNodeManager {
   renderStagingArea() {
     this.preview.stagingArea.render(
       this.stateApi.getStagingAreaState(),
+      this.stateApi.getBbox(),
       this.stateApi.getShouldShowStagedImage(),
-      this.stateApi.getLastProgressEvent()
+      this.stateApi.getLastProgressEvent(),
+      this.stateApi.resetLastProgressEvent
     );
   }
 

invokeai/frontend/web/src/features/controlLayers/konva/renderers/preview.ts

@@ -18,18 +18,18 @@ export class CanvasPreview {
     documentSizeOverlay: CanvasDocumentSizeOverlay,
     stagingArea: CanvasStagingArea
   ) {
-    this.layer = new Konva.Layer({ listening: true });
-
-    this.bbox = bbox;
-    this.layer.add(this.bbox.group);
-
-    this.tool = tool;
-    this.layer.add(this.tool.group);
+    this.layer = new Konva.Layer({ listening: true, imageSmoothingEnabled: false });
 
     this.documentSizeOverlay = documentSizeOverlay;
     this.layer.add(this.documentSizeOverlay.group);
 
     this.stagingArea = stagingArea;
     this.layer.add(this.stagingArea.group);
+
+    this.bbox = bbox;
+    this.layer.add(this.bbox.group);
+
+    this.tool = tool;
+    this.layer.add(this.tool.group);
   }
 }

invokeai/frontend/web/src/features/controlLayers/konva/renderers/renderer.ts

@@ -307,6 +307,9 @@ export const initializeRenderer = (
     getStagingAreaState,
     getShouldShowStagedImage: $shouldShowStagedImage.get,
     getLastProgressEvent: $lastProgressEvent.get,
+    resetLastProgressEvent: () => {
+      $lastProgressEvent.set(null);
+    },
 
     // Read-write state
     setTool,

invokeai/frontend/web/src/features/controlLayers/konva/renderers/stagingArea.ts

@@ -2,7 +2,6 @@ import { KonvaImage, KonvaProgressImage } from 'features/controlLayers/konva/ren
 import type { CanvasV2State } from 'features/controlLayers/store/types';
 import Konva from 'konva';
 import type { InvocationDenoiseProgressEvent } from 'services/events/types';
-import { assert } from 'tsafe';
 
 export class CanvasStagingArea {
   group: Konva.Group;
@@ -17,13 +16,54 @@ export class CanvasStagingArea {
 
   async render(
     stagingArea: CanvasV2State['stagingArea'],
+    bbox: CanvasV2State['bbox'],
     shouldShowStagedImage: boolean,
-    lastProgressEvent: InvocationDenoiseProgressEvent | null
+    lastProgressEvent: InvocationDenoiseProgressEvent | null,
+    resetLastProgressEvent: () => void
   ) {
-    if (stagingArea && lastProgressEvent) {
+    const imageDTO = stagingArea.images[stagingArea.selectedImageIndex];
+
+    if (imageDTO) {
+      if (this.image) {
+        if (!this.image.isLoading && !this.image.isError && this.image.imageName !== imageDTO.image_name) {
+          await this.image.updateImageSource(imageDTO.image_name);
+        }
+        this.image.konvaImageGroup.x(bbox.x);
+        this.image.konvaImageGroup.y(bbox.y);
+        this.image.konvaImageGroup.visible(shouldShowStagedImage);
+        this.progressImage?.konvaImageGroup.visible(false);
+      } else {
+        const { image_name, width, height } = imageDTO;
+        this.image = new KonvaImage({
+          imageObject: {
+            id: 'staging-area-image',
+            type: 'image',
+            x: bbox.x,
+            y: bbox.y,
+            width,
+            height,
+            filters: [],
+            image: {
+              name: image_name,
+              width,
+              height,
+            },
+          },
+          onLoad: () => {
+            resetLastProgressEvent();
+          },
+        });
+        this.group.add(this.image.konvaImageGroup);
+        await this.image.updateImageSource(imageDTO.image_name);
+        this.image.konvaImageGroup.visible(shouldShowStagedImage);
+        this.progressImage?.konvaImageGroup.visible(false);
+      }
+    }
+
+    if (stagingArea.isStaging && lastProgressEvent) {
       const { invocation, step, progress_image } = lastProgressEvent;
       const { dataURL } = progress_image;
-      const { x, y, width, height } = stagingArea.bbox;
+      const { x, y, width, height } = bbox;
       const progressImageId = `${invocation.id}_${step}`;
       if (this.progressImage) {
         if (
@@ -42,47 +82,16 @@ export class CanvasStagingArea {
         this.image?.konvaImageGroup.visible(false);
         this.progressImage.konvaImageGroup.visible(true);
       }
-    } else if (stagingArea && stagingArea.selectedImageIndex !== null) {
-      const imageDTO = stagingArea.images[stagingArea.selectedImageIndex];
-      assert(imageDTO, 'Image must exist');
-      if (this.image) {
-        if (!this.image.isLoading && !this.image.isError && this.image.imageName !== imageDTO.image_name) {
-          await this.image.updateImageSource(imageDTO.image_name);
-        }
-        this.image.konvaImageGroup.x(stagingArea.bbox.x);
-        this.image.konvaImageGroup.y(stagingArea.bbox.y);
-        this.image.konvaImageGroup.visible(shouldShowStagedImage);
-        this.progressImage?.konvaImageGroup.visible(false);
-      } else {
-        const { image_name, width, height } = imageDTO;
-        this.image = new KonvaImage({
-          imageObject: {
-            id: 'staging-area-image',
-            type: 'image',
-            x: stagingArea.bbox.x,
-            y: stagingArea.bbox.y,
-            width,
-            height,
-            filters: [],
-            image: {
-              name: image_name,
-              width,
-              height,
-            },
-          },
-        });
-        this.group.add(this.image.konvaImageGroup);
-        await this.image.updateImageSource(imageDTO.image_name);
-        this.image.konvaImageGroup.visible(shouldShowStagedImage);
-        this.progressImage?.konvaImageGroup.visible(false);
-      }
-    } else {
+    }
+
+    if (!imageDTO && !lastProgressEvent) {
       if (this.image) {
         this.image.konvaImageGroup.visible(false);
       }
       if (this.progressImage) {
         this.progressImage.konvaImageGroup.visible(false);
       }
+      resetLastProgressEvent();
     }
   }
 }

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

@@ -121,7 +121,11 @@ const initialState: CanvasV2State = {
     refinerNegativeAestheticScore: 2.5,
     refinerStart: 0.8,
   },
-  stagingArea: null,
+  stagingArea: {
+    isStaging: false,
+    images: [],
+    selectedImageIndex: 0,
+  },
 };
 
 export const canvasV2Slice = createSlice({
@@ -332,12 +336,11 @@ export const {
   imLinePointAdded,
   imRectAdded,
   // Staging
-  stagingAreaInitialized,
+  stagingAreaStartedStaging,
   stagingAreaImageAdded,
-  stagingAreaBatchIdAdded,
   stagingAreaImageDiscarded,
   stagingAreaImageAccepted,
-  stagingAreaReset,
+  stagingAreaCanceledStaging,
   stagingAreaNextImageSelected,
   stagingAreaPreviousImageSelected,
 } = canvasV2Slice.actions;

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

@@ -1,16 +1,11 @@
 import type { PayloadAction, SliceCaseReducers } from '@reduxjs/toolkit';
-import type { CanvasV2State, Rect } from 'features/controlLayers/store/types';
+import type { CanvasV2State } from 'features/controlLayers/store/types';
 import type { ImageDTO } from 'services/api/types';
 
 export const stagingAreaReducers = {
-  stagingAreaInitialized: (state, action: PayloadAction<{ bbox: Rect; batchIds: string[] }>) => {
-    const { bbox, batchIds } = action.payload;
-    state.stagingArea = {
-      bbox,
-      batchIds,
-      selectedImageIndex: null,
-      images: [],
-    };
+  stagingAreaStartedStaging: (state) => {
+    state.stagingArea.isStaging = true;
+    state.stagingArea.selectedImageIndex = 0;
     // When we start staging, the user should not be interacting with the stage except to move it around. Set the tool
     // to view.
     state.tool.selectedBuffer = state.tool.selected;
@@ -18,67 +13,41 @@ export const stagingAreaReducers = {
   },
   stagingAreaImageAdded: (state, action: PayloadAction<{ imageDTO: ImageDTO }>) => {
     const { imageDTO } = action.payload;
-    if (!state.stagingArea) {
-      // Should not happen
-      return;
-    }
     state.stagingArea.images.push(imageDTO);
-    if (!state.stagingArea.selectedImageIndex) {
-      state.stagingArea.selectedImageIndex = state.stagingArea.images.length - 1;
-    }
+    state.stagingArea.selectedImageIndex = state.stagingArea.images.length - 1;
   },
   stagingAreaNextImageSelected: (state) => {
-    if (!state.stagingArea) {
-      // Should not happen
-      return;
-    }
-    if (state.stagingArea.selectedImageIndex === null) {
-      if (state.stagingArea.images.length > 0) {
-        state.stagingArea.selectedImageIndex = 0;
-      }
-      return;
-    }
     state.stagingArea.selectedImageIndex = (state.stagingArea.selectedImageIndex + 1) % state.stagingArea.images.length;
   },
   stagingAreaPreviousImageSelected: (state) => {
-    if (!state.stagingArea) {
-      // Should not happen
-      return;
-    }
-    if (state.stagingArea.selectedImageIndex === null) {
-      if (state.stagingArea.images.length > 0) {
-        state.stagingArea.selectedImageIndex = 0;
-      }
-      return;
-    }
     state.stagingArea.selectedImageIndex =
       (state.stagingArea.selectedImageIndex - 1 + state.stagingArea.images.length) % state.stagingArea.images.length;
   },
-  stagingAreaBatchIdAdded: (state, action: PayloadAction<{ batchId: string }>) => {
-    const { batchId } = action.payload;
-    if (!state.stagingArea) {
-      // Should not happen
-      return;
-    }
-    state.stagingArea.batchIds.push(batchId);
-  },
   stagingAreaImageDiscarded: (state, action: PayloadAction<{ imageDTO: ImageDTO }>) => {
     const { imageDTO } = action.payload;
-    if (!state.stagingArea) {
-      // Should not happen
-      return;
-    }
     state.stagingArea.images = state.stagingArea.images.filter((image) => image.image_name !== imageDTO.image_name);
+    state.stagingArea.selectedImageIndex = Math.min(
+      state.stagingArea.selectedImageIndex,
+      state.stagingArea.images.length - 1
+    );
+    if (state.stagingArea.images.length === 0) {
+      state.stagingArea.isStaging = false;
+    }
   },
   stagingAreaImageAccepted: (state, _: PayloadAction<{ imageDTO: ImageDTO }>) => {
     // When we finish staging, reset the tool back to the previous selection.
+    state.stagingArea.isStaging = false;
+    state.stagingArea.images = [];
+    state.stagingArea.selectedImageIndex = 0;
     if (state.tool.selectedBuffer) {
       state.tool.selected = state.tool.selectedBuffer;
       state.tool.selectedBuffer = null;
     }
   },
-  stagingAreaReset: (state) => {
-    state.stagingArea = null;
+  stagingAreaCanceledStaging: (state) => {
+    state.stagingArea.isStaging = false;
+    state.stagingArea.images = [];
+    state.stagingArea.selectedImageIndex = 0;
     // When we finish staging, reset the tool back to the previous selection.
     if (state.tool.selectedBuffer) {
       state.tool.selected = state.tool.selectedBuffer;

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

@@ -883,11 +883,10 @@ export type CanvasV2State = {
     refinerStart: number;
   };
   stagingArea: {
-    bbox: Rect;
+    isStaging: boolean;
     images: ImageDTO[];
-    selectedImageIndex: number | null;
-    batchIds: string[];
-  } | null;
+    selectedImageIndex: number;
+  };
 };
 
 export type StageAttrs = { x: number; y: number; width: number; height: number; scale: number };

invokeai/frontend/web/src/features/nodes/util/graph/buildLinearBatchConfig.ts

@@ -107,6 +107,7 @@ export const prepareLinearUIBatch = (state: RootState, g: Graph, prepend: boolea
       graph: g.getGraph(),
       runs: 1,
       data,
+      origin: 'canvas',
     },
   };
 

invokeai/frontend/web/src/services/api/endpoints/queue.ts

@@ -276,6 +276,26 @@ export const queueApi = api.injectEndpoints({
       },
       invalidatesTags: ['SessionQueueStatus', 'BatchStatus'],
     }),
+    cancelByBatchOrigin: build.mutation<
+      paths['/api/v1/queue/{queue_id}/cancel_by_origin']['put']['responses']['200']['content']['application/json'],
+      paths['/api/v1/queue/{queue_id}/cancel_by_origin']['put']['parameters']['query']
+    >({
+      query: (params) => ({
+        url: buildQueueUrl('cancel_by_origin'),
+        method: 'PUT',
+        params,
+      }),
+      onQueryStarted: async (arg, api) => {
+        const { dispatch, queryFulfilled } = api;
+        try {
+          await queryFulfilled;
+          resetListQueryData(dispatch);
+        } catch {
+          // no-op
+        }
+      },
+      invalidatesTags: ['SessionQueueStatus', 'BatchStatus'],
+    }),
     listQueueItems: build.query<
       EntityState<components['schemas']['SessionQueueItemDTO'], string> & {
         has_more: boolean;