refactor(ui): do not rely on konva internal canvas cache for layer previews

Instead of pulling the preview canvas from the konva internals, use the canvas created for bbox calculations as the preview canvas. This doesn't change perf characteristics, because we were already creating this canvas. It just means we don't need to dip into the konva internals. It fixes an issue where the layer preview didn't update or show when a layer is disabled or otherwise hidden.
2026-04-23 03:00:31 -04:00 · 2024-10-07 19:55:54 +10:00
parent ad76399702
commit 883beb90eb
9 changed files with 70 additions and 104 deletions
--- a/invokeai/frontend/web/src/features/controlLayers/components/common/CanvasEntityPreviewImage.tsx
+++ b/invokeai/frontend/web/src/features/controlLayers/components/common/CanvasEntityPreviewImage.tsx
@@ -6,13 +6,12 @@ import { useEntityAdapter } from 'features/controlLayers/contexts/EntityAdapterC
 import { useEntityIdentifierContext } from 'features/controlLayers/contexts/EntityIdentifierContext';
 import { TRANSPARENCY_CHECKERBOARD_PATTERN_DARK_DATAURL } from 'features/controlLayers/konva/patterns/transparency-checkerboard-pattern';
 import { selectCanvasSlice, selectEntity } from 'features/controlLayers/store/selectors';
+import { debounce } from 'lodash-es';
 import { memo, useEffect, useMemo, useRef } from 'react';
 import { useSelector } from 'react-redux';

 const ChakraCanvas = chakra.canvas;

-const PADDING = 2;
-
 const CONTAINER_WIDTH = 36; // this is size 12 in our theme - need it in px for the canvas
 const CONTAINER_WIDTH_PX = `${CONTAINER_WIDTH}px`;

@@ -35,48 +34,57 @@ export const CanvasEntityPreviewImage = memo(() => {
  );
  const maskColor = useSelector(selectMaskColor);
  const canvasRef = useRef<HTMLCanvasElement>(null);
-  const cache = useStore(adapter.renderer.$canvasCache);
-  useEffect(() => {
-    if (!canvasRef.current) {
-      return;
-    }
-    const ctx = canvasRef.current.getContext('2d');
-    if (!ctx) {
-      return;
-    }
+  const pixelRect = useStore(adapter.transformer.$pixelRect);
+  const nodeRect = useStore(adapter.transformer.$nodeRect);
+  const canvasCache = useStore(adapter.$canvasCache);

-    if (!cache) {
-      // Draw an empty canvas
-      ctx.clearRect(0, 0, canvasRef.current.width, canvasRef.current.height);
-      return;
-    }
+  const updatePreview = useMemo(
+    () =>
+      debounce(() => {
+        if (!canvasRef.current) {
+          return;
+        }
+        const ctx = canvasRef.current.getContext('2d');
+        if (!ctx) {
+          return;
+        }

-    const { rect, canvas } = cache;
+        const pixelRect = adapter.transformer.$pixelRect.get();
+        const nodeRect = adapter.transformer.$nodeRect.get();
+        const canvasCache = adapter.$canvasCache.get();

-    ctx.clearRect(0, 0, canvasRef.current.width, canvasRef.current.height);
+        if (!canvasCache || canvasCache.width === 0 || canvasCache.height === 0) {
+          // Draw an empty canvas
+          ctx.clearRect(0, 0, canvasRef.current.width, canvasRef.current.height);
+          return;
+        }

-    canvasRef.current.width = rect.width;
-    canvasRef.current.height = rect.height;
+        ctx.clearRect(0, 0, canvasRef.current.width, canvasRef.current.height);

-    const scale = CONTAINER_WIDTH / rect.width;
+        canvasRef.current.width = pixelRect.width;
+        canvasRef.current.height = pixelRect.height;

-    const sx = rect.x;
-    const sy = rect.y;
-    const sWidth = rect.width;
-    const sHeight = rect.height;
-    const dx = PADDING / scale;
-    const dy = PADDING / scale;
-    const dWidth = rect.width - (PADDING * 2) / scale;
-    const dHeight = rect.height - (PADDING * 2) / scale;
+        const sx = pixelRect.x - nodeRect.x;
+        const sy = pixelRect.y - nodeRect.y;
+        const sWidth = pixelRect.width;
+        const sHeight = pixelRect.height;
+        const dx = 0;
+        const dy = 0;
+        const dWidth = pixelRect.width;
+        const dHeight = pixelRect.height;

-    ctx.drawImage(canvas, sx, sy, sWidth, sHeight, dx, dy, dWidth, dHeight);
+        ctx.drawImage(canvasCache, sx, sy, sWidth, sHeight, dx, dy, dWidth, dHeight);

-    if (maskColor) {
-      ctx.fillStyle = maskColor;
-      ctx.globalCompositeOperation = 'source-in';
-      ctx.fillRect(0, 0, rect.width, rect.height);
-    }
-  }, [cache, maskColor]);
+        if (maskColor) {
+          ctx.fillStyle = maskColor;
+          ctx.globalCompositeOperation = 'source-in';
+          ctx.fillRect(0, 0, pixelRect.width, pixelRect.height);
+        }
+      }, 300),
+    [adapter.$canvasCache, adapter.transformer.$nodeRect, adapter.transformer.$pixelRect, maskColor]
+  );
+
+  useEffect(updatePreview, [updatePreview, canvasCache, nodeRect, pixelRect]);

  return (
    <Flex
--- a/invokeai/frontend/web/src/features/controlLayers/konva/CanvasEntity/CanvasEntityAdapterBase.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/konva/CanvasEntity/CanvasEntityAdapterBase.ts
@@ -124,6 +124,10 @@ export abstract class CanvasEntityAdapterBase<
  $isInteractable = computed([this.$isLocked, this.$isDisabled, this.$isHidden], (isLocked, isDisabled, isHidden) => {
    return !isLocked && !isDisabled && !isHidden;
  });
+  /**
+   * A cache of the entity's canvas element. This is generated from a clone of the entity's Konva layer.
+   */
+  $canvasCache = atom<HTMLCanvasElement | null>(null);

  constructor(entityIdentifier: CanvasEntityIdentifier<T['type']>, manager: CanvasManager, adapterType: U) {
    super();
@@ -333,6 +337,7 @@ export abstract class CanvasEntityAdapterBase<
      renderer: this.renderer.repr(),
      bufferRenderer: this.bufferRenderer.repr(),
      filterer: this.filterer?.repr(),
+      hasCache: this.$canvasCache.get() !== null,
    };
  };
 }
--- a/invokeai/frontend/web/src/features/controlLayers/konva/CanvasEntity/CanvasEntityAdapterControlLayer.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/konva/CanvasEntity/CanvasEntityAdapterControlLayer.ts
@@ -69,7 +69,7 @@ export class CanvasEntityAdapterControlLayer extends CanvasEntityAdapterBase<
    this.log.trace({ rect }, 'Getting canvas');
    // The opacity may have been changed in response to user selecting a different entity category, so we must restore
    // the original opacity before rendering the canvas
-    const attrs: GroupConfig = { opacity: this.state.opacity };
+    const attrs: GroupConfig = { opacity: this.state.opacity, filters: [] };
    const canvas = this.renderer.getCanvas({ rect, attrs });
    return canvas;
  };
--- a/invokeai/frontend/web/src/features/controlLayers/konva/CanvasEntity/CanvasEntityAdapterInpaintMask.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/konva/CanvasEntity/CanvasEntityAdapterInpaintMask.ts
@@ -78,7 +78,7 @@ export class CanvasEntityAdapterInpaintMask extends CanvasEntityAdapterBase<
  getCanvas = (rect?: Rect): HTMLCanvasElement => {
    // The opacity may have been changed in response to user selecting a different entity category, and the mask regions
    // should be fully opaque - set opacity to 1 before rendering the canvas
-    const attrs: GroupConfig = { opacity: 1 };
+    const attrs: GroupConfig = { opacity: 1, filters: [] };
    const canvas = this.renderer.getCanvas({ rect, attrs });
    return canvas;
  };
--- a/invokeai/frontend/web/src/features/controlLayers/konva/CanvasEntity/CanvasEntityAdapterRasterLayer.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/konva/CanvasEntity/CanvasEntityAdapterRasterLayer.ts
@@ -62,7 +62,7 @@ export class CanvasEntityAdapterRasterLayer extends CanvasEntityAdapterBase<
    this.log.trace({ rect }, 'Getting canvas');
    // The opacity may have been changed in response to user selecting a different entity category, so we must restore
    // the original opacity before rendering the canvas
-    const attrs: GroupConfig = { opacity: this.state.opacity };
+    const attrs: GroupConfig = { opacity: this.state.opacity, filters: [] };
    const canvas = this.renderer.getCanvas({ rect, attrs });
    return canvas;
  };
--- a/invokeai/frontend/web/src/features/controlLayers/konva/CanvasEntity/CanvasEntityAdapterRegionalGuidance.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/konva/CanvasEntity/CanvasEntityAdapterRegionalGuidance.ts
@@ -78,7 +78,7 @@ export class CanvasEntityAdapterRegionalGuidance extends CanvasEntityAdapterBase
  getCanvas = (rect?: Rect): HTMLCanvasElement => {
    // The opacity may have been changed in response to user selecting a different entity category, and the mask regions
    // should be fully opaque - set opacity to 1 before rendering the canvas
-    const attrs: GroupConfig = { opacity: 1 };
+    const attrs: GroupConfig = { opacity: 1, filters: [] };
    const canvas = this.renderer.getCanvas({ rect, attrs });
    return canvas;
  };
--- a/invokeai/frontend/web/src/features/controlLayers/konva/CanvasEntity/CanvasEntityObjectRenderer.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/konva/CanvasEntity/CanvasEntityObjectRenderer.ts
@@ -1,6 +1,5 @@
 import { $authToken } from 'app/store/nanostores/authToken';
 import { rgbColorToString } from 'common/util/colorCodeTransformers';
-import { withResult } from 'common/util/result';
 import { SyncableMap } from 'common/util/SyncableMap/SyncableMap';
 import type { CanvasEntityAdapter } from 'features/controlLayers/konva/CanvasEntity/types';
 import type { CanvasManager } from 'features/controlLayers/konva/CanvasManager';
@@ -26,10 +25,7 @@ import type { Rect } from 'features/controlLayers/store/types';
 import { imageDTOToImageObject } from 'features/controlLayers/store/util';
 import Konva from 'konva';
 import type { GroupConfig } from 'konva/lib/Group';
-import { debounce } from 'lodash-es';
-import { atom } from 'nanostores';
 import type { Logger } from 'roarr';
-import { serializeError } from 'serialize-error';
 import { getImageDTOSafe, uploadImage } from 'services/api/endpoints/images';
 import type { ImageDTO } from 'services/api/types';
 import { assert } from 'tsafe';
@@ -96,17 +92,6 @@ export class CanvasEntityObjectRenderer extends CanvasModuleBase {
    } | null;
  };

-  /**
-   * The entity's object group as a canvas element along with the pixel rect of the entity at the time the canvas was
-   * drawn.
-   *
-   * Technically, this is an internal Konva object, created when a Konva node's `.cache()` method is called. We cache
-   * the object group after every update, so we get this as a "free" side effect.
-   *
-   * This is used to render the entity's preview in the control layer.
-   */
-  $canvasCache = atom<{ canvas: HTMLCanvasElement; rect: Rect } | null>(null);
-
  constructor(parent: CanvasEntityAdapter) {
    super();
    this.id = getPrefixedId(this.type);
@@ -198,12 +183,10 @@ export class CanvasEntityObjectRenderer extends CanvasModuleBase {
    if (this.renderers.size === 0) {
      this.log.trace('Clearing object group cache');
      this.konva.objectGroup.clearCache();
-      this.$canvasCache.set(null);
    } else if (force || !this.konva.objectGroup.isCached()) {
      this.log.trace('Caching object group');
      this.konva.objectGroup.clearCache();
      this.konva.objectGroup.cache({ pixelRatio: 1, imageSmoothingEnabled: false });
-      this.parent.renderer.updatePreviewCanvas();
    }
  };

@@ -459,52 +442,15 @@ export class CanvasEntityObjectRenderer extends CanvasModuleBase {
    return imageDTO;
  };

-  updatePreviewCanvas = debounce(() => {
-    if (this.parent.transformer.$isPendingRectCalculation.get()) {
-      return;
-    }
-
-    const pixelRect = this.parent.transformer.$pixelRect.get();
-    if (pixelRect.width === 0 || pixelRect.height === 0) {
-      return;
-    }
-
-    /**
-     * TODO(psyche): This is an internal Konva method, so it may break in the future. Can we make this API public?
-     *
-     * This method's API is unknown. It has been experimentally determined that it may throw, so we need to handle
-     * errors.
-     */
-    const getCacheCanvasResult = withResult(
-      () => this.konva.objectGroup._getCachedSceneCanvas()._canvas as HTMLCanvasElement | undefined | null
-    );
-    if (getCacheCanvasResult.isErr()) {
-      // We are using an internal Konva method, so we need to catch any errors that may occur.
-      this.log.warn({ error: serializeError(getCacheCanvasResult.error) }, 'Failed to update preview canvas');
-      return;
-    }
-
-    const canvas = getCacheCanvasResult.value;
-
-    if (canvas) {
-      const nodeRect = this.parent.transformer.$nodeRect.get();
-      const rect = {
-        x: pixelRect.x - nodeRect.x,
-        y: pixelRect.y - nodeRect.y,
-        width: pixelRect.width,
-        height: pixelRect.height,
-      };
-      this.$canvasCache.set({ rect, canvas });
-    }
-  }, 300);
-
  cloneObjectGroup = (arg: { attrs?: GroupConfig } = {}): Konva.Group => {
    const { attrs } = arg;
    const clone = this.konva.objectGroup.clone();
    if (attrs) {
      clone.setAttrs(attrs);
    }
-    clone.cache({ pixelRatio: 1, imageSmoothingEnabled: false });
+    if (clone.hasChildren()) {
+      clone.cache({ pixelRatio: 1, imageSmoothingEnabled: false });
+    }
    return clone;
  };

@@ -551,7 +497,6 @@ export class CanvasEntityObjectRenderer extends CanvasModuleBase {
      konva: {
        objectGroup: getKonvaNodeDebugAttrs(this.konva.objectGroup),
      },
-      hasCache: this.$canvasCache.get() !== null,
    };
  };
 }
--- a/invokeai/frontend/web/src/features/controlLayers/konva/CanvasEntity/CanvasEntityTransformer.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/konva/CanvasEntity/CanvasEntityTransformer.ts
@@ -96,14 +96,21 @@ export class CanvasEntityTransformer extends CanvasModuleBase {
  config: CanvasEntityTransformerConfig = DEFAULT_CONFIG;

  /**
-   * The rect of the parent, _including_ transparent regions.
+   * The rect of the parent, _including_ transparent regions, **relative to the parent's position**. To get the rect
+   * relative to the _stage_, add the parent's position.
+   *
   * It is calculated via Konva's getClientRect method, which is fast but includes transparent regions.
+   *
+   * This rect is relative _to the parent's position_, not the stage.
   */
  $nodeRect = atom<Rect>(getEmptyRect());

  /**
-   * The rect of the parent, _excluding_ transparent regions.
+   * The rect of the parent, _excluding_ transparent regions, **relative to the parent's position**. To get the rect
+   * relative to the _stage_, add the parent's position.
+   *
   * If the parent's nodes have no possibility of transparent regions, this will be calculated the same way as nodeRect.
+   *
   * If the parent's nodes may have transparent regions, this will be calculated manually by rasterizing the parent and
   * checking the pixel data.
   */
@@ -795,19 +802,19 @@ export class CanvasEntityTransformer extends CanvasModuleBase {
      this.parent.renderer.konva.objectGroup.setAttrs(groupAttrs);
      this.parent.bufferRenderer.konva.group.setAttrs(groupAttrs);
    }
-
-    this.parent.renderer.updatePreviewCanvas();
  };

  calculateRect = debounce(() => {
    this.log.debug('Calculating bbox');

    this.$isPendingRectCalculation.set(true);
+    const canvas = this.parent.getCanvas();

    if (!this.parent.renderer.hasObjects()) {
      this.log.trace('No objects, resetting bbox');
      this.$nodeRect.set(getEmptyRect());
      this.$pixelRect.set(getEmptyRect());
+      this.parent.$canvasCache.set(canvas);
      this.$isPendingRectCalculation.set(false);
      this.updateBbox();
      return;
@@ -819,13 +826,13 @@ export class CanvasEntityTransformer extends CanvasModuleBase {
      this.$nodeRect.set({ ...rect });
      this.$pixelRect.set({ ...rect });
      this.log.trace({ nodeRect: this.$nodeRect.get(), pixelRect: this.$pixelRect.get() }, 'Got bbox from client rect');
+      this.parent.$canvasCache.set(canvas);
      this.$isPendingRectCalculation.set(false);
      this.updateBbox();
      return;
    }

    // We have eraser strokes - we must calculate the bbox using pixel data
-    const canvas = this.parent.renderer.getCanvas({ attrs: { opacity: 1, filters: [] } });
    const imageData = canvasToImageData(canvas);
    this.manager.worker.requestBbox(
      { buffer: imageData.data.buffer, width: imageData.width, height: imageData.height },
@@ -847,6 +854,7 @@ export class CanvasEntityTransformer extends CanvasModuleBase {
          { nodeRect: this.$nodeRect.get(), pixelRect: this.$pixelRect.get(), extents },
          `Got bbox from worker`
        );
+        this.parent.$canvasCache.set(canvas);
        this.$isPendingRectCalculation.set(false);
        this.updateBbox();
      }
--- a/invokeai/frontend/web/src/features/controlLayers/konva/util.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/konva/util.ts
@@ -333,7 +333,7 @@ export const dataURLToImageData = (dataURL: string, width: number, height: numbe

 export const konvaNodeToCanvas = (arg: { node: Konva.Node; rect?: Rect; bg?: string }): HTMLCanvasElement => {
  const { node, rect, bg } = arg;
-  const canvas = node.toCanvas({ ...(rect ?? {}), imageSmoothingEnabled: false });
+  const canvas = node.toCanvas({ ...(rect ?? {}), imageSmoothingEnabled: false, pixelRatio: 1 });

  if (!bg) {
    return canvas;