feat(ui): iterate on dynamic gallery limit

- Simplify the gallery layout
- Set an initial gallery limit to load _some_ images immediately.
- Refactor the resize observer to use the actual rendered image component to calculate the number of images per row/col. This prevents inaccuracies caused by image padding that could result in the wrong number of images.
- Debounce the limit update to not thrash teh API
- Use absolute positioning trick to ensure the gallery container is always exactly the right size
- Minimum of `imagesPerRow` images loaded at all times

invokeai/frontend/web/src/features/gallery/components/ImageGalleryContent.tsx

@@ -1,9 +1,9 @@
-import { Box, Button, ButtonGroup, Flex, Tab, TabList, Tabs, useDisclosure, VStack } from '@invoke-ai/ui-library';
+import { Box, Button, ButtonGroup, Flex, Tab, TabList, Tabs, useDisclosure } from '@invoke-ai/ui-library';
 import { useStore } from '@nanostores/react';
 import { $galleryHeader } from 'app/store/nanostores/galleryHeader';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import { galleryViewChanged } from 'features/gallery/store/gallerySlice';
-import { memo, useCallback, useRef } from 'react';
+import { memo, useCallback } from 'react';
 import { useTranslation } from 'react-i18next';
 import { PiImagesBold } from 'react-icons/pi';
 import { RiServerLine } from 'react-icons/ri';
@@ -12,13 +12,10 @@ import BoardsList from './Boards/BoardsList/BoardsList';
 import GalleryBoardName from './GalleryBoardName';
 import GallerySettingsPopover from './GallerySettingsPopover';
 import GalleryImageGrid from './ImageGrid/GalleryImageGrid';
-import { GalleryImageGridContainer } from './ImageGrid/GalleryImageGridContainer';
 import { GalleryPagination } from './ImageGrid/GalleryPagination';
 
 const ImageGalleryContent = () => {
   const { t } = useTranslation();
-  const resizeObserverRef = useRef<HTMLDivElement>(null);
-  const galleryGridRef = useRef<HTMLDivElement>(null);
   const galleryView = useAppSelector((s) => s.gallery.galleryView);
   const dispatch = useAppDispatch();
   const galleryHeader = useStore($galleryHeader);
@@ -33,10 +30,10 @@ const ImageGalleryContent = () => {
   }, [dispatch]);
 
   return (
-    <Flex layerStyle="first" flexDirection="column" h="full" w="full" borderRadius="base" p={2}>
+    <Flex layerStyle="first" flexDirection="column" h="full" w="full" borderRadius="base" p={2} gap={2}>
       {galleryHeader}
-      <Box w="full">
+      <Box>
-        <Flex ref={resizeObserverRef} alignItems="center" justifyContent="space-between" gap={2}>
+        <Flex alignItems="center" justifyContent="space-between" gap={2}>
           <GalleryBoardName isOpen={isBoardListOpen} onToggle={onToggleBoardList} />
           <GallerySettingsPopover />
         </Flex>
@@ -44,7 +41,6 @@ const ImageGalleryContent = () => {
           <BoardsList isOpen={isBoardListOpen} />
         </Box>
       </Box>
-      <Flex ref={galleryGridRef} direction="column" gap={2} h="full" w="full" minH={0}>
       <Flex alignItems="center" justifyContent="space-between" gap={2}>
         <Tabs index={galleryView === 'images' ? 0 : 1} variant="unstyled" size="sm" w="full">
           <TabList>
@@ -75,8 +71,7 @@ const ImageGalleryContent = () => {
           </TabList>
         </Tabs>
       </Flex>
-        <GalleryImageGridContainer />
+      <GalleryImageGrid />
-      </Flex>
       <GalleryPagination />
     </Flex>
   );

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

@@ -17,8 +17,10 @@ import { memo, useCallback, useMemo, useState } from 'react';
 import { useTranslation } from 'react-i18next';
 import { PiStarBold, PiStarFill, PiTrashSimpleFill } from 'react-icons/pi';
 import { useStarImagesMutation, useUnstarImagesMutation } from 'services/api/endpoints/images';
-import { ImageDTO } from '../../../../services/api/types';
+import type { ImageDTO } from 'services/api/types';
-import { imageItemContainerTestId } from './ImageGridItemContainer';
+
+// This class name is used to calculate the number of images that fit in the gallery
+export const GALLERY_IMAGE_CLASS_NAME = 'gallery-image';
 
 const imageSx: SystemStyleObject = { w: 'full', h: 'full' };
 const boxSx: SystemStyleObject = {
@@ -135,7 +137,7 @@ const GalleryImage = ({ index, imageDTO }: HoverableImageProps) => {
   }
 
   return (
-    <Box w="full" h="full" p={1.5} className={imageItemContainerTestId} data-testid={dataTestId} sx={boxSx}>
+    <Box w="full" h="full" p={1.5} className={GALLERY_IMAGE_CLASS_NAME} data-testid={dataTestId} sx={boxSx}>
       <Flex
         ref={imageContainerRef}
         userSelect="none"

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

@@ -1,25 +1,23 @@
 import { Box, Flex, Grid } from '@invoke-ai/ui-library';
 import { EMPTY_ARRAY } from 'app/store/constants';
-import { useAppSelector } from 'app/store/storeHooks';
+import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import { IAINoContentFallback } from 'common/components/IAIImageFallback';
 import { useGalleryHotkeys } from 'features/gallery/hooks/useGalleryHotkeys';
 import { selectListImagesQueryArgs } from 'features/gallery/store/gallerySelectors';
-import { memo } from 'react';
+import { limitChanged } from 'features/gallery/store/gallerySlice';
+import { debounce } from 'lodash-es';
+import { memo, useEffect, useMemo, useState } from 'react';
 import { useTranslation } from 'react-i18next';
 import { PiImageBold, PiWarningCircleBold } from 'react-icons/pi';
+import { useListImagesQuery } from 'services/api/endpoints/images';
 
-import GalleryImage from './GalleryImage';
+import GalleryImage, { GALLERY_IMAGE_CLASS_NAME } from './GalleryImage';
-import { useListImagesQuery } from '../../../../services/api/endpoints/images';
-import { GalleryPagination } from './GalleryPagination';
 
-export const imageListContainerTestId = 'image-list-container';
+export const GALLERY_GRID_CLASS_NAME = 'gallery-grid';
-export const imageItemContainerTestId = 'image-item-container';
 
 const GalleryImageGrid = () => {
   useGalleryHotkeys();
   const { t } = useTranslation();
-
-  const { galleryImageMinimumWidth, limit } = useAppSelector((s) => s.gallery);
   const queryArgs = useAppSelector(selectListImagesQueryArgs);
   const { imageDTOs, isLoading, isSuccess, isError } = useListImagesQuery(queryArgs, {
     selectFromResult: ({ data, isLoading, isSuccess, isError }) => ({
@@ -38,7 +36,7 @@ const GalleryImageGrid = () => {
     );
   }
 
-  if (isLoading || !limit) {
+  if (isLoading) {
     return (
       <Flex w="full" h="full" alignItems="center" justifyContent="center">
         <IAINoContentFallback label={t('gallery.loading')} icon={PiImageBold} />
@@ -54,22 +52,99 @@ const GalleryImageGrid = () => {
     );
   }
 
+  return <Content />;
+};
+
+export default memo(GalleryImageGrid);
+
+const Content = () => {
+  const dispatch = useAppDispatch();
+  const galleryImageMinimumWidth = useAppSelector((s) => s.gallery.galleryImageMinimumWidth);
+
+  const queryArgs = useAppSelector(selectListImagesQueryArgs);
+  const { imageDTOs } = useListImagesQuery(queryArgs, {
+    selectFromResult: ({ data }) => ({ imageDTOs: data?.items ?? EMPTY_ARRAY }),
+  });
+  // Use a callback ref to get reactivity on the container element because it is conditionally rendered
+  const [container, containerRef] = useState<HTMLDivElement | null>(null);
+
+  const calculateNewLimit = useMemo(() => {
+    // Debounce this to not thrash the API
+    return debounce(() => {
+      if (!container) {
+        // Container not rendered yet
+        return;
+      }
+      // Managing refs for dynamically rendered components is a bit tedious:
+      // - https://react.dev/learn/manipulating-the-dom-with-refs#how-to-manage-a-list-of-refs-using-a-ref-callback
+      // As a easy workaround, we can just grab the first gallery image element directly.
+      const galleryImageEl = document.querySelector(`.${GALLERY_IMAGE_CLASS_NAME}`);
+      if (!galleryImageEl) {
+        // No images in gallery?
+        return;
+      }
+
+      const galleryImageRect = galleryImageEl.getBoundingClientRect();
+      const containerRect = container.getBoundingClientRect();
+
+      if (!galleryImageRect.width || !galleryImageRect.height || !containerRect.width || !containerRect.height) {
+        // Gallery is too small to fit images or not rendered yet
+        return;
+      }
+
+      // Floating-point precision requires we round to get the correct number of images per row
+      const imagesPerRow = Math.round(containerRect.width / galleryImageRect.width);
+      // However, when calculating the number of images per column, we want to floor the value to not overflow the container
+      const imagesPerColumn = Math.floor(containerRect.height / galleryImageRect.height);
+      // Always load at least 1 row of images
+      const limit = Math.max(imagesPerRow, imagesPerRow * imagesPerColumn);
+      dispatch(limitChanged(limit));
+    }, 300);
+  }, [container, dispatch]);
+
+  useEffect(() => {
+    // We want to recalculate the limit when image size changes
+    calculateNewLimit();
+  }, [calculateNewLimit, galleryImageMinimumWidth]);
+
+  useEffect(() => {
+    if (!container) {
+      return;
+    }
+
+    const resizeObserver = new ResizeObserver(calculateNewLimit);
+    resizeObserver.observe(container);
+
+    // First render
+    calculateNewLimit();
+
+    return () => {
+      resizeObserver.disconnect();
+    };
+  }, [calculateNewLimit, container, dispatch]);
+
   return (
-    <>
+    <Box position="relative" w="full" h="full">
-      <Box data-overlayscrollbars="" h="100%" id="gallery-grid">
+      <Box
+        ref={containerRef}
+        position="absolute"
+        top={0}
+        right={0}
+        bottom={0}
+        left={0}
+        w="full"
+        h="full"
+        overflow="hidden"
+      >
         <Grid
-          className="list-container"
+          className={GALLERY_GRID_CLASS_NAME}
           gridTemplateColumns={`repeat(auto-fill, minmax(${galleryImageMinimumWidth}px, 1fr))`}
-          data-testid={imageListContainerTestId}
         >
           {imageDTOs.map((imageDTO, index) => (
             <GalleryImage key={imageDTO.image_name} imageDTO={imageDTO} index={index} />
           ))}
         </Grid>
       </Box>
-      {/* <GalleryPagination /> */}
+    </Box>
-    </>
   );
 };
-
-export default memo(GalleryImageGrid);

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

@@ -1,58 +0,0 @@
-import { useCallback, useEffect, useRef, useState } from 'react';
-import { Box, Flex } from '@invoke-ai/ui-library';
-import GalleryImageGrid from './GalleryImageGrid';
-import { useAppSelector, useAppDispatch } from '../../../../app/store/storeHooks';
-import { limitChanged } from '../../store/gallerySlice';
-
-export const GalleryImageGridContainer = () => {
-  const { galleryImageMinimumWidth, limit } = useAppSelector((s) => s.gallery);
-  const dispatch = useAppDispatch();
-  const containerRef = useRef<HTMLDivElement>(null);
-
-  const calculateItemsPerPage = useCallback(() => {
-    const containerWidth = containerRef.current?.clientWidth;
-    const containerHeight = containerRef.current?.clientHeight;
-    console.log({ containerWidth, containerHeight, galleryImageMinimumWidth });
-    if (containerHeight && containerWidth) {
-      const numberHorizontal = Math.floor(containerWidth / galleryImageMinimumWidth);
-      const imageWidth = containerWidth / numberHorizontal;
-      const numberAllowedVertical = Math.floor(containerHeight / imageWidth);
-      console.log({ numberAllowedVertical, numberHorizontal });
-      dispatch(limitChanged(numberAllowedVertical * numberHorizontal));
-    }
-  }, [containerRef, galleryImageMinimumWidth]);
-
-  useEffect(() => {
-    dispatch(limitChanged(undefined));
-    calculateItemsPerPage();
-  }, [galleryImageMinimumWidth]);
-
-  useEffect(() => {
-    if (!containerRef.current) {
-      return;
-    }
-
-    const resizeObserver = new ResizeObserver(() => {
-      console.log('resize');
-      if (!containerRef.current) {
-        return;
-      }
-      dispatch(limitChanged(undefined));
-      calculateItemsPerPage();
-    });
-
-    resizeObserver.observe(containerRef.current);
-    dispatch(limitChanged(undefined));
-    calculateItemsPerPage();
-
-    return () => {
-      resizeObserver.disconnect();
-    };
-  }, []);
-
-  return (
-    <Flex flexDir="column" w="full" h="full" overflow="hidden" ref={containerRef}>
-      {limit && <GalleryImageGrid />}
-    </Flex>
-  );
-};

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

@@ -1,15 +0,0 @@
-import type { FlexProps } from '@invoke-ai/ui-library';
-import { Box, forwardRef } from '@invoke-ai/ui-library';
-import type { PropsWithChildren } from 'react';
-import { memo } from 'react';
-
-export const imageItemContainerTestId = 'image-item-container';
-
-type ItemContainerProps = PropsWithChildren & FlexProps;
-const ItemContainer = forwardRef((props: ItemContainerProps, ref) => (
-  <Box className="item-container" ref={ref} p={1.5} data-testid={imageItemContainerTestId}>
-    {props.children}
-  </Box>
-));
-
-export default memo(ItemContainer);

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

@@ -1,26 +0,0 @@
-import type { FlexProps } from '@invoke-ai/ui-library';
-import { forwardRef, Grid } from '@invoke-ai/ui-library';
-import { useAppSelector } from 'app/store/storeHooks';
-import type { PropsWithChildren } from 'react';
-import { memo } from 'react';
-
-export const imageListContainerTestId = 'image-list-container';
-
-type ListContainerProps = PropsWithChildren & FlexProps;
-const ListContainer = forwardRef((props: ListContainerProps, ref) => {
-  const galleryImageMinimumWidth = useAppSelector((s) => s.gallery.galleryImageMinimumWidth);
-
-  return (
-    <Grid
-      {...props}
-      className="list-container"
-      ref={ref}
-      gridTemplateColumns={`repeat(auto-fill, minmax(${galleryImageMinimumWidth}px, 1fr))`}
-      data-testid={imageListContainerTestId}
-    >
-      {props.children}
-    </Grid>
-  );
-});
-
-export default memo(ListContainer);

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

@@ -1,8 +1,8 @@
 import { useAltModifier } from '@invoke-ai/ui-library';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
+import { GALLERY_IMAGE_CLASS_NAME } from 'features/gallery/components/ImageGrid/GalleryImage';
+import { GALLERY_GRID_CLASS_NAME } from 'features/gallery/components/ImageGrid/GalleryImageGrid';
 import { getGalleryImageDataTestId } from 'features/gallery/components/ImageGrid/getGalleryImageDataTestId';
-import { imageItemContainerTestId } from 'features/gallery/components/ImageGrid/ImageGridItemContainer';
-import { imageListContainerTestId } from 'features/gallery/components/ImageGrid/ImageGridListContainer';
 import { virtuosoGridRefs } from 'features/gallery/components/ImageGrid/types';
 import { useGalleryImages } from 'features/gallery/hooks/useGalleryImages';
 import { imageSelected, imageToCompareChanged } from 'features/gallery/store/gallerySlice';
@@ -11,7 +11,6 @@ import { getScrollToIndexAlign } from 'features/gallery/util/getScrollToIndexAli
 import { clamp } from 'lodash-es';
 import { useCallback, useMemo } from 'react';
 import type { ImageDTO } from 'services/api/types';
-import { imagesSelectors } from 'services/api/util';
 
 /**
  * This hook is used to navigate the gallery using the arrow keys.
@@ -29,10 +28,9 @@ import { imagesSelectors } from 'services/api/util';
  */
 const getImagesPerRow = (): number => {
   const widthOfGalleryImage =
-    document.querySelector(`.${imageItemContainerTestId}`)?.getBoundingClientRect().width ?? 1;
+    document.querySelector(`.${GALLERY_IMAGE_CLASS_NAME}`)?.getBoundingClientRect().width ?? 1;
 
-  const widthOfGalleryGrid =
+  const widthOfGalleryGrid = document.querySelector(`.${GALLERY_GRID_CLASS_NAME}`)?.getBoundingClientRect().width ?? 0;
-    document.querySelector(`[data-testid="${imageListContainerTestId}"]`)?.getBoundingClientRect().width ?? 0;
 
   const imagesPerRow = Math.round(widthOfGalleryGrid / widthOfGalleryImage);
 

invokeai/frontend/web/src/features/gallery/store/gallerySlice.ts

@@ -19,7 +19,7 @@ const initialGalleryState: GalleryState = {
   selectedBoardId: 'none',
   galleryView: 'images',
   boardSearchText: '',
-  limit: undefined,
+  limit: 20,
   offset: 0,
   isImageViewerOpen: true,
   imageToCompare: null,
@@ -153,7 +153,7 @@ export const {
   comparisonFitChanged,
   comparisonModeCycled,
   offsetChanged,
-  limitChanged
+  limitChanged,
 } = gallerySlice.actions;
 
 const isAnyBoardDeleted = isAnyOf(

invokeai/frontend/web/src/features/gallery/store/types.ts

@@ -19,7 +19,7 @@ export type GalleryState = {
   galleryView: GalleryView;
   boardSearchText: string;
   offset: number;
-  limit: number | undefined;
+  limit: number;
   alwaysShowImageSizeBadge: boolean;
   imageToCompare: ImageDTO | null;
   comparisonMode: ComparisonMode;