feat(ui): replace automask apply w/ save as menu

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

@@ -1876,7 +1876,7 @@
             "exclude": "Exclude",
             "neutral": "Neutral",
             "reset": "Reset",
-            "apply": "Apply",
+            "saveAs": "Save As",
             "cancel": "Cancel",
             "process": "Process"
         },

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

@@ -1,4 +1,14 @@
-import { Button, ButtonGroup, Flex, Heading, Spacer } from '@invoke-ai/ui-library';
+import {
+  Button,
+  ButtonGroup,
+  Flex,
+  Heading,
+  Menu,
+  MenuButton,
+  MenuItem,
+  MenuList,
+  Spacer,
+} from '@invoke-ai/ui-library';
 import { useStore } from '@nanostores/react';
 import { useAppSelector } from 'app/store/storeHooks';
 import { useFocusRegion, useIsRegionFocused } from 'common/hooks/focus';
@@ -10,7 +20,7 @@ import type { CanvasEntityAdapterControlLayer } from 'features/controlLayers/kon
 import type { CanvasEntityAdapterRasterLayer } from 'features/controlLayers/konva/CanvasEntity/CanvasEntityAdapterRasterLayer';
 import { selectAutoProcess } from 'features/controlLayers/store/canvasSettingsSlice';
 import { useRegisteredHotkeys } from 'features/system/components/HotkeysModal/useHotkeyData';
-import { memo, useRef } from 'react';
+import { memo, useCallback, useRef } from 'react';
 import { useTranslation } from 'react-i18next';
 import { PiArrowsCounterClockwiseBold, PiCheckBold, PiStarBold, PiXBold } from 'react-icons/pi';
 
@@ -22,8 +32,25 @@ const SegmentAnythingContent = memo(
     const isCanvasFocused = useIsRegionFocused('canvas');
     const isProcessing = useStore(adapter.segmentAnything.$isProcessing);
     const hasPoints = useStore(adapter.segmentAnything.$hasPoints);
+    const hasImageState = useStore(adapter.segmentAnything.$hasImageState);
     const autoProcess = useAppSelector(selectAutoProcess);
 
+    const saveAsInpaintMask = useCallback(() => {
+      adapter.segmentAnything.saveAs('inpaint_mask');
+    }, [adapter.segmentAnything]);
+
+    const saveAsRegionalGuidance = useCallback(() => {
+      adapter.segmentAnything.saveAs('regional_guidance');
+    }, [adapter.segmentAnything]);
+
+    const saveAsRasterLayer = useCallback(() => {
+      adapter.segmentAnything.saveAs('raster_layer');
+    }, [adapter.segmentAnything]);
+
+    const saveAsControlLayer = useCallback(() => {
+      adapter.segmentAnything.saveAs('control_layer');
+    }, [adapter.segmentAnything]);
+
     useRegisteredHotkeys({
       id: 'applySegmentAnything',
       category: 'canvas',
@@ -86,15 +113,32 @@ const SegmentAnythingContent = memo(
           >
             {t('controlLayers.segment.reset')}
           </Button>
-          <Button
-            leftIcon={<PiCheckBold />}
-            onClick={adapter.segmentAnything.apply}
-            isLoading={isProcessing}
-            loadingText={t('controlLayers.segment.apply')}
-            variant="ghost"
-          >
-            {t('controlLayers.segment.apply')}
-          </Button>
+          <Menu>
+            <MenuButton
+              as={Button}
+              leftIcon={<PiCheckBold />}
+              isLoading={isProcessing}
+              loadingText={t('controlLayers.segment.saveAs')}
+              variant="ghost"
+              isDisabled={!hasImageState}
+            >
+              {t('controlLayers.segment.saveAs')}
+            </MenuButton>
+            <MenuList>
+              <MenuItem isDisabled={!hasImageState} onClick={saveAsInpaintMask}>
+                {t('controlLayers.inpaintMask')}
+              </MenuItem>
+              <MenuItem isDisabled={!hasImageState} onClick={saveAsRegionalGuidance}>
+                {t('controlLayers.regionalGuidance')}
+              </MenuItem>
+              <MenuItem isDisabled={!hasImageState} onClick={saveAsControlLayer}>
+                {t('controlLayers.controlLayer')}
+              </MenuItem>
+              <MenuItem isDisabled={!hasImageState} onClick={saveAsRasterLayer}>
+                {t('controlLayers.rasterLayer')}
+              </MenuItem>
+            </MenuList>
+          </Menu>
           <Button
             leftIcon={<PiXBold />}
             onClick={adapter.segmentAnything.cancel}

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

@@ -15,6 +15,7 @@ import {
 } from 'features/controlLayers/konva/util';
 import { selectAutoProcess } from 'features/controlLayers/store/canvasSettingsSlice';
 import type {
+  CanvasEntityType,
   CanvasImageState,
   Coordinate,
   RgbaColor,
@@ -34,6 +35,8 @@ import type { Logger } from 'roarr';
 import { serializeError } from 'serialize-error';
 import type { ImageDTO } from 'services/api/types';
 import stableHash from 'stable-hash';
+import type { Equals } from 'tsafe';
+import { assert } from 'tsafe';
 
 type CanvasSegmentAnythingModuleConfig = {
   /**
@@ -152,7 +155,12 @@ export class CanvasSegmentAnythingModule extends CanvasModuleBase {
   /**
    * The ephemeral image state of the processed image. Only used while segmenting.
    */
-  imageState: CanvasImageState | null = null;
+  $imageState = atom<CanvasImageState | null>(null);
+
+  /**
+   * Whether the module has an image state. This is a computed value based on $imageState.
+   */
+  $hasImageState = computed(this.$imageState, (imageState) => imageState !== null);
 
   /**
    * The current input points. A listener is added to this atom to process the points when they change.
@@ -567,21 +575,23 @@ export class CanvasSegmentAnythingModule extends CanvasModuleBase {
     this.log.trace({ imageDTO: segmentResult.value }, 'Segmented');
 
     // Prepare the ephemeral image state
-    this.imageState = imageDTOToImageObject(segmentResult.value);
+    const imageState = imageDTOToImageObject(segmentResult.value);
+    this.$imageState.set(imageState);
 
     // Destroy any existing masked image and create a new one
     if (this.maskedImage) {
       this.maskedImage.destroy();
     }
-    this.maskedImage = new CanvasObjectImage(this.imageState, this);
+
+    this.maskedImage = new CanvasObjectImage(imageState, this);
 
     // Force update the masked image - after awaiting, the image will be rendered (in memory)
-    await this.maskedImage.update(this.imageState, true);
+    await this.maskedImage.update(imageState, true);
 
     // Update the compositing rect to match the image size
     this.konva.compositingRect.setAttrs({
-      width: this.imageState.image.width,
-      height: this.imageState.image.height,
+      width: imageState.image.width,
+      height: imageState.image.height,
       visible: true,
     });
 
@@ -614,7 +624,7 @@ export class CanvasSegmentAnythingModule extends CanvasModuleBase {
    * Applies the segmented image to the entity.
    */
   apply = () => {
-    const imageState = this.imageState;
+    const imageState = this.$imageState.get();
     if (!imageState) {
       this.log.error('No image state to apply');
       return;
@@ -641,6 +651,55 @@ export class CanvasSegmentAnythingModule extends CanvasModuleBase {
     this.teardown();
   };
 
+  /**
+   * Applies the segmented image to the entity.
+   */
+  saveAs = (type: Exclude<CanvasEntityType, 'reference_image'>) => {
+    const imageState = this.$imageState.get();
+    if (!imageState) {
+      this.log.error('No image state to save as');
+      return;
+    }
+    this.log.trace(`Saving as ${type}`);
+
+    // Clear the buffer - we are creating a new entity, so we don't want to keep the old one
+    this.parent.bufferRenderer.clearBuffer();
+
+    // Create the new entity with the masked image as its only object
+    const rect = this.parent.transformer.getRelativeRect();
+    const arg = {
+      overrides: {
+        objects: [imageState],
+        position: {
+          x: Math.round(rect.x),
+          y: Math.round(rect.y),
+        },
+      },
+      isSelected: true,
+    };
+
+    switch (type) {
+      case 'raster_layer':
+        this.manager.stateApi.addRasterLayer(arg);
+        break;
+      case 'control_layer':
+        this.manager.stateApi.addControlLayer(arg);
+        break;
+      case 'inpaint_mask':
+        this.manager.stateApi.addInpaintMask(arg);
+        break;
+      case 'regional_guidance':
+        this.manager.stateApi.addRegionalGuidance(arg);
+        break;
+      default:
+        assert<Equals<typeof type, never>>(false);
+    }
+
+    // Final cleanup and teardown, returning user to main canvas UI
+    this.resetEphemeralState();
+    this.teardown();
+  };
+
   /**
    * Resets the module (e.g. remove all points and the mask image).
    *
@@ -703,7 +762,7 @@ export class CanvasSegmentAnythingModule extends CanvasModuleBase {
 
     // Empty internal module state
     this.$points.set([]);
-    this.imageState = null;
+    this.$imageState.set(null);
     this.$pointType.set(1);
     this.$lastProcessedHash.set('');
     this.$isProcessing.set(false);
@@ -773,7 +832,7 @@ export class CanvasSegmentAnythingModule extends CanvasModuleBase {
         label,
         circle: getKonvaNodeDebugAttrs(konva.circle),
       })),
-      imageState: deepClone(this.imageState),
+      imageState: deepClone(this.$imageState.get()),
       maskedImage: this.maskedImage?.repr(),
       config: deepClone(this.config),
       $isSegmenting: this.$isSegmenting.get(),

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

@@ -17,12 +17,16 @@ import {
 } from 'features/controlLayers/store/canvasSettingsSlice';
 import {
   bboxChangedFromCanvas,
+  controlLayerAdded,
   entityBrushLineAdded,
   entityEraserLineAdded,
   entityMoved,
   entityRasterized,
   entityRectAdded,
   entityReset,
+  inpaintMaskAdded,
+  rasterLayerAdded,
+  rgAdded,
 } from 'features/controlLayers/store/canvasSlice';
 import { selectCanvasStagingAreaSlice } from 'features/controlLayers/store/canvasStagingAreaSlice';
 import {
@@ -51,6 +55,7 @@ import { getImageDTO } from 'services/api/endpoints/images';
 import { queueApi } from 'services/api/endpoints/queue';
 import type { BatchConfig, ImageDTO, S } from 'services/api/types';
 import { QueueError } from 'services/events/errors';
+import type { Param0 } from 'tsafe';
 import { assert } from 'tsafe';
 
 import type { CanvasEntityAdapter } from './CanvasEntity/types';
@@ -160,6 +165,34 @@ export class CanvasStateApiModule extends CanvasModuleBase {
     this.store.dispatch(entityRectAdded(arg));
   };
 
+  /**
+   * Adds a raster layer to the canvas, pushing state to redux.
+   */
+  addRasterLayer = (arg: Param0<typeof rasterLayerAdded>) => {
+    this.store.dispatch(rasterLayerAdded(arg));
+  };
+
+  /**
+   * Adds a control layer to the canvas, pushing state to redux.
+   */
+  addControlLayer = (arg: Param0<typeof controlLayerAdded>) => {
+    this.store.dispatch(controlLayerAdded(arg));
+  };
+
+  /**
+   * Adds an inpaint mask to the canvas, pushing state to redux.
+   */
+  addInpaintMask = (arg: Param0<typeof inpaintMaskAdded>) => {
+    this.store.dispatch(inpaintMaskAdded(arg));
+  };
+
+  /**
+   * Adds regional guidance to the canvas, pushing state to redux.
+   */
+  addRegionalGuidance = (arg: Param0<typeof rgAdded>) => {
+    this.store.dispatch(rgAdded(arg));
+  };
+
   /**
    * Rasterizes an entity, pushing state to redux.
    */