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

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

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

invokeai/frontend/web/src/common/hooks/useHandleOldUrls.ts

@@ -0,0 +1,49 @@
+import { useAppToaster } from 'app/components/Toaster';
+import { RootState } from 'app/store/store';
+import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
+import { useCallback, useState } from 'react';
+import { ImageDTO } from 'services/api';
+import { imageUrlsReceived } from 'services/thunks/image';
+
+export const useHandleOldUrls = () => {
+  const [didGetUrls, setDidGetUrls] = useState(false);
+  const dispatch = useAppDispatch();
+  const toaster = useAppToaster();
+  const shouldFetchImages = useAppSelector(
+    (state: RootState) => state.config.shouldFetchImages
+  );
+
+  return useCallback(
+    (image: ImageDTO | undefined, failedCallback?: () => void) => {
+      if (shouldFetchImages && image) {
+        if (didGetUrls) {
+          toaster({
+            title: 'Something went wrong, please refresh',
+            status: 'error',
+            isClosable: true,
+          });
+          if (failedCallback) {
+            failedCallback();
+          }
+          return;
+        }
+
+        const { image_origin, image_name } = image;
+
+        dispatch(
+          imageUrlsReceived({
+            imageOrigin: image_origin,
+            imageName: image_name,
+          })
+        );
+        setDidGetUrls(true);
+        return;
+      }
+
+      if (failedCallback) {
+        failedCallback();
+      }
+    },
+    [didGetUrls, dispatch, shouldFetchImages, toaster]
+  );
+};

invokeai/frontend/web/src/features/canvas/components/IAICanvasImage.tsx

@@ -1,15 +1,32 @@
+import { useHandleOldUrls } from 'common/hooks/useHandleOldUrls';
+import { useEffect } from 'react';
 import { Image } from 'react-konva';
+import { ImageDTO } from 'services/api';
 import useImage from 'use-image';
 
 type IAICanvasImageProps = {
-  url: string;
+  image: ImageDTO;
   x: number;
   y: number;
 };
+
 const IAICanvasImage = (props: IAICanvasImageProps) => {
-  const { url, x, y } = props;
+  const { image, x, y } = props;
-  const [image] = useImage(url, 'anonymous');
+  const handleOldUrls = useHandleOldUrls();
-  return <Image x={x} y={y} image={image} listening={false} />;
+  const [imageElement] = useImage(image.image_url, 'anonymous');
+
+  useEffect(() => {
+    if (imageElement) {
+      imageElement.onerror = () => handleOldUrls(image);
+    }
+
+    return () => {
+      if (imageElement) {
+        imageElement.onerror = () => undefined;
+      }
+    };
+  }, [imageElement, image, handleOldUrls]);
+  return <Image x={x} y={y} image={imageElement} listening={false} />;
 };
 
 export default IAICanvasImage;

invokeai/frontend/web/src/features/canvas/components/IAICanvasObjectRenderer.tsx

@@ -33,7 +33,6 @@ const selector = createSelector(
 
 const IAICanvasObjectRenderer = () => {
   const { objects } = useAppSelector(selector);
-  const { getUrl } = useGetUrl();
 
   if (!objects) return null;
 
@@ -42,12 +41,7 @@ const IAICanvasObjectRenderer = () => {
       {objects.map((obj, i) => {
         if (isCanvasBaseImage(obj)) {
           return (
-            <IAICanvasImage
+            <IAICanvasImage key={i} x={obj.x} y={obj.y} image={obj.image} />
-              key={i}
-              x={obj.x}
-              y={obj.y}
-              url={getUrl(obj.image.image_url)}
-            />
           );
         } else if (isCanvasBaseLine(obj)) {
           const line = (

invokeai/frontend/web/src/features/canvas/components/IAICanvasStagingArea.tsx

@@ -1,6 +1,5 @@
 import { createSelector } from '@reduxjs/toolkit';
 import { useAppSelector } from 'app/store/storeHooks';
-import { useGetUrl } from 'common/util/getUrl';
 import { canvasSelector } from 'features/canvas/store/canvasSelectors';
 import { GroupConfig } from 'konva/lib/Group';
 import { isEqual } from 'lodash-es';
@@ -56,16 +55,11 @@ const IAICanvasStagingArea = (props: Props) => {
     width,
     height,
   } = useAppSelector(selector);
-  const { getUrl } = useGetUrl();
 
   return (
     <Group {...rest}>
       {shouldShowStagingImage && currentStagingAreaImage && (
-        <IAICanvasImage
+        <IAICanvasImage image={currentStagingAreaImage.image} x={x} y={y} />
-          url={getUrl(currentStagingAreaImage.image.image_url) ?? ''}
-          x={x}
-          y={y}
-        />
       )}
       {shouldShowStagingOutline && (
         <Group>

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

@@ -30,6 +30,7 @@ import {
 } from './canvasTypes';
 import { ImageDTO } from 'services/api';
 import { sessionCanceled } from 'services/thunks/session';
+import { imageUrlsUpdated } from 'features/gallery/store/imagesSlice';
 
 export const initialLayerState: CanvasLayerState = {
   objects: [],
@@ -851,6 +852,27 @@ export const canvasSlice = createSlice({
         state.layerState.stagingArea = initialLayerState.stagingArea;
       }
     });
+
+    builder.addCase(imageUrlsUpdated, (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/canvas/util/konvaNodeToBlob.ts

@@ -12,5 +12,6 @@ export const konvaNodeToBlob = async (
   node: Konva.Node,
   boundingBox: IRect
 ): Promise<Blob> => {
+  console.log(node);
   return await canvasToBlob(node.toCanvas(boundingBox));
 };

invokeai/frontend/web/src/features/gallery/components/CurrentImagePreview.tsx

@@ -12,9 +12,8 @@ import { DragEvent, memo, useCallback } from 'react';
 import { systemSelector } from 'features/system/store/systemSelectors';
 import ImageFallbackSpinner from './ImageFallbackSpinner';
 import ImageMetadataOverlay from 'common/components/ImageMetadataOverlay';
-import { configSelector } from '../../system/store/configSelectors';
-import { useAppToaster } from 'app/components/Toaster';
 import { imageSelected } from '../store/gallerySlice';
+import { useHandleOldUrls } from 'common/hooks/useHandleOldUrls';
 
 export const imagesSelector = createSelector(
   [uiSelector, gallerySelector, systemSelector],
@@ -51,10 +50,9 @@ const CurrentImagePreview = () => {
     shouldShowProgressInViewer,
     shouldAntialiasProgressImage,
   } = useAppSelector(imagesSelector);
-  const { shouldFetchImages } = useAppSelector(configSelector);
   const { getUrl } = useGetUrl();
-  const toaster = useAppToaster();
   const dispatch = useAppDispatch();
+  const handleOldUrls = useHandleOldUrls();
 
   const handleDragStart = useCallback(
     (e: DragEvent<HTMLDivElement>) => {
@@ -68,15 +66,10 @@ const CurrentImagePreview = () => {
   );
 
   const handleError = useCallback(() => {
-    dispatch(imageSelected());
+    handleOldUrls(image, () => {
-    if (shouldFetchImages) {
+      dispatch(imageSelected());
-      toaster({
+    });
-        title: 'Something went wrong, please refresh',
+  }, [dispatch, handleOldUrls, image]);
-        status: 'error',
-        isClosable: true,
-      });
-    }
-  }, [dispatch, toaster, shouldFetchImages]);
 
   return (
     <Flex

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 { imageUpserted } from './imagesSlice';
+import { imageUrlsUpdated, imageUpserted } from './imagesSlice';
 
 type GalleryImageObjectFitType = 'contain' | 'cover';
 
@@ -54,6 +54,18 @@ export const gallerySlice = createSlice({
         state.selectedImage = action.payload;
       }
     });
+    builder.addCase(imageUrlsUpdated, (state, action) => {
+      const { image_name, image_origin, image_url, thumbnail_url } =
+        action.payload;
+
+      if (
+        state.selectedImage &&
+        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

@@ -5,7 +5,7 @@ import {
   createSlice,
 } from '@reduxjs/toolkit';
 import { RootState } from 'app/store/store';
-import { ImageCategory, ImageDTO } from 'services/api';
+import { ImageCategory, ImageDTO, ImageUrlsDTO } from 'services/api';
 import { dateComparator } from 'common/util/dateComparator';
 import { isString, keyBy } from 'lodash-es';
 import { receivedPageOfImages } from 'services/thunks/image';
@@ -57,6 +57,12 @@ const imagesSlice = createSlice({
 
       imagesAdapter.removeOne(state, action.payload.image_name);
     },
+    imageUrlsUpdated: (state, action: PayloadAction<ImageUrlsDTO>) => {
+      imagesAdapter.updateOne(state, {
+        id: action.payload.image_name,
+        changes: action.payload,
+      });
+    },
     imageCategoriesChanged: (state, action: PayloadAction<ImageCategory[]>) => {
       state.categories = action.payload;
     },
@@ -87,8 +93,12 @@ export const {
   selectTotal: selectImagesTotal,
 } = imagesAdapter.getSelectors<RootState>((state) => state.images);
 
-export const { imageUpserted, imageRemoved, imageCategoriesChanged } =
+export const {
-  imagesSlice.actions;
+  imageUpserted,
+  imageUrlsUpdated,
+  imageRemoved,
+  imageCategoriesChanged,
+} = imagesSlice.actions;
 
 export default imagesSlice.reducer;
 

invokeai/frontend/web/src/features/parameters/components/Parameters/ImageToImage/InitialImagePreview.tsx

@@ -4,15 +4,13 @@ import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import { useGetUrl } from 'common/util/getUrl';
 import { clearInitialImage } from 'features/parameters/store/generationSlice';
 import { DragEvent, useCallback } from 'react';
-import { useTranslation } from 'react-i18next';
 import ImageMetadataOverlay from 'common/components/ImageMetadataOverlay';
 import { generationSelector } from 'features/parameters/store/generationSelectors';
 import { initialImageSelected } from 'features/parameters/store/actions';
 import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
 import ImageFallbackSpinner from 'features/gallery/components/ImageFallbackSpinner';
 import { FaImage } from 'react-icons/fa';
-import { configSelector } from '../../../../system/store/configSelectors';
+import { useHandleOldUrls } from 'common/hooks/useHandleOldUrls';
-import { useAppToaster } from 'app/components/Toaster';
 
 const selector = createSelector(
   [generationSelector],
@@ -27,29 +25,15 @@ const selector = createSelector(
 
 const InitialImagePreview = () => {
   const { initialImage } = useAppSelector(selector);
-  const { shouldFetchImages } = useAppSelector(configSelector);
   const { getUrl } = useGetUrl();
   const dispatch = useAppDispatch();
-  const { t } = useTranslation();
+  const handleOldUrls = useHandleOldUrls();
-  const toaster = useAppToaster();
 
   const handleError = useCallback(() => {
-    dispatch(clearInitialImage());
+    handleOldUrls(initialImage, () => {
-    if (shouldFetchImages) {
+      dispatch(clearInitialImage());
-      toaster({
+    });
-        title: 'Something went wrong, please refresh',
+  }, [dispatch, handleOldUrls, initialImage]);
-        status: 'error',
-        isClosable: true,
-      });
-    } else {
-      toaster({
-        title: t('toast.parametersFailed'),
-        description: t('toast.parametersFailedDesc'),
-        status: 'error',
-        isClosable: true,
-      });
-    }
-  }, [dispatch, t, toaster, shouldFetchImages]);
 
   const handleDrop = useCallback(
     (e: DragEvent<HTMLDivElement>) => {

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;
+      }
+    });
   },
 });
 

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

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