feat(ui): switch to view tool when staging

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

@@ -1,7 +1,11 @@
 import { enqueueRequested } from 'app/store/actions';
 import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
 import { getNodeManager } from 'features/controlLayers/konva/nodeManager';
-import { stagingAreaBatchIdAdded, stagingAreaInitialized } from 'features/controlLayers/store/canvasV2Slice';
+import {
+  stagingAreaBatchIdAdded,
+  stagingAreaInitialized,
+  stagingAreaReset,
+} from 'features/controlLayers/store/canvasV2Slice';
 import { isImageViewerOpenChanged } from 'features/gallery/store/gallerySlice';
 import { prepareLinearUIBatch } from 'features/nodes/util/graph/buildLinearBatchConfig';
 import { buildSD1Graph } from 'features/nodes/util/graph/generation/buildSD1Graph';
@@ -19,47 +23,58 @@ export const addEnqueueRequestedLinear = (startAppListening: AppStartListening)
       const model = state.canvasV2.params.model;
       const { prepend } = action.payload;
 
-      let g;
+      let didInitializeStagingArea = false;
 
-      const manager = getNodeManager();
+      if (state.canvasV2.stagingArea === null) {
-      assert(model, 'No model found in state');
+        dispatch(
-      const base = model.base;
+          stagingAreaInitialized({
-
+            batchIds: [],
-      if (base === 'sdxl') {
+            bbox: state.canvasV2.bbox,
-        g = await buildSDXLGraph(state, manager);
+          })
-      } else if (base === 'sd-1' || base === 'sd-2') {
+        );
-        g = await buildSD1Graph(state, manager);
+        didInitializeStagingArea = true;
-      } else {
-        assert(false, `No graph builders for base ${base}`);
       }
 
-      const batchConfig = prepareLinearUIBatch(state, g, prepend);
-
-      const req = dispatch(
-        queueApi.endpoints.enqueueBatch.initiate(batchConfig, {
-          fixedCacheKey: 'enqueueBatch',
-        })
-      );
       try {
+        let g;
+
+        const manager = getNodeManager();
+        assert(model, 'No model found in state');
+        const base = model.base;
+
+        if (base === 'sdxl') {
+          g = await buildSDXLGraph(state, manager);
+        } else if (base === 'sd-1' || base === 'sd-2') {
+          g = await buildSD1Graph(state, manager);
+        } else {
+          assert(false, `No graph builders for base ${base}`);
+        }
+
+        const batchConfig = prepareLinearUIBatch(state, g, prepend);
+
+        const req = dispatch(
+          queueApi.endpoints.enqueueBatch.initiate(batchConfig, {
+            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');
-        if (!state.canvasV2.stagingArea) {
+        dispatch(stagingAreaBatchIdAdded({ batchId }));
-          dispatch(
+      } catch {
-            stagingAreaInitialized({
+        if (didInitializeStagingArea) {
-              batchIds: [batchId],
+          // We initialized the staging area in this listener, and there was a problem at some point. This means
-              bbox: state.canvasV2.bbox,
+          // 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());
-        } else {
-          dispatch(stagingAreaBatchIdAdded({ batchId }));
         }
-      } finally {
-        req.reset();
       }
     },
   });

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

@@ -43,6 +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 isDrawingToolDisabled = useMemo(
     () => !getIsDrawingToolEnabled(selectedEntityIdentifier),
     [selectedEntityIdentifier]
@@ -53,19 +54,35 @@ export const ToolChooser: React.FC = () => {
   const setToolToBrush = useCallback(() => {
     dispatch(toolChanged('brush'));
   }, [dispatch]);
-  useHotkeys('b', setToolToBrush, { enabled: !isDrawingToolDisabled }, [isDrawingToolDisabled, setToolToBrush]);
+  useHotkeys('b', setToolToBrush, { enabled: !isDrawingToolDisabled && !isStaging }, [
+    isDrawingToolDisabled,
+    isStaging,
+    setToolToBrush,
+  ]);
   const setToolToEraser = useCallback(() => {
     dispatch(toolChanged('eraser'));
   }, [dispatch]);
-  useHotkeys('e', setToolToEraser, { enabled: !isDrawingToolDisabled }, [isDrawingToolDisabled, setToolToEraser]);
+  useHotkeys('e', setToolToEraser, { enabled: !isDrawingToolDisabled && !isStaging }, [
+    isDrawingToolDisabled,
+    isStaging,
+    setToolToEraser,
+  ]);
   const setToolToRect = useCallback(() => {
     dispatch(toolChanged('rect'));
   }, [dispatch]);
-  useHotkeys('u', setToolToRect, { enabled: !isDrawingToolDisabled }, [isDrawingToolDisabled, setToolToRect]);
+  useHotkeys('u', setToolToRect, { enabled: !isDrawingToolDisabled && !isStaging }, [
+    isDrawingToolDisabled,
+    isStaging,
+    setToolToRect,
+  ]);
   const setToolToMove = useCallback(() => {
     dispatch(toolChanged('move'));
   }, [dispatch]);
-  useHotkeys('v', setToolToMove, { enabled: !isMoveToolDisabled }, [isMoveToolDisabled, setToolToMove]);
+  useHotkeys('v', setToolToMove, { enabled: !isMoveToolDisabled && !isStaging }, [
+    isMoveToolDisabled,
+    isStaging,
+    setToolToMove,
+  ]);
   const setToolToView = useCallback(() => {
     dispatch(toolChanged('view'));
   }, [dispatch]);
@@ -92,12 +109,16 @@ export const ToolChooser: React.FC = () => {
   }, [dispatch, selectedEntityIdentifier]);
   const isResetEnabled = useMemo(
     () =>
-      selectedEntityIdentifier?.type === 'layer' ||
+      (!isStaging && selectedEntityIdentifier?.type === 'layer') ||
       selectedEntityIdentifier?.type === 'regional_guidance' ||
       selectedEntityIdentifier?.type === 'inpaint_mask',
-    [selectedEntityIdentifier]
+    [isStaging, selectedEntityIdentifier?.type]
   );
-  useHotkeys('shift+c', resetSelectedLayer, { enabled: isResetEnabled }, [isResetEnabled, resetSelectedLayer]);
+  useHotkeys('shift+c', resetSelectedLayer, { enabled: isResetEnabled }, [
+    isResetEnabled,
+    isStaging,
+    resetSelectedLayer,
+  ]);
 
   const deleteSelectedLayer = useCallback(() => {
     if (selectedEntityIdentifier === null) {
@@ -117,7 +138,10 @@ export const ToolChooser: React.FC = () => {
       dispatch(ipaDeleted({ id }));
     }
   }, [dispatch, selectedEntityIdentifier]);
-  const isDeleteEnabled = useMemo(() => selectedEntityIdentifier !== null, [selectedEntityIdentifier]);
+  const isDeleteEnabled = useMemo(
+    () => selectedEntityIdentifier !== null && !isStaging,
+    [selectedEntityIdentifier, isStaging]
+  );
   useHotkeys('shift+d', deleteSelectedLayer, { enabled: isDeleteEnabled }, [isDeleteEnabled, deleteSelectedLayer]);
 
   return (
@@ -128,7 +152,7 @@ export const ToolChooser: React.FC = () => {
         icon={<PiPaintBrushBold />}
         variant={tool === 'brush' ? 'solid' : 'outline'}
         onClick={setToolToBrush}
-        isDisabled={isDrawingToolDisabled}
+        isDisabled={isDrawingToolDisabled || isStaging}
       />
       <IconButton
         aria-label={`${t('unifiedCanvas.eraser')} (E)`}
@@ -136,7 +160,7 @@ export const ToolChooser: React.FC = () => {
         icon={<PiEraserBold />}
         variant={tool === 'eraser' ? 'solid' : 'outline'}
         onClick={setToolToEraser}
-        isDisabled={isDrawingToolDisabled}
+        isDisabled={isDrawingToolDisabled || isStaging}
       />
       <IconButton
         aria-label={`${t('controlLayers.rectangle')} (U)`}
@@ -144,7 +168,7 @@ export const ToolChooser: React.FC = () => {
         icon={<PiRectangleBold />}
         variant={tool === 'rect' ? 'solid' : 'outline'}
         onClick={setToolToRect}
-        isDisabled={isDrawingToolDisabled}
+        isDisabled={isDrawingToolDisabled || isStaging}
       />
       <IconButton
         aria-label={`${t('unifiedCanvas.move')} (V)`}
@@ -152,7 +176,7 @@ export const ToolChooser: React.FC = () => {
         icon={<PiArrowsOutCardinalBold />}
         variant={tool === 'move' ? 'solid' : 'outline'}
         onClick={setToolToMove}
-        isDisabled={isMoveToolDisabled}
+        isDisabled={isMoveToolDisabled || isStaging}
       />
       <IconButton
         aria-label={`${t('unifiedCanvas.view')} (H)`}
@@ -160,6 +184,7 @@ export const ToolChooser: React.FC = () => {
         icon={<PiHandBold />}
         variant={tool === 'view' ? 'solid' : 'outline'}
         onClick={setToolToView}
+        isDisabled={isStaging}
       />
       <IconButton
         aria-label={`${t('controlLayers.bbox')} (Q)`}
@@ -167,6 +192,7 @@ export const ToolChooser: React.FC = () => {
         icon={<PiBoundingBoxBold />}
         variant={tool === 'bbox' ? 'solid' : 'outline'}
         onClick={setToolToBbox}
+        isDisabled={isStaging}
       />
     </ButtonGroup>
   );

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

@@ -128,8 +128,6 @@ export const setStageEventHandlers = (manager: KonvaNodeManager): (() => void) =
 
   //#region mouseenter
   stage.on('mouseenter', () => {
-    const tool = getToolState().selected;
-    stage.findOne<Konva.Layer>(`#${PREVIEW_TOOL_GROUP_ID}`)?.visible(tool === 'brush' || tool === 'eraser');
     manager.renderToolPreview();
   });
 

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

@@ -446,7 +446,12 @@ export const initializeRenderer = (
   const unsubscribeRenderer = subscribe(renderCanvas);
 
   // When we this flag, we need to render the staging area
-  $shouldShowStagedImage.subscribe(manager.renderStagingArea.bind(manager));
+  $shouldShowStagedImage.subscribe((shouldShowStagedImage, prevShouldShowStagedImage) => {
+    logIfDebugging('Rendering staging area');
+    if (shouldShowStagedImage !== prevShouldShowStagedImage) {
+      manager.renderStagingArea();
+    }
+  });
 
   logIfDebugging('First render of konva stage');
   // On first render, the document should be fit to the stage.

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

@@ -2,13 +2,12 @@ import { rgbaColorToString } from 'common/util/colorCodeTransformers';
 import {
   BRUSH_BORDER_INNER_COLOR,
   BRUSH_BORDER_OUTER_COLOR,
-  BRUSH_ERASER_BORDER_WIDTH
+  BRUSH_ERASER_BORDER_WIDTH,
 } from 'features/controlLayers/konva/constants';
 import { PREVIEW_RECT_ID } from 'features/controlLayers/konva/naming';
 import type { CanvasEntity, CanvasV2State, Position, RgbaColor } from 'features/controlLayers/store/types';
 import Konva from 'konva';
 
-
 export class CanvasTool {
   group: Konva.Group;
   brush: {
@@ -125,7 +124,8 @@ export class CanvasTool {
     isMouseDown: boolean
   ) {
     const tool = toolState.selected;
-    const isDrawableEntity = selectedEntity?.type === 'regional_guidance' ||
+    const isDrawableEntity =
+      selectedEntity?.type === 'regional_guidance' ||
       selectedEntity?.type === 'layer' ||
       selectedEntity?.type === 'inpaint_mask';
 

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

@@ -11,6 +11,10 @@ export const stagingAreaReducers = {
       selectedImageIndex: null,
       images: [],
     };
+    // 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;
+    state.tool.selected = 'view';
   },
   stagingAreaImageAdded: (state, action: PayloadAction<{ imageDTO: ImageDTO }>) => {
     const { imageDTO } = action.payload;
@@ -66,8 +70,19 @@ export const stagingAreaReducers = {
     }
     state.stagingArea.images = state.stagingArea.images.filter((image) => image.image_name !== imageDTO.image_name);
   },
-  stagingAreaImageAccepted: (state, _: PayloadAction<{ imageDTO: ImageDTO }>) => state,
+  stagingAreaImageAccepted: (state, _: PayloadAction<{ imageDTO: ImageDTO }>) => {
+    // When we finish staging, reset the tool back to the previous selection.
+    if (state.tool.selectedBuffer) {
+      state.tool.selected = state.tool.selectedBuffer;
+      state.tool.selectedBuffer = null;
+    }
+  },
   stagingAreaReset: (state) => {
     state.stagingArea = null;
+    // When we finish staging, reset the tool back to the previous selection.
+    if (state.tool.selectedBuffer) {
+      state.tool.selected = state.tool.selectedBuffer;
+      state.tool.selectedBuffer = null;
+    }
   },
 } satisfies SliceCaseReducers<CanvasV2State>;