From 68a231afeae4c2f8b5b7eadfd3f9f18b24ddcffe Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Tue, 9 Jan 2024 18:24:50 +1100
Subject: [PATCH] feat(ui): move canvas stage and base layer to nanostores

---
 .../listeners/canvasMerged.ts                 |  4 +-
 .../features/canvas/components/IAICanvas.tsx  | 27 ++++-----
 .../components/IAICanvasToolPreview.tsx       |  3 +-
 .../IAICanvasToolbar/IAICanvasToolbar.tsx     | 16 +++--
 .../features/canvas/hooks/useCanvasHotkeys.ts | 59 ++++++++-----------
 .../canvas/hooks/useColorUnderCursor.ts       | 14 ++---
 .../features/canvas/store/canvasNanostore.ts  |  5 +-
 .../features/canvas/util/getBaseLayerBlob.ts  |  4 +-
 .../src/features/canvas/util/getCanvasData.ts |  8 +--
 .../canvas/util/getFullBaseLayerBlob.ts       |  5 +-
 .../canvas/util/konvaInstanceProvider.ts      | 16 -----
 11 files changed, 67 insertions(+), 94 deletions(-)
 delete mode 100644 invokeai/frontend/web/src/features/canvas/util/konvaInstanceProvider.ts

diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/canvasMerged.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/canvasMerged.ts
index defbb04402..946f413d34 100644
--- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/canvasMerged.ts
+++ b/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');
diff --git a/invokeai/frontend/web/src/features/canvas/components/IAICanvas.tsx b/invokeai/frontend/web/src/features/canvas/components/IAICanvas.tsx
index 7d08a14766..0c77e27d07 100644
--- a/invokeai/frontend/web/src/features/canvas/components/IAICanvas.tsx
+++ b/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 });
 
diff --git a/invokeai/frontend/web/src/features/canvas/components/IAICanvasToolPreview.tsx b/invokeai/frontend/web/src/features/canvas/components/IAICanvasToolPreview.tsx
index 09e332caa3..97ad0c6073 100644
--- a/invokeai/frontend/web/src/features/canvas/components/IAICanvasToolPreview.tsx
+++ b/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);
diff --git a/invokeai/frontend/web/src/features/canvas/components/IAICanvasToolbar/IAICanvasToolbar.tsx b/invokeai/frontend/web/src/features/canvas/components/IAICanvasToolbar/IAICanvasToolbar.tsx
index ce2c0ca9ee..157e6ae9c1 100644
--- a/invokeai/frontend/web/src/features/canvas/components/IAICanvasToolbar/IAICanvasToolbar.tsx
+++ b/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;
     }
diff --git a/invokeai/frontend/web/src/features/canvas/hooks/useCanvasHotkeys.ts b/invokeai/frontend/web/src/features/canvas/hooks/useCanvasHotkeys.ts
index a79043a1a6..e29213cf7a 100644
--- a/invokeai/frontend/web/src/features/canvas/hooks/useCanvasHotkeys.ts
+++ b/invokeai/frontend/web/src/features/canvas/hooks/useCanvasHotkeys.ts
@@ -1,10 +1,9 @@
 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,
@@ -12,7 +11,6 @@ import {
   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);
diff --git a/invokeai/frontend/web/src/features/canvas/hooks/useColorUnderCursor.ts b/invokeai/frontend/web/src/features/canvas/hooks/useColorUnderCursor.ts
index 47076c56b5..bc531d7f3a 100644
--- a/invokeai/frontend/web/src/features/canvas/hooks/useColorUnderCursor.ts
+++ b/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 };
diff --git a/invokeai/frontend/web/src/features/canvas/store/canvasNanostore.ts b/invokeai/frontend/web/src/features/canvas/store/canvasNanostore.ts
index f1d0ab8fc5..12d2f386b6 100644
--- a/invokeai/frontend/web/src/features/canvas/store/canvasNanostore.ts
+++ b/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);
+
diff --git a/invokeai/frontend/web/src/features/canvas/util/getBaseLayerBlob.ts b/invokeai/frontend/web/src/features/canvas/util/getBaseLayerBlob.ts
index 047d74d380..ac6bed10a7 100644
--- a/invokeai/frontend/web/src/features/canvas/util/getBaseLayerBlob.ts
+++ b/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');
diff --git a/invokeai/frontend/web/src/features/canvas/util/getCanvasData.ts b/invokeai/frontend/web/src/features/canvas/util/getCanvasData.ts
index c23d2f9429..35dac03e01 100644
--- a/invokeai/frontend/web/src/features/canvas/util/getCanvasData.ts
+++ b/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');
diff --git a/invokeai/frontend/web/src/features/canvas/util/getFullBaseLayerBlob.ts b/invokeai/frontend/web/src/features/canvas/util/getFullBaseLayerBlob.ts
index ba855723fb..a5fbc99922 100644
--- a/invokeai/frontend/web/src/features/canvas/util/getFullBaseLayerBlob.ts
+++ b/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;
diff --git a/invokeai/frontend/web/src/features/canvas/util/konvaInstanceProvider.ts b/invokeai/frontend/web/src/features/canvas/util/konvaInstanceProvider.ts
deleted file mode 100644
index 350f23b2c7..0000000000
--- a/invokeai/frontend/web/src/features/canvas/util/konvaInstanceProvider.ts
+++ /dev/null
@@ -1,16 +0,0 @@
-import type Konva from 'konva';
-
-let canvasBaseLayer: Konva.Layer | null = null;
-let canvasStage: Konva.Stage | null = null;
-
-export const setCanvasBaseLayer = (layer: Konva.Layer) => {
-  canvasBaseLayer = layer;
-};
-
-export const getCanvasBaseLayer = () => canvasBaseLayer;
-
-export const setCanvasStage = (stage: Konva.Stage) => {
-  canvasStage = stage;
-};
-
-export const getCanvasStage = () => canvasStage;