From fc6d244071693a7bc3385482480ee1e545ee6762 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Sat, 24 Aug 2024 10:10:21 +1000
Subject: [PATCH] fix(ui): calculate renderable entities correctly in tool
 module

---
 .../controlLayers/konva/CanvasStateApiModule.ts | 12 ++++++++++++
 .../controlLayers/konva/CanvasToolModule.ts     |  2 +-
 .../controlLayers/store/canvasV2Slice.ts        | 17 ++++++++++++++++-
 3 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/invokeai/frontend/web/src/features/controlLayers/konva/CanvasStateApiModule.ts b/invokeai/frontend/web/src/features/controlLayers/konva/CanvasStateApiModule.ts
index 0f77f60713..41e30d017e 100644
--- a/invokeai/frontend/web/src/features/controlLayers/konva/CanvasStateApiModule.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/konva/CanvasStateApiModule.ts
@@ -25,6 +25,7 @@ import {
   entitySelected,
   eraserWidthChanged,
   fillChanged,
+  selectAllRenderableEntities,
   toolBufferChanged,
   toolChanged,
 } from 'features/controlLayers/store/canvasV2Slice';
@@ -198,6 +199,17 @@ export class CanvasStateApiModule {
     return null;
   }
 
+  getRenderedEntityCount = () => {
+    const renderableEntities = selectAllRenderableEntities(this.getState());
+    let count = 0;
+    for (const entity of renderableEntities) {
+      if (entity.isEnabled) {
+        count++;
+      }
+    }
+    return count;
+  };
+
   getSelectedEntity = () => {
     const state = this.getState();
     if (state.selectedEntityIdentifier) {
diff --git a/invokeai/frontend/web/src/features/controlLayers/konva/CanvasToolModule.ts b/invokeai/frontend/web/src/features/controlLayers/konva/CanvasToolModule.ts
index 275ebe4ea0..9e3d2aedaa 100644
--- a/invokeai/frontend/web/src/features/controlLayers/konva/CanvasToolModule.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/konva/CanvasToolModule.ts
@@ -219,7 +219,7 @@ export class CanvasToolModule {
 
   render() {
     const stage = this.manager.stage;
-    const renderedEntityCount: number = 1; // TODO(psyche): this.manager should be renderable entity count
+    const renderedEntityCount = this.manager.stateApi.getRenderedEntityCount();
     const toolState = this.manager.stateApi.getToolState();
     const selectedEntity = this.manager.stateApi.getSelectedEntity();
     const cursorPos = this.manager.stateApi.$lastCursorPos.get();
diff --git a/invokeai/frontend/web/src/features/controlLayers/store/canvasV2Slice.ts b/invokeai/frontend/web/src/features/controlLayers/store/canvasV2Slice.ts
index fa6391c1d8..a9a04b163c 100644
--- a/invokeai/frontend/web/src/features/controlLayers/store/canvasV2Slice.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/store/canvasV2Slice.ts
@@ -24,8 +24,12 @@ import { atom } from 'nanostores';
 import { assert } from 'tsafe';
 
 import type {
+  CanvasControlLayerState,
   CanvasEntityIdentifier,
   CanvasEntityState,
+  CanvasInpaintMaskState,
+  CanvasRasterLayerState,
+  CanvasRegionalGuidanceState,
   CanvasV2State,
   Coordinate,
   EntityBrushLineAddedPayload,
@@ -170,7 +174,7 @@ function selectAllEntitiesOfType(state: CanvasV2State, type: CanvasEntityState['
   }
 }
 
-function selectAllEntities(state: CanvasV2State): CanvasEntityState[] {
+export function selectAllEntities(state: CanvasV2State): CanvasEntityState[] {
   // These are in the same order as they are displayed in the list!
   return [
     ...state.inpaintMasks.entities.toReversed(),
@@ -181,6 +185,17 @@ function selectAllEntities(state: CanvasV2State): CanvasEntityState[] {
   ];
 }
 
+export function selectAllRenderableEntities(
+  state: CanvasV2State
+): (CanvasRasterLayerState | CanvasControlLayerState | CanvasInpaintMaskState | CanvasRegionalGuidanceState)[] {
+  return [
+    ...state.rasterLayers.entities,
+    ...state.controlLayers.entities,
+    ...state.inpaintMasks.entities,
+    ...state.regions.entities,
+  ];
+}
+
 export const canvasV2Slice = createSlice({
   name: 'canvasV2',
   initialState,