From 421c23d3ea6200e30e58cc5e44603a807d8a5794 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Wed, 21 Jun 2023 21:27:15 +1000
Subject: [PATCH] fix(ui): fix gallery image fetching for board categories

---
 .../listeners/boardIdSelected.ts              | 27 ++++++++++++++++---
 .../components/ImageGalleryContent.tsx        | 18 ++++++++-----
 2 files changed, 36 insertions(+), 9 deletions(-)

diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/boardIdSelected.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/boardIdSelected.ts
index 3889130a9c..30fb696ac3 100644
--- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/boardIdSelected.ts
+++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/boardIdSelected.ts
@@ -2,7 +2,8 @@ import { log } from 'app/logging/useLogger';
 import { startAppListening } from '..';
 import { boardIdSelected } from 'features/gallery/store/boardSlice';
 import { selectImagesAll } from 'features/gallery/store/imagesSlice';
-import { receivedPageOfImages } from 'services/thunks/image';
+import { IMAGES_PER_PAGE, receivedPageOfImages } from 'services/thunks/image';
+import { api } from 'services/apiSlice';
 
 const moduleLog = log.child({ namespace: 'boards' });
 
@@ -11,16 +12,36 @@ export const addBoardIdSelectedListener = () => {
     actionCreator: boardIdSelected,
     effect: (action, { getState, dispatch }) => {
       const boardId = action.payload;
+
+      // 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 images = selectImagesAll(state).filter((i) => {
+      const filteredImages = selectImagesAll(state).filter((i) => {
         const isInCategory = categories.includes(i.image_category);
         const isInSelectedBoard = boardId ? i.board_id === boardId : true;
         return isInCategory && isInSelectedBoard;
       });
 
-      if (images.length === 0) {
+      // 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...
+        return;
+      }
+
+      // 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 }));
       }
     },
diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageGalleryContent.tsx b/invokeai/frontend/web/src/features/gallery/components/ImageGalleryContent.tsx
index 866a68f0b6..c2f8a30409 100644
--- a/invokeai/frontend/web/src/features/gallery/components/ImageGalleryContent.tsx
+++ b/invokeai/frontend/web/src/features/gallery/components/ImageGalleryContent.tsx
@@ -69,7 +69,7 @@ import { useListAllBoardsQuery } from 'services/apiSlice';
 const itemSelector = createSelector(
   [(state: RootState) => state],
   (state) => {
-    const { categories, total, isLoading } = state.images;
+    const { categories, total: allImagesTotal, isLoading } = state.images;
     const { selectedBoardId } = state.boards;
 
     const allImages = selectImagesAll(state);
@@ -82,12 +82,10 @@ const itemSelector = createSelector(
       return isInCategory && isInSelectedBoard;
     });
 
-    const areMoreAvailable = images.length < total;
-
     return {
       images,
+      allImagesTotal,
       isLoading,
-      areMoreAvailable,
       categories,
     };
   },
@@ -151,8 +149,7 @@ const ImageGalleryContent = () => {
     selectedBoardId,
   } = useAppSelector(mainSelector);
 
-  const { images, areMoreAvailable, isLoading, categories } =
-    useAppSelector(itemSelector);
+  const { images, isLoading, allImagesTotal } = useAppSelector(itemSelector);
 
   const { selectedBoard } = useListAllBoardsQuery(undefined, {
     selectFromResult: ({ data }) => ({
@@ -160,6 +157,15 @@ const ImageGalleryContent = () => {
     }),
   });
 
+  const filteredImagesTotal = useMemo(
+    () => selectedBoard?.image_count ?? allImagesTotal,
+    [allImagesTotal, selectedBoard?.image_count]
+  );
+
+  const areMoreAvailable = useMemo(() => {
+    return images.length < filteredImagesTotal;
+  }, [images.length, filteredImagesTotal]);
+
   const handleLoadMoreImages = useCallback(() => {
     dispatch(receivedPageOfImages({}));
   }, [dispatch]);