fix(ui): attempt to fix fetch images

- Change the name of the redux action from `imagesUpdatedOne` to `imageUrlsUpdated` which is much clearer
- Add handling for initial image and canvas staging images (I had missed this initially)

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 { imageUpdatedOne } from 'features/gallery/store/imagesSlice';
+import { imageUrlsUpdated } from 'features/gallery/store/imagesSlice';
 
 const moduleLog = log.child({ namespace: 'image' });
 
@@ -11,7 +11,7 @@ export const addImageUrlsReceivedFulfilledListener = () => {
     effect: (action, { getState, dispatch }) => {
       const image = action.payload;
       moduleLog.debug({ data: { image } }, 'Image URLs received');
-      dispatch(imageUpdatedOne(image));
+      dispatch(imageUrlsUpdated(image));
     },
   });
 };

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

@@ -30,7 +30,7 @@ import {
 } from './canvasTypes';
 import { ImageDTO } from 'services/api';
 import { sessionCanceled } from 'services/thunks/session';
-import { imageUpdatedOne } from 'features/gallery/store/imagesSlice';
+import { imageUrlsUpdated } from 'features/gallery/store/imagesSlice';
 
 export const initialLayerState: CanvasLayerState = {
   objects: [],
@@ -853,7 +853,7 @@ export const canvasSlice = createSlice({
       }
     });
 
-    builder.addCase(imageUpdatedOne, (state, action) => {
+    builder.addCase(imageUrlsUpdated, (state, action) => {
       const { image_name, image_origin, image_url, thumbnail_url } =
         action.payload;
 
@@ -865,6 +865,13 @@ export const canvasSlice = createSlice({
           }
         }
       });
+
+      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/gallery/store/gallerySlice.ts

@@ -1,7 +1,7 @@
 import type { PayloadAction } from '@reduxjs/toolkit';
 import { createSlice } from '@reduxjs/toolkit';
 import { ImageDTO } from 'services/api';
-import { imageUpdatedOne, imageUpserted } from './imagesSlice';
+import { imageUrlsUpdated, imageUpserted } from './imagesSlice';
 
 type GalleryImageObjectFitType = 'contain' | 'cover';
 
@@ -54,7 +54,7 @@ export const gallerySlice = createSlice({
         state.selectedImage = action.payload;
       }
     });
-    builder.addCase(imageUpdatedOne, (state, action) => {
+    builder.addCase(imageUrlsUpdated, (state, action) => {
       const { image_name, image_origin, image_url, thumbnail_url } =
         action.payload;
 

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

@@ -57,7 +57,7 @@ const imagesSlice = createSlice({
 
       imagesAdapter.removeOne(state, action.payload.image_name);
     },
-    imageUpdatedOne: (state, action: PayloadAction<ImageUrlsDTO>) => {
+    imageUrlsUpdated: (state, action: PayloadAction<ImageUrlsDTO>) => {
       imagesAdapter.updateOne(state, {
         id: action.payload.image_name,
         changes: action.payload,
@@ -95,7 +95,7 @@ export const {
 
 export const {
   imageUpserted,
-  imageUpdatedOne,
+  imageUrlsUpdated,
   imageRemoved,
   imageCategoriesChanged,
 } = imagesSlice.actions;

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

@@ -8,6 +8,7 @@ import { receivedModels } from 'services/thunks/model';
 import { Scheduler } from 'app/constants';
 import { ImageDTO } from 'services/api';
 import { configChanged } from 'features/system/store/configSlice';
+import { imageUrlsUpdated } from 'features/gallery/store/imagesSlice';
 
 export interface GenerationState {
   cfgScale: number;
@@ -239,6 +240,15 @@ export const generationSlice = createSlice({
         state.model = defaultModel;
       }
     });
+    builder.addCase(imageUrlsUpdated, (state, action) => {
+      const { image_name, image_origin, image_url, thumbnail_url } =
+        action.payload;
+
+      if (state.initialImage && state.initialImage.image_name === image_name) {
+        state.initialImage.image_url = image_url;
+        state.initialImage.thumbnail_url = thumbnail_url;
+      }
+    });
   },
 });