Selected in View within Gallery (#5240)

* selector added

* ref and useeffect added

* scrolling done using useeffect

* fixed scroll and changed the ref name

* fixed scroll again

* created hook for scroll logic

* feat(ui): debounce metadata fetch by 300ms

This vastly reduces the network requests when using the arrow keys to quickly skim through images.

* feat(ui): extract logic to determine virtuoso scrollToIndex align

This needs to be used in `useNextPrevImage()` to ensure the scrolling puts the image at the top or bottom appropriately

* feat(ui): add debounce to image workflow hook

This was spamming network requests like the metadata query

---------

Co-authored-by: psychedelicious <4822129+psychedelicious@users.noreply.github.com>

invokeai/frontend/web/src/features/gallery/components/ImageGrid/GalleryImage.tsx

@@ -3,6 +3,7 @@ import { useStore } from '@nanostores/react';
 import { $customStarUI } from 'app/store/nanostores/customStarUI';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import IAIDndImage from 'common/components/IAIDndImage';
+import IAIDndImageIcon from 'common/components/IAIDndImageIcon';
 import IAIFillSkeleton from 'common/components/IAIFillSkeleton';
 import { imagesToDeleteSelected } from 'features/deleteImageModal/store/slice';
 import {
@@ -10,7 +11,9 @@ import {
   ImageDraggableData,
   TypesafeDraggableData,
 } from 'features/dnd/types';
+import { VirtuosoGalleryContext } from 'features/gallery/components/ImageGrid/types';
 import { useMultiselect } from 'features/gallery/hooks/useMultiselect';
+import { useScrollToVisible } from 'features/gallery/hooks/useScrollToVisible';
 import { MouseEvent, memo, useCallback, useMemo, useState } from 'react';
 import { useTranslation } from 'react-i18next';
 import { FaTrash } from 'react-icons/fa';
@@ -20,15 +23,16 @@ import {
   useStarImagesMutation,
   useUnstarImagesMutation,
 } from 'services/api/endpoints/images';
-import IAIDndImageIcon from 'common/components/IAIDndImageIcon';
 
 interface HoverableImageProps {
   imageName: string;
+  index: number;
+  virtuosoContext: VirtuosoGalleryContext;
 }
 
 const GalleryImage = (props: HoverableImageProps) => {
   const dispatch = useAppDispatch();
-  const { imageName } = props;
+  const { imageName, virtuosoContext } = props;
   const { currentData: imageDTO } = useGetImageDTOQuery(imageName);
   const shift = useAppSelector((state) => state.hotkeys.shift);
   const { t } = useTranslation();
@@ -38,6 +42,13 @@ const GalleryImage = (props: HoverableImageProps) => {
 
   const customStarUi = useStore($customStarUI);
 
+  const imageContainerRef = useScrollToVisible(
+    isSelected,
+    props.index,
+    selectionCount,
+    virtuosoContext
+  );
+
   const handleDelete = useCallback(
     (e: MouseEvent<HTMLButtonElement>) => {
       e.stopPropagation();
@@ -122,6 +133,7 @@ const GalleryImage = (props: HoverableImageProps) => {
       data-testid={`image-${imageDTO.image_name}`}
     >
       <Flex
+        ref={imageContainerRef}
         userSelect="none"
         sx={{
           position: 'relative',

invokeai/frontend/web/src/features/gallery/components/ImageGrid/GalleryImageGrid.tsx

@@ -1,7 +1,10 @@
 import { Box, Flex } from '@chakra-ui/react';
+import { EntityId } from '@reduxjs/toolkit';
 import { useAppSelector } from 'app/store/storeHooks';
 import IAIButton from 'common/components/IAIButton';
 import { IAINoContentFallback } from 'common/components/IAIImageFallback';
+import { VirtuosoGalleryContext } from 'features/gallery/components/ImageGrid/types';
+import { $useNextPrevImageState } from 'features/gallery/hooks/useNextPrevImage';
 import { selectListImagesBaseQueryArgs } from 'features/gallery/store/gallerySelectors';
 import { IMAGE_LIMIT } from 'features/gallery/store/types';
 import {
@@ -11,7 +14,12 @@ import {
 import { memo, useCallback, useEffect, useMemo, useRef, useState } from 'react';
 import { useTranslation } from 'react-i18next';
 import { FaExclamationCircle, FaImage } from 'react-icons/fa';
-import { VirtuosoGrid } from 'react-virtuoso';
+import {
+  ItemContent,
+  ListRange,
+  VirtuosoGrid,
+  VirtuosoGridHandle,
+} from 'react-virtuoso';
 import {
   useLazyListImagesQuery,
   useListImagesQuery,
@@ -20,7 +28,6 @@ import { useBoardTotal } from 'services/api/hooks/useBoardTotal';
 import GalleryImage from './GalleryImage';
 import ImageGridItemContainer from './ImageGridItemContainer';
 import ImageGridListContainer from './ImageGridListContainer';
-import { EntityId } from '@reduxjs/toolkit';
 
 const overlayScrollbarsConfig: UseOverlayScrollbarsParams = {
   defer: true,
@@ -48,6 +55,10 @@ const GalleryImageGrid = () => {
   const { currentViewTotal } = useBoardTotal(selectedBoardId);
   const queryArgs = useAppSelector(selectListImagesBaseQueryArgs);
 
+  const virtuosoRangeRef = useRef<ListRange | null>(null);
+
+  const virtuosoRef = useRef<VirtuosoGridHandle>(null);
+
   const { currentData, isFetching, isSuccess, isError } =
     useListImagesQuery(queryArgs);
 
@@ -72,12 +83,26 @@ const GalleryImageGrid = () => {
     });
   }, [areMoreAvailable, listImages, queryArgs, currentData?.ids.length]);
 
-  const itemContentFunc = useCallback(
-    (index: number, imageName: EntityId) => (
-      <GalleryImage key={imageName} imageName={imageName as string} />
-    ),
-    []
-  );
+  const virtuosoContext = useMemo<VirtuosoGalleryContext>(() => {
+    return {
+      virtuosoRef,
+      rootRef,
+      virtuosoRangeRef,
+    };
+  }, []);
+
+  const itemContentFunc: ItemContent<EntityId, VirtuosoGalleryContext> =
+    useCallback(
+      (index, imageName, virtuosoContext) => (
+        <GalleryImage
+          key={imageName}
+          index={index}
+          imageName={imageName as string}
+          virtuosoContext={virtuosoContext}
+        />
+      ),
+      []
+    );
 
   useEffect(() => {
     // Initialize the gallery's custom scrollbar
@@ -93,6 +118,15 @@ const GalleryImageGrid = () => {
     return () => osInstance()?.destroy();
   }, [scroller, initialize, osInstance]);
 
+  const onRangeChanged = useCallback((range: ListRange) => {
+    virtuosoRangeRef.current = range;
+  }, []);
+
+  useEffect(() => {
+    $useNextPrevImageState.setKey('virtuosoRef', virtuosoRef);
+    $useNextPrevImageState.setKey('virtuosoRangeRef', virtuosoRangeRef);
+  }, []);
+
   if (!currentData) {
     return (
       <Flex
@@ -140,6 +174,10 @@ const GalleryImageGrid = () => {
             }}
             scrollerRef={setScroller}
             itemContent={itemContentFunc}
+            ref={virtuosoRef}
+            rangeChanged={onRangeChanged}
+            context={virtuosoContext}
+            overscan={10}
           />
         </Box>
         <IAIButton

invokeai/frontend/web/src/features/gallery/components/ImageGrid/types.ts

@@ -0,0 +1,8 @@
+import { RefObject } from 'react';
+import { ListRange, VirtuosoGridHandle } from 'react-virtuoso';
+
+export type VirtuosoGalleryContext = {
+  virtuosoRef: RefObject<VirtuosoGridHandle>;
+  rootRef: RefObject<HTMLDivElement>;
+  virtuosoRangeRef: RefObject<ListRange>;
+};

invokeai/frontend/web/src/features/gallery/components/ImageMetadataViewer/ImageMetadataWorkflowTabContent.tsx

@@ -1,7 +1,7 @@
 import { IAINoContentFallback } from 'common/components/IAIImageFallback';
 import { memo } from 'react';
 import { useTranslation } from 'react-i18next';
-import { useGetImageWorkflowQuery } from 'services/api/endpoints/images';
+import { useDebouncedImageWorkflow } from 'services/api/hooks/useDebouncedImageWorkflow';
 import { ImageDTO } from 'services/api/types';
 import DataViewer from './DataViewer';
 
@@ -11,7 +11,7 @@ type Props = {
 
 const ImageMetadataWorkflowTabContent = ({ image }: Props) => {
   const { t } = useTranslation();
-  const { currentData: workflow } = useGetImageWorkflowQuery(image.image_name);
+  const { workflow } = useDebouncedImageWorkflow(image);
 
   if (!workflow) {
     return <IAINoContentFallback label={t('nodes.noWorkflow')} />;

invokeai/frontend/web/src/features/gallery/hooks/useNextPrevImage.ts

@@ -4,8 +4,11 @@ import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import { selectListImagesBaseQueryArgs } from 'features/gallery/store/gallerySelectors';
 import { imageSelected } from 'features/gallery/store/gallerySlice';
 import { IMAGE_LIMIT } from 'features/gallery/store/types';
+import { getScrollToIndexAlign } from 'features/gallery/util/getScrollToIndexAlign';
 import { clamp } from 'lodash-es';
-import { useCallback } from 'react';
+import { map } from 'nanostores';
+import { RefObject, useCallback } from 'react';
+import { ListRange, VirtuosoGridHandle } from 'react-virtuoso';
 import { boardsApi } from 'services/api/endpoints/boards';
 import {
   imagesApi,
@@ -14,6 +17,16 @@ import {
 import { ListImagesArgs } from 'services/api/types';
 import { imagesAdapter } from 'services/api/util';
 
+export type UseNextPrevImageState = {
+  virtuosoRef: RefObject<VirtuosoGridHandle> | undefined;
+  virtuosoRangeRef: RefObject<ListRange> | undefined;
+};
+
+export const $useNextPrevImageState = map<UseNextPrevImageState>({
+  virtuosoRef: undefined,
+  virtuosoRangeRef: undefined,
+});
+
 export const nextPrevImageButtonsSelector = createMemoizedSelector(
   [stateSelector, selectListImagesBaseQueryArgs],
   (state, baseQueryArgs) => {
@@ -78,6 +91,8 @@ export const nextPrevImageButtonsSelector = createMemoizedSelector(
       isFetching: status === 'pending',
       nextImage,
       prevImage,
+      nextImageIndex,
+      prevImageIndex,
       queryArgs,
     };
   }
@@ -88,7 +103,9 @@ export const useNextPrevImage = () => {
 
   const {
     nextImage,
+    nextImageIndex,
     prevImage,
+    prevImageIndex,
     areMoreImagesAvailable,
     isFetching,
     queryArgs,
@@ -98,11 +115,43 @@ export const useNextPrevImage = () => {
 
   const handlePrevImage = useCallback(() => {
     prevImage && dispatch(imageSelected(prevImage));
-  }, [dispatch, prevImage]);
+    const range = $useNextPrevImageState.get().virtuosoRangeRef?.current;
+    const virtuoso = $useNextPrevImageState.get().virtuosoRef?.current;
+    if (!range || !virtuoso) {
+      return;
+    }
+
+    if (
+      prevImageIndex !== undefined &&
+      (prevImageIndex < range.startIndex || prevImageIndex > range.endIndex)
+    ) {
+      virtuoso.scrollToIndex({
+        index: prevImageIndex,
+        behavior: 'smooth',
+        align: getScrollToIndexAlign(prevImageIndex, range),
+      });
+    }
+  }, [dispatch, prevImage, prevImageIndex]);
 
   const handleNextImage = useCallback(() => {
     nextImage && dispatch(imageSelected(nextImage));
-  }, [dispatch, nextImage]);
+    const range = $useNextPrevImageState.get().virtuosoRangeRef?.current;
+    const virtuoso = $useNextPrevImageState.get().virtuosoRef?.current;
+    if (!range || !virtuoso) {
+      return;
+    }
+
+    if (
+      nextImageIndex !== undefined &&
+      (nextImageIndex < range.startIndex || nextImageIndex > range.endIndex)
+    ) {
+      virtuoso.scrollToIndex({
+        index: nextImageIndex,
+        behavior: 'smooth',
+        align: getScrollToIndexAlign(nextImageIndex, range),
+      });
+    }
+  }, [dispatch, nextImage, nextImageIndex]);
 
   const [listImages] = useLazyListImagesQuery();
 

invokeai/frontend/web/src/features/gallery/hooks/useScrollToVisible.ts

@@ -0,0 +1,46 @@
+import { VirtuosoGalleryContext } from 'features/gallery/components/ImageGrid/types';
+import { getScrollToIndexAlign } from 'features/gallery/util/getScrollToIndexAlign';
+import { useEffect, useRef } from 'react';
+
+export const useScrollToVisible = (
+  isSelected: boolean,
+  index: number,
+  selectionCount: number,
+  virtuosoContext: VirtuosoGalleryContext
+) => {
+  const imageContainerRef = useRef<HTMLDivElement>(null);
+
+  useEffect(() => {
+    if (
+      !isSelected ||
+      selectionCount !== 1 ||
+      !virtuosoContext.rootRef.current ||
+      !virtuosoContext.virtuosoRef.current ||
+      !virtuosoContext.virtuosoRangeRef.current ||
+      !imageContainerRef.current
+    ) {
+      return;
+    }
+
+    const itemRect = imageContainerRef.current.getBoundingClientRect();
+    const rootRect = virtuosoContext.rootRef.current.getBoundingClientRect();
+    const itemIsVisible =
+      itemRect.top >= rootRect.top &&
+      itemRect.bottom <= rootRect.bottom &&
+      itemRect.left >= rootRect.left &&
+      itemRect.right <= rootRect.right;
+
+    if (!itemIsVisible) {
+      virtuosoContext.virtuosoRef.current.scrollToIndex({
+        index,
+        behavior: 'smooth',
+        align: getScrollToIndexAlign(
+          index,
+          virtuosoContext.virtuosoRangeRef.current
+        ),
+      });
+    }
+  }, [isSelected, index, selectionCount, virtuosoContext]);
+
+  return imageContainerRef;
+};

invokeai/frontend/web/src/features/gallery/util/getScrollToIndexAlign.ts

@@ -0,0 +1,17 @@
+import { ListRange } from 'react-virtuoso';
+
+/**
+ * Gets the alignment for react-virtuoso's scrollToIndex function.
+ * @param index The index of the item.
+ * @param range The range of items currently visible.
+ * @returns
+ */
+export const getScrollToIndexAlign = (
+  index: number,
+  range: ListRange
+): 'start' | 'end' => {
+  if (index > (range.endIndex - range.startIndex) / 2 + range.startIndex) {
+    return 'end';
+  }
+  return 'start';
+};

invokeai/frontend/web/src/services/api/hooks/useDebouncedImageWorkflow.ts

@@ -0,0 +1,22 @@
+import { skipToken } from '@reduxjs/toolkit/query';
+import { useAppSelector } from 'app/store/storeHooks';
+import { useGetImageWorkflowQuery } from 'services/api/endpoints/images';
+import { ImageDTO } from 'services/api/types';
+import { useDebounce } from 'use-debounce';
+
+export const useDebouncedImageWorkflow = (imageDTO?: ImageDTO | null) => {
+  const workflowFetchDebounce = useAppSelector(
+    (state) => state.config.workflowFetchDebounce ?? 300
+  );
+
+  const [debouncedImageName] = useDebounce(
+    imageDTO?.has_workflow ? imageDTO.image_name : null,
+    workflowFetchDebounce
+  );
+
+  const { data: workflow, isLoading } = useGetImageWorkflowQuery(
+    debouncedImageName ?? skipToken
+  );
+
+  return { workflow, isLoading };
+};

invokeai/frontend/web/src/services/api/hooks/useDebouncedMetadata.ts

@@ -1,17 +1,14 @@
 import { skipToken } from '@reduxjs/toolkit/query';
-import { useDebounce } from 'use-debounce';
-import { useGetImageMetadataQuery } from 'services/api/endpoints/images';
 import { useAppSelector } from 'app/store/storeHooks';
+import { useGetImageMetadataQuery } from 'services/api/endpoints/images';
+import { useDebounce } from 'use-debounce';
 
 export const useDebouncedMetadata = (imageName?: string | null) => {
   const metadataFetchDebounce = useAppSelector(
-    (state) => state.config.metadataFetchDebounce
+    (state) => state.config.metadataFetchDebounce ?? 300
   );
 
-  const [debouncedImageName] = useDebounce(
-    imageName,
-    metadataFetchDebounce ?? 0
-  );
+  const [debouncedImageName] = useDebounce(imageName, metadataFetchDebounce);
 
   const { data: metadata, isLoading } = useGetImageMetadataQuery(
     debouncedImageName ?? skipToken