feat(ui): improve handling for urls/metadata received

Update images everywhere when urls or metadata is received:
- control images
- init images
- canvas
- nodes
- init image

Also renamed the variable.

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/imageDeleted.ts

@@ -5,6 +5,7 @@ import { log } from 'app/logging/useLogger';
 import { clamp } from 'lodash-es';
 import { imageSelected } from 'features/gallery/store/gallerySlice';
 import {
+  imageRemoved,
   selectImagesEntities,
   selectImagesIds,
 } from 'features/gallery/store/imagesSlice';
@@ -56,6 +57,10 @@ export const addRequestedImageDeletionListener = () => {
         }
       }
 
+      // Preemptively remove from gallery
+      dispatch(imageRemoved(image_name));
+
+      // Delete from server
       dispatch(
         imageDeleted({ imageName: image_name, imageOrigin: image_origin })
       );

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/imageMetadataReceived.ts

@@ -26,6 +26,10 @@ export const addImageMetadataReceivedFulfilledListener = () => {
         );
       } else if (image.is_intermediate) {
         // No further actions needed for intermediate images
+        moduleLog.trace(
+          { data: { image } },
+          'Image metadata received (intermediate), skipping'
+        );
         return;
       }
 

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/imageUrlsReceived.ts

@@ -1,7 +1,7 @@
 import { log } from 'app/logging/useLogger';
 import { startAppListening } from '..';
 import { imageUrlsReceived } from 'services/thunks/image';
-import { imagesAdapter } from 'features/gallery/store/imagesSlice';
+import { imageUpdatedOne } from 'features/gallery/store/imagesSlice';
 
 const moduleLog = log.child({ namespace: 'image' });
 
@@ -14,13 +14,12 @@ export const addImageUrlsReceivedFulfilledListener = () => {
 
       const { image_name, image_url, thumbnail_url } = image;
 
-      imagesAdapter.updateOne(getState().images, {
-        id: image_name,
-        changes: {
-          image_url,
-          thumbnail_url,
-        },
-      });
+      dispatch(
+        imageUpdatedOne({
+          id: image_name,
+          changes: { image_url, thumbnail_url },
+        })
+      );
     },
   });
 };

invokeai/frontend/web/src/app/types/invokeai.ts

@@ -114,6 +114,7 @@ export type AppConfig = {
   /**
    * Whether or not we need to re-fetch images
    */
+  shouldUpdateImageUrlsOnError: boolean;
   disabledTabs: InvokeTabName[];
   disabledFeatures: AppFeature[];
   disabledSDFeatures: SDFeature[];

invokeai/frontend/web/src/features/canvas/store/canvasSlice.ts

@@ -31,6 +31,7 @@ import {
 import { ImageDTO } from 'services/api';
 import { sessionCanceled } from 'services/thunks/session';
 import { setShouldUseCanvasBetaLayout } from 'features/ui/store/uiSlice';
+import { imageUrlsReceived } from 'services/thunks/image';
 
 export const initialLayerState: CanvasLayerState = {
   objects: [],
@@ -856,6 +857,26 @@ export const canvasSlice = createSlice({
     builder.addCase(setShouldUseCanvasBetaLayout, (state, action) => {
       state.doesCanvasNeedScaling = true;
     });
+    builder.addCase(imageUrlsReceived.fulfilled, (state, action) => {
+      const { image_name, image_origin, image_url, thumbnail_url } =
+        action.payload;
+
+      state.layerState.objects.forEach((object) => {
+        if (object.kind === 'image') {
+          if (object.image.image_name === image_name) {
+            object.image.image_url = image_url;
+            object.image.thumbnail_url = thumbnail_url;
+          }
+        }
+      });
+
+      state.layerState.stagingArea.images.forEach((stagedImage) => {
+        if (stagedImage.image.image_name === image_name) {
+          stagedImage.image.image_url = image_url;
+          stagedImage.image.thumbnail_url = thumbnail_url;
+        }
+      });
+    });
   },
 });
 

invokeai/frontend/web/src/features/controlNet/store/controlNetSlice.ts

@@ -13,7 +13,7 @@ import {
   ControlNetModel,
 } from './constants';
 import { controlNetImageProcessed } from './actions';
-import { imageDeleted } from 'services/thunks/image';
+import { imageDeleted, imageUrlsReceived } from 'services/thunks/image';
 import { forEach } from 'lodash-es';
 
 export const initialControlNet: Omit<ControlNetConfig, 'controlNetId'> = {
@@ -210,6 +210,22 @@ export const controlNetSlice = createSlice({
         }
       });
     });
+
+    builder.addCase(imageUrlsReceived.fulfilled, (state, action) => {
+      const { image_name, image_origin, image_url, thumbnail_url } =
+        action.payload;
+
+      forEach(state.controlNets, (c) => {
+        if (c.controlImage?.image_name === image_name) {
+          c.controlImage.image_url = image_url;
+          c.controlImage.thumbnail_url = thumbnail_url;
+        }
+        if (c.processedControlImage?.image_name === image_name) {
+          c.processedControlImage.image_url = image_url;
+          c.processedControlImage.thumbnail_url = thumbnail_url;
+        }
+      });
+    });
   },
 });
 

invokeai/frontend/web/src/features/gallery/store/gallerySlice.ts

@@ -2,6 +2,7 @@ import type { PayloadAction } from '@reduxjs/toolkit';
 import { createSlice } from '@reduxjs/toolkit';
 import { ImageDTO } from 'services/api';
 import { imageUpserted } from './imagesSlice';
+import { imageUrlsReceived } from 'services/thunks/image';
 
 type GalleryImageObjectFitType = 'contain' | 'cover';
 
@@ -57,6 +58,15 @@ export const gallerySlice = createSlice({
         state.selectedImage = action.payload;
       }
     });
+    builder.addCase(imageUrlsReceived.fulfilled, (state, action) => {
+      const { image_name, image_origin, image_url, thumbnail_url } =
+        action.payload;
+
+      if (state.selectedImage?.image_name === image_name) {
+        state.selectedImage.image_url = image_url;
+        state.selectedImage.thumbnail_url = thumbnail_url;
+      }
+    });
   },
 });
 

invokeai/frontend/web/src/features/gallery/store/imagesSlice.ts

@@ -1,5 +1,6 @@
 import {
   PayloadAction,
+  Update,
   createEntityAdapter,
   createSelector,
   createSlice,
@@ -8,7 +9,12 @@ import { RootState } from 'app/store/store';
 import { ImageCategory, ImageDTO } from 'services/api';
 import { dateComparator } from 'common/util/dateComparator';
 import { keyBy } from 'lodash-es';
-import { imageDeleted, receivedPageOfImages } from 'services/thunks/image';
+import {
+  imageDeleted,
+  imageMetadataReceived,
+  imageUrlsReceived,
+  receivedPageOfImages,
+} from 'services/thunks/image';
 
 export const imagesAdapter = createEntityAdapter<ImageDTO>({
   selectId: (image) => image.image_name,
@@ -49,6 +55,12 @@ const imagesSlice = createSlice({
     imageUpserted: (state, action: PayloadAction<ImageDTO>) => {
       imagesAdapter.upsertOne(state, action.payload);
     },
+    imageUpdatedOne: (state, action: PayloadAction<Update<ImageDTO>>) => {
+      imagesAdapter.updateOne(state, action.payload);
+    },
+    imageRemoved: (state, action: PayloadAction<string>) => {
+      imagesAdapter.removeOne(state, action.payload);
+    },
     imageCategoriesChanged: (state, action: PayloadAction<ImageCategory[]>) => {
       state.categories = action.payload;
     },
@@ -69,10 +81,19 @@ const imagesSlice = createSlice({
       imagesAdapter.upsertMany(state, items);
     });
     builder.addCase(imageDeleted.pending, (state, action) => {
-      // Preemptively remove the image from the gallery
+      // Image deleted
       const { imageName } = action.meta.arg;
       imagesAdapter.removeOne(state, imageName);
     });
+    builder.addCase(imageUrlsReceived.fulfilled, (state, action) => {
+      const { image_name, image_origin, image_url, thumbnail_url } =
+        action.payload;
+
+      imagesAdapter.updateOne(state, {
+        id: image_name,
+        changes: { image_url, thumbnail_url },
+      });
+    });
   },
 });
 
@@ -84,7 +105,12 @@ export const {
   selectTotal: selectImagesTotal,
 } = imagesAdapter.getSelectors<RootState>((state) => state.images);
 
-export const { imageUpserted, imageCategoriesChanged } = imagesSlice.actions;
+export const {
+  imageUpserted,
+  imageUpdatedOne,
+  imageRemoved,
+  imageCategoriesChanged,
+} = imagesSlice.actions;
 
 export default imagesSlice.reducer;
 

invokeai/frontend/web/src/features/nodes/store/nodesSlice.ts

@@ -16,9 +16,10 @@ import { receivedOpenAPISchema } from 'services/thunks/schema';
 import { InvocationTemplate, InvocationValue } from '../types/types';
 import { parseSchema } from '../util/parseSchema';
 import { log } from 'app/logging/useLogger';
-import { size } from 'lodash-es';
+import { forEach, size } from 'lodash-es';
 import { isAnyGraphBuilt } from './actions';
 import { RgbaColor } from 'react-colorful';
+import { imageUrlsReceived } from 'services/thunks/image';
 
 export type NodesState = {
   nodes: Node<InvocationValue>[];
@@ -98,9 +99,20 @@ const nodesSlice = createSlice({
       state.schema = action.payload;
     });
 
-    builder.addMatcher(isAnyGraphBuilt, (state, action) => {
-      // TODO: Achtung! Side effect in a reducer!
-      log.info({ namespace: 'nodes', data: action.payload }, 'Graph built');
+    builder.addCase(imageUrlsReceived.fulfilled, (state, action) => {
+      const { image_name, image_origin, image_url, thumbnail_url } =
+        action.payload;
+
+      state.nodes.forEach((node) => {
+        forEach(node.data.inputs, (input) => {
+          if (input.type === 'image') {
+            if (input.value?.image_name === image_name) {
+              input.value.image_url = image_url;
+              input.value.thumbnail_url = thumbnail_url;
+            }
+          }
+        });
+      });
     });
   },
 });

invokeai/frontend/web/src/features/parameters/store/generationSlice.ts

@@ -17,6 +17,7 @@ import {
   StrengthParam,
   WidthParam,
 } from './parameterZodSchemas';
+import { imageUrlsReceived } from 'services/thunks/image';
 
 export interface GenerationState {
   cfgScale: CfgScaleParam;
@@ -231,6 +232,16 @@ export const generationSlice = createSlice({
         state.model = defaultModel;
       }
     });
+
+    builder.addCase(imageUrlsReceived.fulfilled, (state, action) => {
+      const { image_name, image_origin, image_url, thumbnail_url } =
+        action.payload;
+
+      if (state.initialImage?.image_name === image_name) {
+        state.initialImage.image_url = image_url;
+        state.initialImage.thumbnail_url = thumbnail_url;
+      }
+    });
   },
 });
 

invokeai/frontend/web/src/features/system/store/configSlice.ts

@@ -5,6 +5,7 @@ import { merge } from 'lodash-es';
 
 export const initialConfigState: AppConfig = {
   shouldTransformUrls: false,
+  shouldUpdateImageUrlsOnError: false,
   disabledTabs: [],
   disabledFeatures: [],
   disabledSDFeatures: [],