From 4e29a751d892ba50e3a5466ea7d3829ca0cc3995 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Mon, 22 May 2023 16:11:28 +1000
Subject: [PATCH] feat(ui): add POC image record fetching

---
 .../listeners/invocationComplete.ts           | 32 +++++++++++++++++--
 .../frontend/web/src/services/thunks/image.ts | 32 +++++++++++++++++++
 2 files changed, 62 insertions(+), 2 deletions(-)

diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/invocationComplete.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/invocationComplete.ts
index 9d84b2cbf0..b5c391afe4 100644
--- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/invocationComplete.ts
+++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/invocationComplete.ts
@@ -6,7 +6,12 @@ import {
 } from 'services/util/deserializeImageField';
 import { Image } from 'app/types/invokeai';
 import { resultAdded } from 'features/gallery/store/resultsSlice';
-import { imageReceived, thumbnailReceived } from 'services/thunks/image';
+import {
+  imageReceived,
+  imageRecordReceived,
+  imageUrlsReceived,
+  thumbnailReceived,
+} from 'services/thunks/image';
 import { startAppListening } from '..';
 import { imageSelected } from 'features/gallery/store/gallerySlice';
 import { addImageToStagingArea } from 'features/canvas/store/canvasSlice';
@@ -24,7 +29,7 @@ export const addImageResultReceivedListener = () => {
       }
       return false;
     },
-    effect: (action, { getState, dispatch }) => {
+    effect: async (action, { getState, dispatch, take }) => {
       if (!invocationComplete.match(action)) {
         return;
       }
@@ -35,6 +40,29 @@ export const addImageResultReceivedListener = () => {
       if (isImageOutput(result) && !nodeDenylist.includes(node.type)) {
         const name = result.image.image_name;
         const type = result.image.image_type;
+
+        dispatch(imageUrlsReceived({ imageName: name, imageType: type }));
+
+        const [{ payload }] = await take(
+          (action): action is ReturnType<typeof imageUrlsReceived.fulfilled> =>
+            imageUrlsReceived.fulfilled.match(action) &&
+            action.payload.image_name === name
+        );
+
+        console.log(payload);
+
+        dispatch(imageRecordReceived({ imageName: name, imageType: type }));
+
+        const [x] = await take(
+          (
+            action
+          ): action is ReturnType<typeof imageRecordReceived.fulfilled> =>
+            imageRecordReceived.fulfilled.match(action) &&
+            action.payload.image_name === name
+        );
+
+        console.log(x);
+
         const state = getState();
 
         // if we need to refetch, set URLs to placeholder for now
diff --git a/invokeai/frontend/web/src/services/thunks/image.ts b/invokeai/frontend/web/src/services/thunks/image.ts
index ec2533b61b..5528e41bfc 100644
--- a/invokeai/frontend/web/src/services/thunks/image.ts
+++ b/invokeai/frontend/web/src/services/thunks/image.ts
@@ -6,6 +6,38 @@ import { getHeaders } from 'services/util/getHeaders';
 
 const imagesLog = log.child({ namespace: 'image' });
 
+type imageUrlsReceivedArg = Parameters<
+  (typeof ImagesService)['getImageUrls']
+>[0];
+
+/**
+ * `ImagesService.getImageUrls()` thunk
+ */
+export const imageUrlsReceived = createAppAsyncThunk(
+  'api/imageUrlsReceived',
+  async (arg: imageUrlsReceivedArg) => {
+    const response = await ImagesService.getImageUrls(arg);
+    imagesLog.info({ arg, response }, 'Received image urls');
+    return response;
+  }
+);
+
+type imageRecordReceivedArg = Parameters<
+  (typeof ImagesService)['getImageUrls']
+>[0];
+
+/**
+ * `ImagesService.getImageUrls()` thunk
+ */
+export const imageRecordReceived = createAppAsyncThunk(
+  'api/imageUrlsReceived',
+  async (arg: imageRecordReceivedArg) => {
+    const response = await ImagesService.getImageRecord(arg);
+    imagesLog.info({ arg, response }, 'Received image record');
+    return response;
+  }
+);
+
 type ImageReceivedArg = Parameters<(typeof ImagesService)['getImage']>[0];
 
 /**