feat(ui): optimized empty mask logic

Turns out, it's more efficient to just use the bbox logic for empty mask calculations. We already track if if the bbox needs updating, so this calculation does minimal work. The dedicated calculation wasn't able to use the bbox tracking so it ran far more often than the bbox calculation. Removed the "fast" bbox calculation logic, bc the new logic means we are continually updating the bbox in the background - not only when the user switches to the move tool and/or selects a layer. The bbox calculation logic is split out from the bbox rendering logic to support this. Result - better perf overall, with the empty mask handling retained.
2024-08-30 20:32:17 +00:00 · 2024-05-09 11:54:21 +10:00 · 2024-05-09 11:54:21 +10:00 · 1533429e54
commit 1533429e54
parent fc000214a5
5 changed files with 69 additions and 135 deletions
--- a/invokeai/frontend/web/src/features/controlLayers/components/StageComponent.tsx
+++ b/invokeai/frontend/web/src/features/controlLayers/components/StageComponent.tsx
@ -178,7 +178,7 @@ const useStageRenderer = (
      // Preview should not display bboxes
      return;
    }
-    renderers.renderBbox(stage, state.layers, tool, onBboxChanged);
+    renderers.renderBboxes(stage, state.layers, tool);
  }, [stage, asPreview, state.layers, tool, onBboxChanged, renderers]);

  useLayoutEffect(() => {
@ -186,8 +186,8 @@ const useStageRenderer = (
      // Preview should not check for transparency
      return;
    }
-    log.trace('Checking for transparency');
-    debouncedRenderers.checkForTransparency(stage, state.layers, onBboxChanged);
+    log.trace('Updating bboxes');
+    debouncedRenderers.updateBboxes(stage, state.layers, onBboxChanged);
  }, [stage, asPreview, state.layers, onBboxChanged]);

  useLayoutEffect(() => {
--- a/invokeai/frontend/web/src/features/controlLayers/store/controlLayersSlice.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/store/controlLayersSlice.ts
@ -137,13 +137,11 @@ export const controlLayersSlice = createSlice({
  reducers: {
    //#region Any Layer Type
    layerSelected: (state, action: PayloadAction<string>) => {
-      for (const layer of state.layers.filter(isRenderableLayer)) {
-        if (layer.id === action.payload) {
-          layer.isSelected = true;
-          state.selectedLayerId = action.payload;
-        } else {
-          layer.isSelected = false;
-        }
+      deselectAllLayers(state);
+      const layer = state.layers.find((l) => l.id === action.payload);
+      if (isRenderableLayer(layer)) {
+        layer.isSelected = true;
+        state.selectedLayerId = layer.id;
      }
    },
    layerVisibilityToggled: (state, action: PayloadAction<string>) => {
@ -173,7 +171,6 @@ export const controlLayersSlice = createSlice({
          // The layer was fully erased, empty its objects to prevent accumulation of invisible objects
          layer.maskObjects = [];
          layer.uploadedMaskImage = null;
-          layer.needsPixelBbox = false;
        }
      }
    },
@ -184,7 +181,6 @@ export const controlLayersSlice = createSlice({
        layer.maskObjects = [];
        layer.bbox = null;
        layer.isEnabled = true;
-        layer.needsPixelBbox = false;
        layer.bboxNeedsUpdate = false;
        layer.uploadedMaskImage = null;
      }
@ -236,6 +232,7 @@ export const controlLayersSlice = createSlice({
        action: PayloadAction<{ layerId: string; controlAdapter: ControlNetConfigV2 | T2IAdapterConfigV2 }>
      ) => {
        const { layerId, controlAdapter } = action.payload;
+        deselectAllLayers(state);
        const layer: ControlAdapterLayer = {
          id: getCALayerId(layerId),
          type: 'control_adapter_layer',
@ -251,11 +248,6 @@ export const controlLayersSlice = createSlice({
        };
        state.layers.push(layer);
        state.selectedLayerId = layer.id;
-        for (const layer of state.layers.filter(isRenderableLayer)) {
-          if (layer.id !== layerId) {
-            layer.isSelected = false;
-          }
-        }
      },
      prepare: (controlAdapter: ControlNetConfigV2 | T2IAdapterConfigV2) => ({
        payload: { layerId: uuidv4(), controlAdapter },
@ -439,6 +431,7 @@ export const controlLayersSlice = createSlice({
    rgLayerAdded: {
      reducer: (state, action: PayloadAction<{ layerId: string }>) => {
        const { layerId } = action.payload;
+        deselectAllLayers(state);
        const layer: RegionalGuidanceLayer = {
          id: getRGLayerId(layerId),
          type: 'regional_guidance_layer',
@ -450,7 +443,6 @@ export const controlLayersSlice = createSlice({
          x: 0,
          y: 0,
          autoNegative: 'invert',
-          needsPixelBbox: false,
          positivePrompt: '',
          negativePrompt: null,
          ipAdapters: [],
@ -459,11 +451,6 @@ export const controlLayersSlice = createSlice({
        };
        state.layers.push(layer);
        state.selectedLayerId = layer.id;
-        for (const layer of state.layers.filter(isRenderableLayer)) {
-          if (layer.id !== layerId) {
-            layer.isSelected = false;
-          }
-        }
      },
      prepare: () => ({ payload: { layerId: uuidv4() } }),
    },
@ -511,9 +498,6 @@ export const controlLayersSlice = createSlice({
        });
        layer.bboxNeedsUpdate = true;
        layer.uploadedMaskImage = null;
-        if (!layer.needsPixelBbox && tool === 'eraser') {
-          layer.needsPixelBbox = true;
-        }
      },
      prepare: (payload: { layerId: string; points: [number, number, number, number]; tool: DrawingTool }) => ({
        payload: { ...payload, lineUuid: uuidv4() },
@ -638,6 +622,7 @@ export const controlLayersSlice = createSlice({
    //#region Initial Image Layer
    iiLayerAdded: {
      reducer: (state, action: PayloadAction<{ layerId: string; imageDTO: ImageDTO | null }>) => {
+        deselectAllLayers(state);
        const { layerId, imageDTO } = action.payload;
        // Highlander! There can be only one!
        state.layers = state.layers.filter((l) => (isInitialImageLayer(l) ? false : true));
@ -656,11 +641,6 @@ export const controlLayersSlice = createSlice({
        };
        state.layers.push(layer);
        state.selectedLayerId = layer.id;
-        for (const layer of state.layers.filter(isRenderableLayer)) {
-          if (layer.id !== layerId) {
-            layer.isSelected = false;
-          }
-        }
      },
      prepare: (imageDTO: ImageDTO | null) => ({ payload: { layerId: INITIAL_IMAGE_LAYER_ID, imageDTO } }),
    },
--- a/invokeai/frontend/web/src/features/controlLayers/store/types.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/store/types.ts
@ -92,11 +92,6 @@ const zRegionalGuidanceLayer = zRenderableLayerBase.extend({
  ipAdapters: z.array(zIPAdapterConfigV2),
  previewColor: zRgbColor,
  autoNegative: zAutoNegative,
-  needsPixelBbox: z
-    .boolean()
-    .describe(
-      'Whether the layer needs the slower pixel-based bbox calculation. Set to true when an there is an eraser object.'
-    ),
  uploadedMaskImage: zImageWithDims.nullable(),
 });
 export type RegionalGuidanceLayer = z.infer<typeof zRegionalGuidanceLayer>;
--- a/invokeai/frontend/web/src/features/controlLayers/util/bbox.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/util/bbox.ts
@ -52,25 +52,6 @@ const getImageDataBbox = (imageData: ImageData): Extents | null => {
  return isEmpty ? null : { minX, minY, maxX, maxY };
 };

-/**
- * Check if an image is fully transparent.
- * @param imageData The ImageData object to check for transparency.
- * @returns Whether the image is fully transparent.
- */
-const getIsFullyTransparent = (imageData: ImageData) => {
-  if (!imageData.height || !imageData.width || imageData.data.length === 0) {
-    return true;
-  }
-  const data = imageData.data;
-  const len = data.length / 4;
-  for (let i = 0; i < len; i++) {
-    if (data[i * 4 + 3] ?? 0 > 0) {
-      return false;
-    }
-  }
-  return true;
-};
-
 /**
 * Clones a regional guidance konva layer onto an offscreen stage/canvas. This allows the pixel data for a given layer
 * to be captured, manipulated or analyzed without interference from other layers.
@ -172,22 +153,3 @@ export const getLayerBboxFast = (layer: Konva.Layer): IRect => {
    height: Math.floor(bbox.height),
  };
 };
-
-export const getIsLayerTransparent = (layer: Konva.Layer): boolean => {
-  const { stageClone, layerClone } = getIsolatedRGLayerClone(layer);
-
-  // Get a worst-case rect using the relatively fast `getClientRect`.
-  const layerRect = layerClone.getClientRect();
-  if (layerRect.width === 0 || layerRect.height === 0) {
-    return true;
-  }
-
-  // Capture the image data with the above rect.
-  const layerImageData = stageClone
-    .toCanvas(layerRect)
-    .getContext('2d')
-    ?.getImageData(0, 0, layerRect.width, layerRect.height);
-  assert(layerImageData, "Unable to get layer's image data");
-
-  return getIsFullyTransparent(layerImageData);
-};
--- a/invokeai/frontend/web/src/features/controlLayers/util/renderers.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/util/renderers.ts
@ -40,7 +40,7 @@ import type {
  VectorMaskLine,
  VectorMaskRect,
 } from 'features/controlLayers/store/types';
-import { getIsLayerTransparent, getLayerBboxFast, getLayerBboxPixels } from 'features/controlLayers/util/bbox';
+import { getLayerBboxFast, getLayerBboxPixels } from 'features/controlLayers/util/bbox';
 import { t } from 'i18next';
 import Konva from 'konva';
 import type { IRect, Vector2d } from 'konva/lib/types';
@ -437,8 +437,8 @@ const renderRegionalGuidanceLayer = (
    konvaObjectGroup.opacity(1);

    compositingRect.setAttrs({
-      // The rect should be the size of the layer - use the fast method bc it's OK if the rect is larger
-      ...getLayerBboxFast(konvaLayer),
+      // The rect should be the size of the layer - use the fast method if we don't have a pixel-perfect bbox already
+      ...(!reduxLayer.bboxNeedsUpdate && reduxLayer.bbox ? reduxLayer.bbox : getLayerBboxFast(konvaLayer)),
      fill: rgbColor,
      opacity: globalMaskLayerOpacity,
      // Draw this rect only where there are non-transparent pixels under it (e.g. the mask shapes)
@ -716,6 +716,7 @@ const createBboxRect = (reduxLayer: Layer, konvaLayer: Konva.Layer) => {
    id: getLayerBboxId(reduxLayer.id),
    name: LAYER_BBOX_NAME,
    strokeWidth: 1,
+    visible: false,
  });
  konvaLayer.add(rect);
  return rect;
@ -725,18 +726,10 @@ const createBboxRect = (reduxLayer: Layer, konvaLayer: Konva.Layer) => {
 * Renders the bounding boxes for the layers.
 * @param stage The konva stage to render on
 * @param reduxLayers An array of all redux layers to draw bboxes for
- * @param selectedLayerId The selected layer's id
 * @param tool The current tool
- * @param onBboxChanged Callback for when the bbox is changed
- * @param onBboxMouseDown Callback for when the bbox is clicked
 * @returns
 */
-const renderBbox = (
-  stage: Konva.Stage,
-  reduxLayers: Layer[],
-  tool: Tool,
-  onBboxChanged: (layerId: string, bbox: IRect | null) => void
-) => {
+const renderBboxes = (stage: Konva.Stage, reduxLayers: Layer[], tool: Tool) => {
  // Hide all bboxes so they don't interfere with getClientRect
  for (const bboxRect of stage.find<Konva.Rect>(`.${LAYER_BBOX_NAME}`)) {
    bboxRect.visible(false);
@ -747,36 +740,59 @@ const renderBbox = (
    return;
  }

-  for (const reduxLayer of reduxLayers) {
-    if (reduxLayer.type === 'regional_guidance_layer') {
-      const konvaLayer = stage.findOne<Konva.Layer>(`#${reduxLayer.id}`);
-      assert(konvaLayer, `Layer ${reduxLayer.id} not found in stage`);
+  for (const reduxLayer of reduxLayers.filter(isRegionalGuidanceLayer)) {
+    if (!reduxLayer.bbox) {
+      continue;
+    }
+    const konvaLayer = stage.findOne<Konva.Layer>(`#${reduxLayer.id}`);
+    assert(konvaLayer, `Layer ${reduxLayer.id} not found in stage`);

-      let bbox = reduxLayer.bbox;
+    const bboxRect = konvaLayer.findOne<Konva.Rect>(`.${LAYER_BBOX_NAME}`) ?? createBboxRect(reduxLayer, konvaLayer);

-      // We only need to recalculate the bbox if the layer has changed and it has objects
-      if (reduxLayer.bboxNeedsUpdate && reduxLayer.maskObjects.length) {
-        // We only need to use the pixel-perfect bounding box if the layer has eraser strokes
-        bbox = reduxLayer.needsPixelBbox ? getLayerBboxPixels(konvaLayer) : getLayerBboxFast(konvaLayer);
-        // Update the layer's bbox in the redux store
-        onBboxChanged(reduxLayer.id, bbox);
+    bboxRect.setAttrs({
+      visible: !reduxLayer.bboxNeedsUpdate,
+      listening: reduxLayer.isSelected,
+      x: reduxLayer.bbox.x,
+      y: reduxLayer.bbox.y,
+      width: reduxLayer.bbox.width,
+      height: reduxLayer.bbox.height,
+      stroke: reduxLayer.isSelected ? BBOX_SELECTED_STROKE : '',
+    });
+  }
+};
+
+/**
+ * Calculates the bbox of each regional guidance layer. Only calculates if the mask has changed.
+ * @param stage The konva stage to render on.
+ * @param reduxLayers An array of redux layers to calculate bboxes for
+ * @param onBboxChanged Callback for when the bounding box changes
+ */
+const updateBboxes = (
+  stage: Konva.Stage,
+  reduxLayers: Layer[],
+  onBboxChanged: (layerId: string, bbox: IRect | null) => void
+) => {
+  for (const rgLayer of reduxLayers.filter(isRegionalGuidanceLayer)) {
+    const konvaLayer = stage.findOne<Konva.Layer>(`#${rgLayer.id}`);
+    assert(konvaLayer, `Layer ${rgLayer.id} not found in stage`);
+    // We only need to recalculate the bbox if the layer has changed
+    if (rgLayer.bboxNeedsUpdate) {
+      const bboxRect = konvaLayer.findOne<Konva.Rect>(`.${LAYER_BBOX_NAME}`) ?? createBboxRect(rgLayer, konvaLayer);
+
+      // Hide the bbox while we calculate the new bbox, else the bbox will be included in the calculation
+      const visible = bboxRect.visible();
+      bboxRect.visible(false);
+
+      if (rgLayer.maskObjects.length === 0) {
+        // No objects - no bbox to calculate
+        onBboxChanged(rgLayer.id, null);
+      } else {
+        // Calculate the bbox by rendering the layer and checking its pixels
+        onBboxChanged(rgLayer.id, getLayerBboxPixels(konvaLayer));
      }

-      if (!bbox) {
-        continue;
-      }
-
-      const rect = konvaLayer.findOne<Konva.Rect>(`.${LAYER_BBOX_NAME}`) ?? createBboxRect(reduxLayer, konvaLayer);
-
-      rect.setAttrs({
-        visible: true,
-        listening: reduxLayer.isSelected,
-        x: bbox.x,
-        y: bbox.y,
-        width: bbox.width,
-        height: bbox.height,
-        stroke: reduxLayer.isSelected ? BBOX_SELECTED_STROKE : '',
-      });
+      // Restore the visibility of the bbox
+      bboxRect.visible(visible);
    }
  }
 };
@ -890,33 +906,14 @@ const renderNoLayersMessage = (stage: Konva.Stage, layerCount: number, width: nu
  }
 };

-const checkForTransparency = (
-  stage: Konva.Stage,
-  reduxLayers: Layer[],
-  onBboxChanged: (layerId: string, bbox: IRect | null) => void
-) => {
-  for (const reduxLayer of reduxLayers.filter(isRegionalGuidanceLayer)) {
-    if (!reduxLayer.needsPixelBbox) {
-      continue;
-    }
-    const konvaLayer = stage.findOne<Konva.Layer>(`#${reduxLayer.id}`);
-    if (!konvaLayer) {
-      continue;
-    }
-    if (getIsLayerTransparent(konvaLayer)) {
-      onBboxChanged(reduxLayer.id, null);
-    }
-  }
-};
-
 export const renderers = {
  renderToolPreview,
  renderLayers,
-  renderBbox,
+  renderBboxes,
  renderBackground,
  renderNoLayersMessage,
  arrangeLayers,
-  checkForTransparency,
+  updateBboxes,
 };

 const DEBOUNCE_MS = 300;
@ -924,11 +921,11 @@ const DEBOUNCE_MS = 300;
 export const debouncedRenderers = {
  renderToolPreview: debounce(renderToolPreview, DEBOUNCE_MS),
  renderLayers: debounce(renderLayers, DEBOUNCE_MS),
-  renderBbox: debounce(renderBbox, DEBOUNCE_MS),
+  renderBboxes: debounce(renderBboxes, DEBOUNCE_MS),
  renderBackground: debounce(renderBackground, DEBOUNCE_MS),
  renderNoLayersMessage: debounce(renderNoLayersMessage, DEBOUNCE_MS),
  arrangeLayers: debounce(arrangeLayers, DEBOUNCE_MS),
-  checkForTransparency: debounce(checkForTransparency, DEBOUNCE_MS),
+  updateBboxes: debounce(updateBboxes, DEBOUNCE_MS),
 };

 /**