From 0a71d6baa16a929484ace8eae6356269058f708e Mon Sep 17 00:00:00 2001 From: Mary Hipp Date: Fri, 11 Aug 2023 14:41:48 -0400 Subject: [PATCH] (ui) update cache to render pinned images alongside unpinned correctly, as well as changes in pinnedness --- .../SingleSelectionMenuItems.tsx | 8 +- .../web/src/services/api/endpoints/images.ts | 92 +++++++++++++++++-- .../frontend/web/src/services/api/util.ts | 31 +++++-- 3 files changed, 108 insertions(+), 23 deletions(-) diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageContextMenu/SingleSelectionMenuItems.tsx b/invokeai/frontend/web/src/features/gallery/components/ImageContextMenu/SingleSelectionMenuItems.tsx index f6c5a16f0a..b24146c798 100644 --- a/invokeai/frontend/web/src/features/gallery/components/ImageContextMenu/SingleSelectionMenuItems.tsx +++ b/invokeai/frontend/web/src/features/gallery/components/ImageContextMenu/SingleSelectionMenuItems.tsx @@ -134,12 +134,12 @@ const SingleSelectionMenuItems = (props: SingleSelectionMenuItemsProps) => { }, [copyImageToClipboard, imageDTO.image_url]); const handlePinImage = useCallback(() => { - togglePin({ imageName: imageDTO.image_name, pinned: true }); - }, [togglePin, imageDTO.image_name]); + togglePin({ imageDTO, pinned: true }); + }, [togglePin, imageDTO]); const handleUnpinImage = useCallback(() => { - togglePin({ imageName: imageDTO.image_name, pinned: false }); - }, [togglePin, imageDTO.image_name]); + togglePin({ imageDTO, pinned: false }); + }, [togglePin, imageDTO]); return ( <> diff --git a/invokeai/frontend/web/src/services/api/endpoints/images.ts b/invokeai/frontend/web/src/services/api/endpoints/images.ts index e9d26c40d8..9320e7a137 100644 --- a/invokeai/frontend/web/src/services/api/endpoints/images.ts +++ b/invokeai/frontend/web/src/services/api/endpoints/images.ts @@ -392,23 +392,28 @@ export const imagesApi = api.injectEndpoints({ */ changeImagePinned: build.mutation< ImageDTO, - { imageName: string; pinned: boolean } + { imageDTO: ImageDTO; pinned: boolean } >({ - query: ({ imageName, pinned }) => ({ - url: `images/i/${imageName}`, + query: ({ imageDTO, pinned }) => ({ + url: `images/i/${imageDTO.image_name}`, method: 'PATCH', body: { pinned }, }), - // invalidatesTags: (result, error, { imageDTO }) => [ - // { type: 'BoardImagesTotal', id: imageDTO.board_id ?? 'none' }, - // { type: 'BoardAssetsTotal', id: imageDTO.board_id ?? 'none' }, - // ], + invalidatesTags: (result, error, { imageDTO }) => [ + { + type: 'ImageList', + id: getListImagesUrl({ + board_id: imageDTO.board_id, + categories: IMAGE_CATEGORIES, + }), + }, + ], async onQueryStarted( - { imageName, pinned }, + { imageDTO, pinned }, { dispatch, queryFulfilled, getState } ) { /** - * Cache changes for `changeImageSessionId`: + * Cache changes for `changeImagePinned`: * - *update* getImageDTO */ @@ -420,7 +425,7 @@ export const imagesApi = api.injectEndpoints({ dispatch( imagesApi.util.updateQueryData( 'getImageDTO', - imageName, + imageDTO.image_name, (draft) => { Object.assign(draft, { pinned }); } @@ -428,6 +433,73 @@ export const imagesApi = api.injectEndpoints({ ) ); + const categories = getCategories(imageDTO); + + const queryArgs = { + board_id: imageDTO.board_id ?? 'none', + categories, + }; + + const currentCache = imagesApi.endpoints.listImages.select(queryArgs)( + getState() + ); + + const { data: total } = IMAGE_CATEGORIES.includes( + imageDTO.image_category + ) + ? boardsApi.endpoints.getBoardImagesTotal.select( + imageDTO.board_id ?? 'none' + )(getState()) + : boardsApi.endpoints.getBoardAssetsTotal.select( + imageDTO.board_id ?? 'none' + )(getState()); + + // IF it eligible for insertion into existing $cache + // "eligible" means either: + // - The cache is fully populated, with all images in the db cached + // OR + // - The image's `created_at` is within the range of the cached images within that pinned state + + const updatedImage: ImageDTO = { ...imageDTO, pinned } + + const isCacheFullyPopulated = + currentCache.data && currentCache.data.ids.length >= (total ?? 0); + + const isInDateRangeForPinnedState = getIsImageInDateRange( + currentCache.data, + updatedImage + ); + + if (!isInDateRangeForPinnedState) { + // if newly pinned or unpinned image is not in date range for its new state, remove from cache + patches.push( + dispatch( + imagesApi.util.updateQueryData( + 'listImages', + queryArgs, + (draft) => { + imagesAdapter.removeOne(draft, updatedImage.image_name); + } + ) + ) + ); + } + + if (isCacheFullyPopulated || isInDateRangeForPinnedState) { + // *upsert* to $cache + patches.push( + dispatch( + imagesApi.util.updateQueryData( + 'listImages', + queryArgs, + (draft) => { + imagesAdapter.upsertOne(draft, updatedImage); + } + ) + ) + ); + } + try { await queryFulfilled; } catch { diff --git a/invokeai/frontend/web/src/services/api/util.ts b/invokeai/frontend/web/src/services/api/util.ts index 20c9baedbb..10c4557eec 100644 --- a/invokeai/frontend/web/src/services/api/util.ts +++ b/invokeai/frontend/web/src/services/api/util.ts @@ -14,22 +14,26 @@ export const getIsImageInDateRange = ( if (!data) { return false; } - const cacheImageDTOS = imagesSelectors.selectAll(data); - if (cacheImageDTOS.length > 1) { - // Images are sorted by `created_at` DESC - // check if the image is newer than the oldest image in the cache + const totalCachedImageDtos = imagesSelectors.selectAll(data); + + if (totalCachedImageDtos.length <= 1) { + return true; + } + + const cacheImageDTOSForPinnedState = totalCachedImageDtos.filter((image) => image.pinned === imageDTO.pinned); + + if (cacheImageDTOSForPinnedState.length > 1) { + // Images are sorted by `pinned` DESC and then `created_at` DESC + // check if the image is newer than the oldest image in the cache for either the pinned group or unpinned group const createdDate = new Date(imageDTO.created_at); - const oldestImage = cacheImageDTOS[cacheImageDTOS.length - 1]; + const oldestImage = cacheImageDTOSForPinnedState[cacheImageDTOSForPinnedState.length - 1]; if (!oldestImage) { // satisfy TS gods, we already confirmed the array has more than one image return false; } const oldestDate = new Date(oldestImage.created_at); return createdDate >= oldestDate; - } else if ([0, 1].includes(cacheImageDTOS.length)) { - // if there are only 1 or 0 images in the cache, we consider the image to be in the date range - return true; } return false; }; @@ -45,7 +49,16 @@ export const getCategories = (imageDTO: ImageDTO) => { // with some other store of data. We will use the RTK Query cache as that store. export const imagesAdapter = createEntityAdapter({ selectId: (image) => image.image_name, - sortComparer: (a, b) => dateComparator(b.updated_at, a.updated_at), + sortComparer: (a, b) => { + // Compare pinned images first + if (a.pinned && !b.pinned) { + return -1; + } + if (!a.pinned && b.pinned) { + return 1; + } + return dateComparator(b.created_at, a.created_at) + }, }); // Create selectors for the adapter.