fix(ui): fix issue with gallery not letting you load more images

To determine whether the Load More button should work, we need to keep track of how many images are left to load for a given board or category.

The Assets tab doesn't work, though. Need to figure out a better way to handle this.

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

@@ -4,6 +4,7 @@ import { boardIdSelected } from 'features/gallery/store/boardSlice';
 import { selectImagesAll } from 'features/gallery/store/imagesSlice';
 import { IMAGES_PER_PAGE, receivedPageOfImages } from 'services/thunks/image';
 import { api } from 'services/apiSlice';
+import { imageSelected } from 'features/gallery/store/gallerySlice';
 
 const moduleLog = log.child({ namespace: 'boards' });
 
@@ -15,12 +16,62 @@ export const addBoardIdSelectedListener = () => {
 
       // we need to check if we need to fetch more images
 
+      const state = getState();
+      const allImages = selectImagesAll(state);
+
+      if (!boardId) {
+        // a board was unselected
+        dispatch(imageSelected(allImages[0]?.image_name));
+        return;
+      }
+
+      const { categories } = state.images;
+
+      const filteredImages = allImages.filter((i) => {
+        const isInCategory = categories.includes(i.image_category);
+        const isInSelectedBoard = boardId ? i.board_id === boardId : true;
+        return isInCategory && isInSelectedBoard;
+      });
+
+      // get the board from the cache
+      const { data: boards } = api.endpoints.listAllBoards.select()(state);
+      const board = boards?.find((b) => b.board_id === boardId);
+
+      if (!board) {
+        // can't find the board in cache...
+        dispatch(imageSelected(allImages[0]?.image_name));
+        return;
+      }
+
+      console.log('setting image');
+      dispatch(imageSelected(board.cover_image_name));
+
+      // if we haven't loaded one full page of images from this board, load more
+      if (
+        filteredImages.length < board.image_count &&
+        filteredImages.length < IMAGES_PER_PAGE
+      ) {
+        dispatch(receivedPageOfImages({ categories, boardId }));
+      }
+    },
+  });
+};
+
+export const addBoardIdSelected_changeSelectedImage_listener = () => {
+  startAppListening({
+    actionCreator: boardIdSelected,
+    effect: (action, { getState, dispatch }) => {
+      const boardId = action.payload;
+
+      const state = getState();
+
+      // we need to check if we need to fetch more images
+
       if (!boardId) {
         // a board was unselected - we don't need to do anything
         return;
       }
 
-      const state = getState();
       const { categories } = state.images;
 
       const filteredImages = selectImagesAll(state).filter((i) => {

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

@@ -12,12 +12,16 @@ export const addImageCategoriesChangedListener = () => {
   startAppListening({
     actionCreator: imageCategoriesChanged,
     effect: (action, { getState, dispatch }) => {
-      const filteredImagesCount = selectFilteredImagesAsArray(
-        getState()
-      ).length;
+      const state = getState();
+      const filteredImagesCount = selectFilteredImagesAsArray(state).length;
 
       if (!filteredImagesCount) {
-        dispatch(receivedPageOfImages());
+        dispatch(
+          receivedPageOfImages({
+            categories: action.payload,
+            boardId: state.boards.selectedBoardId,
+          })
+        );
       }
     },
   });

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

@@ -20,7 +20,12 @@ export const addSocketConnectedEventListener = () => {
       const { disabledTabs } = config;
 
       if (!images.ids.length) {
-        dispatch(receivedPageOfImages());
+        dispatch(
+          receivedPageOfImages({
+            categories: ['general'],
+            isIntermediate: false,
+          })
+        );
       }
 
       if (!models.ids.length) {

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

@@ -9,14 +9,10 @@ import ImageMetadataViewer from './ImageMetaDataViewer/ImageMetadataViewer';
 import NextPrevImageButtons from './NextPrevImageButtons';
 import { memo, useCallback } from 'react';
 import { systemSelector } from 'features/system/store/systemSelectors';
-import { configSelector } from '../../system/store/configSelectors';
-import { useAppToaster } from 'app/components/Toaster';
 import { imageSelected } from '../store/gallerySlice';
 import IAIDndImage from 'common/components/IAIDndImage';
 import { ImageDTO } from 'services/api';
 import { IAIImageLoadingFallback } from 'common/components/IAIImageFallback';
-import { RootState } from 'app/store/store';
-import { selectImagesById } from '../store/imagesSlice';
 import { useGetImageDTOQuery } from 'services/apiSlice';
 import { skipToken } from '@reduxjs/toolkit/dist/query';
 
@@ -114,14 +110,14 @@ const CurrentImagePreview = () => {
           }}
         >
           <IAIDndImage
-            image={image}
+            image={selectedImage && image ? image : undefined}
             onDrop={handleDrop}
             fallback={<IAIImageLoadingFallback sx={{ bg: 'none' }} />}
             isUploadDisabled={true}
           />
         </Flex>
       )}
-      {shouldShowImageDetails && image && (
+      {shouldShowImageDetails && image && selectedImage && (
         <Box
           sx={{
             position: 'absolute',
@@ -135,7 +131,7 @@ const CurrentImagePreview = () => {
           <ImageMetadataViewer image={image} />
         </Box>
       )}
-      {!shouldShowImageDetails && image && (
+      {!shouldShowImageDetails && image && selectedImage && (
         <Box
           sx={{
             position: 'absolute',

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

@@ -56,11 +56,7 @@ import {
   imageCategoriesChanged,
   selectImagesAll,
 } from '../store/imagesSlice';
-import {
-  IMAGES_PER_PAGE,
-  receivedImages,
-  receivedPageOfImages,
-} from 'services/thunks/image';
+import { receivedPageOfImages } from 'services/thunks/image';
 import BoardsList from './Boards/BoardsList';
 import { boardsSelector } from '../store/boardSlice';
 import { ChevronUpIcon } from '@chakra-ui/icons';
@@ -87,6 +83,7 @@ const itemSelector = createSelector(
       allImagesTotal,
       isLoading,
       categories,
+      selectedBoardId,
     };
   },
   defaultSelectorOptions
@@ -146,10 +143,10 @@ const ImageGalleryContent = () => {
     shouldUseSingleGalleryColumn,
     selectedImage,
     galleryView,
-    selectedBoardId,
   } = useAppSelector(mainSelector);
 
-  const { images, isLoading, allImagesTotal } = useAppSelector(itemSelector);
+  const { images, isLoading, allImagesTotal, categories, selectedBoardId } =
+    useAppSelector(itemSelector);
 
   const { selectedBoard } = useListAllBoardsQuery(undefined, {
     selectFromResult: ({ data }) => ({
@@ -167,8 +164,13 @@ const ImageGalleryContent = () => {
   }, [images.length, filteredImagesTotal]);
 
   const handleLoadMoreImages = useCallback(() => {
-    dispatch(receivedPageOfImages({}));
-  }, [dispatch]);
+    dispatch(
+      receivedPageOfImages({
+        categories,
+        boardId: selectedBoardId,
+      })
+    );
+  }, [categories, dispatch, selectedBoardId]);
 
   const handleEndReached = useMemo(() => {
     if (areMoreAvailable && !isLoading) {

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

@@ -1,8 +1,6 @@
 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';
 

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

@@ -11,7 +11,6 @@ import { dateComparator } from 'common/util/dateComparator';
 import { keyBy } from 'lodash-es';
 import {
   imageDeleted,
-  imageMetadataReceived,
   imageUrlsReceived,
   receivedPageOfImages,
 } from 'services/thunks/image';
@@ -74,11 +73,21 @@ const imagesSlice = createSlice({
     });
     builder.addCase(receivedPageOfImages.fulfilled, (state, action) => {
       state.isLoading = false;
+      const { boardId, categories, imageOrigin, isIntermediate } =
+        action.meta.arg;
+
       const { items, offset, limit, total } = action.payload;
+      imagesAdapter.upsertMany(state, items);
+
+      if (!categories?.includes('general') || boardId) {
+        // need to skip updating the total images count if the images recieved were for a specific board
+        // TODO: this doesn't work when on the Asset tab/category...
+        return;
+      }
+
       state.offset = offset;
       state.limit = limit;
       state.total = total;
-      imagesAdapter.upsertMany(state, items);
     });
     builder.addCase(imageDeleted.pending, (state, action) => {
       // Image deleted

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

@@ -148,7 +148,11 @@ export const receivedPageOfImages = createAppAsyncThunk(
     let queryArg: ReceivedImagesArg = {};
 
     if (size(arg)) {
-      queryArg = { ...DEFAULT_IMAGES_LISTED_ARG, ...arg };
+      queryArg = {
+        ...DEFAULT_IMAGES_LISTED_ARG,
+        offset: images.length,
+        ...arg,
+      };
     } else {
       queryArg = {
         ...DEFAULT_IMAGES_LISTED_ARG,