feat(ui): tidy stateApi atoms & add docstrings

invokeai/frontend/web/src/features/controlLayers/components/StagingArea/StagingAreaToolbar.tsx

@@ -45,7 +45,7 @@ export const StagingAreaToolbar = memo(() => {
   const index = useAppSelector(selectStagedImageIndex);
   const selectedImage = useAppSelector(selectSelectedImage);
   const imageCount = useAppSelector(selectImageCount);
-  const shouldShowStagedImage = useStore(canvasManager.stateApi.$shouldShowStagedImage);
+  const shouldShowStagedImage = useStore(canvasManager.stagingArea.$shouldShowStagedImage);
   const isCanvasActive = useStore(INTERACTION_SCOPES.canvas.$isActive);
   const [changeIsImageIntermediate] = useChangeImageIsIntermediateMutation();
   useScopeOnMount('stagingArea');
@@ -83,8 +83,8 @@ export const StagingAreaToolbar = memo(() => {
   }, [dispatch]);
 
   const onToggleShouldShowStagedImage = useCallback(() => {
-    canvasManager.stateApi.$shouldShowStagedImage.set(!shouldShowStagedImage);
-  }, [canvasManager.stateApi.$shouldShowStagedImage, shouldShowStagedImage]);
+    canvasManager.stagingArea.$shouldShowStagedImage.set(!shouldShowStagedImage);
+  }, [canvasManager.stagingArea.$shouldShowStagedImage, shouldShowStagedImage]);
 
   const onSaveStagingImage = useCallback(() => {
     if (!selectedImage) {

invokeai/frontend/web/src/features/controlLayers/components/Tool/ToolSettings.tsx

@@ -6,7 +6,7 @@ import { memo } from 'react';
 
 export const ToolSettings = memo(() => {
   const canvasManager = useCanvasManager();
-  const tool = useStore(canvasManager.stateApi.$tool);
+  const tool = useStore(canvasManager.tool.$tool);
   if (tool === 'brush') {
     return <ToolBrushWidth />;
   }

invokeai/frontend/web/src/features/controlLayers/components/Tool/hooks.ts

@@ -6,14 +6,14 @@ import { useCallback } from 'react';
 
 export const useToolIsSelected = (tool: Tool) => {
   const canvasManager = useCanvasManager();
-  const isSelected = useStore(computed(canvasManager.stateApi.$tool, (t) => t === tool));
+  const isSelected = useStore(computed(canvasManager.tool.$tool, (t) => t === tool));
   return isSelected;
 };
 
 export const useSelectTool = (tool: Tool) => {
   const canvasManager = useCanvasManager();
   const setTool = useCallback(() => {
-    canvasManager.stateApi.$tool.set(tool);
-  }, [canvasManager.stateApi.$tool, tool]);
+    canvasManager.tool.$tool.set(tool);
+  }, [canvasManager.tool.$tool, tool]);
   return setTool;
 };

invokeai/frontend/web/src/features/controlLayers/components/Toolbar/CanvasToolbarScale.tsx

@@ -73,7 +73,7 @@ const snapCandidates = marks.slice(1, marks.length - 1);
 export const CanvasToolbarScale = memo(() => {
   const { t } = useTranslation();
   const canvasManager = useCanvasManager();
-  const scale = useStore(computed(canvasManager.stateApi.$stageAttrs, (attrs) => attrs.scale));
+  const scale = useStore(computed(canvasManager.stage.$stageAttrs, (attrs) => attrs.scale));
   const [localScale, setLocalScale] = useState(scale * 100);
 
   const onChangeSlider = useCallback(

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

@@ -60,7 +60,7 @@ export class CanvasBackgroundModule extends CanvasModuleBase {
      * - position
      * - size
      */
-    this.subscriptions.add(this.manager.stateApi.$stageAttrs.listen(this.render));
+    this.subscriptions.add(this.manager.stage.$stageAttrs.listen(this.render));
   }
 
   /**

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

@@ -112,7 +112,7 @@ export class CanvasBboxModule extends CanvasModuleBase {
     this.konva.group.add(this.konva.transformer);
 
     // We will listen to the tool state to determine if the bbox should be visible or not.
-    this.subscriptions.add(this.manager.stateApi.$tool.listen(this.render));
+    this.subscriptions.add(this.manager.tool.$tool.listen(this.render));
   }
 
   /**
@@ -122,7 +122,7 @@ export class CanvasBboxModule extends CanvasModuleBase {
     this.log.trace('Rendering');
 
     const { x, y, width, height } = this.manager.stateApi.getBbox().rect;
-    const tool = this.manager.stateApi.$tool.get();
+    const tool = this.manager.tool.$tool.get();
 
     this.konva.group.visible(true);
 

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

@@ -88,7 +88,7 @@ export class CanvasBrushToolPreview extends CanvasModuleBase {
   }
 
   render = () => {
-    const cursorPos = this.manager.stateApi.$lastCursorPos.get();
+    const cursorPos = this.manager.tool.$lastCursorPos.get();
 
     // If the cursor position is not available, do not update the brush preview. The tool module will handle visiblity.
     if (!cursorPos) {

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

@@ -192,7 +192,7 @@ export class CanvasColorPickerToolPreview extends CanvasModuleBase {
    * Renders the color picker tool preview on the canvas.
    */
   render = () => {
-    const cursorPos = this.manager.stateApi.$lastCursorPos.get();
+    const cursorPos = this.manager.tool.$lastCursorPos.get();
 
     // If the cursor position is not available, do not render the preview. The tool module will handle visibility.
     if (!cursorPos) {
@@ -200,7 +200,7 @@ export class CanvasColorPickerToolPreview extends CanvasModuleBase {
     }
 
     const toolState = this.manager.stateApi.getToolState();
-    const colorUnderCursor = this.manager.stateApi.$colorUnderCursor.get();
+    const colorUnderCursor = this.parent.$colorUnderCursor.get();
     const colorPickerInnerRadius = this.manager.stage.getScaledPixels(this.config.RING_INNER_RADIUS);
     const colorPickerOuterRadius = this.manager.stage.getScaledPixels(this.config.RING_OUTER_RADIUS);
     const onePixel = this.manager.stage.getScaledPixels(1);

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

@@ -175,7 +175,7 @@ export class CanvasEntityRenderer extends CanvasModuleBase {
     // user switches tool mid-drawing, for example by pressing space to pan the stage. It's easy to press space
     // to pan _before_ releasing the mouse button, which would cause the buffer to be lost if we didn't commit it.
     this.subscriptions.add(
-      this.manager.stateApi.$tool.listen(() => {
+      this.manager.tool.$tool.listen(() => {
         this.commitBuffer();
       })
     );
@@ -183,7 +183,7 @@ export class CanvasEntityRenderer extends CanvasModuleBase {
     // The compositing rect must cover the whole stage at all times. When the stage is scaled, moved or resized, we
     // need to update the compositing rect to match the stage.
     this.subscriptions.add(
-      this.manager.stateApi.$stageAttrs.listen(() => {
+      this.manager.stage.$stageAttrs.listen(() => {
         if (this.konva.compositing && this.parent.type === 'entity_mask_adapter') {
           this.updateCompositingRectSize();
         }
@@ -256,7 +256,7 @@ export class CanvasEntityRenderer extends CanvasModuleBase {
     this.log.trace('Updating compositing rect size');
     assert(this.konva.compositing, 'Missing compositing rect');
 
-    const { x, y, width, height, scale } = this.manager.stateApi.$stageAttrs.get();
+    const { x, y, width, height, scale } = this.manager.stage.$stageAttrs.get();
 
     this.konva.compositing.rect.setAttrs({
       x: -x / scale,

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

@@ -220,7 +220,7 @@ export class CanvasEntityTransformer extends CanvasModuleBase {
     // When the stage scale changes, we may need to re-scale some of the transformer's components. For example,
     // the bbox outline should always be 1 screen pixel wide, so we need to update its stroke width.
     this.subscriptions.add(
-      this.manager.stateApi.$stageAttrs.listen((newVal, oldVal) => {
+      this.manager.stage.$stageAttrs.listen((newVal, oldVal) => {
         if (newVal.scale !== oldVal.scale) {
           this.syncScale();
         }
@@ -236,7 +236,7 @@ export class CanvasEntityTransformer extends CanvasModuleBase {
     );
 
     // When the selected tool changes, we need to update the transformer's interaction state.
-    this.subscriptions.add(this.manager.stateApi.$tool.listen(this.syncInteractionState));
+    this.subscriptions.add(this.manager.tool.$tool.listen(this.syncInteractionState));
 
     // When the selected entity changes, we need to update the transformer's interaction state.
     this.subscriptions.add(this.manager.stateApi.$selectedEntityIdentifier.listen(this.syncInteractionState));
@@ -511,7 +511,7 @@ export class CanvasEntityTransformer extends CanvasModuleBase {
       return;
     }
 
-    const tool = this.manager.stateApi.$tool.get();
+    const tool = this.manager.tool.$tool.get();
     const isSelected = this.manager.stateApi.getIsSelected(this.parent.id);
 
     if (!this.parent.renderer.hasObjects() || this.parent.state.isLocked || !this.parent.state.isEnabled) {
@@ -565,12 +565,12 @@ export class CanvasEntityTransformer extends CanvasModuleBase {
   startTransform = () => {
     this.log.debug('Starting transform');
     this.$isTransforming.set(true);
-    this.manager.stateApi.$tool.set('move');
+    this.manager.tool.$tool.set('move');
     // When transforming, we want the stage to still be movable if the view tool is selected. If the transformer or
     // interaction rect are listening, it will interrupt the stage's drag events. So we should disable listening
     // when the view tool is selected
     // TODO(psyche): We just set the tool to 'move', why would it be 'view'? Investigate and figure out if this is needed
-    const shouldListen = this.manager.stateApi.$tool.get() !== 'view';
+    const shouldListen = this.manager.tool.$tool.get() !== 'view';
     this.parent.konva.layer.listening(shouldListen);
     this.setInteractionMode('all');
     this.manager.stateApi.$transformingAdapter.set(this.parent);

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

@@ -79,7 +79,7 @@ export class CanvasEraserToolPreview extends CanvasModuleBase {
   }
 
   render = () => {
-    const cursorPos = this.manager.stateApi.$lastCursorPos.get();
+    const cursorPos = this.manager.tool.$lastCursorPos.get();
 
     if (!cursorPos) {
       return;

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

@@ -50,7 +50,7 @@ export class CanvasFilterModule extends CanvasModuleBase {
       return;
     }
     this.$adapter.set(entity.adapter);
-    this.manager.stateApi.$tool.set('view');
+    this.manager.tool.$tool.set('view');
   };
 
   previewFilter = async () => {

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

@@ -111,16 +111,15 @@ export class CanvasManager extends CanvasModuleBase {
     };
     this.stage.addLayer(this.konva.previewLayer);
 
-    this.stagingArea = new CanvasStagingAreaModule(this);
-    this.konva.previewLayer.add(this.stagingArea.konva.group);
-
-    this.progressImage = new CanvasProgressImageModule(this);
-    this.konva.previewLayer.add(this.progressImage.konva.group);
-
-    this.bbox = new CanvasBboxModule(this);
-    this.konva.previewLayer.add(this.bbox.konva.group);
-
     this.tool = new CanvasToolModule(this);
+    this.stagingArea = new CanvasStagingAreaModule(this);
+    this.progressImage = new CanvasProgressImageModule(this);
+    this.bbox = new CanvasBboxModule(this);
+
+    // Must add in this order for correct z-index
+    this.konva.previewLayer.add(this.stagingArea.konva.group);
+    this.konva.previewLayer.add(this.progressImage.konva.group);
+    this.konva.previewLayer.add(this.bbox.konva.group);
     this.konva.previewLayer.add(this.tool.konva.group);
   }
 

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

@@ -1,10 +1,17 @@
 import type { CanvasManager } from 'features/controlLayers/konva/CanvasManager';
 import { CanvasModuleBase } from 'features/controlLayers/konva/CanvasModuleBase';
 import { getPrefixedId, getRectUnion } from 'features/controlLayers/konva/util';
-import type { CanvasEntityIdentifier, Coordinate, Dimensions, Rect } from 'features/controlLayers/store/types';
+import type {
+  CanvasEntityIdentifier,
+  Coordinate,
+  Dimensions,
+  Rect,
+  StageAttrs,
+} from 'features/controlLayers/store/types';
 import type Konva from 'konva';
 import type { KonvaEventObject } from 'konva/lib/Node';
 import { clamp } from 'lodash-es';
+import { atom } from 'nanostores';
 import type { Logger } from 'roarr';
 
 type CanvasStageModuleConfig = {
@@ -42,6 +49,14 @@ export class CanvasStageModule extends CanvasModuleBase {
 
   config: CanvasStageModuleConfig = DEFAULT_CONFIG;
 
+  $stageAttrs = atom<StageAttrs>({
+    x: 0,
+    y: 0,
+    width: 0,
+    height: 0,
+    scale: 0,
+  });
+
   subscriptions = new Set<() => void>();
 
   constructor(stage: Konva.Stage, container: HTMLDivElement, manager: CanvasManager) {
@@ -89,7 +104,7 @@ export class CanvasStageModule extends CanvasModuleBase {
     this.log.trace('Fitting stage to container');
     this.konva.stage.width(this.konva.stage.container().offsetWidth);
     this.konva.stage.height(this.konva.stage.container().offsetHeight);
-    this.manager.stateApi.$stageAttrs.set({
+    this.$stageAttrs.set({
       x: this.konva.stage.x(),
       y: this.konva.stage.y(),
       width: this.konva.stage.width(),
@@ -149,8 +164,8 @@ export class CanvasStageModule extends CanvasModuleBase {
       scaleY: scale,
     });
 
-    this.manager.stateApi.$stageAttrs.set({
-      ...this.manager.stateApi.$stageAttrs.get(),
+    this.$stageAttrs.set({
+      ...this.$stageAttrs.get(),
       x,
       y,
       scale,
@@ -203,7 +218,7 @@ export class CanvasStageModule extends CanvasModuleBase {
       scaleY: newScale,
     });
 
-    this.manager.stateApi.$stageAttrs.set({
+    this.$stageAttrs.set({
       x: Math.floor(this.konva.stage.x()),
       y: Math.floor(this.konva.stage.y()),
       width: this.konva.stage.width(),
@@ -235,7 +250,7 @@ export class CanvasStageModule extends CanvasModuleBase {
       return;
     }
 
-    this.manager.stateApi.$stageAttrs.set({
+    this.$stageAttrs.set({
       // Stage position should always be an integer, else we get fractional pixels which are blurry
       x: Math.floor(this.konva.stage.x()),
       y: Math.floor(this.konva.stage.y()),
@@ -250,7 +265,7 @@ export class CanvasStageModule extends CanvasModuleBase {
       return;
     }
 
-    this.manager.stateApi.$stageAttrs.set({
+    this.$stageAttrs.set({
       // Stage position should always be an integer, else we get fractional pixels which are blurry
       x: Math.floor(this.konva.stage.x()),
       y: Math.floor(this.konva.stage.y()),

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

@@ -4,6 +4,7 @@ import { CanvasObjectImageRenderer } from 'features/controlLayers/konva/CanvasOb
 import { getPrefixedId } from 'features/controlLayers/konva/util';
 import { imageDTOToImageWithDims, type StagingAreaImage } from 'features/controlLayers/store/types';
 import Konva from 'konva';
+import { atom } from 'nanostores';
 import type { Logger } from 'roarr';
 
 export class CanvasStagingAreaModule extends CanvasModuleBase {
@@ -20,6 +21,8 @@ export class CanvasStagingAreaModule extends CanvasModuleBase {
   image: CanvasObjectImageRenderer | null;
   selectedImage: StagingAreaImage | null;
 
+  $shouldShowStagedImage = atom<boolean>(true);
+
   constructor(manager: CanvasManager) {
     super();
     this.id = getPrefixedId(this.type);
@@ -34,14 +37,14 @@ export class CanvasStagingAreaModule extends CanvasModuleBase {
     this.image = null;
     this.selectedImage = null;
 
-    this.subscriptions.add(this.manager.stateApi.$shouldShowStagedImage.listen(this.render));
+    this.subscriptions.add(this.$shouldShowStagedImage.listen(this.render));
   }
 
   render = async () => {
     this.log.trace('Rendering staging area');
     const session = this.manager.stateApi.getSession();
     const { x, y, width, height } = this.manager.stateApi.getBbox().rect;
-    const shouldShowStagedImage = this.manager.stateApi.$shouldShowStagedImage.get();
+    const shouldShowStagedImage = this.$shouldShowStagedImage.get();
 
     this.selectedImage = session.stagedImages[session.selectedStagedImageIndex] ?? null;
     this.konva.group.position({ x, y });

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

@@ -13,7 +13,6 @@ import {
   entityRasterized,
   entityRectAdded,
   entityReset,
-  entitySelected,
 } from 'features/controlLayers/store/canvasSlice';
 import { selectAllRenderableEntities, selectCanvasSlice } from 'features/controlLayers/store/selectors';
 import {
@@ -28,7 +27,6 @@ import type {
   CanvasInpaintMaskState,
   CanvasRasterLayerState,
   CanvasRegionalGuidanceState,
-  Coordinate,
   EntityBrushLineAddedPayload,
   EntityEraserLineAddedPayload,
   EntityIdentifierPayload,
@@ -37,9 +35,6 @@ import type {
   EntityRectAddedPayload,
   Rect,
   RgbaColor,
-  RgbColor,
-  StageAttrs,
-  Tool,
 } from 'features/controlLayers/store/types';
 import { RGBA_BLACK } from 'features/controlLayers/store/types';
 import type { WritableAtom } from 'nanostores';
@@ -84,6 +79,9 @@ export class CanvasStateApiModule extends CanvasModuleBase {
   manager: CanvasManager;
   log: Logger;
 
+  /**
+   * The redux store.
+   */
   store: AppStore;
 
   constructor(store: AppStore, manager: CanvasManager) {
@@ -99,43 +97,88 @@ export class CanvasStateApiModule extends CanvasModuleBase {
     this.store = store;
   }
 
-  // Reminder - use arrow functions to avoid binding issues
+  /**
+   * Gets the canvas slice.
+   *
+   * The state is stored in redux.
+   */
   getCanvasState = () => {
     return selectCanvasSlice(this.store.getState());
   };
+
+  /**
+   * Resets an entity, pushing state to redux.
+   */
   resetEntity = (arg: EntityIdentifierPayload) => {
     this.store.dispatch(entityReset(arg));
   };
+
+  /**
+   * Updates an entity's position, pushing state to redux.
+   */
   setEntityPosition = (arg: EntityMovedPayload) => {
     this.store.dispatch(entityMoved(arg));
   };
+
+  /**
+   * Adds a brush line to an entity, pushing state to redux.
+   */
   addBrushLine = (arg: EntityBrushLineAddedPayload) => {
     this.store.dispatch(entityBrushLineAdded(arg));
   };
+
+  /**
+   * Adds an eraser line to an entity, pushing state to redux.
+   */
   addEraserLine = (arg: EntityEraserLineAddedPayload) => {
     this.store.dispatch(entityEraserLineAdded(arg));
   };
+
+  /**
+   * Adds a rectangle to an entity, pushing state to redux.
+   */
   addRect = (arg: EntityRectAddedPayload) => {
     this.store.dispatch(entityRectAdded(arg));
   };
+
+  /**
+   * Rasterizes an entity, pushing state to redux.
+   */
   rasterizeEntity = (arg: EntityRasterizedPayload) => {
     this.store.dispatch(entityRasterized(arg));
   };
-  setSelectedEntity = (arg: EntityIdentifierPayload) => {
-    this.store.dispatch(entitySelected(arg));
-  };
-  setGenerationBbox = (bbox: Rect) => {
-    this.store.dispatch(bboxChanged(bbox));
+
+  /**
+   * Sets the generation bbox rect, pushing state to redux.
+   */
+  setGenerationBbox = (rect: Rect) => {
+    this.store.dispatch(bboxChanged(rect));
   };
+
+  /**
+   * Sets the brush width, pushing state to redux.
+   */
   setBrushWidth = (width: number) => {
     this.store.dispatch(brushWidthChanged(width));
   };
+
+  /**
+   * Sets the eraser width, pushing state to redux.
+   */
   setEraserWidth = (width: number) => {
     this.store.dispatch(eraserWidthChanged(width));
   };
+
+  /**
+   * Sets the fill color, pushing state to redux.
+   */
   setFill = (fill: RgbaColor) => {
     return this.store.dispatch(fillChanged(fill));
   };
+
+  /**
+   * Enqueues a batch, pushing state to redux.
+   */
   enqueueBatch = (batch: BatchConfig) => {
     this.store.dispatch(
       queueApi.endpoints.enqueueBatch.initiate(batch, {
@@ -143,35 +186,76 @@ export class CanvasStateApiModule extends CanvasModuleBase {
       })
     );
   };
+
+  /**
+   * Gets the generation bbox state from redux.
+   */
   getBbox = () => {
     return this.getCanvasState().bbox;
   };
 
+  /**
+   * Gets the tool state from redux.
+   */
   getToolState = () => {
     return this.store.getState().tool;
   };
+
+  /**
+   * Gets the canvas settings from redux.
+   */
   getSettings = () => {
     return this.store.getState().canvasSettings;
   };
+
+  /**
+   * Gets the regions state from redux.
+   */
   getRegionsState = () => {
     return this.getCanvasState().regions;
   };
+
+  /**
+   * Gets the raster layers state from redux.
+   */
   getRasterLayersState = () => {
     return this.getCanvasState().rasterLayers;
   };
+
+  /**
+   * Gets the control layers state from redux.
+   */
   getControlLayersState = () => {
     return this.getCanvasState().controlLayers;
   };
+
+  /**
+   * Gets the inpaint masks state from redux.
+   */
   getInpaintMasksState = () => {
     return this.getCanvasState().inpaintMasks;
   };
+
+  /**
+   * Gets the canvas session state from redux.
+   */
   getSession = () => {
     return this.store.getState().canvasSession;
   };
-  getIsSelected = (id: string) => {
+
+  /**
+   * Checks if an entity is selected.
+   */
+  getIsSelected = (id: string): boolean => {
     return this.getCanvasState().selectedEntityIdentifier?.id === id;
   };
 
+  /**
+   * Gets an entity by its identifier. The entity's state is retrieved from the redux store, and its adapter is
+   * retrieved from the canvas manager.
+   *
+   * Both state and adapter must exist for the entity to be returned.
+   */
   getEntity(identifier: CanvasEntityIdentifier): EntityStateAndAdapter | null {
     const state = this.getCanvasState();
 
@@ -204,6 +288,9 @@ export class CanvasStateApiModule extends CanvasModuleBase {
     return null;
   }
 
+  /**
+   * Gets the number of entities that are currently rendered on the canvas.
+   */
   getRenderedEntityCount = () => {
     const renderableEntities = selectAllRenderableEntities(this.getCanvasState());
     let count = 0;
@@ -215,6 +302,10 @@ export class CanvasStateApiModule extends CanvasModuleBase {
     return count;
   };
 
+  /**
+   * Gets the currently selected entity, if any. The entity's state is retrieved from the redux store, and its adapter
+   * is retrieved from the canvas manager.
+   */
   getSelectedEntity = () => {
     const state = this.getCanvasState();
     if (state.selectedEntityIdentifier) {
@@ -223,6 +314,16 @@ export class CanvasStateApiModule extends CanvasModuleBase {
     return null;
   };
 
+  /**
+   * Gets the current fill color. The fill color is determined by the tool state and the selected entity.
+   *
+   * The fill color is determined by the tool state, except when the selected entity is a regional guidance or inpaint
+   * mask. In that case, the fill color is always black.
+   *
+   * Regional guidance and inpaint mask entities use a compositing rect to draw with their selected color and texture,
+   * so the fill color for lines and rects doesn't matter - it is never seen. The only requirement is that it is opaque.
+   * For consistency with conventional black and white mask images, we use black as the fill color for these entities.
+   */
   getCurrentFill = () => {
     let currentFill: RgbaColor = this.getToolState().fill;
     const selectedEntity = this.getSelectedEntity();
@@ -235,45 +336,86 @@ export class CanvasStateApiModule extends CanvasModuleBase {
     return currentFill;
   };
 
+  /**
+   * Gets the brush preview fill color. The brush preview fill color is determined by the tool state and the selected
+   * entity.
+   *
+   * The color is the tool state's fill color, except when the selected entity is a regional guidance or inpaint mask.
+   *
+   * These entities have their own fill color and texture, so the brush preview should use those instead of the tool
+   * state's fill color.
+   */
   getBrushPreviewFill = (): RgbaColor => {
     const selectedEntity = this.getSelectedEntity();
     if (selectedEntity?.state.type === 'regional_guidance' || selectedEntity?.state.type === 'inpaint_mask') {
-      // The brush should use the mask opacity for these enktity types
+      // TODO(psyche): If we move the brush preview's Konva nodes to the selected entity renderer, we can draw them
+      // under the entity's compositing rect, so they would use selected entity's selected color and texture. As a
+      // temporary workaround to improve the UX when using a brush on a regional guidance or inpaint mask, we use the
+      // selected entity's fill color with 50% opacity.
       return { ...selectedEntity.state.fill.color, a: 0.5 };
     } else {
       return this.getToolState().fill;
     }
   };
 
+  /**
+   * The entity adapter being transformed, if any.
+   */
   $transformingAdapter = atom<CanvasEntityLayerAdapter | CanvasEntityMaskAdapter | null>(null);
+
+  /**
+   * Whether an entity is currently being transformed. Derived from `$transformingAdapter`.
+   */
   $isTranforming = computed(this.$transformingAdapter, (transformingAdapter) => Boolean(transformingAdapter));
 
+  /**
+   * A nanostores atom, kept in sync with the redux store's tool state.
+   */
   $toolState: WritableAtom<ToolState> = atom();
-  $currentFill: WritableAtom<RgbaColor> = atom();
-  $selectedEntity: WritableAtom<EntityStateAndAdapter | null> = atom();
-  $selectedEntityIdentifier: WritableAtom<CanvasEntityIdentifier | null> = atom();
-  $colorUnderCursor: WritableAtom<RgbColor> = atom(RGBA_BLACK);
 
-  // Read-write state, ephemeral interaction state
-  $tool = atom<Tool>('brush');
-  $toolBuffer = atom<Tool | null>(null);
-  $isDrawing = atom<boolean>(false);
-  $isMouseDown = atom<boolean>(false);
-  $lastAddedPoint = atom<Coordinate | null>(null);
-  $lastMouseDownPos = atom<Coordinate | null>(null);
-  $lastCursorPos = atom<Coordinate | null>(null);
+  /**
+   * The current fill color, derived from the tool state and the selected entity.
+   */
+  $currentFill: WritableAtom<RgbaColor> = atom();
+
+  /**
+   * The currently selected entity, if any. Includes the entity latest state and its adapter.
+   */
+  $selectedEntity: WritableAtom<EntityStateAndAdapter | null> = atom();
+
+  /**
+   * The currently selected entity's identifier, if an entity is selected.
+   */
+  $selectedEntityIdentifier: WritableAtom<CanvasEntityIdentifier | null> = atom();
+
+  /**
+   * The last canvas progress event. This is set in a global event listener. The staging area may set it to null when it
+   * consumes the event.
+   */
   $lastCanvasProgressEvent = $lastCanvasProgressEvent;
+
+  /**
+   * Whether the space key is currently pressed.
+   */
   $spaceKey = atom<boolean>(false);
+
+  /**
+   * Whether the alt key is currently pressed.
+   */
   $altKey = $alt;
+
+  /**
+   * Whether the ctrl key is currently pressed.
+   */
   $ctrlKey = $ctrl;
+
+  /**
+   * Whether the meta key is currently pressed.
+   */
   $metaKey = $meta;
+
+  /**
+   * Whether the shift key is currently pressed.
+   */
   $shiftKey = $shift;
-  $shouldShowStagedImage = atom(true);
-  $stageAttrs = atom<StageAttrs>({
-    x: 0,
-    y: 0,
-    width: 0,
-    height: 0,
-    scale: 0,
-  });
 }

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

@@ -22,9 +22,10 @@ import type {
   RgbColor,
   Tool,
 } from 'features/controlLayers/store/types';
-import { isDrawableEntity } from 'features/controlLayers/store/types';
+import { isDrawableEntity, RGBA_BLACK } from 'features/controlLayers/store/types';
 import Konva from 'konva';
 import type { KonvaEventObject } from 'konva/lib/Node';
+import { atom } from 'nanostores';
 import type { Logger } from 'roarr';
 
 type CanvasToolModuleConfig = {
@@ -51,6 +52,36 @@ export class CanvasToolModule extends CanvasModuleBase {
   eraserToolPreview: CanvasEraserToolPreview;
   colorPickerToolPreview: CanvasColorPickerToolPreview;
 
+  /**
+   * The currently selected tool.
+   */
+  $tool = atom<Tool>('brush');
+  /**
+   * A buffer for the currently selected tool. This is used to temporarily store the tool while the user is using any
+   * hold-to-activate tools, like the view or color picker tools.
+   */
+  $toolBuffer = atom<Tool | null>(null);
+  /**
+   * The last point added to the current entity.
+   */
+  $lastAddedPoint = atom<Coordinate | null>(null);
+  /**
+   * Whether the mouse is currently down.
+   */
+  $isMouseDown = atom<boolean>(false);
+  /**
+   * The last position where the mouse was down.
+   */
+  $lastMouseDownPos = atom<Coordinate | null>(null);
+  /**
+   * The last cursor position.
+   */
+  $lastCursorPos = atom<Coordinate | null>(null);
+  /**
+   * The color currently under the cursor. Only has a value when the color picker tool is active.
+   */
+  $colorUnderCursor = atom<RgbColor>(RGBA_BLACK);
+
   konva: {
     stage: Konva.Stage;
     group: Konva.Group;
@@ -79,7 +110,7 @@ export class CanvasToolModule extends CanvasModuleBase {
     this.konva.group.add(this.eraserToolPreview.konva.group);
     this.konva.group.add(this.colorPickerToolPreview.konva.group);
 
-    this.subscriptions.add(this.manager.stateApi.$stageAttrs.listen(this.render));
+    this.subscriptions.add(this.manager.stage.$stageAttrs.listen(this.render));
     this.subscriptions.add(
       this.manager.stateApi.$toolState.listen((value, oldValue) => {
         if (
@@ -92,7 +123,7 @@ export class CanvasToolModule extends CanvasModuleBase {
         }
       })
     );
-    this.subscriptions.add(this.manager.stateApi.$tool.listen(this.render));
+    this.subscriptions.add(this.$tool.listen(this.render));
 
     const cleanupListeners = this.setEventListeners();
 
@@ -109,8 +140,8 @@ export class CanvasToolModule extends CanvasModuleBase {
     const stage = this.manager.stage;
     const renderedEntityCount = this.manager.stateApi.getRenderedEntityCount();
     const selectedEntity = this.manager.stateApi.getSelectedEntity();
-    const isMouseDown = this.manager.stateApi.$isMouseDown.get();
-    const tool = this.manager.stateApi.$tool.get();
+    const isMouseDown = this.$isMouseDown.get();
+    const tool = this.$tool.get();
 
     const isDrawable =
       !!selectedEntity &&
@@ -153,8 +184,8 @@ export class CanvasToolModule extends CanvasModuleBase {
     const stage = this.manager.stage;
     const renderedEntityCount = this.manager.stateApi.getRenderedEntityCount();
     const selectedEntity = this.manager.stateApi.getSelectedEntity();
-    const cursorPos = this.manager.stateApi.$lastCursorPos.get();
-    const tool = this.manager.stateApi.$tool.get();
+    const cursorPos = this.$lastCursorPos.get();
+    const tool = this.$tool.get();
 
     const isDrawable =
       !!selectedEntity &&
@@ -187,7 +218,7 @@ export class CanvasToolModule extends CanvasModuleBase {
 
   syncLastCursorPos = (): Coordinate | null => {
     const pos = getScaledCursorPosition(this.konva.stage);
-    this.manager.stateApi.$lastCursorPos.set(pos);
+    this.$lastCursorPos.set(pos);
     return pos;
   };
 
@@ -268,16 +299,16 @@ export class CanvasToolModule extends CanvasModuleBase {
   };
 
   onStageMouseDown = async (e: KonvaEventObject<MouseEvent>) => {
-    this.manager.stateApi.$isMouseDown.set(true);
+    this.$isMouseDown.set(true);
     const toolState = this.manager.stateApi.getToolState();
-    const tool = this.manager.stateApi.$tool.get();
+    const tool = this.$tool.get();
     const pos = this.syncLastCursorPos();
     const selectedEntity = this.manager.stateApi.getSelectedEntity();
 
     if (tool === 'colorPicker') {
       const color = this.getColorUnderCursor();
       if (color) {
-        this.manager.stateApi.$colorUnderCursor.set(color);
+        this.$colorUnderCursor.set(color);
       }
       if (color) {
         this.manager.stateApi.setFill({ ...toolState.fill, ...color });
@@ -286,7 +317,7 @@ export class CanvasToolModule extends CanvasModuleBase {
     } else {
       const isDrawable = selectedEntity?.state.isEnabled && !selectedEntity.state.isLocked;
       if (pos && isDrawable && !this.manager.stateApi.$spaceKey.get() && getIsPrimaryMouseDown(e)) {
-        this.manager.stateApi.$lastMouseDownPos.set(pos);
+        this.$lastMouseDownPos.set(pos);
         const normalizedPoint = offsetCoord(pos, selectedEntity.state.position);
 
         if (tool === 'brush') {
@@ -325,7 +356,7 @@ export class CanvasToolModule extends CanvasModuleBase {
               clip: this.getClip(selectedEntity.state),
             });
           }
-          this.manager.stateApi.$lastAddedPoint.set(alignedPoint);
+          this.$lastAddedPoint.set(alignedPoint);
         }
 
         if (tool === 'eraser') {
@@ -361,7 +392,7 @@ export class CanvasToolModule extends CanvasModuleBase {
               clip: this.getClip(selectedEntity.state),
             });
           }
-          this.manager.stateApi.$lastAddedPoint.set(alignedPoint);
+          this.$lastAddedPoint.set(alignedPoint);
         }
 
         if (tool === 'rect') {
@@ -380,11 +411,11 @@ export class CanvasToolModule extends CanvasModuleBase {
   };
 
   onStageMouseUp = (_: KonvaEventObject<MouseEvent>) => {
-    this.manager.stateApi.$isMouseDown.set(false);
-    const pos = this.manager.stateApi.$lastCursorPos.get();
+    this.$isMouseDown.set(false);
+    const pos = this.$lastCursorPos.get();
     const selectedEntity = this.manager.stateApi.getSelectedEntity();
     const isDrawable = selectedEntity?.state.isEnabled && !selectedEntity.state.isLocked;
-    const tool = this.manager.stateApi.$tool.get();
+    const tool = this.$tool.get();
 
     if (pos && isDrawable && !this.manager.stateApi.$spaceKey.get()) {
       if (tool === 'brush') {
@@ -414,7 +445,7 @@ export class CanvasToolModule extends CanvasModuleBase {
         }
       }
 
-      this.manager.stateApi.$lastMouseDownPos.set(null);
+      this.$lastMouseDownPos.set(null);
     }
     this.render();
   };
@@ -423,12 +454,12 @@ export class CanvasToolModule extends CanvasModuleBase {
     const toolState = this.manager.stateApi.getToolState();
     const pos = this.syncLastCursorPos();
     const selectedEntity = this.manager.stateApi.getSelectedEntity();
-    const tool = this.manager.stateApi.$tool.get();
+    const tool = this.$tool.get();
 
     if (tool === 'colorPicker') {
       const color = this.getColorUnderCursor();
       if (color) {
-        this.manager.stateApi.$colorUnderCursor.set(color);
+        this.$colorUnderCursor.set(color);
       }
     } else {
       const isDrawable = selectedEntity?.state.isEnabled && !selectedEntity.state.isLocked;
@@ -446,7 +477,7 @@ export class CanvasToolModule extends CanvasModuleBase {
                 if (lastPoint.x !== alignedPoint.x || lastPoint.y !== alignedPoint.y) {
                   drawingBuffer.points.push(alignedPoint.x, alignedPoint.y);
                   await selectedEntity.adapter.renderer.setBuffer(drawingBuffer);
-                  this.manager.stateApi.$lastAddedPoint.set(alignedPoint);
+                  this.$lastAddedPoint.set(alignedPoint);
                 }
               }
             } else {
@@ -466,7 +497,7 @@ export class CanvasToolModule extends CanvasModuleBase {
               color: this.manager.stateApi.getCurrentFill(),
               clip: this.getClip(selectedEntity.state),
             });
-            this.manager.stateApi.$lastAddedPoint.set(alignedPoint);
+            this.$lastAddedPoint.set(alignedPoint);
           }
         }
 
@@ -483,7 +514,7 @@ export class CanvasToolModule extends CanvasModuleBase {
                 if (lastPoint.x !== alignedPoint.x || lastPoint.y !== alignedPoint.y) {
                   drawingBuffer.points.push(alignedPoint.x, alignedPoint.y);
                   await selectedEntity.adapter.renderer.setBuffer(drawingBuffer);
-                  this.manager.stateApi.$lastAddedPoint.set(alignedPoint);
+                  this.$lastAddedPoint.set(alignedPoint);
                 }
               }
             } else {
@@ -502,7 +533,7 @@ export class CanvasToolModule extends CanvasModuleBase {
               strokeWidth: toolState.eraser.width,
               clip: this.getClip(selectedEntity.state),
             });
-            this.manager.stateApi.$lastAddedPoint.set(alignedPoint);
+            this.$lastAddedPoint.set(alignedPoint);
           }
         }
 
@@ -527,12 +558,12 @@ export class CanvasToolModule extends CanvasModuleBase {
 
   onStageMouseLeave = async (e: KonvaEventObject<MouseEvent>) => {
     const pos = this.syncLastCursorPos();
-    this.manager.stateApi.$lastCursorPos.set(null);
-    this.manager.stateApi.$lastMouseDownPos.set(null);
+    this.$lastCursorPos.set(null);
+    this.$lastMouseDownPos.set(null);
     const selectedEntity = this.manager.stateApi.getSelectedEntity();
     const toolState = this.manager.stateApi.getToolState();
     const isDrawable = selectedEntity?.state.isEnabled && !selectedEntity.state.isLocked;
-    const tool = this.manager.stateApi.$tool.get();
+    const tool = this.$tool.get();
 
     if (pos && isDrawable && !this.manager.stateApi.$spaceKey.get() && getIsPrimaryMouseDown(e)) {
       const drawingBuffer = selectedEntity.adapter.renderer.bufferState;
@@ -566,7 +597,7 @@ export class CanvasToolModule extends CanvasModuleBase {
     }
 
     const toolState = this.manager.stateApi.getToolState();
-    const tool = this.manager.stateApi.$tool.get();
+    const tool = this.$tool.get();
 
     let delta = e.evt.deltaY;
 
@@ -596,19 +627,19 @@ export class CanvasToolModule extends CanvasModuleBase {
       const selectedEntity = this.manager.stateApi.getSelectedEntity();
       if (selectedEntity) {
         selectedEntity.adapter.renderer.clearBuffer();
-        this.manager.stateApi.$lastMouseDownPos.set(null);
+        this.$lastMouseDownPos.set(null);
       }
     } else if (e.key === ' ') {
       // Select the view tool on space key down
-      this.manager.stateApi.$toolBuffer.set(this.manager.stateApi.$tool.get());
-      this.manager.stateApi.$tool.set('view');
+      this.$toolBuffer.set(this.$tool.get());
+      this.$tool.set('view');
       this.manager.stateApi.$spaceKey.set(true);
-      this.manager.stateApi.$lastCursorPos.set(null);
-      this.manager.stateApi.$lastMouseDownPos.set(null);
+      this.$lastCursorPos.set(null);
+      this.$lastMouseDownPos.set(null);
     } else if (e.key === 'Alt') {
       // Select the color picker on alt key down
-      this.manager.stateApi.$toolBuffer.set(this.manager.stateApi.$tool.get());
-      this.manager.stateApi.$tool.set('colorPicker');
+      this.$toolBuffer.set(this.$tool.get());
+      this.$tool.set('colorPicker');
     }
   };
 
@@ -621,15 +652,15 @@ export class CanvasToolModule extends CanvasModuleBase {
     }
     if (e.key === ' ') {
       // Revert the tool to the previous tool on space key up
-      const toolBuffer = this.manager.stateApi.$toolBuffer.get();
-      this.manager.stateApi.$tool.set(toolBuffer ?? 'move');
-      this.manager.stateApi.$toolBuffer.set(null);
+      const toolBuffer = this.$toolBuffer.get();
+      this.$tool.set(toolBuffer ?? 'move');
+      this.$toolBuffer.set(null);
       this.manager.stateApi.$spaceKey.set(false);
     } else if (e.key === 'Alt') {
       // Revert the tool to the previous tool on alt key up
-      const toolBuffer = this.manager.stateApi.$toolBuffer.get();
-      this.manager.stateApi.$tool.set(toolBuffer ?? 'move');
-      this.manager.stateApi.$toolBuffer.set(null);
+      const toolBuffer = this.$toolBuffer.get();
+      this.$tool.set(toolBuffer ?? 'move');
+      this.$toolBuffer.set(null);
     }
   };