feat(ui): rough out canvas mode

invokeai/frontend/web/public/locales/en.json

@@ -1645,6 +1645,11 @@
         "storeNotInitialized": "Store is not initialized"
     },
     "controlLayers": {
+        "generateMode": "Generate",
+        "generateModeDesc": "Create individual images. Generated images are added directly to the gallery.",
+        "composeMode": "Compose",
+        "composeModeDesc": "Compose your work iterative. Generated images are added back to the canvas.",
+        "autoSave": "Auto-save to Gallery",
         "resetCanvas": "Reset Canvas",
         "resetAll": "Reset All",
         "clearCaches": "Clear Caches",

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

@@ -21,7 +21,7 @@ export const addEnqueueRequestedLinear = (startAppListening: AppStartListening)
       assert(manager, 'No model found in state');
 
       let didStartStaging = false;
-      if (!state.canvasV2.session.isStaging) {
+      if (!state.canvasV2.session.isStaging && state.canvasV2.session.mode === 'compose') {
         dispatch(sessionStartedStaging());
         didStartStaging = true;
       }
@@ -50,7 +50,6 @@ export const addEnqueueRequestedLinear = (startAppListening: AppStartListening)
         req.reset();
         await req.unwrap();
       } catch (error) {
-        console.log('Error in enqueueRequestedLinear', error);
         if (didStartStaging && getState().canvasV2.session.isStaging) {
           dispatch(sessionStagingAreaReset());
         }

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

@@ -0,0 +1,26 @@
+import { Button, ButtonGroup } from '@invoke-ai/ui-library';
+import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
+import { sessionModeChanged } from 'features/controlLayers/store/canvasV2Slice';
+import { memo, useCallback } from 'react';
+import { useTranslation } from 'react-i18next';
+
+export const CanvasModeSwitcher = memo(() => {
+  const { t } = useTranslation();
+  const dispatch = useAppDispatch();
+  const mode = useAppSelector((s) => s.canvasV2.session.mode);
+  const onClickGenerate = useCallback(() => dispatch(sessionModeChanged({ mode: 'generate' })), [dispatch]);
+  const onClickCompose = useCallback(() => dispatch(sessionModeChanged({ mode: 'compose' })), [dispatch]);
+
+  return (
+    <ButtonGroup variant="outline">
+      <Button onClick={onClickGenerate} colorScheme={mode === 'generate' ? 'invokeBlue' : 'base'}>
+        {t('controlLayers.generateMode')}
+      </Button>
+      <Button onClick={onClickCompose} colorScheme={mode === 'compose' ? 'invokeBlue' : 'base'}>
+        {t('controlLayers.composeMode')}
+      </Button>
+    </ButtonGroup>
+  );
+});
+
+CanvasModeSwitcher.displayName = 'CanvasModeSwitcher';

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

@@ -23,7 +23,7 @@ export const CanvasResetViewButton = memo(() => {
     if (!canvasManager) {
       return;
     }
-    canvasManager.stage.resetView();
+    canvasManager.stage.fitLayersToStage();
   }, [canvasManager]);
 
   const onReset = useCallback(() => {

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

@@ -1,6 +1,7 @@
 /* eslint-disable i18next/no-literal-string */
 import { Flex, Spacer } from '@invoke-ai/ui-library';
 import { useAppSelector } from 'app/store/storeHooks';
+import { CanvasModeSwitcher } from 'features/controlLayers/components/CanvasModeSwitcher';
 import { CanvasResetViewButton } from 'features/controlLayers/components/CanvasResetViewButton';
 import { CanvasScale } from 'features/controlLayers/components/CanvasScale';
 import { CanvasSettingsPopover } from 'features/controlLayers/components/Settings/CanvasSettingsPopover';
@@ -23,11 +24,12 @@ export const ControlLayersToolbar = memo(() => {
         <ToolChooser />
         {tool === 'brush' && <ToolBrushWidth />}
         {tool === 'eraser' && <ToolEraserWidth />}
+        <ToolFillColorPicker />
         <Spacer />
         <CanvasScale />
         <CanvasResetViewButton />
         <Spacer />
-        <ToolFillColorPicker />
+        <CanvasModeSwitcher />
         <UndoRedoButtonGroup />
         <CanvasSettingsPopover />
         <ViewerToggleMenu />

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

@@ -70,13 +70,6 @@ export class CanvasProgressImageModule {
       return;
     }
 
-    const { isStaging } = this.manager.stateApi.getSession();
-
-    if (!isStaging) {
-      release();
-      return;
-    }
-
     this.isLoading = true;
 
     const { x, y, width, height } = this.manager.stateApi.getBbox().rect;

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

@@ -33,6 +33,7 @@ export class CanvasStageModule {
     const resizeObserver = new ResizeObserver(this.fitStageToContainer);
     resizeObserver.observe(this.container);
     this.fitStageToContainer();
+    this.fitLayersToStage();
 
     return () => {
       this.log.debug('Destroying stage');
@@ -91,10 +92,20 @@ export class CanvasStageModule {
     }
   };
 
-  resetView() {
+  fitBboxToStage = () => {
-    this.log.trace('Resetting view');
+    this.log.trace('Fitting bbox to stage');
-    const { width, height } = this.getSize();
+    const bbox = this.manager.stateApi.getBbox();
+    this.fitRect(bbox.rect);
+  };
+
+  fitLayersToStage() {
+    this.log.trace('Fitting layers to stage');
     const rect = this.getVisibleRect();
+    this.fitRect(rect);
+  }
+
+  fitRect = (rect: Rect) => {
+    const { width, height } = this.getSize();
 
     const padding = 20; // Padding in absolute pixels
 
@@ -118,7 +129,7 @@ export class CanvasStageModule {
       y,
       scale,
     });
-  }
+  };
 
   /**
    * Gets the center of the stage in either absolute or relative coordinates

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

@@ -132,6 +132,7 @@ const initialState: CanvasV2State = {
     refinerStart: 0.8,
   },
   session: {
+    mode: 'generate',
     isStaging: false,
     stagedImages: [],
     selectedStagedImageIndex: 0,
@@ -474,6 +475,7 @@ export const {
   clipToBboxChanged,
   canvasReset,
   settingsDynamicGridToggled,
+  settingsAutoSaveToggled,
   // All entities
   entitySelected,
   entityNameChanged,
@@ -601,6 +603,7 @@ export const {
   sessionStagingAreaReset,
   sessionNextStagedImageSelected,
   sessionPrevStagedImageSelected,
+  sessionModeChanged,
 } = canvasV2Slice.actions;
 
 export const selectCanvasV2Slice = (state: RootState) => state.canvasV2;

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

@@ -1,5 +1,5 @@
 import type { PayloadAction, SliceCaseReducers } from '@reduxjs/toolkit';
-import type { CanvasV2State, StagingAreaImage } from 'features/controlLayers/store/types';
+import type { CanvasV2State, SessionMode, StagingAreaImage } from 'features/controlLayers/store/types';
 
 export const sessionReducers = {
   sessionStartedStaging: (state) => {
@@ -45,4 +45,8 @@ export const sessionReducers = {
       state.tool.selectedBuffer = null;
     }
   },
+  sessionModeChanged: (state, action: PayloadAction<{ mode: SessionMode }>) => {
+    const { mode } = action.payload;
+    state.session.mode = mode;
+  },
 } satisfies SliceCaseReducers<CanvasV2State>;

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

@@ -8,4 +8,7 @@ export const settingsReducers = {
   settingsDynamicGridToggled: (state) => {
     state.settings.dynamicGrid = !state.settings.dynamicGrid;
   },
+  settingsAutoSaveToggled: (state) => {
+    state.settings.autoSave = !state.settings.autoSave;
+  },
 } satisfies SliceCaseReducers<CanvasV2State>;

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

@@ -834,6 +834,8 @@ export type StagingAreaImage = {
   offsetY: number;
 };
 
+export type SessionMode = 'generate' | 'compose';
+
 export type CanvasV2State = {
   _version: 3;
   selectedEntityIdentifier: CanvasEntityIdentifier | null;
@@ -929,6 +931,7 @@ export type CanvasV2State = {
     refinerStart: number;
   };
   session: {
+    mode: SessionMode;
     isStaging: boolean;
     stagedImages: StagingAreaImage[];
     selectedStagedImageIndex: number;

invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSD1Graph.ts

@@ -41,7 +41,7 @@ export const buildSD1Graph = async (state: RootState, manager: CanvasManager): P
   const generationMode = manager.compositor.getGenerationMode();
   log.debug({ generationMode }, 'Building SD1/SD2 graph');
 
-  const { bbox, params } = state.canvasV2;
+  const { bbox, params, session, settings } = state.canvasV2;
 
   const {
     model,
@@ -249,10 +249,11 @@ export const buildSD1Graph = async (state: RootState, manager: CanvasManager): P
     canvasOutput = addWatermarker(g, canvasOutput);
   }
 
-  // This is the terminal node and must always save to gallery.
+  const shouldSaveToGallery = session.mode === 'generate' || settings.autoSave;
+
   g.updateNode(canvasOutput, {
     id: CANVAS_OUTPUT,
-    is_intermediate: false,
+    is_intermediate: !shouldSaveToGallery,
     use_cache: false,
     board: getBoardField(state),
   });

invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSDXLGraph.ts

@@ -40,7 +40,7 @@ export const buildSDXLGraph = async (state: RootState, manager: CanvasManager):
   const generationMode = manager.compositor.getGenerationMode();
   log.debug({ generationMode }, 'Building SDXL graph');
 
-  const { bbox, params } = state.canvasV2;
+  const { bbox, params, session, settings } = state.canvasV2;
 
   const {
     model,
@@ -246,10 +246,11 @@ export const buildSDXLGraph = async (state: RootState, manager: CanvasManager):
     canvasOutput = addWatermarker(g, canvasOutput);
   }
 
-  // This is the terminal node and must always save to gallery.
+  const shouldSaveToGallery = session.mode === 'generate' || settings.autoSave;
+
   g.updateNode(canvasOutput, {
     id: CANVAS_OUTPUT,
-    is_intermediate: false,
+    is_intermediate: !shouldSaveToGallery,
     use_cache: false,
     board: getBoardField(state),
   });

invokeai/frontend/web/src/services/events/onInvocationComplete.ts

@@ -0,0 +1,145 @@
+import { logger } from 'app/logging/logger';
+import type { AppDispatch, RootState } from 'app/store/store';
+import type { SerializableObject } from 'common/types';
+import { deepClone } from 'common/util/deepClone';
+import { sessionImageStaged } from 'features/controlLayers/store/canvasV2Slice';
+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 { boardsApi } from 'services/api/endpoints/boards';
+import { getImageDTO, imagesApi } from 'services/api/endpoints/images';
+import type { ImageDTO } from 'services/api/types';
+import { getCategories, getListImagesUrl } from 'services/api/util';
+import type { InvocationCompleteEvent, InvocationDenoiseProgressEvent } from 'services/events/types';
+
+const log = logger('events');
+
+export const buildOnInvocationComplete = (
+  getState: () => RootState,
+  dispatch: AppDispatch,
+  nodeTypeDenylist: string[],
+  setLastProgressEvent: (event: InvocationDenoiseProgressEvent | null) => void,
+  setLastCanvasProgressEvent: (event: InvocationDenoiseProgressEvent | null) => void
+) => {
+  const addImageToGallery = (imageDTO: ImageDTO) => {
+    if (imageDTO.is_intermediate) {
+      return;
+    }
+
+    // update the total images for the board
+    dispatch(
+      boardsApi.util.updateQueryData('getBoardImagesTotal', imageDTO.board_id ?? 'none', (draft) => {
+        // eslint-disable-next-line @typescript-eslint/no-unused-vars
+        draft.total += 1;
+      })
+    );
+
+    dispatch(
+      imagesApi.util.invalidateTags([
+        { type: 'Board', id: imageDTO.board_id ?? 'none' },
+        {
+          type: 'ImageList',
+          id: getListImagesUrl({
+            board_id: imageDTO.board_id ?? 'none',
+            categories: getCategories(imageDTO),
+          }),
+        },
+      ])
+    );
+
+    const { shouldAutoSwitch, galleryView, selectedBoardId } = getState().gallery;
+
+    // If auto-switch is enabled, select the new image
+    if (shouldAutoSwitch) {
+      // if auto-add is enabled, switch the gallery view and board if needed as the image comes in
+      if (galleryView !== 'images') {
+        dispatch(galleryViewChanged('images'));
+      }
+
+      if (imageDTO.board_id && imageDTO.board_id !== selectedBoardId) {
+        dispatch(
+          boardIdSelected({
+            boardId: imageDTO.board_id,
+            selectedImageName: imageDTO.image_name,
+          })
+        );
+      }
+
+      dispatch(offsetChanged({ offset: 0 }));
+
+      if (!imageDTO.board_id && selectedBoardId !== 'none') {
+        dispatch(
+          boardIdSelected({
+            boardId: 'none',
+            selectedImageName: imageDTO.image_name,
+          })
+        );
+      }
+
+      dispatch(imageSelected(imageDTO));
+    }
+  };
+
+  return async (data: InvocationCompleteEvent) => {
+    log.debug(
+      { data } as SerializableObject,
+      `Invocation complete (${data.invocation.type}, ${data.invocation_source_id})`
+    );
+
+    const { result, invocation_source_id } = data;
+
+    // Update the node execution states - the image output is handled below
+    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);
+      }
+    }
+
+    // This complete event has an associated image output
+    if (
+      (data.result.type === 'image_output' || data.result.type === 'canvas_v2_mask_and_crop_output') &&
+      !nodeTypeDenylist.includes(data.invocation.type)
+    ) {
+      const { image_name } = data.result.image;
+      const { session } = getState().canvasV2;
+
+      const imageDTO = await getImageDTO(image_name);
+
+      if (!imageDTO) {
+        log.error({ data } as SerializableObject, 'Failed to fetch image DTO after generation');
+        return;
+      }
+
+      if (data.origin === 'canvas') {
+        if (data.invocation_source_id !== 'canvas_output') {
+          // Not a canvas output image - ignore
+          return;
+        }
+        if (session.mode === 'compose' && session.isStaging) {
+          if (data.result.type === 'canvas_v2_mask_and_crop_output') {
+            const { offset_x, offset_y } = data.result;
+            if (session.isStaging) {
+              dispatch(sessionImageStaged({ stagingAreaImage: { imageDTO, offsetX: offset_x, offsetY: offset_y } }));
+            }
+          } else if (data.result.type === 'image_output') {
+            if (session.isStaging) {
+              dispatch(sessionImageStaged({ stagingAreaImage: { imageDTO, offsetX: 0, offsetY: 0 } }));
+            }
+          }
+          addImageToGallery(imageDTO);
+        } else {
+          addImageToGallery(imageDTO);
+          setLastCanvasProgressEvent(null);
+        }
+      }
+    }
+
+    setLastProgressEvent(null);
+  };
+};

invokeai/frontend/web/src/services/events/setEventListeners.tsx

@@ -7,8 +7,6 @@ import { $queueId } from 'app/store/nanostores/queueId';
 import type { AppDispatch, RootState } from 'app/store/store';
 import type { SerializableObject } from 'common/types';
 import { deepClone } from 'common/util/deepClone';
-import { sessionImageStaged } from 'features/controlLayers/store/canvasV2Slice';
-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 ErrorToastDescription, { getTitleFromErrorType } from 'features/toast/ErrorToastDescription';
@@ -17,11 +15,9 @@ import { t } from 'i18next';
 import { forEach } from 'lodash-es';
 import { atom, computed } from 'nanostores';
 import { api, LIST_TAG } from 'services/api';
-import { boardsApi } from 'services/api/endpoints/boards';
-import { imagesApi } from 'services/api/endpoints/images';
 import { modelsApi } from 'services/api/endpoints/models';
 import { queueApi, queueItemsAdapter } from 'services/api/endpoints/queue';
-import { getCategories, getListImagesUrl } from 'services/api/util';
+import { buildOnInvocationComplete } from 'services/events/onInvocationComplete';
 import type { ClientToServerEvents, InvocationDenoiseProgressEvent, ServerToClientEvents } from 'services/events/types';
 import type { Socket } from 'socket.io-client';
 
@@ -147,116 +143,14 @@ export const setEventListeners = ({ socket, dispatch, getState, setIsConnected }
     }
   });
 
-  socket.on('invocation_complete', async (data) => {
+  const onInvocationComplete = buildOnInvocationComplete(
-    log.debug(
+    getState,
-      { data } as SerializableObject,
+    dispatch,
-      `Invocation complete (${data.invocation.type}, ${data.invocation_source_id})`
+    nodeTypeDenylist,
-    );
+    $lastProgressEvent.set,
-
+    $lastCanvasProgressEvent.set
-    const { result, invocation_source_id } = data;
+  );
-
+  socket.on('invocation_complete', onInvocationComplete);
-    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);
-      }
-    }
-
-    // This complete event has an associated image output
-    if (
-      (data.result.type === 'image_output' || data.result.type === 'canvas_v2_mask_and_crop_output') &&
-      !nodeTypeDenylist.includes(data.invocation.type)
-    ) {
-      const { image_name } = data.result.image;
-      const { gallery, canvasV2 } = getState();
-
-      // This populates the `getImageDTO` cache
-      const imageDTORequest = dispatch(
-        imagesApi.endpoints.getImageDTO.initiate(image_name, {
-          forceRefetch: true,
-        })
-      );
-
-      const imageDTO = await imageDTORequest.unwrap();
-      imageDTORequest.unsubscribe();
-
-      // handle tab-specific logic
-      if (data.origin === 'canvas' && data.invocation_source_id === 'canvas_output') {
-        if (data.result.type === 'canvas_v2_mask_and_crop_output') {
-          const { offset_x, offset_y } = data.result;
-          if (canvasV2.session.isStaging) {
-            dispatch(sessionImageStaged({ stagingAreaImage: { imageDTO, offsetX: offset_x, offsetY: offset_y } }));
-          }
-        } else if (data.result.type === 'image_output') {
-          if (canvasV2.session.isStaging) {
-            dispatch(sessionImageStaged({ stagingAreaImage: { imageDTO, offsetX: 0, offsetY: 0 } }));
-          }
-        }
-      }
-
-      if (!imageDTO.is_intermediate) {
-        // update the total images for the board
-        dispatch(
-          boardsApi.util.updateQueryData('getBoardImagesTotal', imageDTO.board_id ?? 'none', (draft) => {
-            // eslint-disable-next-line @typescript-eslint/no-unused-vars
-            draft.total += 1;
-          })
-        );
-
-        dispatch(
-          imagesApi.util.invalidateTags([
-            { type: 'Board', id: imageDTO.board_id ?? 'none' },
-            {
-              type: 'ImageList',
-              id: getListImagesUrl({
-                board_id: imageDTO.board_id ?? 'none',
-                categories: getCategories(imageDTO),
-              }),
-            },
-          ])
-        );
-
-        const { shouldAutoSwitch } = gallery;
-
-        // If auto-switch is enabled, select the new image
-        if (shouldAutoSwitch) {
-          // if auto-add is enabled, switch the gallery view and board if needed as the image comes in
-          if (gallery.galleryView !== 'images') {
-            dispatch(galleryViewChanged('images'));
-          }
-
-          if (imageDTO.board_id && imageDTO.board_id !== gallery.selectedBoardId) {
-            dispatch(
-              boardIdSelected({
-                boardId: imageDTO.board_id,
-                selectedImageName: imageDTO.image_name,
-              })
-            );
-          }
-
-          dispatch(offsetChanged({ offset: 0 }));
-
-          if (!imageDTO.board_id && gallery.selectedBoardId !== 'none') {
-            dispatch(
-              boardIdSelected({
-                boardId: 'none',
-                selectedImageName: imageDTO.image_name,
-              })
-            );
-          }
-
-          dispatch(imageSelected(imageDTO));
-        }
-      }
-    }
-
-    $lastProgressEvent.set(null);
-  });
 
   socket.on('model_load_started', (data) => {
     const { config, submodel_type } = data;