From 454683e6eb058b4f83ee2045bb86464a6558645f Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Wed, 7 Jun 2023 00:23:51 +1000
Subject: [PATCH] feat(ui): update image urls on connect (#3507)

* feat(ui): update image urls on connect

Add `updateImageUrlsOnConnect` RTK listener:
- requests URLs for *every* image the app knows about, on connect: gallery, selectedImage, initialImage, canvas images, nodes images, controlnet images
- only fires when `shouldUpdateImagesOnConnect` config is enabled

* remove prop

---------

Co-authored-by: Mary Hipp <maryhipp@Marys-MacBook-Air.local>
---
 .../middleware/listenerMiddleware/index.ts    |  4 +
 .../listeners/updateImageUrlsOnConnect.ts     | 93 +++++++++++++++++++
 .../frontend/web/src/app/types/invokeai.ts    |  2 +-
 .../src/features/system/store/configSlice.ts  |  2 +-
 4 files changed, 99 insertions(+), 2 deletions(-)
 create mode 100644 invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/updateImageUrlsOnConnect.ts

diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/index.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/index.ts
index a9349dc863..8c073e81d6 100644
--- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/index.ts
+++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/index.ts
@@ -72,6 +72,7 @@ import { addCommitStagingAreaImageListener } from './listeners/addCommitStagingA
 import { addImageCategoriesChangedListener } from './listeners/imageCategoriesChanged';
 import { addControlNetImageProcessedListener } from './listeners/controlNetImageProcessed';
 import { addControlNetAutoProcessListener } from './listeners/controlNetAutoProcess';
+import { addUpdateImageUrlsOnConnectListener } from './listeners/updateImageUrlsOnConnect';
 
 export const listenerMiddleware = createListenerMiddleware();
 
@@ -179,3 +180,6 @@ addImageCategoriesChangedListener();
 // ControlNet
 addControlNetImageProcessedListener();
 addControlNetAutoProcessListener();
+
+// Update image URLs on connect
+addUpdateImageUrlsOnConnectListener();
diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/updateImageUrlsOnConnect.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/updateImageUrlsOnConnect.ts
new file mode 100644
index 0000000000..d02ffbe931
--- /dev/null
+++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/updateImageUrlsOnConnect.ts
@@ -0,0 +1,93 @@
+import { socketConnected } from 'services/events/actions';
+import { startAppListening } from '..';
+import { createSelector } from '@reduxjs/toolkit';
+import { generationSelector } from 'features/parameters/store/generationSelectors';
+import { canvasSelector } from 'features/canvas/store/canvasSelectors';
+import { nodesSelecter } from 'features/nodes/store/nodesSlice';
+import { controlNetSelector } from 'features/controlNet/store/controlNetSlice';
+import { ImageDTO } from 'services/api';
+import { forEach, uniqBy } from 'lodash-es';
+import { imageUrlsReceived } from 'services/thunks/image';
+import { log } from 'app/logging/useLogger';
+import { selectImagesEntities } from 'features/gallery/store/imagesSlice';
+
+const moduleLog = log.child({ namespace: 'images' });
+
+const selectAllUsedImages = createSelector(
+  [
+    generationSelector,
+    canvasSelector,
+    nodesSelecter,
+    controlNetSelector,
+    selectImagesEntities,
+  ],
+  (generation, canvas, nodes, controlNet, imageEntities) => {
+    const allUsedImages: ImageDTO[] = [];
+
+    if (generation.initialImage) {
+      allUsedImages.push(generation.initialImage);
+    }
+
+    canvas.layerState.objects.forEach((obj) => {
+      if (obj.kind === 'image') {
+        allUsedImages.push(obj.image);
+      }
+    });
+
+    nodes.nodes.forEach((node) => {
+      forEach(node.data.inputs, (input) => {
+        if (input.type === 'image' && input.value) {
+          allUsedImages.push(input.value);
+        }
+      });
+    });
+
+    forEach(controlNet.controlNets, (c) => {
+      if (c.controlImage) {
+        allUsedImages.push(c.controlImage);
+      }
+      if (c.processedControlImage) {
+        allUsedImages.push(c.processedControlImage);
+      }
+    });
+
+    forEach(imageEntities, (image) => {
+      if (image) {
+        allUsedImages.push(image);
+      }
+    });
+
+    const uniqueImages = uniqBy(allUsedImages, 'image_name');
+
+    return uniqueImages;
+  }
+);
+
+export const addUpdateImageUrlsOnConnectListener = () => {
+  startAppListening({
+    actionCreator: socketConnected,
+    effect: async (action, { dispatch, getState, take }) => {
+      const state = getState();
+
+      if (!state.config.shouldUpdateImagesOnConnect) {
+        return;
+      }
+
+      const allUsedImages = selectAllUsedImages(state);
+
+      moduleLog.trace(
+        { data: allUsedImages },
+        `Fetching new image URLs for ${allUsedImages.length} images`
+      );
+
+      allUsedImages.forEach(({ image_name, image_origin }) => {
+        dispatch(
+          imageUrlsReceived({
+            imageName: image_name,
+            imageOrigin: image_origin,
+          })
+        );
+      });
+    },
+  });
+};
diff --git a/invokeai/frontend/web/src/app/types/invokeai.ts b/invokeai/frontend/web/src/app/types/invokeai.ts
index f202b66ca2..4931c498bf 100644
--- a/invokeai/frontend/web/src/app/types/invokeai.ts
+++ b/invokeai/frontend/web/src/app/types/invokeai.ts
@@ -110,7 +110,7 @@ export type AppConfig = {
   /**
    * Whether or not we should update image urls when image loading errors
    */
-  shouldUpdateImageUrlsOnError: boolean;
+  shouldUpdateImagesOnConnect: boolean;
   disabledTabs: InvokeTabName[];
   disabledFeatures: AppFeature[];
   disabledSDFeatures: SDFeature[];
diff --git a/invokeai/frontend/web/src/features/system/store/configSlice.ts b/invokeai/frontend/web/src/features/system/store/configSlice.ts
index fb00a7a5d4..5f4dd68959 100644
--- a/invokeai/frontend/web/src/features/system/store/configSlice.ts
+++ b/invokeai/frontend/web/src/features/system/store/configSlice.ts
@@ -4,7 +4,7 @@ import { AppConfig, PartialAppConfig } from 'app/types/invokeai';
 import { merge } from 'lodash-es';
 
 export const initialConfigState: AppConfig = {
-  shouldUpdateImageUrlsOnError: false,
+  shouldUpdateImagesOnConnect: false,
   disabledTabs: [],
   disabledFeatures: [],
   disabledSDFeatures: [],