gallery autoscroll attempt 2

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

@@ -10,8 +10,17 @@ import {
   ImageDraggableData,
   TypesafeDraggableData,
 } from 'features/dnd/types';
+import { VirtuosoGalleryContext } from 'features/gallery/components/ImageGrid/types';
 import { useMultiselect } from 'features/gallery/hooks/useMultiselect';
-import { MouseEvent, memo, useCallback, useMemo, useState } from 'react';
+import {
+  MouseEvent,
+  memo,
+  useCallback,
+  useEffect,
+  useMemo,
+  useRef,
+  useState,
+} from 'react';
 import { useTranslation } from 'react-i18next';
 import { FaTrash } from 'react-icons/fa';
 import { MdStar, MdStarBorder } from 'react-icons/md';
@@ -24,20 +33,66 @@ 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();
-
   const { handleClick, isSelected, selection, selectionCount } =
     useMultiselect(imageDTO);
+  const imageContainerRef = useRef<HTMLDivElement>(null);
 
   const customStarUi = useStore($customStarUI);
 
+  useEffect(() => {
+    if (
+      !isSelected ||
+      selectionCount !== 1 ||
+      !virtuosoContext.rootRef.current ||
+      !virtuosoContext.virtuosoRef.current ||
+      !virtuosoContext.virtuosoRangeRef.current ||
+      !imageContainerRef.current
+    ) {
+      return;
+    }
+
+    const root = virtuosoContext.rootRef.current;
+    const virtuoso = virtuosoContext.virtuosoRef.current;
+    const item = imageContainerRef.current;
+    const range = virtuosoContext.virtuosoRangeRef.current;
+    const itemRect = item.getBoundingClientRect();
+    const rootRect = root.getBoundingClientRect();
+    const itemIsVisible =
+      itemRect.top >= rootRect.top &&
+      itemRect.bottom <= rootRect.bottom &&
+      itemRect.left >= rootRect.left &&
+      itemRect.right <= rootRect.right;
+
+    if (!itemIsVisible) {
+      virtuoso.scrollToIndex({
+        index: props.index,
+        behavior: 'smooth',
+        align:
+          props.index >
+          (range.endIndex - range.startIndex) / 2 + range.startIndex
+            ? 'end'
+            : 'start',
+      });
+    }
+  }, [
+    isSelected,
+    props.index,
+    selectionCount,
+    virtuosoContext.virtuosoRangeRef,
+    virtuosoContext.rootRef,
+    virtuosoContext.virtuosoRef,
+  ]);
+
   const handleDelete = useCallback(
     (e: MouseEvent<HTMLButtonElement>) => {
       e.stopPropagation();
@@ -122,6 +177,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,
@@ -35,6 +42,15 @@ const overlayScrollbarsConfig: UseOverlayScrollbarsParams = {
   },
 };
 
+// const selector = createSelector(stateSelector, (state) => {
+//   const selection = state.gallery.selection;
+
+//   if (selection.length !== 1) {
+//     return undefined;
+//   }
+//   return selection[0]?.image_name;
+// });
+
 const GalleryImageGrid = () => {
   const { t } = useTranslation();
   const rootRef = useRef<HTMLDivElement>(null);
@@ -47,6 +63,9 @@ 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 +91,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 +126,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 +182,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/hooks/useNextPrevImage.ts

@@ -3,16 +3,28 @@ import { stateSelector } from 'app/store/store';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import { imageSelected } from 'features/gallery/store/gallerySlice';
 import { clamp, isEqual } 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,
   useLazyListImagesQuery,
 } from 'services/api/endpoints/images';
-import { selectListImagesBaseQueryArgs } from '../store/gallerySelectors';
-import { IMAGE_LIMIT } from '../store/types';
 import { ListImagesArgs } from 'services/api/types';
 import { imagesAdapter } from 'services/api/util';
+import { selectListImagesBaseQueryArgs } from '../store/gallerySelectors';
+import { IMAGE_LIMIT } from '../store/types';
+
+export type UseNextPrevImageState = {
+  virtuosoRef: RefObject<VirtuosoGridHandle> | undefined;
+  virtuosoRangeRef: RefObject<ListRange> | undefined;
+};
+
+export const $useNextPrevImageState = map<UseNextPrevImageState>({
+  virtuosoRef: undefined,
+  virtuosoRangeRef: undefined,
+});
 
 export const nextPrevImageButtonsSelector = createSelector(
   [stateSelector, selectListImagesBaseQueryArgs],
@@ -78,6 +90,8 @@ export const nextPrevImageButtonsSelector = createSelector(
       isFetching: status === 'pending',
       nextImage,
       prevImage,
+      nextImageIndex,
+      prevImageIndex,
       queryArgs,
     };
   },
@@ -93,7 +107,9 @@ export const useNextPrevImage = () => {
 
   const {
     nextImage,
+    nextImageIndex,
     prevImage,
+    prevImageIndex,
     areMoreImagesAvailable,
     isFetching,
     queryArgs,
@@ -103,11 +119,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: 'start',
+      });
+    }
+  }, [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: 'end',
+      });
+    }
+  }, [dispatch, nextImage, nextImageIndex]);
 
   const [listImages] = useLazyListImagesQuery();