fix(ui): auto image selection on invocation complete, board change

invokeai/frontend/web/src/app/store/middleware/debugLoggerMiddleware.ts

@@ -8,10 +8,13 @@ import { diff } from 'jsondiffpatch';
  * Super simple logger middleware. Useful for debugging when the redux devtools are awkward.
  */
 export const getDebugLoggerMiddleware =
-  (options?: { withDiff?: boolean; withNextState?: boolean }): Middleware =>
+  (options?: { filter?: (action: unknown) => boolean; withDiff?: boolean; withNextState?: boolean }): Middleware =>
   (api: MiddlewareAPI) =>
   (next) =>
   (action) => {
+    if (options?.filter?.(action)) {
+      return next(action);
+    }
     const originalState = api.getState();
     console.log('REDUX: dispatching', action);
     const result = next(action);

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/index.ts

@@ -9,7 +9,6 @@ import { addDeleteBoardAndImagesFulfilledListener } from 'app/store/middleware/l
 import { addBoardIdSelectedListener } from 'app/store/middleware/listenerMiddleware/listeners/boardIdSelected';
 import { addBulkDownloadListeners } from 'app/store/middleware/listenerMiddleware/listeners/bulkDownload';
 import { addEnqueueRequestedLinear } from 'app/store/middleware/listenerMiddleware/listeners/enqueueRequestedLinear';
-import { addEnsureImageIsSelectedListener } from 'app/store/middleware/listenerMiddleware/listeners/ensureImageIsSelectedListener';
 import { addGetOpenAPISchemaListener } from 'app/store/middleware/listenerMiddleware/listeners/getOpenAPISchema';
 import { addImageAddedToBoardFulfilledListener } from 'app/store/middleware/listenerMiddleware/listeners/imageAddedToBoard';
 import { addImageRemovedFromBoardFulfilledListener } from 'app/store/middleware/listenerMiddleware/listeners/imageRemovedFromBoard';
@@ -76,5 +75,3 @@ addAppConfigReceivedListener(startAppListening);
 addAdHocPostProcessingRequestedListener(startAppListening);
 
 addSetDefaultSettingsListener(startAppListening);
-
-addEnsureImageIsSelectedListener(startAppListening);

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

@@ -1,46 +1,59 @@
 import { isAnyOf } from '@reduxjs/toolkit';
 import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
-import { selectListImagesBaseQueryArgs } from 'features/gallery/store/gallerySelectors';
+import {
+  selectLastSelectedImage,
+  selectListImageNamesQueryArgs,
+  selectSelectedBoardId,
+} from 'features/gallery/store/gallerySelectors';
 import { boardIdSelected, galleryViewChanged, imageSelected } from 'features/gallery/store/gallerySlice';
 import { imagesApi } from 'services/api/endpoints/images';
 
 export const addBoardIdSelectedListener = (startAppListening: AppStartListening) => {
   startAppListening({
-    matcher: isAnyOf(boardIdSelected, galleryViewChanged),
+    matcher: isAnyOf(boardIdSelected, galleryViewChanged, imagesApi.endpoints.getImageNames.matchFulfilled),
     effect: async (action, { getState, dispatch, condition, cancelActiveListeners }) => {
       // Cancel any in-progress instances of this listener, we don't want to select an image from a previous board
       cancelActiveListeners();
 
+      if (boardIdSelected.match(action) && action.payload.selectedImageName) {
+        // This action already has a selected image name, we trust it is valid
+        return;
+      }
+
       const state = getState();
 
-      const queryArgs = { ...selectListImagesBaseQueryArgs(state), offset: 0 };
+      const board_id = selectSelectedBoardId(state);
+      const lastSelectedImage = selectLastSelectedImage(state);
+
+      if (
+        imagesApi.endpoints.getImageNames.matchFulfilled(action) &&
+        lastSelectedImage &&
+        action.meta.arg.originalArgs.board_id === board_id
+      ) {
+        // We just loaded image names for the current board, and we have a last selected image
+        return;
+      }
+
+      const queryArgs = { ...selectListImageNamesQueryArgs(state), board_id };
 
       // wait until the board has some images - maybe it already has some from a previous fetch
       // must use getState() to ensure we do not have stale state
       const isSuccess = await condition(
-        () => imagesApi.endpoints.listImages.select(queryArgs)(getState()).isSuccess,
+        () => imagesApi.endpoints.getImageNames.select(queryArgs)(getState()).isSuccess,
         5000
       );
 
-      if (isSuccess) {
-        // the board was just changed - we can select the first image
-        const { data: boardImagesData } = imagesApi.endpoints.listImages.select(queryArgs)(getState());
-
-        if (boardImagesData && boardIdSelected.match(action) && action.payload.selectedImageName) {
-          const selectedImage = boardImagesData.items.find(
-            (item) => item.image_name === action.payload.selectedImageName
-          );
-          dispatch(imageSelected(selectedImage?.image_name ?? null));
-        } else if (boardImagesData) {
-          dispatch(imageSelected(boardImagesData.items[0]?.image_name ?? null));
-        } else {
-          // board has no images - deselect
-          dispatch(imageSelected(null));
-        }
-      } else {
-        // fallback - deselect
+      if (!isSuccess) {
         dispatch(imageSelected(null));
+        return;
       }
+
+      // the board was just changed - we can select the first image
+      const imageNames = imagesApi.endpoints.getImageNames.select(queryArgs)(getState()).data?.image_names;
+
+      const imageToSelect = imageNames?.at(0) ?? null;
+
+      dispatch(imageSelected(imageToSelect));
     },
   });
 };

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

@@ -1,16 +0,0 @@
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
-import { imageSelected } from 'features/gallery/store/gallerySlice';
-import { imagesApi } from 'services/api/endpoints/images';
-
-export const addEnsureImageIsSelectedListener = (startAppListening: AppStartListening) => {
-  // When we list images, if no images is selected, select the first one.
-  startAppListening({
-    matcher: imagesApi.endpoints.listImages.matchFulfilled,
-    effect: (action, { dispatch, getState }) => {
-      const selection = getState().gallery.selection;
-      if (selection.length === 0) {
-        dispatch(imageSelected(action.payload.items[0]?.image_name ?? null));
-      }
-    },
-  });
-};

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

@@ -39,6 +39,7 @@ import { authToastMiddleware } from 'services/api/authToastMiddleware';
 import type { JsonObject } from 'type-fest';
 
 import { STORAGE_PREFIX } from './constants';
+import { getDebugLoggerMiddleware } from './middleware/debugLoggerMiddleware';
 import { actionSanitizer } from './middleware/devtools/actionSanitizer';
 import { actionsDenylist } from './middleware/devtools/actionsDenylist';
 import { stateSanitizer } from './middleware/devtools/stateSanitizer';
@@ -176,7 +177,17 @@ export const createStore = (uniqueStoreKey?: string, persist = true) =>
         .concat(api.middleware)
         .concat(dynamicMiddlewares)
         .concat(authToastMiddleware)
-        // .concat(getDebugLoggerMiddleware())
+        .concat(
+          getDebugLoggerMiddleware({
+            filter: (action) => {
+              try {
+                return (action as UnknownAction).type.startsWith('api');
+              } catch {
+                return false;
+              }
+            },
+          })
+        )
         .prepend(listenerMiddleware.middleware),
     enhancers: (getDefaultEnhancers) => {
       const _enhancers = getDefaultEnhancers().concat(autoBatchEnhancer());

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

@@ -10,6 +10,7 @@ import type { ProgressImage as ProgressImageType } from 'features/nodes/types/co
 import { selectShouldShowImageDetails, selectShouldShowProgressInViewer } from 'features/ui/store/uiSelectors';
 import type { AnimationProps } from 'framer-motion';
 import { AnimatePresence, motion } from 'framer-motion';
+import { atom } from 'nanostores';
 import { memo, useCallback, useEffect, useRef, useState } from 'react';
 import type { ImageDTO, S } from 'services/api/types';
 import { $socket } from 'services/events/stores';
@@ -25,8 +26,10 @@ export const CurrentImagePreview = memo(({ imageDTO }: { imageDTO: ImageDTO | nu
   const autoSwitch = useAppSelector(selectAutoSwitch);
 
   const socket = useStore($socket);
-  const [progressEvent, setProgressEvent] = useState<S['InvocationProgressEvent'] | null>(null);
-  const [progressImage, setProgressImage] = useState<ProgressImageType | null>(null);
+  const $progressEvent = useState(() => atom<S['InvocationProgressEvent'] | null>(null))[0];
+  const $progressImage = useState(() => atom<ProgressImageType | null>(null))[0];
+  const progressImage = useStore($progressImage);
+  const progressEvent = useStore($progressEvent);
 
   // Show and hide the next/prev buttons on mouse move
   const [shouldShowNextPrevButtons, setShouldShowNextPrevButtons] = useState<boolean>(false);
@@ -47,9 +50,9 @@ export const CurrentImagePreview = memo(({ imageDTO }: { imageDTO: ImageDTO | nu
     }
 
     const onInvocationProgress = (data: S['InvocationProgressEvent']) => {
-      setProgressEvent(data);
+      $progressEvent.set(data);
       if (data.image) {
-        setProgressImage(data.image);
+        $progressImage.set(data.image);
       }
     };
 
@@ -58,7 +61,7 @@ export const CurrentImagePreview = memo(({ imageDTO }: { imageDTO: ImageDTO | nu
     return () => {
       socket.off('invocation_progress', onInvocationProgress);
     };
-  }, [socket]);
+  }, [$progressEvent, $progressImage, socket]);
 
   useEffect(() => {
     if (!socket) {
@@ -72,8 +75,8 @@ export const CurrentImagePreview = memo(({ imageDTO }: { imageDTO: ImageDTO | nu
     // creating the illusion of the progress image turning into the new image.
     // But when auto-switch is disabled, we won't get that load event, so we need to clear the progress image manually.
     const onQueueItemStatusChanged = () => {
-      setProgressEvent(null);
-      setProgressImage(null);
+      $progressEvent.set(null);
+      $progressImage.set(null);
     };
 
     socket.on('queue_item_status_changed', onQueueItemStatusChanged);
@@ -81,17 +84,12 @@ export const CurrentImagePreview = memo(({ imageDTO }: { imageDTO: ImageDTO | nu
     return () => {
       socket.off('queue_item_status_changed', onQueueItemStatusChanged);
     };
-  }, [autoSwitch, socket]);
+  }, [$progressEvent, $progressImage, autoSwitch, socket]);
 
   const onLoadImage = useCallback(() => {
-    if (!progressEvent || !imageDTO) {
-      return;
-    }
-    if (progressEvent.session_id === imageDTO.session_id) {
-      setProgressEvent(null);
-      setProgressImage(null);
-    }
-  }, [imageDTO, progressEvent]);
+    $progressEvent.set(null);
+    $progressImage.set(null);
+  }, [$progressEvent, $progressImage]);
 
   const withProgress = shouldShowProgressInViewer && progressEvent && progressImage;
 

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

@@ -111,8 +111,12 @@ export const gallerySlice = createSlice({
       state.autoAssignBoardOnClick = action.payload;
     },
     boardIdSelected: (state, action: PayloadAction<{ boardId: BoardId; selectedImageName?: string }>) => {
-      state.selectedBoardId = action.payload.boardId;
+      const { boardId, selectedImageName } = action.payload;
+      state.selectedBoardId = boardId;
       state.galleryView = 'images';
+      if (selectedImageName) {
+        state.selection = [selectedImageName];
+      }
     },
     autoAddBoardIdChanged: (state, action: PayloadAction<BoardId>) => {
       if (!action.payload) {

invokeai/frontend/web/src/services/api/endpoints/images.ts

@@ -609,7 +609,7 @@ export const imageDTOToFile = async (imageDTO: ImageDTO): Promise<File> => {
 };
 
 export const useImageDTO = (imageName: string | null | undefined) => {
-  const { data: imageDTO } = useGetImageDTOQuery(imageName ?? skipToken);
+  const { currentData: imageDTO } = useGetImageDTOQuery(imageName ?? skipToken);
   return imageDTO ?? null;
 };
 

invokeai/frontend/web/src/services/events/onInvocationComplete.tsx

@@ -127,6 +127,7 @@ export const buildOnInvocationComplete = (getState: AppGetState, dispatch: AppDi
     // If the image is from a different board, switch to that board & select the image - otherwise just select the
     // image. This implicitly changes the view to 'images' if it was not already.
     if (board_id !== selectedBoardId) {
+      console.log('boardIdSelected');
       dispatch(
         boardIdSelected({
           boardId: board_id,
@@ -140,6 +141,7 @@ export const buildOnInvocationComplete = (getState: AppGetState, dispatch: AppDi
         dispatch(galleryViewChanged('images'));
       }
       // Select the image immediately since we've optimistically updated the cache
+      console.log('imageSelected');
       dispatch(imageSelected(lastImageDTO.image_name));
     }
   };