feat(ui): efficient canvas compositing

Also solves issue of exporting layers at different opacities than what is visible

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

@@ -4,13 +4,14 @@ import { CanvasFilter } from 'features/controlLayers/konva/CanvasFilter';
 import type { CanvasManager } from 'features/controlLayers/konva/CanvasManager';
 import { CanvasObjectRenderer } from 'features/controlLayers/konva/CanvasObjectRenderer';
 import { CanvasTransformer } from 'features/controlLayers/konva/CanvasTransformer';
-import {
-  type CanvasControlLayerState,
-  type CanvasEntityIdentifier,
-  type CanvasRasterLayerState,
-  type CanvasV2State,
-  getEntityIdentifier,
+import type {
+  CanvasControlLayerState,
+  CanvasEntityIdentifier,
+  CanvasRasterLayerState,
+  CanvasV2State,
+  Rect,
 } from 'features/controlLayers/store/types';
+import { getEntityIdentifier } from 'features/controlLayers/store/types';
 import Konva from 'konva';
 import { get } from 'lodash-es';
 import type { Logger } from 'roarr';
@@ -149,6 +150,13 @@ export class CanvasLayerAdapter {
     return { ...this.manager.getLoggingContext(), path: this.path.join('.') };
   };
 
+  getCanvas = (rect: Rect): HTMLCanvasElement => {
+    // TODO(psyche) - cache this - maybe with package `memoizee`? Would require careful review of cache invalidation
+    this.log.trace({ rect }, 'Getting canvas');
+    const canvas = this.renderer.getCanvas(rect);
+    return canvas;
+  };
+
   logDebugInfo(msg = 'Debug info') {
     const info = {
       repr: this.repr(),

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

@@ -4,11 +4,11 @@ import type { AppStore } from 'app/store/store';
 import type { JSONObject } from 'common/types';
 import { MAX_CANVAS_SCALE, MIN_CANVAS_SCALE } from 'features/controlLayers/konva/constants';
 import {
+  canvasToBlob,
+  canvasToImageData,
   getImageDataTransparency,
   getPrefixedId,
   getRectUnion,
-  konvaNodeToBlob,
-  konvaNodeToImageData,
   nanoid,
   previewBlob,
 } from 'features/controlLayers/konva/util';
@@ -27,6 +27,7 @@ import { atom } from 'nanostores';
 import type { Logger } from 'roarr';
 import { getImageDTO, uploadImage } from 'services/api/endpoints/images';
 import type { ImageDTO } from 'services/api/types';
+import { assert } from 'tsafe';
 
 import { CanvasBackground } from './CanvasBackground';
 import { CanvasLayerAdapter } from './CanvasLayerAdapter';
@@ -576,48 +577,54 @@ export class CanvasManager {
     return pixels / this.getStageScale();
   }
 
-  getCompositeRasterLayerStageClone = (): Konva.Stage => {
-    const layersState = this.stateApi.getRasterLayersState();
-    const stageClone = this.stage.clone();
+  getCompositeRasterLayerCanvas = (rect: Rect): HTMLCanvasElement => {
+    this.log.trace({ rect }, 'Building composite raster layer canvas');
 
-    stageClone.scaleX(1);
-    stageClone.scaleY(1);
-    stageClone.x(0);
-    stageClone.y(0);
+    const canvas = document.createElement('canvas');
+    canvas.width = rect.width;
+    canvas.height = rect.height;
 
-    const validLayers = layersState.entities.filter((entity) => entity.isEnabled && entity.objects.length > 0);
+    const ctx = canvas.getContext('2d');
+    assert(ctx !== null);
 
-    // getLayers() returns the internal `children` array of the stage directly - calling destroy on a layer will
-    // mutate that array. We need to clone the array to avoid mutating the original.
-    for (const konvaLayer of stageClone.getLayers().slice()) {
-      if (!validLayers.find((l) => l.id === konvaLayer.id())) {
-        konvaLayer.destroy();
+    for (const { id } of this.stateApi.getRasterLayersState().entities) {
+      const adapter = this.rasterLayerAdapters.get(id);
+      if (!adapter) {
+        this.log.warn({ id }, 'Raster layer adapter not found');
+        continue;
+      }
+      if (adapter.state.isEnabled && adapter.renderer.hasObjects()) {
+        this.log.trace({ id }, 'Drawing raster layer to composite canvas');
+        const adapterCanvas = adapter.getCanvas(rect);
+        ctx.drawImage(adapterCanvas, 0, 0);
       }
     }
-
-    return stageClone;
+    return canvas;
   };
 
-  getCompositeInpaintMaskStageClone = (): Konva.Stage => {
-    const entities = this.stateApi.getInpaintMasksState().entities;
-    const validEntities = entities.filter((entity) => entity.isEnabled && entity.objects.length > 0);
+  getCompositeInpaintMaskCanvas = (rect: Rect): HTMLCanvasElement => {
+    this.log.trace({ rect }, 'Building composite inpaint mask canvas');
 
-    const stageClone = this.stage.clone();
+    const canvas = document.createElement('canvas');
+    canvas.width = rect.width;
+    canvas.height = rect.height;
 
-    stageClone.scaleX(1);
-    stageClone.scaleY(1);
-    stageClone.x(0);
-    stageClone.y(0);
+    const ctx = canvas.getContext('2d');
+    assert(ctx !== null);
 
-    // getLayers() returns the internal `children` array of the stage directly - calling destroy on a layer will
-    // mutate that array. We need to clone the array to avoid mutating the original.
-    for (const konvaLayer of stageClone.getLayers().slice()) {
-      if (!validEntities.find((l) => l.id === konvaLayer.id())) {
-        konvaLayer.destroy();
+    for (const { id } of this.stateApi.getInpaintMasksState().entities) {
+      const adapter = this.inpaintMaskAdapters.get(id);
+      if (!adapter) {
+        this.log.warn({ id }, 'Inpaint mask adapter not found');
+        continue;
+      }
+      if (adapter.state.isEnabled && adapter.renderer.hasObjects()) {
+        this.log.trace({ id }, 'Drawing inpaint mask to composite canvas');
+        const adapterCanvas = adapter.getCanvas(rect);
+        ctx.drawImage(adapterCanvas, 0, 0);
       }
     }
-
-    return stageClone;
+    return canvas;
   };
 
   getCompositeInpaintMaskImageCache = (rect: Rect): ImageCache | null => {
@@ -646,10 +653,10 @@ export class CanvasManager {
 
     this.log.trace({ rect }, 'Rasterizing composite raster layer');
 
-    const blob = await konvaNodeToBlob(this.getCompositeRasterLayerStageClone(), rect);
-
+    const canvas = this.getCompositeRasterLayerCanvas(rect);
+    const blob = await canvasToBlob(canvas);
     if (this._isDebugging) {
-      previewBlob(blob, 'Composite raster layer');
+      previewBlob(blob, 'Composite raster layer canvas');
     }
 
     imageDTO = await uploadImage(blob, 'composite-raster-layer.png', 'general', true);
@@ -671,10 +678,10 @@ export class CanvasManager {
 
     this.log.trace({ rect }, 'Rasterizing composite inpaint mask');
 
-    const blob = await konvaNodeToBlob(this.getCompositeInpaintMaskStageClone(), rect);
-
+    const canvas = this.getCompositeInpaintMaskCanvas(rect);
+    const blob = await canvasToBlob(canvas);
     if (this._isDebugging) {
-      previewBlob(blob, 'Composite inpaint mask');
+      previewBlob(blob, 'Composite inpaint mask canvas');
     }
 
     imageDTO = await uploadImage(blob, 'composite-inpaint-mask.png', 'general', true);
@@ -684,9 +691,9 @@ export class CanvasManager {
 
   getGenerationMode(): GenerationMode {
     const { rect } = this.stateApi.getBbox();
-    const inpaintMaskImageData = konvaNodeToImageData(this.getCompositeInpaintMaskStageClone(), rect);
+    const inpaintMaskImageData = canvasToImageData(this.getCompositeInpaintMaskCanvas(rect));
     const inpaintMaskTransparency = getImageDataTransparency(inpaintMaskImageData);
-    const compositeLayerImageData = konvaNodeToImageData(this.getCompositeRasterLayerStageClone(), rect);
+    const compositeLayerImageData = canvasToImageData(this.getCompositeRasterLayerCanvas(rect));
     const compositeLayerTransparency = getImageDataTransparency(compositeLayerImageData);
     if (compositeLayerTransparency === 'FULLY_TRANSPARENT') {
       // When the initial image is fully transparent, we are always doing txt2img

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

@@ -3,14 +3,16 @@ 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 {
-  type CanvasEntityIdentifier,
-  type CanvasInpaintMaskState,
-  type CanvasRegionalGuidanceState,
-  type CanvasV2State,
-  getEntityIdentifier,
+import type {
+  CanvasEntityIdentifier,
+  CanvasInpaintMaskState,
+  CanvasRegionalGuidanceState,
+  CanvasV2State,
+  Rect,
 } from 'features/controlLayers/store/types';
+import { getEntityIdentifier } from 'features/controlLayers/store/types';
 import Konva from 'konva';
+import type { GroupConfig } from 'konva/lib/Group';
 import { get } from 'lodash-es';
 import type { Logger } from 'roarr';
 
@@ -141,6 +143,13 @@ export class CanvasMaskAdapter {
     };
   };
 
+  getCanvas = (rect: Rect): HTMLCanvasElement => {
+    // TODO(psyche): Cache this?
+    // Backend expects masks to be fully opaque
+    const attrs: GroupConfig = { opacity: 1 };
+    const canvas = this.renderer.getCanvas(rect, attrs);
+    return canvas;
+  };
   getLoggingContext = (): JSONObject => {
     return { ...this.manager.getLoggingContext(), path: this.path.join('.') };
   };

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

@@ -9,7 +9,13 @@ import type { CanvasMaskAdapter } from 'features/controlLayers/konva/CanvasMaskA
 import { CanvasRectRenderer } from 'features/controlLayers/konva/CanvasRect';
 import { LightnessToAlphaFilter } from 'features/controlLayers/konva/filters';
 import { getPatternSVG } from 'features/controlLayers/konva/patterns/getPatternSVG';
-import { getPrefixedId, konvaNodeToBlob, konvaNodeToImageData, previewBlob } from 'features/controlLayers/konva/util';
+import {
+  getPrefixedId,
+  konvaNodeToBlob,
+  konvaNodeToCanvas,
+  konvaNodeToImageData,
+  previewBlob,
+} from 'features/controlLayers/konva/util';
 import type {
   CanvasBrushLineState,
   CanvasEraserLineState,
@@ -21,6 +27,7 @@ import type {
 } from 'features/controlLayers/store/types';
 import { imageDTOToImageObject } from 'features/controlLayers/store/types';
 import Konva from 'konva';
+import type { GroupConfig } from 'konva/lib/Group';
 import { isEqual } from 'lodash-es';
 import type { Logger } from 'roarr';
 import { getImageDTO, uploadImage } from 'services/api/endpoints/images';
@@ -527,12 +534,28 @@ export class CanvasObjectRenderer {
     return imageDTO;
   };
 
-  getBlob = (rect?: Rect): Promise<Blob> => {
-    return konvaNodeToBlob(this.konva.objectGroup.clone(), rect);
+  cloneObjectGroup = (attrs?: GroupConfig): Konva.Group => {
+    const clone = this.konva.objectGroup.clone();
+    clone.cache();
+    if (attrs) {
+      clone.setAttrs(attrs);
+    }
+    return clone;
   };
 
-  getImageData = (rect?: Rect): ImageData => {
-    return konvaNodeToImageData(this.konva.objectGroup.clone(), rect);
+  getCanvas = (rect?: Rect, attrs?: GroupConfig): HTMLCanvasElement => {
+    const clone = this.cloneObjectGroup(attrs);
+    return konvaNodeToCanvas(clone, rect);
+  };
+
+  getBlob = (rect?: Rect, attrs?: GroupConfig): Promise<Blob> => {
+    const clone = this.cloneObjectGroup(attrs);
+    return konvaNodeToBlob(clone, rect);
+  };
+
+  getImageData = (rect?: Rect, attrs?: GroupConfig): ImageData => {
+    const clone = this.cloneObjectGroup(attrs);
+    return konvaNodeToImageData(clone, rect);
   };
 
   /**

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

@@ -49,7 +49,7 @@ import type {
   RgbaColor,
   Tool,
 } from 'features/controlLayers/store/types';
-import { RGBA_RED } from 'features/controlLayers/store/types';
+import { RGBA_BLACK } from 'features/controlLayers/store/types';
 import type { WritableAtom } from 'nanostores';
 import { atom } from 'nanostores';
 import { $lastCanvasProgressEvent } from 'services/events/setEventListeners';
@@ -167,9 +167,6 @@ export class CanvasStateApi {
   getIsSelected = (id: string) => {
     return this.getState().selectedEntityIdentifier?.id === id;
   };
-  getFilterState = () => {
-    return this._store.getState().canvasV2.filter;
-  };
 
   getEntity(identifier: CanvasEntityIdentifier): EntityStateAndAdapter | null {
     const state = this.getState();
@@ -218,7 +215,7 @@ export class CanvasStateApi {
     if (selectedEntity) {
       // These two entity types use a compositing rect for opacity. Their fill is always a solid color.
       if (selectedEntity.state.type === 'regional_guidance' || selectedEntity.state.type === 'inpaint_mask') {
-        currentFill = RGBA_RED;
+        currentFill = RGBA_BLACK;
       }
     }
     return currentFill;

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

@@ -509,6 +509,7 @@ const zRgbaColor = zRgbColor.extend({
 });
 export type RgbaColor = z.infer<typeof zRgbaColor>;
 export const RGBA_RED: RgbaColor = { r: 255, g: 0, b: 0, a: 1 };
+export const RGBA_BLACK: RgbaColor = { r: 0, g: 0, b: 0, a: 1 };
 export const RGBA_WHITE: RgbaColor = { r: 255, g: 255, b: 255, a: 1 };
 
 const zOpacity = z.number().gte(0).lte(1);