feat(ui): move canvas stage and base layer to nanostores

Author: psychedelicious

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/canvasMerged.ts

@@ -1,8 +1,8 @@
 import { $logger } from 'app/logging/logger';
 import { canvasMerged } from 'features/canvas/store/actions';
+import { $canvasBaseLayer } from 'features/canvas/store/canvasNanostore';
 import { setMergedCanvas } from 'features/canvas/store/canvasSlice';
 import { getFullBaseLayerBlob } from 'features/canvas/util/getFullBaseLayerBlob';
-import { getCanvasBaseLayer } from 'features/canvas/util/konvaInstanceProvider';
 import { addToast } from 'features/system/store/systemSlice';
 import { t } from 'i18next';
 import { imagesApi } from 'services/api/endpoints/images';
@@ -30,7 +30,7 @@ export const addCanvasMergedListener = () => {
         return;
       }
 
-      const canvasBaseLayer = getCanvasBaseLayer();
+      const canvasBaseLayer = $canvasBaseLayer.get();
 
       if (!canvasBaseLayer) {
         moduleLog.error('Problem getting canvas base layer');

invokeai/frontend/web/src/features/canvas/components/IAICanvas.tsx

@@ -10,21 +10,16 @@ import useCanvasMouseOut from 'features/canvas/hooks/useCanvasMouseOut';
 import useCanvasMouseUp from 'features/canvas/hooks/useCanvasMouseUp';
 import useCanvasWheel from 'features/canvas/hooks/useCanvasZoom';
 import {
-  $isModifyingBoundingBox,
+ $canvasBaseLayer,  $canvasStage,$isModifyingBoundingBox,
   $isMouseOverBoundingBox,
   $isMovingStage,
   $isTransformingBoundingBox,
-  $tool,
-} from 'features/canvas/store/canvasNanostore';
+  $tool  } from 'features/canvas/store/canvasNanostore';
 import { isStagingSelector } from 'features/canvas/store/canvasSelectors';
 import {
   canvasResized,
   selectCanvasSlice,
 } from 'features/canvas/store/canvasSlice';
-import {
-  setCanvasBaseLayer,
-  setCanvasStage,
-} from 'features/canvas/util/konvaInstanceProvider';
 import type Konva from 'konva';
 import type { KonvaEventObject } from 'konva/lib/Node';
 import type { Vector2d } from 'konva/lib/types';
@@ -80,9 +75,9 @@ const IAICanvas = () => {
   const isMouseOverBoundingBox = useStore($isMouseOverBoundingBox);
   const tool = useStore($tool);
   useCanvasHotkeys();
-  const canvasStageRefCallback = useCallback((el: Konva.Stage) => {
-    setCanvasStage(el as Konva.Stage);
-    stageRef.current = el;
+  const canvasStageRefCallback = useCallback((stageElement: Konva.Stage) => {
+    $canvasStage.set(stageElement);
+    stageRef.current = stageElement;
   }, []);
   const stageCursor = useMemo(() => {
     if (tool === 'move' || isStaging) {
@@ -105,10 +100,14 @@ const IAICanvas = () => {
     shouldRestrictStrokesToBox,
     tool,
   ]);
-  const canvasBaseLayerRefCallback = useCallback((el: Konva.Layer) => {
-    setCanvasBaseLayer(el as Konva.Layer);
-    canvasBaseLayerRef.current = el;
-  }, []);
+
+  const canvasBaseLayerRefCallback = useCallback(
+    (layerElement: Konva.Layer) => {
+      $canvasBaseLayer.set(layerElement);
+      canvasBaseLayerRef.current = layerElement;
+    },
+    []
+  );
 
   const lastCursorPositionRef = useRef<Vector2d>({ x: 0, y: 0 });
 

invokeai/frontend/web/src/features/canvas/components/IAICanvasToolPreview.tsx

@@ -5,6 +5,7 @@ import {
   $cursorPosition,
   $isMovingBoundingBox,
   $isTransformingBoundingBox,
+  $tool,
 } from 'features/canvas/store/canvasNanostore';
 import { selectCanvasSlice } from 'features/canvas/store/canvasSlice';
 import { rgbaColorToString } from 'features/canvas/util/colorToString';
@@ -89,7 +90,7 @@ const IAICanvasToolPreview = (props: GroupConfig) => {
   const maskColorString = useAppSelector((s) =>
     rgbaColorToString({ ...s.canvas.maskColor, a: 0.5 })
   );
-  const tool = useAppSelector((s) => s.canvas.tool);
+  const tool = useStore($tool);
   const layer = useAppSelector((s) => s.canvas.layer);
   const dotRadius = useAppSelector((s) => 1.5 / s.canvas.stageScale);
   const strokeWidth = useAppSelector((s) => 1.5 / s.canvas.stageScale);

invokeai/frontend/web/src/features/canvas/components/IAICanvasToolbar/IAICanvasToolbar.tsx

@@ -15,7 +15,7 @@ import {
   canvasMerged,
   canvasSavedToGallery,
 } from 'features/canvas/store/actions';
-import { $tool } from 'features/canvas/store/canvasNanostore';
+import { $canvasBaseLayer,$tool  } from 'features/canvas/store/canvasNanostore';
 import { isStagingSelector } from 'features/canvas/store/canvasSelectors';
 import {
   resetCanvas,
@@ -25,7 +25,6 @@ import {
 } from 'features/canvas/store/canvasSlice';
 import type { CanvasLayer } from 'features/canvas/store/canvasTypes';
 import { LAYER_NAMES_DICT } from 'features/canvas/store/canvasTypes';
-import { getCanvasBaseLayer } from 'features/canvas/util/konvaInstanceProvider';
 import { InvIconButton } from 'index';
 import { memo, useCallback, useMemo } from 'react';
 import { useHotkeys } from 'react-hotkeys-hook';
@@ -53,7 +52,6 @@ const IAICanvasToolbar = () => {
   const layer = useAppSelector((s) => s.canvas.layer);
   const tool = useStore($tool);
   const isStaging = useAppSelector(isStagingSelector);
-  const canvasBaseLayer = getCanvasBaseLayer();
   const { t } = useTranslation();
   const { isClipboardAPIAvailable } = useCopyImageToClipboard();
 
@@ -82,7 +80,7 @@ const IAICanvasToolbar = () => {
       enabled: () => true,
       preventDefault: true,
     },
-    [canvasBaseLayer]
+    []
   );
 
   useHotkeys(
@@ -94,7 +92,7 @@ const IAICanvasToolbar = () => {
       enabled: () => !isStaging,
       preventDefault: true,
     },
-    [canvasBaseLayer]
+    []
   );
 
   useHotkeys(
@@ -106,7 +104,7 @@ const IAICanvasToolbar = () => {
       enabled: () => !isStaging,
       preventDefault: true,
     },
-    [canvasBaseLayer]
+    []
   );
 
   useHotkeys(
@@ -118,7 +116,7 @@ const IAICanvasToolbar = () => {
       enabled: () => !isStaging && isClipboardAPIAvailable,
       preventDefault: true,
     },
-    [canvasBaseLayer, isClipboardAPIAvailable]
+    [isClipboardAPIAvailable]
   );
 
   useHotkeys(
@@ -130,7 +128,7 @@ const IAICanvasToolbar = () => {
       enabled: () => !isStaging,
       preventDefault: true,
     },
-    [canvasBaseLayer]
+    []
   );
 
   const handleSelectMoveTool = useCallback(() => {
@@ -143,7 +141,7 @@ const IAICanvasToolbar = () => {
   );
 
   const handleResetCanvasView = (shouldScaleTo1 = false) => {
-    const canvasBaseLayer = getCanvasBaseLayer();
+    const canvasBaseLayer = $canvasBaseLayer.get();
     if (!canvasBaseLayer) {
       return;
     }

invokeai/frontend/web/src/features/canvas/hooks/useCanvasHotkeys.ts

@@ -1,18 +1,16 @@
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import {
-  $tool,
+ $canvasStage,  $tool,
   $toolStash,
   resetCanvasInteractionState,
-  resetToolInteractionState,
-} from 'features/canvas/store/canvasNanostore';
+  resetToolInteractionState } from 'features/canvas/store/canvasNanostore';
 import { isStagingSelector } from 'features/canvas/store/canvasSelectors';
 import {
   clearMask,
   setIsMaskEnabled,
   setShouldShowBoundingBox,
   setShouldSnapToGrid,
 } from 'features/canvas/store/canvasSlice';
-import { getCanvasStage } from 'features/canvas/util/konvaInstanceProvider';
 import { activeTabNameSelector } from 'features/ui/store/uiSelectors';
 import { useCallback, useEffect } from 'react';
 import { useHotkeys } from 'react-hotkeys-hook';
@@ -26,7 +24,6 @@ const useInpaintingCanvasHotkeys = () => {
   const isStaging = useAppSelector(isStagingSelector);
   const isMaskEnabled = useAppSelector((s) => s.canvas.isMaskEnabled);
   const shouldSnapToGrid = useAppSelector((s) => s.canvas.shouldSnapToGrid);
-  const canvasStage = getCanvasStage();
 
   // Beta Keys
   const handleClearMask = useCallback(() => dispatch(clearMask()), [dispatch]);
@@ -102,35 +99,29 @@ const useInpaintingCanvasHotkeys = () => {
     });
   }, []);
 
-  const onKeyDown = useCallback(
-    (e: KeyboardEvent) => {
-      if (e.repeat || e.key !== ' ') {
-        return;
-      }
-      if ($toolStash.get() || $tool.get() === 'move') {
-        return;
-      }
-      canvasStage?.container().focus();
-      $toolStash.set($tool.get());
-      $tool.set('move');
-      resetToolInteractionState();
-    },
-    [canvasStage]
-  );
-  const onKeyUp = useCallback(
-    (e: KeyboardEvent) => {
-      if (e.repeat || e.key !== ' ') {
-        return;
-      }
-      if (!$toolStash.get() || $tool.get() !== 'move') {
-        return;
-      }
-      canvasStage?.container().focus();
-      $tool.set($toolStash.get() ?? 'move');
-      $toolStash.set(null);
-    },
-    [canvasStage]
-  );
+  const onKeyDown = useCallback((e: KeyboardEvent) => {
+    if (e.repeat || e.key !== ' ') {
+      return;
+    }
+    if ($toolStash.get() || $tool.get() === 'move') {
+      return;
+    }
+    $canvasStage.get()?.container().focus();
+    $toolStash.set($tool.get());
+    $tool.set('move');
+    resetToolInteractionState();
+  }, []);
+  const onKeyUp = useCallback((e: KeyboardEvent) => {
+    if (e.repeat || e.key !== ' ') {
+      return;
+    }
+    if (!$toolStash.get() || $tool.get() !== 'move') {
+      return;
+    }
+    $canvasStage.get()?.container().focus();
+    $tool.set($toolStash.get() ?? 'move');
+    $toolStash.set(null);
+  }, []);
 
   useEffect(() => {
     window.addEventListener('keydown', onKeyDown);

invokeai/frontend/web/src/features/canvas/hooks/useColorUnderCursor.ts

@@ -1,22 +1,18 @@
 import { useAppDispatch } from 'app/store/storeHooks';
-import { $tool } from 'features/canvas/store/canvasNanostore';
+import { $canvasBaseLayer,$canvasStage,$tool   } from 'features/canvas/store/canvasNanostore';
 import {
   commitColorPickerColor,
   setColorPickerColor,
 } from 'features/canvas/store/canvasSlice';
-import {
-  getCanvasBaseLayer,
-  getCanvasStage,
-} from 'features/canvas/util/konvaInstanceProvider';
 import Konva from 'konva';
 import { useCallback } from 'react';
 
 const useColorPicker = () => {
   const dispatch = useAppDispatch();
-  const canvasBaseLayer = getCanvasBaseLayer();
-  const stage = getCanvasStage();
 
   const updateColorUnderCursor = useCallback(() => {
+    const stage = $canvasStage.get();
+    const canvasBaseLayer = $canvasBaseLayer.get();
     if (!stage || !canvasBaseLayer) {
       return;
     }
@@ -48,11 +44,11 @@ const useColorPicker = () => {
     }
 
     dispatch(setColorPickerColor({ r, g, b, a }));
-  }, [canvasBaseLayer, dispatch, stage]);
+  }, [dispatch]);
 
   const commitColorUnderCursor = useCallback(() => {
     dispatch(commitColorPickerColor());
-    $tool.set('brush')
+    $tool.set('brush');
   }, [dispatch]);
 
   return { updateColorUnderCursor, commitColorUnderCursor };

invokeai/frontend/web/src/features/canvas/store/canvasNanostore.ts

@@ -1,4 +1,5 @@
 import type { CanvasTool } from 'features/canvas/store/canvasTypes';
+import type Konva from "konva";
 import type { Vector2d } from 'konva/lib/types';
 import { atom, computed } from 'nanostores';
 
@@ -38,4 +39,6 @@ export const resetToolInteractionState = () => {
 
 export const setCanvasInteractionStateMouseOut = () => {
   $cursorPosition.set(null);
-};
+};export const $canvasBaseLayer = atom<Konva.Layer | null>(null);
+export const $canvasStage = atom<Konva.Stage | null>(null);
+

invokeai/frontend/web/src/features/canvas/util/getBaseLayerBlob.ts

@@ -1,6 +1,6 @@
 import type { RootState } from 'app/store/store';
+import { $canvasBaseLayer } from 'features/canvas/store/canvasNanostore';
 
-import { getCanvasBaseLayer } from './konvaInstanceProvider';
 import { konvaNodeToBlob } from './konvaNodeToBlob';
 
 /**
@@ -10,7 +10,7 @@ export const getBaseLayerBlob = async (
   state: RootState,
   alwaysUseBoundingBox: boolean = false
 ) => {
-  const canvasBaseLayer = getCanvasBaseLayer();
+  const canvasBaseLayer = $canvasBaseLayer.get();
 
   if (!canvasBaseLayer) {
     throw new Error('Problem getting base layer blob');

invokeai/frontend/web/src/features/canvas/util/getCanvasData.ts

@@ -1,15 +1,15 @@
 import { logger } from 'app/logging/logger';
+import { $canvasBaseLayer , $canvasStage } from 'features/canvas/store/canvasNanostore';
 import type {
   CanvasLayerState,
   Dimensions,
 } from 'features/canvas/store/canvasTypes';
 import { isCanvasMaskLine } from 'features/canvas/store/canvasTypes';
+import { konvaNodeToImageData } from 'features/canvas/util/konvaNodeToImageData';
 import type { Vector2d } from 'konva/lib/types';
 
 import createMaskStage from './createMaskStage';
-import { getCanvasBaseLayer, getCanvasStage } from './konvaInstanceProvider';
 import { konvaNodeToBlob } from './konvaNodeToBlob';
-import { konvaNodeToImageData } from './konvaNodeToImageData';
 
 /**
  * Gets Blob and ImageData objects for the base and mask layers
@@ -23,8 +23,8 @@ export const getCanvasData = async (
 ) => {
   const log = logger('canvas');
 
-  const canvasBaseLayer = getCanvasBaseLayer();
-  const canvasStage = getCanvasStage();
+  const canvasBaseLayer = $canvasBaseLayer.get();
+  const canvasStage = $canvasStage.get();
 
   if (!canvasBaseLayer || !canvasStage) {
     log.error('Unable to find canvas / stage');

invokeai/frontend/web/src/features/canvas/util/getFullBaseLayerBlob.ts

@@ -1,11 +1,12 @@
-import { getCanvasBaseLayer } from './konvaInstanceProvider';
+import { $canvasBaseLayer } from 'features/canvas/store/canvasNanostore';
+
 import { konvaNodeToBlob } from './konvaNodeToBlob';
 
 /**
  * Gets the canvas base layer blob, without bounding box
  */
 export const getFullBaseLayerBlob = async () => {
-  const canvasBaseLayer = getCanvasBaseLayer();
+  const canvasBaseLayer = $canvasBaseLayer.get();
 
   if (!canvasBaseLayer) {
     return;