refactor(ui): canvas flow events (wip)

invokeai/frontend/web/src/services/events/onInvocationComplete.tsx

@@ -7,12 +7,13 @@ import { $nodeExecutionStates, upsertExecutionState } from 'features/nodes/hooks
 import { isImageField, isImageFieldCollection } from 'features/nodes/types/common';
 import { zNodeStatus } from 'features/nodes/types/invocation';
 import { isCanvasOutputEvent } from 'features/nodes/util/graph/graphBuilderUtils';
+import { flushSync } from 'react-dom';
 import type { ApiTagDescription } from 'services/api';
 import { boardsApi } from 'services/api/endpoints/boards';
 import { getImageDTOSafe, imagesApi } from 'services/api/endpoints/images';
 import type { ImageDTO, S } from 'services/api/types';
 import { getCategories, getListImagesUrl } from 'services/api/util';
-import { $lastProgressEvent } from 'services/events/stores';
+import { $lastCanvasProgressImage, $lastProgressEvent } from 'services/events/stores';
 import type { Param0 } from 'tsafe';
 import { objectEntries } from 'tsafe';
 import type { JsonObject } from 'type-fest';
@@ -176,7 +177,11 @@ export const buildOnInvocationComplete = (getState: () => RootState, dispatch: A
       return;
     }
 
-    dispatch(stagingAreaImageStaged({ stagingAreaImage: { imageDTO, offsetX: 0, offsetY: 0 } }));
+    flushSync(() => {
+      dispatch(stagingAreaImageStaged({ stagingAreaImage: { imageDTO, offsetX: 0, offsetY: 0 } }));
+    });
+
+    $lastCanvasProgressImage.set(null);
   };
 
   const handleOriginOther = async (data: S['InvocationCompleteEvent']) => {

invokeai/frontend/web/src/services/events/setEventListeners.tsx

@@ -30,7 +30,7 @@ import type { ClientToServerEvents, ServerToClientEvents } from 'services/events
 import type { Socket } from 'socket.io-client';
 import type { JsonObject } from 'type-fest';
 
-import { $lastProgressEvent } from './stores';
+import { $lastCanvasProgressEvent, $lastProgressEvent } from './stores';
 
 const log = logger('events');
 
@@ -428,6 +428,10 @@ export const setEventListeners = ({ socket, store, setIsConnected }: SetEventLis
       // If the queue item is completed, failed, or cancelled, we want to clear the last progress event
       $lastProgressEvent.set(null);
 
+      if (data.origin === 'canvas') {
+        $lastCanvasProgressEvent.set(null);
+      }
+
       // When a validation run is completed, we want to clear the validation run batch ID & set the workflow as published
       const validationRunData = $validationRunData.get();
       if (!validationRunData || batch_status.batch_id !== validationRunData.batchId || status !== 'completed') {

invokeai/frontend/web/src/services/events/stores.ts

@@ -1,3 +1,4 @@
+import type { ProgressImage } from 'features/nodes/types/common';
 import { round } from 'lodash-es';
 import { atom, computed, map } from 'nanostores';
 import type { S } from 'services/api/types';
@@ -15,18 +16,33 @@ $lastProgressEvent.subscribe((event) => {
   switch (event.destination) {
     case 'workflows':
       $lastWorkflowsProgressEvent.set(event);
+      if (event.image) {
+        $lastWorkflowsProgressImage.set({ sessionId: event.session_id, image: event.image });
+      }
       break;
     case 'upscaling':
       $lastUpscalingProgressEvent.set(event);
+      if (event.image) {
+        $lastUpscalingProgressImage.set({ sessionId: event.session_id, image: event.image });
+      }
       break;
     case 'canvas':
       $lastCanvasProgressEvent.set(event);
+      if (event.image) {
+        $lastCanvasProgressImage.set({ sessionId: event.session_id, image: event.image });
+      }
       break;
   }
 });
+
+type EphemeralProgressImage = { sessionId: string; image: ProgressImage };
+
 export const $lastCanvasProgressEvent = atom<S['InvocationProgressEvent'] | null>(null);
+export const $lastCanvasProgressImage = atom<EphemeralProgressImage | null>(null);
 export const $lastWorkflowsProgressEvent = atom<S['InvocationProgressEvent'] | null>(null);
+export const $lastWorkflowsProgressImage = atom<EphemeralProgressImage | null>(null);
 export const $lastUpscalingProgressEvent = atom<S['InvocationProgressEvent'] | null>(null);
+export const $lastUpscalingProgressImage = atom<EphemeralProgressImage | null>(null);
 
 export const $progressImage = computed($lastProgressEvent, (val) => val?.image ?? null);
 export const $hasProgressImage = computed($lastProgressEvent, (val) => Boolean(val?.image));