From 475b6bef53501a91b9f96d70a944e9a7eb2e41f1 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Mon, 1 May 2023 16:04:25 +1000
Subject: [PATCH] feat(ui): use windowing for gallery

vastly improves the gallery performance when many images are loaded.

- `react-virtuoso` to do the virtualized list
- `overlayscrollbars` for a scrollbar
---
 invokeai/frontend/web/package.json            |   3 +
 .../app/components/ThemeLocaleProvider.tsx    |   2 +
 .../gallery/components/HoverableImage.tsx     |  61 ++---
 .../components/ImageGalleryContent.tsx        | 239 ++++++++++++------
 .../web/src/theme/css/overlayscrollbars.css   |  48 ++++
 invokeai/frontend/web/yarn.lock               |  15 ++
 6 files changed, 261 insertions(+), 107 deletions(-)
 create mode 100644 invokeai/frontend/web/src/theme/css/overlayscrollbars.css

diff --git a/invokeai/frontend/web/package.json b/invokeai/frontend/web/package.json
index d9e4f0af5c..feb2ea1200 100644
--- a/invokeai/frontend/web/package.json
+++ b/invokeai/frontend/web/package.json
@@ -76,6 +76,8 @@
     "i18next-http-backend": "^2.2.0",
     "konva": "^9.0.1",
     "lodash-es": "^4.17.21",
+    "overlayscrollbars": "^2.1.1",
+    "overlayscrollbars-react": "^0.5.0",
     "patch-package": "^7.0.0",
     "re-resizable": "^6.9.9",
     "react": "^18.2.0",
@@ -91,6 +93,7 @@
     "react-rnd": "^10.4.1",
     "react-transition-group": "^4.4.5",
     "react-use": "^17.4.0",
+    "react-virtuoso": "^4.3.3",
     "react-zoom-pan-pinch": "^3.0.7",
     "reactflow": "^11.7.0",
     "redux-deep-persist": "^1.0.7",
diff --git a/invokeai/frontend/web/src/app/components/ThemeLocaleProvider.tsx b/invokeai/frontend/web/src/app/components/ThemeLocaleProvider.tsx
index e574456e37..f0e2e75240 100644
--- a/invokeai/frontend/web/src/app/components/ThemeLocaleProvider.tsx
+++ b/invokeai/frontend/web/src/app/components/ThemeLocaleProvider.tsx
@@ -18,6 +18,8 @@ import '@fontsource/inter/600.css';
 import '@fontsource/inter/700.css';
 import '@fontsource/inter/800.css';
 import '@fontsource/inter/900.css';
+import 'overlayscrollbars/overlayscrollbars.css';
+import 'theme/css/overlayscrollbars.css';
 
 type ThemeLocaleProviderProps = {
   children: ReactNode;
diff --git a/invokeai/frontend/web/src/features/gallery/components/HoverableImage.tsx b/invokeai/frontend/web/src/features/gallery/components/HoverableImage.tsx
index d0ff9aee40..032784fbf9 100644
--- a/invokeai/frontend/web/src/features/gallery/components/HoverableImage.tsx
+++ b/invokeai/frontend/web/src/features/gallery/components/HoverableImage.tsx
@@ -5,6 +5,7 @@ import {
   Image,
   MenuItem,
   MenuList,
+  Skeleton,
   useDisclosure,
   useTheme,
   useToast,
@@ -12,7 +13,7 @@ import {
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import { imageSelected } from 'features/gallery/store/gallerySlice';
 import { DragEvent, memo, useCallback, useState } from 'react';
-import { FaCheck, FaExpand, FaShare, FaTrash } from 'react-icons/fa';
+import { FaCheck, FaExpand, FaImage, FaShare, FaTrash } from 'react-icons/fa';
 import DeleteImageModal from './DeleteImageModal';
 import { ContextMenu } from 'chakra-ui-contextmenu';
 import * as InvokeAI from 'app/types/invokeai';
@@ -268,58 +269,48 @@ const HoverableImage = memo((props: HoverableImageProps) => {
             userSelect="none"
             draggable={true}
             onDragStart={handleDragStart}
+            onClick={handleSelectImage}
             ref={ref}
             sx={{
-              padding: 2,
               display: 'flex',
               justifyContent: 'center',
+              alignItems: 'center',
+              w: 'full',
+              h: 'full',
               transition: 'transform 0.2s ease-out',
-              _hover: {
-                cursor: 'pointer',
-
-                zIndex: 2,
-              },
-              _before: {
-                content: '""',
-                display: 'block',
-                paddingBottom: '100%',
-              },
+              aspectRatio: '1/1',
             }}
           >
             <Image
+              loading="lazy"
               objectFit={
                 shouldUseSingleGalleryColumn ? 'contain' : galleryImageObjectFit
               }
               rounded="md"
               src={getUrl(thumbnail || url)}
-              loading="lazy"
+              fallback={<FaImage />}
               sx={{
-                position: 'absolute',
                 width: '100%',
                 height: '100%',
                 maxWidth: '100%',
                 maxHeight: '100%',
-                top: '50%',
-                transform: 'translate(-50%,-50%)',
-                ...(direction === 'rtl'
-                  ? { insetInlineEnd: '50%' }
-                  : { insetInlineStart: '50%' }),
               }}
             />
-            <Flex
-              onClick={handleSelectImage}
-              sx={{
-                position: 'absolute',
-                top: '0',
-                insetInlineStart: '0',
-                width: '100%',
-                height: '100%',
-                alignItems: 'center',
-                justifyContent: 'center',
-              }}
-            >
-              {isSelected && (
+            {isSelected && (
+              <Flex
+                sx={{
+                  position: 'absolute',
+                  top: '0',
+                  insetInlineStart: '0',
+                  width: '100%',
+                  height: '100%',
+                  alignItems: 'center',
+                  justifyContent: 'center',
+                  pointerEvents: 'none',
+                }}
+              >
                 <Icon
+                  filter={'drop-shadow(0px 0px 1rem black)'}
                   as={FaCheck}
                   sx={{
                     width: '50%',
@@ -327,9 +318,9 @@ const HoverableImage = memo((props: HoverableImageProps) => {
                     fill: 'ok.500',
                   }}
                 />
-              )}
-            </Flex>
-            {isHovered && galleryImageMinimumWidth >= 64 && (
+              </Flex>
+            )}
+            {isHovered && galleryImageMinimumWidth >= 100 && (
               <Box
                 sx={{
                   position: 'absolute',
diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageGalleryContent.tsx b/invokeai/frontend/web/src/features/gallery/components/ImageGalleryContent.tsx
index 6524452e90..1e1e73812f 100644
--- a/invokeai/frontend/web/src/features/gallery/components/ImageGalleryContent.tsx
+++ b/invokeai/frontend/web/src/features/gallery/components/ImageGalleryContent.tsx
@@ -1,5 +1,13 @@
-import { ButtonGroup, Flex, Grid, Icon, Image, Text } from '@chakra-ui/react';
-// import { requestImages } from 'app/socketio/actions';
+import {
+  Box,
+  ButtonGroup,
+  Flex,
+  FlexProps,
+  Grid,
+  Icon,
+  Text,
+  forwardRef,
+} from '@chakra-ui/react';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import IAIButton from 'common/components/IAIButton';
 import IAICheckbox from 'common/components/IAICheckbox';
@@ -15,28 +23,33 @@ import {
   setShouldUseSingleGalleryColumn,
 } from 'features/gallery/store/gallerySlice';
 import { togglePinGalleryPanel } from 'features/ui/store/uiSlice';
+import { useOverlayScrollbars } from 'overlayscrollbars-react';
 
-import { ChangeEvent, useEffect, useRef, useState } from 'react';
+import {
+  ChangeEvent,
+  PropsWithChildren,
+  memo,
+  useCallback,
+  useEffect,
+  useRef,
+  useState,
+} from 'react';
 import { useTranslation } from 'react-i18next';
 import { BsPinAngle, BsPinAngleFill } from 'react-icons/bs';
 import { FaImage, FaUser, FaWrench } from 'react-icons/fa';
 import { MdPhotoLibrary } from 'react-icons/md';
 import HoverableImage from './HoverableImage';
 
-import Scrollable from 'features/ui/components/common/Scrollable';
 import { requestCanvasRescale } from 'features/canvas/store/thunks/requestCanvasScale';
-import {
-  resultsAdapter,
-  selectResultsAll,
-  selectResultsTotal,
-} from '../store/resultsSlice';
+import { resultsAdapter } from '../store/resultsSlice';
 import {
   receivedResultImagesPage,
   receivedUploadImagesPage,
 } from 'services/thunks/gallery';
-import { selectUploadsAll, uploadsAdapter } from '../store/uploadsSlice';
+import { uploadsAdapter } from '../store/uploadsSlice';
 import { createSelector } from '@reduxjs/toolkit';
 import { RootState } from 'app/store/store';
+import { Virtuoso, VirtuosoGrid } from 'react-virtuoso';
 
 const GALLERY_SHOW_BUTTONS_MIN_WIDTH = 290;
 
@@ -68,16 +81,28 @@ const ImageGalleryContent = () => {
   const { t } = useTranslation();
   const resizeObserverRef = useRef<HTMLDivElement>(null);
   const [shouldShouldIconButtons, setShouldShouldIconButtons] = useState(true);
+  const rootRef = useRef(null);
+  const [scroller, setScroller] = useState<HTMLElement | null>(null);
+  const [initialize, osInstance] = useOverlayScrollbars({
+    defer: true,
+    options: {
+      scrollbars: {
+        visibility: 'auto',
+        autoHide: 'leave',
+        autoHideDelay: 1300,
+        theme: 'os-theme-dark',
+      },
+      overflow: { x: 'hidden' },
+    },
+  });
 
   const {
     // images,
     currentCategory,
     shouldPinGallery,
     galleryImageMinimumWidth,
-    galleryGridTemplateColumns,
     galleryImageObjectFit,
     shouldAutoSwitchToNewImages,
-    // areMoreImagesAvailable,
     shouldUseSingleGalleryColumn,
     selectedImage,
   } = useAppSelector(imageGallerySelector);
@@ -85,9 +110,6 @@ const ImageGalleryContent = () => {
   const { images, areMoreImagesAvailable, isLoading } =
     useAppSelector(gallerySelector);
 
-  // const handleClickLoadMore = () => {
-  //   dispatch(requestImages(currentCategory));
-  // };
   const handleClickLoadMore = () => {
     if (currentCategory === 'results') {
       dispatch(receivedResultImagesPage());
@@ -129,6 +151,25 @@ const ImageGalleryContent = () => {
     return () => resizeObserver.disconnect(); // clean up
   }, []);
 
+  useEffect(() => {
+    const { current: root } = rootRef;
+    if (scroller && root) {
+      initialize({
+        target: root,
+        elements: {
+          viewport: scroller,
+        },
+      });
+    }
+    return () => osInstance()?.destroy();
+  }, [scroller, initialize, osInstance]);
+
+  const setScrollerRef = useCallback((ref: HTMLElement | Window | null) => {
+    if (ref instanceof HTMLElement) {
+      setScroller(ref);
+    }
+  }, []);
+
   return (
     <Flex flexDirection="column" w="full" h="full" gap={4}>
       <Flex
@@ -241,65 +282,119 @@ const ImageGalleryContent = () => {
           />
         </Flex>
       </Flex>
-      <Scrollable>
-        <Flex direction="column" gap={2} h="full">
-          {images.length || areMoreImagesAvailable ? (
-            <>
-              <Grid
-                gap={2}
-                style={{ gridTemplateColumns: galleryGridTemplateColumns }}
-              >
-                {images.map((image) => {
-                  const { name } = image;
-                  const isSelected = selectedImage?.name === name;
-                  return (
-                    <HoverableImage
-                      key={`${name}-${image.thumbnail}`}
-                      image={image}
-                      isSelected={isSelected}
-                    />
-                  );
-                })}
-              </Grid>
-              <IAIButton
-                onClick={handleClickLoadMore}
-                isDisabled={!areMoreImagesAvailable}
-                isLoading={isLoading}
-                flexShrink={0}
-              >
-                {areMoreImagesAvailable
-                  ? t('gallery.loadMore')
-                  : t('gallery.allImagesLoaded')}
-              </IAIButton>
-            </>
-          ) : (
-            <Flex
-              sx={{
-                flexDirection: 'column',
-                alignItems: 'center',
-                justifyContent: 'center',
-                gap: 2,
-                padding: 8,
-                h: '100%',
-                w: '100%',
-                color: 'base.500',
-              }}
+      <Flex direction="column" gap={2} h="full">
+        {images.length || areMoreImagesAvailable ? (
+          <>
+            <Box ref={rootRef} data-overlayscrollbars="" h="100%">
+              {shouldUseSingleGalleryColumn ? (
+                <Virtuoso
+                  style={{ height: '100%' }}
+                  data={images}
+                  scrollerRef={(ref) => setScrollerRef(ref)}
+                  itemContent={(index, image) => {
+                    const { name } = image;
+                    const isSelected = selectedImage?.name === name;
+
+                    return (
+                      <Flex sx={{ pb: 2 }}>
+                        <HoverableImage
+                          key={`${name}-${image.thumbnail}`}
+                          image={image}
+                          isSelected={isSelected}
+                        />
+                      </Flex>
+                    );
+                  }}
+                />
+              ) : (
+                <VirtuosoGrid
+                  style={{ height: '100%' }}
+                  data={images}
+                  components={{
+                    Item: ItemContainer,
+                    List: ListContainer,
+                  }}
+                  scrollerRef={setScroller}
+                  itemContent={(index, image) => {
+                    const { name } = image;
+                    const isSelected = selectedImage?.name === name;
+
+                    return (
+                      <HoverableImage
+                        key={`${name}-${image.thumbnail}`}
+                        image={image}
+                        isSelected={isSelected}
+                      />
+                    );
+                  }}
+                />
+              )}
+            </Box>
+            <IAIButton
+              onClick={handleClickLoadMore}
+              isDisabled={!areMoreImagesAvailable}
+              isLoading={isLoading}
+              flexShrink={0}
             >
-              <Icon
-                as={MdPhotoLibrary}
-                sx={{
-                  w: 16,
-                  h: 16,
-                }}
-              />
-              <Text textAlign="center">{t('gallery.noImagesInGallery')}</Text>
-            </Flex>
-          )}
-        </Flex>
-      </Scrollable>
+              {areMoreImagesAvailable
+                ? t('gallery.loadMore')
+                : t('gallery.allImagesLoaded')}
+            </IAIButton>
+          </>
+        ) : (
+          <Flex
+            sx={{
+              flexDirection: 'column',
+              alignItems: 'center',
+              justifyContent: 'center',
+              gap: 2,
+              padding: 8,
+              h: '100%',
+              w: '100%',
+              color: 'base.500',
+            }}
+          >
+            <Icon
+              as={MdPhotoLibrary}
+              sx={{
+                w: 16,
+                h: 16,
+              }}
+            />
+            <Text textAlign="center">{t('gallery.noImagesInGallery')}</Text>
+          </Flex>
+        )}
+      </Flex>
     </Flex>
   );
 };
 
-ImageGalleryContent.displayName = 'ImageGalleryContent';
-export default ImageGalleryContent;
+type ItemContainerProps = PropsWithChildren & FlexProps;
+const ItemContainer = forwardRef((props: ItemContainerProps, ref) => (
+  <Box className="item-container" ref={ref}>
+    {props.children}
+  </Box>
+));
+
+type ListContainerProps = PropsWithChildren & FlexProps;
+const ListContainer = forwardRef((props: ListContainerProps, ref) => {
+  const galleryImageMinimumWidth = useAppSelector(
+    (state: RootState) => state.gallery.galleryImageMinimumWidth
+  );
+
+  return (
+    <Grid
+      {...props}
+      className="list-container"
+      ref={ref}
+      sx={{
+        gap: 2,
+        gridTemplateColumns: `repeat(auto-fit, minmax(${galleryImageMinimumWidth}px, 1fr));`,
+      }}
+    >
+      {props.children}
+    </Grid>
+  );
+});
+
+export default memo(ImageGalleryContent);
diff --git a/invokeai/frontend/web/src/theme/css/overlayscrollbars.css b/invokeai/frontend/web/src/theme/css/overlayscrollbars.css
new file mode 100644
index 0000000000..b5acaca75d
--- /dev/null
+++ b/invokeai/frontend/web/src/theme/css/overlayscrollbars.css
@@ -0,0 +1,48 @@
+.os-scrollbar {
+  /*      The size of the scrollbar */
+  /* --os-size: 0; */
+  /*      The axis-perpedicular padding of the scrollbar (horizontal: padding-y, vertical: padding-x) */
+  /* --os-padding-perpendicular: 0; */
+  /*      The axis padding of the scrollbar (horizontal: padding-x, vertical: padding-y) */
+  /* --os-padding-axis: 0; */
+  /*      The border radius of the scrollbar track */
+  /* --os-track-border-radius: 0; */
+  /*      The background of the scrollbar track */
+  --os-track-bg: rgba(0, 0, 0, 0.3);
+  /*      The :hover background of the scrollbar track */
+  --os-track-bg-hover: rgba(0, 0, 0, 0.3);
+  /*      The :active background of the scrollbar track */
+  --os-track-bg-active: rgba(0, 0, 0, 0.3);
+  /*      The border of the scrollbar track */
+  /* --os-track-border: none; */
+  /*      The :hover background of the scrollbar track */
+  /* --os-track-border-hover: none; */
+  /*      The :active background of the scrollbar track */
+  /* --os-track-border-active: none; */
+  /*      The border radius of the scrollbar handle */
+  /* --os-handle-border-radius: 0; */
+  /*      The background of the scrollbar handle */
+  --os-handle-bg: var(--invokeai-colors-accent-500);
+  /*      The :hover background of the scrollbar handle */
+  --os-handle-bg-hover: var(--invokeai-colors-accent-450);
+  /*      The :active background of the scrollbar handle */
+  --os-handle-bg-active: var(--invokeai-colors-accent-400);
+  /*      The border of the scrollbar handle */
+  /* --os-handle-border: none; */
+  /*      The :hover border of the scrollbar handle */
+  /* --os-handle-border-hover: none; */
+  /*      The :active border of the scrollbar handle */
+  /* --os-handle-border-active: none; */
+  /*      The min size of the scrollbar handle */
+  --os-handle-min-size: 50px;
+  /*      The max size of the scrollbar handle */
+  /* --os-handle-max-size: none; */
+  /*      The axis-perpedicular size of the scrollbar handle (horizontal: height, vertical: width) */
+  /* --os-handle-perpendicular-size: 100%; */
+  /*      The :hover axis-perpedicular size of the scrollbar handle (horizontal: height, vertical: width) */
+  /* --os-handle-perpendicular-size-hover: 100%; */
+  /*      The :active axis-perpedicular size of the scrollbar handle (horizontal: height, vertical: width) */
+  /* --os-handle-perpendicular-size-active: 100%; */
+  /*      Increases the interactive area of the scrollbar handle. */
+  /* --os-handle-interactive-area-offset: 0; */
+}
\ No newline at end of file
diff --git a/invokeai/frontend/web/yarn.lock b/invokeai/frontend/web/yarn.lock
index b60071ab65..872be96040 100644
--- a/invokeai/frontend/web/yarn.lock
+++ b/invokeai/frontend/web/yarn.lock
@@ -5100,6 +5100,16 @@ os-tmpdir@~1.0.2:
   resolved "https://registry.yarnpkg.com/os-tmpdir/-/os-tmpdir-1.0.2.tgz#bbe67406c79aa85c5cfec766fe5734555dfa1274"
   integrity sha512-D2FR03Vir7FIu45XBY20mTb+/ZSWB00sjU9jdQXt83gDrI4Ztz5Fs7/yy74g2N5SVQY4xY1qDr4rNddwYRVX0g==
 
+overlayscrollbars-react@^0.5.0:
+  version "0.5.0"
+  resolved "https://registry.yarnpkg.com/overlayscrollbars-react/-/overlayscrollbars-react-0.5.0.tgz#0272bdc6304c7228a58d30e5b678e97fd5c5d8dd"
+  integrity sha512-uCNTnkfWW74veoiEv3kSwoLelKt4e8gTNv65D771X3il0x5g5Yo0fUbro7SpQzR9yNgi23cvB2mQHTTdQH96pA==
+
+overlayscrollbars@^2.1.1:
+  version "2.1.1"
+  resolved "https://registry.yarnpkg.com/overlayscrollbars/-/overlayscrollbars-2.1.1.tgz#a7414fe9c96cf140dbe4975bbe9312861750388d"
+  integrity sha512-xvs2g8Tcq9+CZDpLEUchN3YUzjJhnTWw9kwqT/qcC53FIkOyP9mqnRMot5sW16tcsPT1KaMyzF0AMXw/7E4a8g==
+
 p-cancelable@^1.0.0:
   version "1.1.0"
   resolved "https://registry.yarnpkg.com/p-cancelable/-/p-cancelable-1.1.0.tgz#d078d15a3af409220c886f1d9a0ca2e441ab26cc"
@@ -5612,6 +5622,11 @@ react-use@^17.4.0:
     ts-easing "^0.2.0"
     tslib "^2.1.0"
 
+react-virtuoso@^4.3.3:
+  version "4.3.3"
+  resolved "https://registry.yarnpkg.com/react-virtuoso/-/react-virtuoso-4.3.3.tgz#04b8d105b5d97365223fb13c9a594f7501f7e975"
+  integrity sha512-x0DeGmVAVOVaTXRMG7jzrHBwK7+dkt7n0G3tNmZXphQUBgkVBYuZoaJltQeZGFN42++3XvrgwStKCtmzgMJ0lA==
+
 react-zoom-pan-pinch@^3.0.7:
   version "3.0.7"
   resolved "https://registry.yarnpkg.com/react-zoom-pan-pinch/-/react-zoom-pan-pinch-3.0.7.tgz#def52f6886bc11e1b160dedf4250aae95470b94d"