feat(ui): inpaint mask rendering (wip)

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

@@ -106,7 +106,12 @@ const maybeAddNextPoint = (
   setLastAddedPoint: Arg['setLastAddedPoint'],
   onPointAddedToLine: Arg['onPointAddedToLine']
 ) => {
-  if (selectedEntity.type !== 'layer' && selectedEntity.type !== 'regional_guidance') {
+  const isDrawableEntity =
+    selectedEntity?.type === 'regional_guidance' ||
+    selectedEntity?.type === 'layer' ||
+    selectedEntity?.type === 'inpaint_mask';
+
+  if (!isDrawableEntity) {
     return;
   }
   // Continue the last line
@@ -189,12 +194,12 @@ export const setStageEventHandlers = ({
     const toolState = getToolState();
     const pos = updateLastCursorPos(stage, setLastCursorPos);
     const selectedEntity = getSelectedEntity();
-    if (
-      pos &&
-      selectedEntity &&
-      (selectedEntity?.type === 'regional_guidance' || selectedEntity?.type === 'layer') &&
-      !getSpaceKey()
-    ) {
+    const isDrawableEntity =
+      selectedEntity?.type === 'regional_guidance' ||
+      selectedEntity?.type === 'layer' ||
+      selectedEntity?.type === 'inpaint_mask';
+
+    if (pos && selectedEntity && isDrawableEntity && !getSpaceKey()) {
       setIsDrawing(true);
       setLastMouseDownPos(pos);
 
@@ -318,13 +323,12 @@ export const setStageEventHandlers = ({
     setIsMouseDown(false);
     const pos = getLastCursorPos();
     const selectedEntity = getSelectedEntity();
+    const isDrawableEntity =
+      selectedEntity?.type === 'regional_guidance' ||
+      selectedEntity?.type === 'layer' ||
+      selectedEntity?.type === 'inpaint_mask';
 
-    if (
-      pos &&
-      selectedEntity &&
-      (selectedEntity?.type === 'regional_guidance' || selectedEntity?.type === 'layer') &&
-      !getSpaceKey()
-    ) {
+    if (pos && selectedEntity && isDrawableEntity && !getSpaceKey()) {
       const toolState = getToolState();
 
       if (toolState.selected === 'rect') {
@@ -372,13 +376,12 @@ export const setStageEventHandlers = ({
       .findOne<Konva.Layer>(`#${PREVIEW_TOOL_GROUP_ID}`)
       ?.visible(toolState.selected === 'brush' || toolState.selected === 'eraser');
 
-    if (
-      pos &&
-      selectedEntity &&
-      (selectedEntity.type === 'regional_guidance' || selectedEntity.type === 'layer') &&
-      !getSpaceKey() &&
-      getIsMouseDown()
-    ) {
+    const isDrawableEntity =
+      selectedEntity?.type === 'regional_guidance' ||
+      selectedEntity?.type === 'layer' ||
+      selectedEntity?.type === 'inpaint_mask';
+
+    if (pos && selectedEntity && isDrawableEntity && !getSpaceKey() && getIsMouseDown()) {
       if (toolState.selected === 'brush') {
         if (getIsDrawing()) {
           // Continue the last line
@@ -489,14 +492,12 @@ export const setStageEventHandlers = ({
     const toolState = getToolState();
 
     stage.findOne<Konva.Layer>(`#${PREVIEW_TOOL_GROUP_ID}`)?.visible(false);
+    const isDrawableEntity =
+      selectedEntity?.type === 'regional_guidance' ||
+      selectedEntity?.type === 'layer' ||
+      selectedEntity?.type === 'inpaint_mask';
 
-    if (
-      pos &&
-      selectedEntity &&
-      (selectedEntity.type === 'regional_guidance' || selectedEntity.type === 'layer') &&
-      !getSpaceKey() &&
-      getIsMouseDown()
-    ) {
+    if (pos && selectedEntity && isDrawableEntity && !getSpaceKey() && getIsMouseDown()) {
       if (getIsMouseDown()) {
         if (toolState.selected === 'brush') {
           onPointAddedToLine({ id: selectedEntity.id, point: [pos.x, pos.y] }, selectedEntity.type);

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

@@ -40,6 +40,10 @@ export const RASTER_LAYER_RECT_SHAPE_NAME = `${RASTER_LAYER_NAME}.rect_shape`;
 export const RASTER_LAYER_IMAGE_NAME = `${RASTER_LAYER_NAME}.image`;
 
 export const INPAINT_MASK_LAYER_NAME = 'inpaint_mask_layer';
+export const INPAINT_MASK_LAYER_OBJECT_GROUP_NAME = `${INPAINT_MASK_LAYER_NAME}.object_group`;
+export const INPAINT_MASK_LAYER_BRUSH_LINE_NAME = `${INPAINT_MASK_LAYER_NAME}.brush_line`;
+export const INPAINT_MASK_LAYER_ERASER_LINE_NAME = `${INPAINT_MASK_LAYER_NAME}.eraser_line`;
+export const INPAINT_MASK_LAYER_RECT_SHAPE_NAME = `${INPAINT_MASK_LAYER_NAME}.rect_shape`;
 
 export const BACKGROUND_LAYER_ID = 'background_layer';
 

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

@@ -18,5 +18,6 @@ export const arrangeEntities = (
   for (const rg of regions) {
     manager.get(rg.id)?.konvaLayer.zIndex(++zIndex);
   }
+  manager.get('inpaint_mask')?.konvaLayer.zIndex(++zIndex);
   manager.preview.layer.zIndex(++zIndex);
 };

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

@@ -0,0 +1,234 @@
+import { rgbColorToString } from 'common/util/colorCodeTransformers';
+import {
+  COMPOSITING_RECT_NAME,
+  INPAINT_MASK_LAYER_BRUSH_LINE_NAME,
+  INPAINT_MASK_LAYER_ERASER_LINE_NAME,
+  INPAINT_MASK_LAYER_NAME,
+  INPAINT_MASK_LAYER_OBJECT_GROUP_NAME,
+  INPAINT_MASK_LAYER_RECT_SHAPE_NAME,
+} from 'features/controlLayers/konva/naming';
+import type { KonvaEntityAdapter, KonvaNodeManager } from 'features/controlLayers/konva/nodeManager';
+import { getLayerBboxFast } from 'features/controlLayers/konva/renderers/bbox';
+import {
+  createObjectGroup,
+  getBrushLine,
+  getEraserLine,
+  getRectShape,
+} from 'features/controlLayers/konva/renderers/objects';
+import { mapId } from 'features/controlLayers/konva/util';
+import type {
+  CanvasEntity,
+  CanvasEntityIdentifier,
+  InpaintMaskEntity,
+  PosChangedArg,
+  Tool,
+} from 'features/controlLayers/store/types';
+import Konva from 'konva';
+
+/**
+ * Logic for creating and rendering regional guidance layers.
+ *
+ * Some special handling is needed to render layer opacity correctly using a "compositing rect". See the comments
+ * in `renderRGLayer`.
+ */
+
+/**
+ * Creates the "compositing rect" for a regional guidance layer.
+ * @param konvaLayer The konva layer
+ */
+const createCompositingRect = (konvaLayer: Konva.Layer): Konva.Rect => {
+  const compositingRect = new Konva.Rect({ name: COMPOSITING_RECT_NAME, listening: false });
+  konvaLayer.add(compositingRect);
+  return compositingRect;
+};
+
+/**
+ * Creates a regional guidance layer.
+ * @param stage The konva stage
+ * @param entity The regional guidance layer state
+ * @param onLayerPosChanged Callback for when the layer's position changes
+ */
+const getInpaintMask = (
+  manager: KonvaNodeManager,
+  entity: InpaintMaskEntity,
+  onPosChanged?: (arg: PosChangedArg, entityType: CanvasEntity['type']) => void
+): KonvaEntityAdapter => {
+  const adapter = manager.get(entity.id);
+  if (adapter) {
+    return adapter;
+  }
+  // This layer hasn't been added to the konva state yet
+  const konvaLayer = new Konva.Layer({
+    id: entity.id,
+    name: INPAINT_MASK_LAYER_NAME,
+    draggable: true,
+    dragDistance: 0,
+  });
+
+  // When a drag on the layer finishes, update the layer's position in state. During the drag, konva handles changing
+  // the position - we do not need to call this on the `dragmove` event.
+  if (onPosChanged) {
+    konvaLayer.on('dragend', function (e) {
+      onPosChanged({ id: entity.id, x: Math.floor(e.target.x()), y: Math.floor(e.target.y()) }, 'inpaint_mask');
+    });
+  }
+
+  const konvaObjectGroup = createObjectGroup(konvaLayer, INPAINT_MASK_LAYER_OBJECT_GROUP_NAME);
+  return manager.add(entity, konvaLayer, konvaObjectGroup);
+};
+
+/**
+ * Renders a raster layer.
+ * @param stage The konva stage
+ * @param entity The regional guidance layer state
+ * @param globalMaskLayerOpacity The global mask layer opacity
+ * @param tool The current tool
+ * @param onPosChanged Callback for when the layer's position changes
+ */
+export const renderInpaintMask = (
+  manager: KonvaNodeManager,
+  entity: InpaintMaskEntity,
+  globalMaskLayerOpacity: number,
+  tool: Tool,
+  selectedEntityIdentifier: CanvasEntityIdentifier | null,
+  onPosChanged?: (arg: PosChangedArg, entityType: CanvasEntity['type']) => void
+): void => {
+  const adapter = getInpaintMask(manager, entity, onPosChanged);
+
+  // Update the layer's position and listening state
+  adapter.konvaLayer.setAttrs({
+    listening: tool === 'move', // The layer only listens when using the move tool - otherwise the stage is handling mouse events
+    x: Math.floor(entity.x),
+    y: Math.floor(entity.y),
+  });
+
+  // Convert the color to a string, stripping the alpha - the object group will handle opacity.
+  const rgbColor = rgbColorToString(entity.fill);
+
+  // We use caching to handle "global" layer opacity, but caching is expensive and we should only do it when required.
+  let groupNeedsCache = false;
+
+  const objectIds = entity.objects.map(mapId);
+  // Destroy any objects that are no longer in state
+  for (const objectRecord of adapter.getAll()) {
+    if (!objectIds.includes(objectRecord.id)) {
+      adapter.destroy(objectRecord.id);
+      groupNeedsCache = true;
+    }
+  }
+
+  for (const obj of entity.objects) {
+    if (obj.type === 'brush_line') {
+      const objectRecord = getBrushLine(adapter, obj, INPAINT_MASK_LAYER_BRUSH_LINE_NAME);
+
+      // Only update the points if they have changed. The point values are never mutated, they are only added to the
+      // array, so checking the length is sufficient to determine if we need to re-cache.
+      if (objectRecord.konvaLine.points().length !== obj.points.length) {
+        objectRecord.konvaLine.points(obj.points);
+        groupNeedsCache = true;
+      }
+      // Only update the color if it has changed.
+      if (objectRecord.konvaLine.stroke() !== rgbColor) {
+        objectRecord.konvaLine.stroke(rgbColor);
+        groupNeedsCache = true;
+      }
+    } else if (obj.type === 'eraser_line') {
+      const objectRecord = getEraserLine(adapter, obj, INPAINT_MASK_LAYER_ERASER_LINE_NAME);
+
+      // Only update the points if they have changed. The point values are never mutated, they are only added to the
+      // array, so checking the length is sufficient to determine if we need to re-cache.
+      if (objectRecord.konvaLine.points().length !== obj.points.length) {
+        objectRecord.konvaLine.points(obj.points);
+        groupNeedsCache = true;
+      }
+      // Only update the color if it has changed.
+      if (objectRecord.konvaLine.stroke() !== rgbColor) {
+        objectRecord.konvaLine.stroke(rgbColor);
+        groupNeedsCache = true;
+      }
+    } else if (obj.type === 'rect_shape') {
+      const objectRecord = getRectShape(adapter, obj, INPAINT_MASK_LAYER_RECT_SHAPE_NAME);
+
+      // Only update the color if it has changed.
+      if (objectRecord.konvaRect.fill() !== rgbColor) {
+        objectRecord.konvaRect.fill(rgbColor);
+        groupNeedsCache = true;
+      }
+    }
+  }
+
+  // Only update layer visibility if it has changed.
+  if (adapter.konvaLayer.visible() !== entity.isEnabled) {
+    adapter.konvaLayer.visible(entity.isEnabled);
+    groupNeedsCache = true;
+  }
+
+  if (adapter.konvaObjectGroup.getChildren().length === 0) {
+    // No objects - clear the cache to reset the previous pixel data
+    adapter.konvaObjectGroup.clearCache();
+    return;
+  }
+
+  const compositingRect =
+    adapter.konvaLayer.findOne<Konva.Rect>(`.${COMPOSITING_RECT_NAME}`) ?? createCompositingRect(adapter.konvaLayer);
+  const isSelected = selectedEntityIdentifier?.id === entity.id;
+
+  /**
+   * When the group is selected, we use a rect of the selected preview color, composited over the shapes. This allows
+   * shapes to render as a "raster" layer with all pixels drawn at the same color and opacity.
+   *
+   * Without this special handling, each shape is drawn individually with the given opacity, atop the other shapes. The
+   * effect is like if you have a Photoshop Group consisting of many shapes, each of which has the given opacity.
+   * Overlapping shapes will have their colors blended together, and the final color is the result of all the shapes.
+   *
+   * Instead, with the special handling, the effect is as if you drew all the shapes at 100% opacity, flattened them to
+   * a single raster image, and _then_ applied the 50% opacity.
+   */
+  if (isSelected && tool !== 'move') {
+    // We must clear the cache first so Konva will re-draw the group with the new compositing rect
+    if (adapter.konvaObjectGroup.isCached()) {
+      adapter.konvaObjectGroup.clearCache();
+    }
+    // The user is allowed to reduce mask opacity to 0, but we need the opacity for the compositing rect to work
+    adapter.konvaObjectGroup.opacity(1);
+
+    compositingRect.setAttrs({
+      // The rect should be the size of the layer - use the fast method if we don't have a pixel-perfect bbox already
+      ...(!entity.bboxNeedsUpdate && entity.bbox ? entity.bbox : getLayerBboxFast(adapter.konvaLayer)),
+      fill: rgbColor,
+      opacity: globalMaskLayerOpacity,
+      // Draw this rect only where there are non-transparent pixels under it (e.g. the mask shapes)
+      globalCompositeOperation: 'source-in',
+      visible: true,
+      // This rect must always be on top of all other shapes
+      zIndex: adapter.konvaObjectGroup.getChildren().length,
+    });
+  } else {
+    // The compositing rect should only be shown when the layer is selected.
+    compositingRect.visible(false);
+    // Cache only if needed - or if we are on this code path and _don't_ have a cache
+    if (groupNeedsCache || !adapter.konvaObjectGroup.isCached()) {
+      adapter.konvaObjectGroup.cache();
+    }
+    // Updating group opacity does not require re-caching
+    adapter.konvaObjectGroup.opacity(globalMaskLayerOpacity);
+  }
+
+  // const bboxRect =
+  //   regionMap.konvaLayer.findOne<Konva.Rect>(`.${LAYER_BBOX_NAME}`) ?? createBboxRect(rg, regionMap.konvaLayer);
+
+  // if (rg.bbox) {
+  //   const active = !rg.bboxNeedsUpdate && isSelected && tool === 'move';
+  //   bboxRect.setAttrs({
+  //     visible: active,
+  //     listening: active,
+  //     x: rg.bbox.x,
+  //     y: rg.bbox.y,
+  //     width: rg.bbox.width,
+  //     height: rg.bbox.height,
+  //     stroke: isSelected ? BBOX_SELECTED_STROKE : '',
+  //   });
+  // } else {
+  //   bboxRect.visible(false);
+  // }
+};

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

@@ -313,6 +313,11 @@ export const renderToolPreview = (
   const stage = manager.stage;
   const layerCount = manager.adapters.size;
   const tool = toolState.selected;
+  const isDrawableEntity =
+    selectedEntity?.type === 'regional_guidance' ||
+    selectedEntity?.type === 'layer' ||
+    selectedEntity?.type === 'inpaint_mask';
+
   // Update the stage's pointer style
   if (tool === 'view') {
     // View gets a hand
@@ -320,8 +325,8 @@ export const renderToolPreview = (
   } else if (layerCount === 0) {
     // We have no layers, so we should not render any tool
     stage.container().style.cursor = 'default';
-  } else if (selectedEntity?.type !== 'regional_guidance' && selectedEntity?.type !== 'layer') {
-    // Non-mask-guidance layers don't have tools
+  } else if (!isDrawableEntity) {
+    // Non-drawable layers don't have tools
     stage.container().style.cursor = 'not-allowed';
   } else if (tool === 'move') {
     // Move tool gets a pointer
@@ -338,11 +343,7 @@ export const renderToolPreview = (
 
   stage.draggable(tool === 'view');
 
-  if (
-    !cursorPos ||
-    layerCount === 0 ||
-    (selectedEntity?.type !== 'regional_guidance' && selectedEntity?.type !== 'layer')
-  ) {
+  if (!cursorPos || layerCount === 0 || !isDrawableEntity) {
     // We can bail early if the mouse isn't over the stage or there are no layers
     manager.preview.tool.group.visible(false);
   } else {

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

@@ -9,6 +9,7 @@ import { arrangeEntities } from 'features/controlLayers/konva/renderers/arrange'
 import { renderBackgroundLayer } from 'features/controlLayers/konva/renderers/background';
 import { updateBboxes } from 'features/controlLayers/konva/renderers/bbox';
 import { renderControlAdapters } from 'features/controlLayers/konva/renderers/controlAdapters';
+import { renderInpaintMask } from 'features/controlLayers/konva/renderers/inpaintMask';
 import { renderLayers } from 'features/controlLayers/konva/renderers/layers';
 import {
   renderBboxPreview,
@@ -24,6 +25,11 @@ import {
   caBboxChanged,
   caTranslated,
   eraserWidthChanged,
+  imBboxChanged,
+  imBrushLineAdded,
+  imEraserLineAdded,
+  imLinePointAdded,
+  imTranslated,
   layerBboxChanged,
   layerBrushLineAdded,
   layerEraserLineAdded,
@@ -99,6 +105,8 @@ export const initializeRenderer = (
       dispatch(caTranslated(arg));
     } else if (entityType === 'regional_guidance') {
       dispatch(rgTranslated(arg));
+    } else if (entityType === 'inpaint_mask') {
+      dispatch(imTranslated(arg));
     }
   };
   const onBboxChanged = (arg: BboxChangedArg, entityType: CanvasEntity['type']) => {
@@ -109,6 +117,8 @@ export const initializeRenderer = (
       dispatch(caBboxChanged(arg));
     } else if (entityType === 'regional_guidance') {
       dispatch(rgBboxChanged(arg));
+    } else if (entityType === 'inpaint_mask') {
+      dispatch(imBboxChanged(arg));
     }
   };
   const onBrushLineAdded = (arg: BrushLineAddedArg, entityType: CanvasEntity['type']) => {
@@ -117,6 +127,8 @@ export const initializeRenderer = (
       dispatch(layerBrushLineAdded(arg));
     } else if (entityType === 'regional_guidance') {
       dispatch(rgBrushLineAdded(arg));
+    } else if (entityType === 'inpaint_mask') {
+      dispatch(imBrushLineAdded(arg));
     }
   };
   const onEraserLineAdded = (arg: EraserLineAddedArg, entityType: CanvasEntity['type']) => {
@@ -125,6 +137,8 @@ export const initializeRenderer = (
       dispatch(layerEraserLineAdded(arg));
     } else if (entityType === 'regional_guidance') {
       dispatch(rgEraserLineAdded(arg));
+    } else if (entityType === 'inpaint_mask') {
+      dispatch(imEraserLineAdded(arg));
     }
   };
   const onPointAddedToLine = (arg: PointAddedToLineArg, entityType: CanvasEntity['type']) => {
@@ -133,6 +147,8 @@ export const initializeRenderer = (
       dispatch(layerLinePointAdded(arg));
     } else if (entityType === 'regional_guidance') {
       dispatch(rgLinePointAdded(arg));
+    } else if (entityType === 'inpaint_mask') {
+      dispatch(imLinePointAdded(arg));
     }
   };
   const onRectShapeAdded = (arg: RectShapeAddedArg, entityType: CanvasEntity['type']) => {
@@ -177,6 +193,8 @@ export const initializeRenderer = (
       selectedEntity = canvasV2.ipAdapters.entities.find((i) => i.id === identifier.id) ?? null;
     } else if (identifier.type === 'regional_guidance') {
       selectedEntity = canvasV2.regions.entities.find((i) => i.id === identifier.id) ?? null;
+    } else if (identifier.type === 'inpaint_mask') {
+      selectedEntity = canvasV2.inpaintMask;
     } else {
       selectedEntity = null;
     }
@@ -328,6 +346,23 @@ export const initializeRenderer = (
       );
     }
 
+    if (
+      isFirstRender ||
+      canvasV2.inpaintMask !== prevCanvasV2.inpaintMask ||
+      canvasV2.settings.maskOpacity !== prevCanvasV2.settings.maskOpacity ||
+      canvasV2.tool.selected !== prevCanvasV2.tool.selected
+    ) {
+      logIfDebugging('Rendering inpaint mask');
+      renderInpaintMask(
+        manager,
+        canvasV2.inpaintMask,
+        canvasV2.settings.maskOpacity,
+        canvasV2.tool.selected,
+        canvasV2.selectedEntityIdentifier,
+        onPosChanged
+      );
+    }
+
     if (isFirstRender || canvasV2.controlAdapters.entities !== prevCanvasV2.controlAdapters.entities) {
       logIfDebugging('Rendering control adapters');
       renderControlAdapters(manager, canvasV2.controlAdapters.entities);

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

@@ -23,7 +23,7 @@ import { DEFAULT_RGBA_COLOR } from './types';
 
 const initialState: CanvasV2State = {
   _version: 3,
-  selectedEntityIdentifier: null,
+  selectedEntityIdentifier: { type: 'inpaint_mask', id: 'inpaint_mask' },
   layers: { entities: [], baseLayerImageCache: null },
   controlAdapters: { entities: [] },
   ipAdapters: { entities: [] },
@@ -36,7 +36,7 @@ const initialState: CanvasV2State = {
     bboxNeedsUpdate: false,
     fill: DEFAULT_RGBA_COLOR,
     imageCache: null,
-    isEnabled: false,
+    isEnabled: true,
     objects: [],
     x: 0,
     y: 0,