From eff5b569908eae323c665b38f80751898e5d2e14 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Tue, 6 Aug 2024 13:43:32 +1000
Subject: [PATCH] fix(ui): brush preview fill for inpaint/region

---
 .../controlLayers/konva/CanvasManager.ts      | 46 ++++++++++++-------
 .../controlLayers/konva/CanvasTool.ts         |  4 +-
 .../src/features/controlLayers/store/types.ts |  1 +
 3 files changed, 32 insertions(+), 19 deletions(-)

diff --git a/invokeai/frontend/web/src/features/controlLayers/konva/CanvasManager.ts b/invokeai/frontend/web/src/features/controlLayers/konva/CanvasManager.ts
index e7e044ce3d..b28197c70b 100644
--- a/invokeai/frontend/web/src/features/controlLayers/konva/CanvasManager.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/konva/CanvasManager.ts
@@ -23,18 +23,19 @@ import {
 } from 'features/controlLayers/konva/util';
 import type { Extents, ExtentsResult, GetBboxTask, WorkerLogMessage } from 'features/controlLayers/konva/worker';
 import { $lastProgressEvent, $shouldShowStagedImage } from 'features/controlLayers/store/canvasV2Slice';
-import type {
-  CanvasControlAdapterState,
-  CanvasEntityIdentifier,
-  CanvasEntityState,
-  CanvasInpaintMaskState,
-  CanvasLayerState,
-  CanvasRegionalGuidanceState,
-  CanvasV2State,
-  Coordinate,
-  GenerationMode,
-  GetLoggingContext,
-  RgbaColor,
+import {
+  type CanvasControlAdapterState,
+  type CanvasEntityIdentifier,
+  type CanvasEntityState,
+  type CanvasInpaintMaskState,
+  type CanvasLayerState,
+  type CanvasRegionalGuidanceState,
+  type CanvasV2State,
+  type Coordinate,
+  type GenerationMode,
+  type GetLoggingContext,
+  RGBA_WHITE,
+  type RgbaColor,
 } from 'features/controlLayers/store/types';
 import type Konva from 'konva';
 import { atom } from 'nanostores';
@@ -366,11 +367,22 @@ export class CanvasManager {
     let currentFill: RgbaColor = state.tool.fill;
     const selectedEntity = this.getSelectedEntity();
     if (selectedEntity) {
-      // These two entity types use a compositing rect for opacity. Their alpha is always 1.
-      if (selectedEntity.state.type === 'regional_guidance') {
-        currentFill = { ...selectedEntity.state.fill, a: 1 };
-      } else if (selectedEntity.state.type === 'inpaint_mask') {
-        currentFill = { ...state.inpaintMask.fill, a: 1 };
+      // These two entity types use a compositing rect for opacity. Their fill is always white.
+      if (selectedEntity.state.type === 'regional_guidance' || selectedEntity.state.type === 'inpaint_mask') {
+        currentFill = RGBA_WHITE;
+      }
+    }
+    return currentFill;
+  };
+
+  getBrushPreviewFill = () => {
+    const state = this.stateApi.getState();
+    let currentFill: RgbaColor = state.tool.fill;
+    const selectedEntity = this.getSelectedEntity();
+    if (selectedEntity) {
+      // The brush should use the mask opacity for these entity types
+      if (selectedEntity.state.type === 'regional_guidance' || selectedEntity.state.type === 'inpaint_mask') {
+        currentFill = { ...selectedEntity.state.fill, a: this.stateApi.getSettings().maskOpacity };
       }
     }
     return currentFill;
diff --git a/invokeai/frontend/web/src/features/controlLayers/konva/CanvasTool.ts b/invokeai/frontend/web/src/features/controlLayers/konva/CanvasTool.ts
index a80a5a02bf..7a20d3579e 100644
--- a/invokeai/frontend/web/src/features/controlLayers/konva/CanvasTool.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/konva/CanvasTool.ts
@@ -127,7 +127,6 @@ export class CanvasTool {
     const stage = this.manager.stage;
     const renderedEntityCount: number = 1; // TODO(psyche): this.manager should be renderable entity count
     const toolState = this.manager.stateApi.getToolState();
-    const currentFill = this.manager.getCurrentFill();
     const selectedEntity = this.manager.getSelectedEntity();
     const cursorPos = this.manager.stateApi.$lastCursorPos.get();
     const isDrawing = this.manager.stateApi.$isDrawing.get();
@@ -172,6 +171,7 @@ export class CanvasTool {
 
       // No need to render the brush preview if the cursor position or color is missing
       if (cursorPos && tool === 'brush') {
+        const brushPreviewFill = this.manager.getBrushPreviewFill();
         const alignedCursorPos = alignCoordForTool(cursorPos, toolState.brush.width);
         const scale = stage.scaleX();
         // Update the fill circle
@@ -181,7 +181,7 @@ export class CanvasTool {
           x: alignedCursorPos.x,
           y: alignedCursorPos.y,
           radius,
-          fill: isDrawing ? '' : rgbaColorToString(currentFill),
+          fill: isDrawing ? '' : rgbaColorToString(brushPreviewFill),
         });
 
         // Update the inner border of the brush preview
diff --git a/invokeai/frontend/web/src/features/controlLayers/store/types.ts b/invokeai/frontend/web/src/features/controlLayers/store/types.ts
index f87a93bd08..520f1086ae 100644
--- a/invokeai/frontend/web/src/features/controlLayers/store/types.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/store/types.ts
@@ -503,6 +503,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_WHITE: RgbaColor = { r: 255, g: 255, b: 255, a: 1 };
 
 const zOpacity = z.number().gte(0).lte(1);