feat(ui): fix a few things that didn't unsubscribe correctly, add helper to manage subscriptions

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

@@ -254,12 +254,6 @@ export class CanvasManager {
     }
   }
 
-  syncStageScale() {
-    for (const layer of this.layers.values()) {
-      layer.syncStageScale();
-    }
-  }
-
   arrangeEntities() {
     const { getLayersState, getControlAdaptersState, getRegionsState } = this.stateApi;
     const layers = getLayersState().entities;
@@ -462,7 +456,7 @@ export class CanvasManager {
     this.log.debug('Initializing renderer');
     this.stage.container(this.container);
 
-    const cleanupListeners = setStageEventHandlers(this);
+    const unsubscribeListeners = setStageEventHandlers(this);
 
     // We can use a resize observer to ensure the stage always fits the container. We also need to re-render the bg and
     // document bounds overlay when the stage is resized.
@@ -473,19 +467,23 @@ export class CanvasManager {
     const unsubscribeRenderer = this._store.subscribe(this.render);
 
     // When we this flag, we need to render the staging area
-    $shouldShowStagedImage.subscribe(async (shouldShowStagedImage, prevShouldShowStagedImage) => {
-      if (shouldShowStagedImage !== prevShouldShowStagedImage) {
-        this.log.debug('Rendering staging area');
-        await this.preview.stagingArea.render();
+    const unsubscribeShouldShowStagedImage = $shouldShowStagedImage.subscribe(
+      async (shouldShowStagedImage, prevShouldShowStagedImage) => {
+        if (shouldShowStagedImage !== prevShouldShowStagedImage) {
+          this.log.debug('Rendering staging area');
+          await this.preview.stagingArea.render();
+        }
       }
-    });
+    );
 
-    $lastProgressEvent.subscribe(async (lastProgressEvent, prevLastProgressEvent) => {
-      if (lastProgressEvent !== prevLastProgressEvent) {
-        this.log.debug('Rendering progress image');
-        await this.preview.progressPreview.render(lastProgressEvent);
+    const unsubscribeLastProgressEvent = $lastProgressEvent.subscribe(
+      async (lastProgressEvent, prevLastProgressEvent) => {
+        if (lastProgressEvent !== prevLastProgressEvent) {
+          this.log.debug('Rendering progress image');
+          await this.preview.progressPreview.render(lastProgressEvent);
+        }
       }
-    });
+    );
 
     this.log.debug('First render of konva stage');
     this.preview.tool.render();
@@ -494,8 +492,9 @@ export class CanvasManager {
     return () => {
       this.log.debug('Cleaning up konva renderer');
       unsubscribeRenderer();
-      cleanupListeners();
-      $shouldShowStagedImage.off();
+      unsubscribeListeners();
+      unsubscribeShouldShowStagedImage();
+      unsubscribeLastProgressEvent();
       resizeObserver.disconnect();
     };
   };

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

@@ -2,6 +2,7 @@ import { $alt, $ctrl, $meta, $shift } from '@invoke-ai/ui-library';
 import type { Store } from '@reduxjs/toolkit';
 import { logger } from 'app/logging/logger';
 import type { RootState } from 'app/store/store';
+import { buildSubscribe } from 'features/controlLayers/konva/util';
 import {
   $isDrawing,
   $isMouseDown,
@@ -293,11 +294,11 @@ export class CanvasStateApi {
   onMetaChanged = $meta.subscribe;
 
   getShiftKey = $shift.get;
-  onShiftChanged = $shift.subscribe;
+  onShiftChanged = buildSubscribe($shift.subscribe, 'onShiftChanged');
 
   getShouldShowStagedImage = $shouldShowStagedImage.get;
   onGetShouldShowStagedImageChanged = $shouldShowStagedImage.subscribe;
 
   setStageAttrs = $stageAttrs.set;
-  onStageAttrsChanged = $stageAttrs.subscribe;
+  onStageAttrsChanged = buildSubscribe($stageAttrs.subscribe, 'onStageAttrsChanged');
 }

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

@@ -494,7 +494,6 @@ export const setStageEventHandlers = (manager: CanvasManager): (() => void) => {
           scale: newScale,
         });
         manager.background.render();
-        manager.syncStageScale();
       }
     }
     manager.preview.tool.render();

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

@@ -7,6 +7,7 @@ import Konva from 'konva';
 import type { KonvaEventObject } from 'konva/lib/Node';
 import type { Vector2d } from 'konva/lib/types';
 import { customAlphabet } from 'nanoid';
+import type { WritableAtom } from 'nanostores';
 import type { ImageDTO } from 'services/api/types';
 import { assert } from 'tsafe';
 
@@ -591,3 +592,23 @@ export function getObjectId(type: RenderableObject['type'], isBuffer?: boolean):
     return getPrefixedId(type);
   }
 }
+
+export type Subscription = {
+  name: string;
+  unsubscribe: () => void;
+};
+
+/**
+ * Builds a subscribe function for a nanostores atom.
+ * @param subscribe The subscribe function of the atom
+ * @param name The name of the atom
+ * @returns A subscribe function that returns an object with the name and unsubscribe function
+ */
+export const buildSubscribe = <T>(subscribe: WritableAtom<T>['subscribe'], name: string) => {
+  return (cb: Parameters<WritableAtom<T>['subscribe']>[0]): Subscription => {
+    return {
+      name,
+      unsubscribe: subscribe(cb),
+    };
+  };
+};