feat(ui): use updated boards data

- Update tooltips to use counts in the DTO - Remove unused `getBoardImagesTotal` and `getBoardAssetsTotal` queries, which were just abusing the list endpoint to get totals... - Remove extraneous optimistic update in invocation complete listener
invoke-ai · Sep 25, 2024 · 9a36b3d · 9a36b3d
1 parent d2f8181
commit 9a36b3d
Show file tree

Hide file tree

Showing 7 changed files with 49 additions and 63 deletions.
diff --git a/...eai/frontend/web/src/features/gallery/components/Boards/BoardsList/BoardTotalsTooltip.tsx b/...eai/frontend/web/src/features/gallery/components/Boards/BoardsList/BoardTotalsTooltip.tsx
@@ -0,0 +1,12 @@
+import { useTranslation } from 'react-i18next';
+
+type Props = {
+  imageCount: number;
+  assetCount: number;
+  isArchived: boolean;
+};
+
+export const BoardTotalsTooltip = ({ imageCount, assetCount, isArchived }: Props) => {
+  const { t } = useTranslation();
+  return `${t('boards.imagesWithCount', { count: imageCount })}, ${t('boards.assetsWithCount', { count: assetCount })}${isArchived ? ` (${t('boards.archived')})` : ''}`;
+};
diff --git a/invokeai/frontend/web/src/features/gallery/components/Boards/BoardsList/GalleryBoard.tsx b/invokeai/frontend/web/src/features/gallery/components/Boards/BoardsList/GalleryBoard.tsx
@@ -17,7 +17,7 @@ import IAIDroppable from 'common/components/IAIDroppable';
 import type { AddToBoardDropData } from 'features/dnd/types';
 import { AutoAddBadge } from 'features/gallery/components/Boards/AutoAddBadge';
 import BoardContextMenu from 'features/gallery/components/Boards/BoardContextMenu';
-import { BoardTooltip } from 'features/gallery/components/Boards/BoardsList/BoardTooltip';
+import { BoardTotalsTooltip } from 'features/gallery/components/Boards/BoardsList/BoardTotalsTooltip';
 import {
   selectAutoAddBoardId,
   selectAutoAssignBoardOnClick,
@@ -119,7 +119,18 @@ const GalleryBoard = ({ board, isSelected }: GalleryBoardProps) => {
   return (
     <BoardContextMenu board={board}>
       {(ref) => (
-        <Tooltip label={<BoardTooltip board={board} />} openDelay={1000} placement="left" closeOnScroll p={2}>
+        <Tooltip
+          label={
+            <BoardTotalsTooltip
+              imageCount={board.image_count}
+              assetCount={board.asset_count}
+              isArchived={Boolean(board.archived)}
+            />
+          }
+          openDelay={1000}
+          placement="left"
+          closeOnScroll
+        >
           <Flex
             position="relative"
             ref={ref}
@@ -167,7 +178,7 @@ const GalleryBoard = ({ board, isSelected }: GalleryBoardProps) => {
             </Editable>
             {autoAddBoardId === board.board_id && !editingDisclosure.isOpen && <AutoAddBadge />}
             {board.archived && !editingDisclosure.isOpen && <Icon as={PiArchiveBold} fill="base.300" />}
-            {!editingDisclosure.isOpen && <Text variant="subtext">{board.image_count}</Text>}
+            {!editingDisclosure.isOpen && <Text variant="subtext">{board.image_count + board.asset_count}</Text>}
 
             <IAIDroppable data={droppableData} dropLabel={t('unifiedCanvas.move')} />
           </Flex>

diff --git a/invokeai/frontend/web/src/features/gallery/components/Boards/BoardsList/NoBoardBoard.tsx b/invokeai/frontend/web/src/features/gallery/components/Boards/BoardsList/NoBoardBoard.tsx
@@ -4,7 +4,7 @@ import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import IAIDroppable from 'common/components/IAIDroppable';
 import type { RemoveFromBoardDropData } from 'features/dnd/types';
 import { AutoAddBadge } from 'features/gallery/components/Boards/AutoAddBadge';
-import { BoardTooltip } from 'features/gallery/components/Boards/BoardsList/BoardTooltip';
+import { BoardTotalsTooltip } from 'features/gallery/components/Boards/BoardsList/BoardTotalsTooltip';
 import NoBoardBoardContextMenu from 'features/gallery/components/Boards/NoBoardBoardContextMenu';
 import {
   selectAutoAddBoardId,
@@ -14,7 +14,7 @@ import {
 import { autoAddBoardIdChanged, boardIdSelected } from 'features/gallery/store/gallerySlice';
 import { memo, useCallback, useMemo } from 'react';
 import { useTranslation } from 'react-i18next';
-import { useGetBoardImagesTotalQuery } from 'services/api/endpoints/boards';
+import { useGetUncategorizedImageCountsQuery } from 'services/api/endpoints/boards';
 import { useBoardName } from 'services/api/hooks/useBoardName';
 
 interface Props {
@@ -27,11 +27,7 @@ const _hover: SystemStyleObject = {
 
 const NoBoardBoard = memo(({ isSelected }: Props) => {
   const dispatch = useAppDispatch();
-  const { imagesTotal } = useGetBoardImagesTotalQuery('none', {
-    selectFromResult: ({ data }) => {
-      return { imagesTotal: data?.total ?? 0 };
-    },
-  });
+  const { data } = useGetUncategorizedImageCountsQuery();
   const autoAddBoardId = useAppSelector(selectAutoAddBoardId);
   const autoAssignBoardOnClick = useAppSelector(selectAutoAssignBoardOnClick);
   const boardSearchText = useAppSelector(selectBoardSearchText);
@@ -60,7 +56,18 @@ const NoBoardBoard = memo(({ isSelected }: Props) => {
   return (
     <NoBoardBoardContextMenu>
       {(ref) => (
-        <Tooltip label={<BoardTooltip board={null} />} openDelay={1000} placement="left" closeOnScroll>
+        <Tooltip
+          label={
+            <BoardTotalsTooltip
+              imageCount={data?.image_count ?? 0}
+              assetCount={data?.asset_count ?? 0}
+              isArchived={false}
+            />
+          }
+          openDelay={1000}
+          placement="left"
+          closeOnScroll
+        >
           <Flex
             position="relative"
             ref={ref}
@@ -91,7 +98,7 @@ const NoBoardBoard = memo(({ isSelected }: Props) => {
               {boardName}
             </Text>
             {autoAddBoardId === 'none' && <AutoAddBadge />}
-            <Text variant="subtext">{imagesTotal}</Text>
+            <Text variant="subtext">{(data?.image_count ?? 0) + (data?.asset_count ?? 0)}</Text>
             <IAIDroppable data={droppableData} dropLabel={t('unifiedCanvas.move')} />
           </Flex>
         </Tooltip>

diff --git a/invokeai/frontend/web/src/services/api/endpoints/boards.ts b/invokeai/frontend/web/src/services/api/endpoints/boards.ts
@@ -1,12 +1,4 @@
-import { ASSETS_CATEGORIES, IMAGE_CATEGORIES } from 'features/gallery/store/types';
-import type {
-  BoardDTO,
-  CreateBoardArg,
-  ListBoardsArgs,
-  OffsetPaginatedResults_ImageDTO_,
-  UpdateBoardArg,
-} from 'services/api/types';
-import { getListImagesUrl } from 'services/api/util';
+import type { BoardDTO, CreateBoardArg, ListBoardsArgs, S, UpdateBoardArg } from 'services/api/types';
 
 import type { ApiTagDescription } from '..';
 import { api, buildV1Url, LIST_TAG } from '..';
@@ -55,38 +47,11 @@ export const boardsApi = api.injectEndpoints({
       keepUnusedDataFor: 0,
     }),
 
-    getBoardImagesTotal: build.query<{ total: number }, string | undefined>({
-      query: (board_id) => ({
-        url: getListImagesUrl({
-          board_id: board_id ?? 'none',
-          categories: IMAGE_CATEGORIES,
-          is_intermediate: false,
-          limit: 0,
-          offset: 0,
-        }),
-        method: 'GET',
+    getUncategorizedImageCounts: build.query<S['UncategorizedImageCounts'], void>({
+      query: () => ({
+        url: buildBoardsUrl('uncategorized/counts'),
       }),
-      providesTags: (result, error, arg) => [{ type: 'BoardImagesTotal', id: arg ?? 'none' }, 'FetchOnReconnect'],
-      transformResponse: (response: OffsetPaginatedResults_ImageDTO_) => {
-        return { total: response.total };
-      },
-    }),
-
-    getBoardAssetsTotal: build.query<{ total: number }, string | undefined>({
-      query: (board_id) => ({
-        url: getListImagesUrl({
-          board_id: board_id ?? 'none',
-          categories: ASSETS_CATEGORIES,
-          is_intermediate: false,
-          limit: 0,
-          offset: 0,
-        }),
-        method: 'GET',
-      }),
-      providesTags: (result, error, arg) => [{ type: 'BoardAssetsTotal', id: arg ?? 'none' }, 'FetchOnReconnect'],
-      transformResponse: (response: OffsetPaginatedResults_ImageDTO_) => {
-        return { total: response.total };
-      },
+      providesTags: ['UncategorizedImageCounts', { type: 'Board', id: LIST_TAG }, { type: 'Board', id: 'none' }],
     }),
 
     /**
@@ -124,9 +89,8 @@ export const boardsApi = api.injectEndpoints({
 
 export const {
   useListAllBoardsQuery,
-  useGetBoardImagesTotalQuery,
-  useGetBoardAssetsTotalQuery,
   useCreateBoardMutation,
   useUpdateBoardMutation,
   useListAllImageNamesForBoardQuery,
+  useGetUncategorizedImageCountsQuery,
 } = boardsApi;
diff --git a/invokeai/frontend/web/src/services/api/index.ts b/invokeai/frontend/web/src/services/api/index.ts
@@ -46,6 +46,7 @@ const tagTypes = [
   // This is invalidated on reconnect. It should be used for queries that have changing data,
   // especially related to the queue and generation.
   'FetchOnReconnect',
+  'UncategorizedImageCounts',
 ] as const;
 export type ApiTagDescription = TagDescription<(typeof tagTypes)[number]>;
 export const LIST_TAG = 'LIST';

diff --git a/invokeai/frontend/web/src/services/api/types.ts b/invokeai/frontend/web/src/services/api/types.ts
@@ -37,7 +37,6 @@ export type AppDependencyVersions = S['AppDependencyVersions'];
 export type ImageDTO = S['ImageDTO'];
 export type BoardDTO = S['BoardDTO'];
 export type ImageCategory = S['ImageCategory'];
-export type OffsetPaginatedResults_ImageDTO_ = S['OffsetPaginatedResults_ImageDTO_'];
 
 // Models
 export type ModelType = S['ModelType'];

diff --git a/invokeai/frontend/web/src/services/events/onInvocationComplete.tsx b/invokeai/frontend/web/src/services/events/onInvocationComplete.tsx
@@ -6,7 +6,6 @@ import { stagingAreaImageStaged } from 'features/controlLayers/store/canvasStagi
 import { boardIdSelected, galleryViewChanged, imageSelected, offsetChanged } from 'features/gallery/store/gallerySlice';
 import { $nodeExecutionStates, upsertExecutionState } from 'features/nodes/hooks/useExecutionState';
 import { zNodeStatus } from 'features/nodes/types/invocation';
-import { boardsApi } from 'services/api/endpoints/boards';
 import { getImageDTOSafe, imagesApi } from 'services/api/endpoints/images';
 import type { ImageDTO, S } from 'services/api/types';
 import { getCategories, getListImagesUrl } from 'services/api/util';
@@ -31,13 +30,6 @@ export const buildOnInvocationComplete = (getState: () => RootState, dispatch: A
       return;
     }
 
-    // update the total images for the board
-    dispatch(
-      boardsApi.util.updateQueryData('getBoardImagesTotal', imageDTO.board_id ?? 'none', (draft) => {
-        draft.total += 1;
-      })
-    );
-
     dispatch(
       imagesApi.util.invalidateTags([
         { type: 'Board', id: imageDTO.board_id ?? 'none' },