feat(ui): inpaint mask transform

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

@@ -19,7 +19,9 @@ export const TransformToolButton = memo(() => {
     if (!canvasManager) {
       return;
     }
-    return canvasManager.isTransforming.subscribe(setIsTransforming);
+    return canvasManager.transformingEntity.subscribe((newValue) => {
+      setIsTransforming(Boolean(newValue));
+    });
   }, [canvasManager]);
 
   const onTransform = useCallback(() => {

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

@@ -5,13 +5,11 @@ import type { CanvasInpaintMaskState, CanvasV2State, GetLoggingContext } from 'f
 import Konva from 'konva';
 import { get } from 'lodash-es';
 import type { Logger } from 'roarr';
-import { assert } from 'tsafe';
 
 export class CanvasInpaintMask {
   static TYPE = 'inpaint_mask' as const;
   static NAME_PREFIX = 'inpaint-mask';
   static KONVA_LAYER_NAME = `${CanvasInpaintMask.NAME_PREFIX}_layer`;
-  static OBJECT_GROUP_NAME = `${CanvasInpaintMask.NAME_PREFIX}_object-group`;
 
   id = CanvasInpaintMask.TYPE;
   type = CanvasInpaintMask.TYPE;
@@ -29,7 +27,6 @@ export class CanvasInpaintMask {
 
   konva: {
     layer: Konva.Layer;
-    objectGroup: Konva.Group;
   };
 
   constructor(state: CanvasInpaintMaskState, manager: CanvasManager) {
@@ -44,16 +41,10 @@ export class CanvasInpaintMask {
         listening: false,
         imageSmoothingEnabled: false,
       }),
-      objectGroup: new Konva.Group({ name: CanvasInpaintMask.OBJECT_GROUP_NAME, listening: false }),
     };
 
     this.transformer = new CanvasTransformer(this);
     this.renderer = new CanvasObjectRenderer(this);
-    assert(this.renderer.konva.compositingRect, 'Compositing rect must be set');
-
-    this.konva.layer.add(this.konva.objectGroup);
-    this.konva.layer.add(this.renderer.konva.compositingRect);
-    this.konva.layer.add(...this.transformer.getNodes());
 
     this.state = state;
     this.maskOpacity = this.manager.stateApi.getMaskOpacity();
@@ -123,12 +114,6 @@ export class CanvasInpaintMask {
     }
   };
 
-  // updateOpacity = (arg?: { opacity: number }) => {
-  //   this.log.trace('Updating opacity');
-  //   const opacity = get(arg, 'opacity', this.state.opacity);
-  //   this.konva.objectGroup.opacity(opacity);
-  // };
-
   updateVisibility = (arg?: { isEnabled: boolean }) => {
     this.log.trace('Updating visibility');
     const isEnabled = get(arg, 'isEnabled', this.state.isEnabled);

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

@@ -2,13 +2,10 @@ import { deepClone } from 'common/util/deepClone';
 import type { CanvasManager } from 'features/controlLayers/konva/CanvasManager';
 import { CanvasObjectRenderer } from 'features/controlLayers/konva/CanvasObjectRenderer';
 import { CanvasTransformer } from 'features/controlLayers/konva/CanvasTransformer';
-import { konvaNodeToBlob, previewBlob } from 'features/controlLayers/konva/util';
 import type { CanvasLayerState, CanvasV2State, GetLoggingContext } from 'features/controlLayers/store/types';
-import { imageDTOToImageObject } from 'features/controlLayers/store/types';
 import Konva from 'konva';
 import { get } from 'lodash-es';
 import type { Logger } from 'roarr';
-import { uploadImage } from 'services/api/endpoints/images';
 
 export class CanvasLayer {
   static TYPE = 'layer' as const;
@@ -25,7 +22,6 @@ export class CanvasLayer {
 
   konva: {
     layer: Konva.Layer;
-    objectGroup: Konva.Group;
   };
   transformer: CanvasTransformer;
   renderer: CanvasObjectRenderer;
@@ -47,15 +43,11 @@ export class CanvasLayer {
         listening: false,
         imageSmoothingEnabled: false,
       }),
-      objectGroup: new Konva.Group({ name: CanvasLayer.KONVA_OBJECT_GROUP_NAME, listening: false }),
     };
 
     this.transformer = new CanvasTransformer(this);
     this.renderer = new CanvasObjectRenderer(this);
 
-    this.konva.layer.add(this.konva.objectGroup);
-    this.konva.layer.add(...this.transformer.getNodes());
-
     this.state = state;
   }
 
@@ -121,26 +113,7 @@ export class CanvasLayer {
   updateOpacity = (arg?: { opacity: number }) => {
     this.log.trace('Updating opacity');
     const opacity = get(arg, 'opacity', this.state.opacity);
-    this.konva.objectGroup.opacity(opacity);
-  };
-
-  rasterize = async () => {
-    this.log.debug('Rasterizing layer');
-
-    const objectGroupClone = this.konva.objectGroup.clone();
-    const interactionRectClone = this.transformer.konva.proxyRect.clone();
-    const rect = interactionRectClone.getClientRect();
-    const blob = await konvaNodeToBlob(objectGroupClone, rect);
-    if (this.manager._isDebugging) {
-      previewBlob(blob, 'Rasterized layer');
-    }
-    const imageDTO = await uploadImage(blob, `${this.id}_rasterized.png`, 'other', true);
-    const imageObject = imageDTOToImageObject(imageDTO);
-    await this.renderer.renderObject(imageObject, true);
-    this.manager.stateApi.rasterizeEntity(
-      { id: this.id, imageObject, position: { x: Math.round(rect.x), y: Math.round(rect.y) } },
-      this.type
-    );
+    this.renderer.konva.objectGroup.opacity(opacity);
   };
 
   repr = () => {
@@ -167,15 +140,15 @@ export class CanvasLayer {
         rotation: this.transformer.konva.proxyRect.rotation(),
       },
       objectGroupAttrs: {
-        x: this.konva.objectGroup.x(),
-        y: this.konva.objectGroup.y(),
-        scaleX: this.konva.objectGroup.scaleX(),
-        scaleY: this.konva.objectGroup.scaleY(),
-        width: this.konva.objectGroup.width(),
-        height: this.konva.objectGroup.height(),
-        rotation: this.konva.objectGroup.rotation(),
-        offsetX: this.konva.objectGroup.offsetX(),
-        offsetY: this.konva.objectGroup.offsetY(),
+        x: this.renderer.konva.objectGroup.x(),
+        y: this.renderer.konva.objectGroup.y(),
+        scaleX: this.renderer.konva.objectGroup.scaleX(),
+        scaleY: this.renderer.konva.objectGroup.scaleY(),
+        width: this.renderer.konva.objectGroup.width(),
+        height: this.renderer.konva.objectGroup.height(),
+        rotation: this.renderer.konva.objectGroup.rotation(),
+        offsetX: this.renderer.konva.objectGroup.offsetX(),
+        offsetY: this.renderer.konva.objectGroup.offsetY(),
       },
     };
     this.log.trace(info, msg);

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

@@ -122,7 +122,7 @@ export class CanvasManager {
   log: Logger;
   workerLog: Logger;
 
-  isTransforming: PubSub<boolean>;
+  transformingEntity: PubSub<CanvasEntityIdentifier | null>;
 
   _store: Store<RootState>;
   _prevState: CanvasV2State;
@@ -208,7 +208,7 @@ export class CanvasManager {
       this.log.error('Worker message error');
     };
 
-    this.isTransforming = new PubSub(false);
+    this.transformingEntity = new PubSub<CanvasEntityIdentifier | null>(null);
     this.toolState = new PubSub(this.stateApi.getToolState());
     this.currentFill = new PubSub(this.getCurrentFill());
     this.selectedEntityIdentifier = new PubSub(
@@ -377,11 +377,24 @@ export class CanvasManager {
   };
 
   getTransformingLayer() {
-    return Array.from(this.layers.values()).find((layer) => layer.transformer.isTransforming);
+    const transformingEntity = this.transformingEntity.getValue();
+    if (!transformingEntity) {
+      return null;
+    }
+
+    const { id, type } = transformingEntity;
+
+    if (type === 'layer') {
+      return this.layers.get(id) ?? null;
+    } else if (type === 'inpaint_mask') {
+      return this.inpaintMask;
+    }
+
+    return null;
   }
 
   getIsTransforming() {
-    return Boolean(this.getTransformingLayer());
+    return Boolean(this.transformingEntity.getValue());
   }
 
   startTransform() {
@@ -395,7 +408,7 @@ export class CanvasManager {
       'No selected layer'
     );
     layer.adapter.transformer.startTransform();
-    this.isTransforming.publish(true);
+    this.transformingEntity.publish({ id: layer.state.id, type: layer.state.type });
   }
 
   async applyTransform() {
@@ -403,7 +416,7 @@ export class CanvasManager {
     if (layer) {
       await layer.transformer.applyTransform();
     }
-    this.isTransforming.publish(false);
+    this.transformingEntity.publish(null);
   }
 
   cancelTransform() {
@@ -411,7 +424,7 @@ export class CanvasManager {
     if (layer) {
       layer.transformer.stopTransform();
     }
-    this.isTransforming.publish(false);
+    this.transformingEntity.publish(null);
   }
 
   render = async () => {

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

@@ -8,16 +8,18 @@ import type { CanvasInpaintMask } from 'features/controlLayers/konva/CanvasInpai
 import type { CanvasLayer } from 'features/controlLayers/konva/CanvasLayer';
 import type { CanvasManager } from 'features/controlLayers/konva/CanvasManager';
 import { CanvasRectRenderer } from 'features/controlLayers/konva/CanvasRect';
-import { getPrefixedId } from 'features/controlLayers/konva/util';
-import type {
-  CanvasBrushLineState,
-  CanvasEraserLineState,
-  CanvasImageState,
-  CanvasRectState,
-  RgbColor,
+import { getPrefixedId, konvaNodeToBlob, previewBlob } from 'features/controlLayers/konva/util';
+import {
+  type CanvasBrushLineState,
+  type CanvasEraserLineState,
+  type CanvasImageState,
+  type CanvasRectState,
+  imageDTOToImageObject,
+  type RgbColor,
 } from 'features/controlLayers/store/types';
 import Konva from 'konva';
 import type { Logger } from 'roarr';
+import { uploadImage } from 'services/api/endpoints/images';
 import { assert } from 'tsafe';
 
 /**
@@ -34,6 +36,7 @@ type AnyObjectState = CanvasBrushLineState | CanvasEraserLineState | CanvasImage
  */
 export class CanvasObjectRenderer {
   static TYPE = 'object_renderer';
+  static KONVA_OBJECT_GROUP_NAME = 'object-group';
   static KONVA_COMPOSITING_RECT_NAME = 'compositing-rect';
 
   id: string;
@@ -63,6 +66,10 @@ export class CanvasObjectRenderer {
    * A object containing singleton Konva nodes.
    */
   konva: {
+    /**
+     * A Konva Group that holds all the object renderers.
+     */
+    objectGroup: Konva.Group;
     /**
      * The compositing rect is used to draw the inpaint mask as a single shape with a given opacity.
      *
@@ -74,6 +81,8 @@ export class CanvasObjectRenderer {
      * of 'source-in'. The shapes effectively become a mask for the "compositing rect".
      *
      * This node is only added when the parent of the renderer is an inpaint mask or region, which require this behavior.
+     *
+     * The compositing rect is not added to the object group.
      */
     compositingRect: Konva.Rect | null;
   };
@@ -87,16 +96,19 @@ export class CanvasObjectRenderer {
     this.log.trace('Creating object renderer');
 
     this.konva = {
+      objectGroup: new Konva.Group({ name: CanvasObjectRenderer.KONVA_OBJECT_GROUP_NAME, listening: false }),
       compositingRect: null,
     };
 
+    this.parent.konva.layer.add(this.konva.objectGroup);
+
     if (this.parent.type === 'inpaint_mask') {
       this.konva.compositingRect = new Konva.Rect({
         name: CanvasObjectRenderer.KONVA_COMPOSITING_RECT_NAME,
         listening: false,
         globalCompositeOperation: 'source-in',
       });
-      this.parent.konva.objectGroup.add(this.konva.compositingRect);
+      this.parent.konva.layer.add(this.konva.compositingRect);
     }
 
     this.subscriptions.add(
@@ -184,7 +196,7 @@ export class CanvasObjectRenderer {
       if (!renderer) {
         renderer = new CanvasBrushLineRenderer(objectState, this);
         this.renderers.set(renderer.id, renderer);
-        this.parent.konva.objectGroup.add(renderer.konva.group);
+        this.konva.objectGroup.add(renderer.konva.group);
       }
 
       didRender = renderer.update(objectState, force || isFirstRender);
@@ -194,7 +206,7 @@ export class CanvasObjectRenderer {
       if (!renderer) {
         renderer = new CanvasEraserLineRenderer(objectState, this);
         this.renderers.set(renderer.id, renderer);
-        this.parent.konva.objectGroup.add(renderer.konva.group);
+        this.konva.objectGroup.add(renderer.konva.group);
       }
 
       didRender = renderer.update(objectState, force || isFirstRender);
@@ -204,7 +216,7 @@ export class CanvasObjectRenderer {
       if (!renderer) {
         renderer = new CanvasRectRenderer(objectState, this);
         this.renderers.set(renderer.id, renderer);
-        this.parent.konva.objectGroup.add(renderer.konva.group);
+        this.konva.objectGroup.add(renderer.konva.group);
       }
 
       didRender = renderer.update(objectState, force || isFirstRender);
@@ -214,7 +226,7 @@ export class CanvasObjectRenderer {
       if (!renderer) {
         renderer = new CanvasImageRenderer(objectState, this);
         this.renderers.set(renderer.id, renderer);
-        this.parent.konva.objectGroup.add(renderer.konva.group);
+        this.konva.objectGroup.add(renderer.konva.group);
       }
       didRender = await renderer.update(objectState, force || isFirstRender);
     }
@@ -311,6 +323,25 @@ export class CanvasObjectRenderer {
     return this.renderers.size > 0 || this.buffer !== null;
   };
 
+  rasterize = async () => {
+    this.log.debug('Rasterizing entity');
+
+    const objectGroupClone = this.konva.objectGroup.clone();
+    const interactionRectClone = this.parent.transformer.konva.proxyRect.clone();
+    const rect = interactionRectClone.getClientRect();
+    const blob = await konvaNodeToBlob(objectGroupClone, rect);
+    if (this.manager._isDebugging) {
+      previewBlob(blob, 'Rasterized layer');
+    }
+    const imageDTO = await uploadImage(blob, `${this.id}_rasterized.png`, 'other', true);
+    const imageObject = imageDTOToImageObject(imageDTO);
+    await this.renderObject(imageObject, true);
+    this.manager.stateApi.rasterizeEntity(
+      { id: this.id, imageObject, position: { x: Math.round(rect.x), y: Math.round(rect.y) } },
+      this.parent.type
+    );
+  };
+
   /**
    * Destroys this renderer and all of its object renderers.
    */

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

@@ -23,6 +23,7 @@ import {
   imImageCacheChanged,
   imRectAdded,
   imTranslated,
+  inpaintMaskRasterized,
   layerBrushLineAdded,
   layerEraserLineAdded,
   layerImageCacheChanged,
@@ -119,6 +120,8 @@ export class CanvasStateApi {
     log.trace({ arg, entityType }, 'Rasterizing entity');
     if (entityType === 'layer') {
       this._store.dispatch(layerRasterized(arg));
+    } else if (entityType === 'inpaint_mask') {
+      this._store.dispatch(inpaintMaskRasterized(arg));
     } else {
       assert(false, 'Rasterizing not supported for this entity type');
     }

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

@@ -149,7 +149,7 @@ export class CanvasTool {
     } else if (!isDrawableEntity) {
       // Non-drawable layers don't have tools
       stage.container().style.cursor = 'not-allowed';
-    } else if (tool === 'move' || this.manager.isTransforming.getValue()) {
+    } else if (tool === 'move' || Boolean(this.manager.transformingEntity.getValue())) {
       // Move tool gets a pointer
       stage.container().style.cursor = 'default';
     } else if (tool === 'rect') {

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

@@ -251,7 +251,7 @@ export class CanvasTransformer {
       // This is called when a transform anchor is dragged. By this time, the transform constraints in the above
       // callbacks have been enforced, and the transformer has updated its nodes' attributes. We need to pass the
       // updated attributes to the object group, propagating the transformation on down.
-      this.parent.konva.objectGroup.setAttrs({
+      this.parent.renderer.konva.objectGroup.setAttrs({
         x: this.konva.proxyRect.x(),
         y: this.konva.proxyRect.y(),
         scaleX: this.konva.proxyRect.scaleX(),
@@ -293,7 +293,7 @@ export class CanvasTransformer {
         scaleX: snappedScaleX,
         scaleY: snappedScaleY,
       });
-      this.parent.konva.objectGroup.setAttrs({
+      this.parent.renderer.konva.objectGroup.setAttrs({
         x: snappedX,
         y: snappedY,
         scaleX: snappedScaleX,
@@ -337,7 +337,7 @@ export class CanvasTransformer {
 
       // The object group is translated by the difference between the interaction rect's new and old positions (which is
       // stored as this.pixelRect)
-      this.parent.konva.objectGroup.setAttrs({
+      this.parent.renderer.konva.objectGroup.setAttrs({
         x: this.konva.proxyRect.x(),
         y: this.konva.proxyRect.y(),
       });
@@ -391,6 +391,10 @@ export class CanvasTransformer {
         this.syncInteractionState();
       })
     );
+
+    this.parent.konva.layer.add(this.konva.bboxOutline);
+    this.parent.konva.layer.add(this.konva.proxyRect);
+    this.parent.konva.layer.add(this.konva.transformer);
   }
 
   /**
@@ -499,7 +503,7 @@ export class CanvasTransformer {
    */
   applyTransform = async () => {
     this.log.debug('Applying transform');
-    await this.parent.rasterize();
+    await this.parent.renderer.rasterize();
     this.requestRectCalculation();
     this.stopTransform();
   };
@@ -534,7 +538,7 @@ export class CanvasTransformer {
       scaleY: 1,
       rotation: 0,
     };
-    this.parent.konva.objectGroup.setAttrs(attrs);
+    this.parent.renderer.konva.objectGroup.setAttrs(attrs);
     this.konva.bboxOutline.setAttrs(attrs);
     this.konva.proxyRect.setAttrs(attrs);
   };
@@ -547,7 +551,7 @@ export class CanvasTransformer {
     this.log.trace('Updating position');
     const position = get(arg, 'position', this.parent.state.position);
 
-    this.parent.konva.objectGroup.setAttrs({
+    this.parent.renderer.konva.objectGroup.setAttrs({
       x: position.x + this.pixelRect.x,
       y: position.y + this.pixelRect.y,
       offsetX: this.pixelRect.x,
@@ -603,7 +607,7 @@ export class CanvasTransformer {
 
     this.syncInteractionState();
     this.update(this.parent.state.position, this.pixelRect);
-    this.parent.konva.objectGroup.setAttrs({
+    this.parent.renderer.konva.objectGroup.setAttrs({
       x: this.parent.state.position.x + this.pixelRect.x,
       y: this.parent.state.position.y + this.pixelRect.y,
       offsetX: this.pixelRect.x,
@@ -625,7 +629,7 @@ export class CanvasTransformer {
       return;
     }
 
-    const rect = this.parent.konva.objectGroup.getClientRect({ skipTransform: true });
+    const rect = this.parent.renderer.konva.objectGroup.getClientRect({ skipTransform: true });
 
     if (!this.parent.renderer.needsPixelBbox()) {
       this.nodeRect = { ...rect };
@@ -638,7 +642,7 @@ export class CanvasTransformer {
 
     // We have eraser strokes - we must calculate the bbox using pixel data
 
-    const clone = this.parent.konva.objectGroup.clone();
+    const clone = this.parent.renderer.konva.objectGroup.clone();
     const canvas = clone.toCanvas();
     const ctx = canvas.getContext('2d');
     if (!ctx) {
@@ -709,12 +713,6 @@ export class CanvasTransformer {
     this.konva.bboxOutline.visible(false);
   };
 
-  /**
-   * Gets the nodes that make up the transformer, in the order they should be added to the layer.
-   * @returns The nodes that make up the transformer.
-   */
-  getNodes = () => [this.konva.bboxOutline, this.konva.proxyRect, this.konva.transformer];
-
   /**
    * Gets a JSON-serializable object that describes the transformer.
    */

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

@@ -345,6 +345,7 @@ export const {
   imBrushLineAdded,
   imEraserLineAdded,
   imRectAdded,
+  inpaintMaskRasterized,
   // Staging
   sessionStarted,
   sessionStartedStaging,

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

@@ -6,6 +6,7 @@ import type {
   CanvasRectState,
   CanvasV2State,
   Coordinate,
+  EntityRasterizedArg,
   ScaleChangedArg,
 } from 'features/controlLayers/store/types';
 import { imageDTOToImageWithDims } from 'features/controlLayers/store/types';
@@ -84,4 +85,10 @@ export const inpaintMaskReducers = {
     state.inpaintMask.bboxNeedsUpdate = true;
     state.layers.imageCache = null;
   },
+  inpaintMaskRasterized: (state, action: PayloadAction<EntityRasterizedArg>) => {
+    const { imageObject, position } = action.payload;
+    state.inpaintMask.objects = [imageObject];
+    state.inpaintMask.position = position;
+    state.inpaintMask.imageCache = null;
+  },
 } satisfies SliceCaseReducers<CanvasV2State>;

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

@@ -683,7 +683,7 @@ const zCanvasInpaintMaskState = z.object({
   position: zCoordinate,
   bbox: zRect.nullable(),
   bboxNeedsUpdate: z.boolean(),
-  objects: z.array(zMaskObject),
+  objects: z.array(zCanvasObjectState),
   fill: zRgbColor,
   imageCache: zImageWithDims.nullable(),
 });