feat(ui): move canvas-specific staging subscriptions to CanvasStagingAreaModule

invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/StagingAreaItemsList.tsx

@@ -4,9 +4,7 @@ import { useStore } from '@nanostores/react';
 import ScrollableContent from 'common/components/OverlayScrollbars/ScrollableContent';
 import { useCanvasSessionContext } from 'features/controlLayers/components/SimpleSession/context';
 import { QueueItemPreviewMini } from 'features/controlLayers/components/SimpleSession/QueueItemPreviewMini';
-import { getOutputImageName } from 'features/controlLayers/components/SimpleSession/shared';
 import { useCanvasManagerSafe } from 'features/controlLayers/contexts/CanvasManagerProviderGate';
-import { effect } from 'nanostores';
 import { memo, useEffect } from 'react';
 
 export const StagingAreaItemsList = memo(() => {
@@ -20,29 +18,8 @@ export const StagingAreaItemsList = memo(() => {
       return;
     }
 
-    return effect([ctx.$selectedItem, ctx.$progressData], (selectedItem, progressData) => {
-      if (!selectedItem) {
-        canvasManager.stagingArea.$imageSrc.set(null);
-        return;
-      }
-
-      const outputImageName = getOutputImageName(selectedItem);
-
-      if (outputImageName) {
-        canvasManager.stagingArea.$imageSrc.set({ type: 'imageName', data: outputImageName });
-        return;
-      }
-
-      const data = progressData[selectedItem.item_id];
-
-      if (data?.progressImage) {
-        canvasManager.stagingArea.$imageSrc.set({ type: 'dataURL', data: data.progressImage.dataURL });
-        return;
-      }
-
-      canvasManager.stagingArea.$imageSrc.set(null);
-    });
-  }, [canvasManager, ctx.$progressData, ctx.$selectedItem]);
+    return canvasManager.stagingArea.connectToSession(ctx.$selectedItemId, ctx.$progressData);
+  }, [canvasManager, ctx.$progressData, ctx.$selectedItemId]);
 
   return (
     <ScrollableContent overflowX="scroll" overflowY="hidden">

invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/context.tsx

@@ -13,18 +13,26 @@ import type { S } from 'services/api/types';
 import { $socket } from 'services/events/stores';
 import { assert } from 'tsafe';
 
-type ProgressData = {
+export type ProgressData = {
   itemId: number;
   progressEvent: S['InvocationProgressEvent'] | null;
   progressImage: ProgressImage | null;
+  outputImageName: string | null;
 };
 
+const getInitialProgressData = (itemId: number): ProgressData => ({
+  itemId,
+  progressEvent: null,
+  progressImage: null,
+  outputImageName: null,
+});
+
 export const useProgressData = (
   $progressData: WritableAtom<Record<number, ProgressData>>,
   itemId: number
 ): ProgressData => {
   const [value, setValue] = useState<ProgressData>(() => {
-    return $progressData.get()[itemId] ?? { itemId, progressEvent: null, progressImage: null };
+    return $progressData.get()[itemId] ?? getInitialProgressData(itemId);
   });
   useEffect(() => {
     const unsub = $progressData.subscribe((data) => {
@@ -62,6 +70,7 @@ const setProgress = ($progressData: WritableAtom<Record<number, ProgressData>>,
         itemId: data.item_id,
         progressEvent: data,
         progressImage: data.image ?? null,
+        outputImageName: null,
       },
     });
   }
@@ -315,18 +324,45 @@ export const CanvasSessionContextProvider = memo(
         const progressData = $progressData.get();
 
         const toDelete: number[] = [];
-        const toClear: number[] = [];
+        const toUpdate: ProgressData[] = [];
 
         for (const datum of Object.values(progressData)) {
           const item = items.find(({ item_id }) => item_id === datum.itemId);
           if (!item) {
             toDelete.push(datum.itemId);
           } else if (item.status === 'canceled' || item.status === 'failed') {
-            toClear.push(datum.itemId);
+            toUpdate[datum.itemId] = {
+              ...datum,
+              progressEvent: null,
+              progressImage: null,
+              outputImageName: null,
+            };
           }
         }
 
-        if (toDelete.length === 0) {
+        for (const item of items) {
+          const datum = progressData[item.item_id];
+
+          if (datum) {
+            if (datum.outputImageName) {
+              continue;
+            }
+            const outputImageName = getOutputImageName(item);
+            if (!outputImageName) {
+              continue;
+            }
+            toUpdate.push({
+              ...datum,
+              outputImageName,
+            });
+          } else {
+            const _datum = getInitialProgressData(item.item_id);
+            _datum.outputImageName = getOutputImageName(item);
+            toUpdate.push(_datum);
+          }
+        }
+
+        if (toDelete.length === 0 && toUpdate.length === 0) {
           return;
         }
 
@@ -336,12 +372,8 @@ export const CanvasSessionContextProvider = memo(
           delete newProgressData[itemId];
         }
 
-        for (const itemId of toClear) {
-          const current = newProgressData[itemId];
-          if (current) {
-            current.progressEvent = null;
-            current.progressImage = null;
-          }
+        for (const datum of toUpdate) {
+          newProgressData[datum.itemId] = datum;
         }
 
         $progressData.set(newProgressData);

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

@@ -1,4 +1,5 @@
 import { Mutex } from 'async-mutex';
+import type { ProgressData } from 'features/controlLayers/components/SimpleSession/context';
 import type { CanvasManager } from 'features/controlLayers/konva/CanvasManager';
 import { CanvasModuleBase } from 'features/controlLayers/konva/CanvasModuleBase';
 import { CanvasObjectImage } from 'features/controlLayers/konva/CanvasObject/CanvasObjectImage';
@@ -6,7 +7,8 @@ import { getPrefixedId } from 'features/controlLayers/konva/util';
 import { selectIsStaging } from 'features/controlLayers/store/canvasStagingAreaSlice';
 import type { CanvasImageState } from 'features/controlLayers/store/types';
 import Konva from 'konva';
-import { atom } from 'nanostores';
+import type { Atom, WritableAtom } from 'nanostores';
+import { atom, effect } from 'nanostores';
 import type { Logger } from 'roarr';
 
 type ImageNameSrc = { type: 'imageName'; data: string };
@@ -133,6 +135,29 @@ export class CanvasStagingAreaModule extends CanvasModuleBase {
     this.$isStaging.set(this.manager.stateApi.runSelector(selectIsStaging));
   };
 
+  connectToSession = (
+    $selectedItemId: Atom<number | null>,
+    $progressData: WritableAtom<Record<string, ProgressData>>
+  ) =>
+    effect([$selectedItemId, $progressData], (selectedItemId, progressData) => {
+      if (!selectedItemId) {
+        this.$imageSrc.set(null);
+        return;
+      }
+
+      const datum = progressData[selectedItemId];
+
+      if (datum?.outputImageName) {
+        this.$imageSrc.set({ type: 'imageName', data: datum.outputImageName });
+        return;
+      } else if (datum?.progressImage) {
+        this.$imageSrc.set({ type: 'dataURL', data: datum.progressImage.dataURL });
+        return;
+      } else {
+        this.$imageSrc.set(null);
+      }
+    });
+
   private _getImageFromSrc = (
     { type, data }: ImageNameSrc | DataURLSrc,
     width: number,