feat(ui): add support for shouldFetchImages if UI needs to re-fetch an image URL (#3250)

* if `shouldFetchImages` is passed in, UI will make an additional
request to get valid image URL when an invocation is complete
* this is necessary in order to have optional authorization for images

invokeai/frontend/web/index.d.ts

@@ -81,6 +81,7 @@ interface InvokeProps extends PropsWithChildren {
   disabledTabs?: InvokeTabName[];
   token?: string;
   shouldTransformUrls?: boolean;
+  shouldFetchImages?: boolean;
 }
 
 declare function Invoke(props: InvokeProps): JSX.Element;

invokeai/frontend/web/src/app/App.tsx

@@ -18,6 +18,7 @@ import { PropsWithChildren, useEffect } from 'react';
 import { setDisabledPanels, setDisabledTabs } from 'features/ui/store/uiSlice';
 import { InvokeTabName } from 'features/ui/store/tabMap';
 import { shouldTransformUrlsChanged } from 'features/system/store/systemSlice';
+import { setShouldFetchImages } from 'features/gallery/store/resultsSlice';
 
 keepGUIAlive();
 
@@ -26,6 +27,7 @@ interface Props extends PropsWithChildren {
     disabledPanels: string[];
     disabledTabs: InvokeTabName[];
     shouldTransformUrls?: boolean;
+    shouldFetchImages: boolean;
   };
 }
 
@@ -50,6 +52,10 @@ const App = (props: Props) => {
     );
   }, [dispatch, props.options.shouldTransformUrls]);
 
+  useEffect(() => {
+    dispatch(setShouldFetchImages(props.options.shouldFetchImages));
+  }, [dispatch, props.options.shouldFetchImages]);
+
   useEffect(() => {
     setColorMode(['light'].includes(currentTheme) ? 'light' : 'dark');
   }, [setColorMode, currentTheme]);

invokeai/frontend/web/src/component.tsx

@@ -30,6 +30,7 @@ interface Props extends PropsWithChildren {
   disabledTabs?: InvokeTabName[];
   token?: string;
   shouldTransformUrls?: boolean;
+  shouldFetchImages?: boolean;
 }
 
 export default function Component({
@@ -39,6 +40,7 @@ export default function Component({
   token,
   children,
   shouldTransformUrls,
+  shouldFetchImages = false,
 }: Props) {
   useEffect(() => {
     // configure API client token
@@ -70,7 +72,12 @@ export default function Component({
           <React.Suspense fallback={<Loading showText />}>
             <ThemeLocaleProvider>
               <App
-                options={{ disabledPanels, disabledTabs, shouldTransformUrls }}
+                options={{
+                  disabledPanels,
+                  disabledTabs,
+                  shouldTransformUrls,
+                  shouldFetchImages,
+                }}
               >
                 {children}
               </App>

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

@@ -254,7 +254,7 @@ const ImageGalleryContent = () => {
                   const isSelected = currentImageUuid === name;
                   return (
                     <HoverableImage
-                      key={name}
+                      key={`${name}-${image.thumbnail}`}
                       image={image}
                       isSelected={isSelected}
                     />

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

@@ -1,4 +1,8 @@
-import { createEntityAdapter, createSlice } from '@reduxjs/toolkit';
+import {
+  PayloadAction,
+  createEntityAdapter,
+  createSlice,
+} from '@reduxjs/toolkit';
 import { Image } from 'app/invokeai';
 import { invocationComplete } from 'services/events/actions';
 
@@ -13,6 +17,7 @@ import {
   extractTimestampFromImageName,
 } from 'services/util/deserializeImageField';
 import { deserializeImageResponse } from 'services/util/deserializeImageResponse';
+import { imageReceived, thumbnailReceived } from 'services/thunks/image';
 
 // use `createEntityAdapter` to create a slice for results images
 // https://redux-toolkit.js.org/api/createEntityAdapter#overview
@@ -34,6 +39,7 @@ type AdditionalResultsState = {
   pages: number; // the total number of pages available
   isLoading: boolean; // whether we are loading more images or not, mostly a placeholder
   nextPage: number; // the next page to request
+  shouldFetchImages: boolean; // whether we need to re-fetch images or not
 };
 
 export const initialResultsState =
@@ -43,6 +49,7 @@ export const initialResultsState =
     pages: 0,
     isLoading: false,
     nextPage: 0,
+    shouldFetchImages: false,
   });
 
 export type ResultsState = typeof initialResultsState;
@@ -57,6 +64,10 @@ const resultsSlice = createSlice({
     // here we just use the function itself as the reducer. we'll call this on `invocation_complete`
     // to add a single result
     resultAdded: resultsAdapter.upsertOne,
+
+    setShouldFetchImages: (state, action: PayloadAction<boolean>) => {
+      state.shouldFetchImages = action.payload;
+    },
   },
   extraReducers: (builder) => {
     // here we can respond to a fulfilled call of the `getNextResultsPage` thunk
@@ -98,7 +109,10 @@ const resultsSlice = createSlice({
       if (isImageOutput(result)) {
         const name = result.image.image_name;
         const type = result.image.image_type;
-        const { url, thumbnail } = buildImageUrls(type, name);
+        // if we need to refetch, set URLs to placeholder for now
+        const { url, thumbnail } = state.shouldFetchImages
+          ? { url: '', thumbnail: '' }
+          : buildImageUrls(type, name);
 
         const timestamp = extractTimestampFromImageName(name);
 
@@ -121,6 +135,30 @@ const resultsSlice = createSlice({
         resultsAdapter.addOne(state, image);
       }
     });
+
+    builder.addCase(imageReceived.fulfilled, (state, action) => {
+      const { imagePath } = action.payload;
+      const { imageName } = action.meta.arg;
+
+      resultsAdapter.updateOne(state, {
+        id: imageName,
+        changes: {
+          url: imagePath,
+        },
+      });
+    });
+
+    builder.addCase(thumbnailReceived.fulfilled, (state, action) => {
+      const { thumbnailPath } = action.payload;
+      const { imageName } = action.meta.arg;
+
+      resultsAdapter.updateOne(state, {
+        id: imageName,
+        changes: {
+          thumbnail: thumbnailPath,
+        },
+      });
+    });
   },
 });
 
@@ -134,6 +172,6 @@ export const {
   selectTotal: selectResultsTotal,
 } = resultsAdapter.getSelectors<RootState>((state) => state.results);
 
-export const { resultAdded } = resultsSlice.actions;
+export const { resultAdded, setShouldFetchImages } = resultsSlice.actions;
 
 export default resultsSlice.reducer;

invokeai/frontend/web/src/services/events/middleware.ts

@@ -30,6 +30,8 @@ import {
 import { OpenAPI } from 'services/api';
 import { receivedModels } from 'services/thunks/model';
 import { receivedOpenAPISchema } from 'services/thunks/schema';
+import { isImageOutput } from 'services/types/guards';
+import { imageReceived, thumbnailReceived } from 'services/thunks/image';
 
 export const socketMiddleware = () => {
   let areListenersSet = false;
@@ -213,6 +215,21 @@ export const socketMiddleware = () => {
         dispatch(sessionInvoked({ sessionId }));
       }
 
+      if (invocationComplete.match(action)) {
+        const { results } = getState();
+
+        if (results.shouldFetchImages) {
+          const { result } = action.payload.data;
+          if (isImageOutput(result)) {
+            const imageName = result.image.image_name;
+            const imageType = result.image.image_type;
+
+            dispatch(imageReceived({ imageName, imageType }));
+            dispatch(thumbnailReceived({ imageName, imageType }));
+          }
+        }
+      }
+
       // Always pass the action on so other middleware and reducers can handle it
       next(action);
     };

invokeai/frontend/web/src/services/thunks/image.ts

@@ -16,6 +16,21 @@ export const imageReceived = createAppAsyncThunk(
   }
 );
 
+type ThumbnailReceivedArg = Parameters<
+  (typeof ImagesService)['getThumbnail']
+>[0];
+
+/**
+ * `ImagesService.getThumbnail()` thunk
+ */
+export const thumbnailReceived = createAppAsyncThunk(
+  'api/thumbnailReceived',
+  async (arg: ThumbnailReceivedArg, _thunkApi) => {
+    const response = await ImagesService.getThumbnail(arg);
+    return response;
+  }
+);
+
 type ImageUploadedArg = Parameters<(typeof ImagesService)['uploadImage']>[0];
 
 /**