From 8d56becf04ea670b641be97fd05bdc92c1780be0 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Tue, 10 Sep 2024 12:54:43 +1000
Subject: [PATCH] fix(ui): retain global canvas manager instance

To prevent losing all ephemeral canvas stage when switching tabs, we will refrain from destroying the canvas manager instance when its tab unmounts, and use the existing canvas manager instance on mount, if there is one.

One small change required in `CanvasStageModule` - a `setContainer` method to update the konva stage DOM element.
---
 .../src/features/controlLayers/hooks/useInvokeCanvas.ts  | 9 ++++++++-
 .../features/controlLayers/konva/CanvasStageModule.ts    | 5 +++++
 .../web/src/features/controlLayers/store/canvasSlice.ts  | 3 +++
 3 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/invokeai/frontend/web/src/features/controlLayers/hooks/useInvokeCanvas.ts b/invokeai/frontend/web/src/features/controlLayers/hooks/useInvokeCanvas.ts
index ca0a8d9d58..286c32125d 100644
--- a/invokeai/frontend/web/src/features/controlLayers/hooks/useInvokeCanvas.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/hooks/useInvokeCanvas.ts
@@ -4,6 +4,7 @@ import { logger } from 'app/logging/logger';
 import { useAppStore } from 'app/store/nanostores/store';
 import { useAssertSingleton } from 'common/hooks/useAssertSingleton';
 import { CanvasManager } from 'features/controlLayers/konva/CanvasManager';
+import { $canvasManager } from 'features/controlLayers/store/canvasSlice';
 import Konva from 'konva';
 import { useLayoutEffect, useState } from 'react';
 import { useDevicePixelRatio } from 'use-device-pixel-ratio';
@@ -24,6 +25,7 @@ const useKonvaPixelRatioWatcher = () => {
 };
 
 export const useInvokeCanvas = (): ((el: HTMLDivElement | null) => void) => {
+  useAssertSingleton('useInvokeCanvas');
   useKonvaPixelRatioWatcher();
   const store = useAppStore();
   const socket = useStore($socket);
@@ -42,9 +44,14 @@ export const useInvokeCanvas = (): ((el: HTMLDivElement | null) => void) => {
       return () => {};
     }
 
+    const currentManager = $canvasManager.get();
+    if (currentManager) {
+      currentManager.stage.setContainer(container);
+      return;
+    }
+
     const manager = new CanvasManager(container, store, socket);
     manager.initialize();
-    return manager.destroy;
   }, [container, socket, store]);
 
   return containerRef;
diff --git a/invokeai/frontend/web/src/features/controlLayers/konva/CanvasStageModule.ts b/invokeai/frontend/web/src/features/controlLayers/konva/CanvasStageModule.ts
index 2ccf22351c..5499fac645 100644
--- a/invokeai/frontend/web/src/features/controlLayers/konva/CanvasStageModule.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/konva/CanvasStageModule.ts
@@ -78,6 +78,11 @@ export class CanvasStageModule extends CanvasModuleBase {
     };
   }
 
+  setContainer = (container: HTMLDivElement) => {
+    this.container = container;
+    this.konva.stage.container(container);
+  };
+
   setEventListeners = () => {
     this.konva.stage.on('wheel', this.onStageMouseWheel);
     this.konva.stage.on('dragmove', this.onStageDragMove);
diff --git a/invokeai/frontend/web/src/features/controlLayers/store/canvasSlice.ts b/invokeai/frontend/web/src/features/controlLayers/store/canvasSlice.ts
index 9855ded807..60d28600fb 100644
--- a/invokeai/frontend/web/src/features/controlLayers/store/canvasSlice.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/store/canvasSlice.ts
@@ -1302,4 +1302,7 @@ function actionsThrottlingFilter(action: UnknownAction) {
 }
 
 export const $lastCanvasProgressEvent = atom<S['InvocationDenoiseProgressEvent'] | null>(null);
+/**
+ * The global canvas manager instance.
+ */
 export const $canvasManager = atom<CanvasManager | null>(null);