feat(ui): virtualized list for staging area

Make the staging area a virtualized list so it doesn't choke when there
are a large number (i.e. more than a few hundred) of queue items.

invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/QueueItemPreviewMini.tsx

@@ -27,6 +27,7 @@ const sx = {
   alignItems: 'center',
   justifyContent: 'center',
   flexShrink: 0,
+  h: 'full',
   aspectRatio: '1/1',
   borderWidth: 2,
   borderRadius: 'base',
@@ -39,11 +40,11 @@ const sx = {
 
 type Props = {
   item: S['SessionQueueItem'];
-  number: number;
+  index: number;
   isSelected: boolean;
 };
 
-export const QueueItemPreviewMini = memo(({ item, isSelected, number }: Props) => {
+export const QueueItemPreviewMini = memo(({ item, isSelected, index }: Props) => {
   const dispatch = useAppDispatch();
   const ctx = useCanvasSessionContext();
   const { imageLoaded } = useProgressData(ctx.$progressData, item.item_id);
@@ -69,7 +70,7 @@ export const QueueItemPreviewMini = memo(({ item, isSelected, number }: Props) =
 
   return (
     <Flex
-      id={getQueueItemElementId(item.item_id)}
+      id={getQueueItemElementId(index)}
       sx={sx}
       data-selected={isSelected}
       onClick={onClick}
@@ -78,7 +79,7 @@ export const QueueItemPreviewMini = memo(({ item, isSelected, number }: Props) =
       <QueueItemStatusLabel item={item} position="absolute" margin="auto" />
       {imageDTO && <DndImage imageDTO={imageDTO} onLoad={onLoad} asThumbnail />}
       {!imageLoaded && <QueueItemProgressImage itemId={item.item_id} position="absolute" />}
-      <QueueItemNumber number={number} position="absolute" top={0} left={1} />
+      <QueueItemNumber number={index + 1} position="absolute" top={0} left={1} />
       <QueueItemCircularProgress itemId={item.item_id} status={item.status} position="absolute" top={1} right={2} />
     </Flex>
   );

invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/StagingAreaItemsList.tsx

@@ -1,17 +1,148 @@
-import { Flex } from '@invoke-ai/ui-library';
+import { Box, Flex, forwardRef } from '@invoke-ai/ui-library';
 import { useStore } from '@nanostores/react';
-import ScrollableContent from 'common/components/OverlayScrollbars/ScrollableContent';
+import { logger } from 'app/logging/logger';
 import { useCanvasSessionContext } from 'features/controlLayers/components/SimpleSession/context';
 import { QueueItemPreviewMini } from 'features/controlLayers/components/SimpleSession/QueueItemPreviewMini';
 import { useCanvasManagerSafe } from 'features/controlLayers/contexts/CanvasManagerProviderGate';
-import { memo, useEffect } from 'react';
+import { useOverlayScrollbars } from 'overlayscrollbars-react';
+import type { CSSProperties, RefObject } from 'react';
+import { memo, useCallback, useEffect, useMemo, useRef, useState } from 'react';
+import type { Components, ItemContent, ListRange, VirtuosoHandle, VirtuosoProps } from 'react-virtuoso';
+import { Virtuoso } from 'react-virtuoso';
+import type { S } from 'services/api/types';
+
+import { getQueueItemElementId } from './shared';
+
+const log = logger('system');
+
+const virtuosoStyles = {
+  width: '100%',
+  height: '72px',
+} satisfies CSSProperties;
+
+type VirtuosoContext = { selectedItemId: number | null };
+
+/**
+ * Scroll the item at the given index into view if it is not currently visible.
+ */
+const scrollIntoView = (
+  targetIndex: number,
+  rootEl: HTMLDivElement,
+  virtuosoHandle: VirtuosoHandle,
+  range: ListRange
+) => {
+  if (range.endIndex === 0) {
+    // No range is rendered; no need to scroll to anything.
+    return;
+  }
+
+  const targetItem = rootEl.querySelector(`#${getQueueItemElementId(targetIndex)}`);
+
+  if (!targetItem) {
+    if (targetIndex > range.endIndex) {
+      virtuosoHandle.scrollToIndex({
+        index: targetIndex,
+        behavior: 'auto',
+        align: 'end',
+      });
+    } else if (targetIndex < range.startIndex) {
+      virtuosoHandle.scrollToIndex({
+        index: targetIndex,
+        behavior: 'auto',
+        align: 'start',
+      });
+    } else {
+      log.debug(
+        `Unable to find queue item at index ${targetIndex} but it is in the rendered range ${range.startIndex}-${range.endIndex}`
+      );
+    }
+    return;
+  }
+
+  // We found the image in the DOM, but it might be in the overscan range - rendered but not in the visible viewport.
+  // Check if it is in the viewport and scroll if necessary.
+
+  const itemRect = targetItem.getBoundingClientRect();
+  const rootRect = rootEl.getBoundingClientRect();
+
+  if (itemRect.left < rootRect.left) {
+    virtuosoHandle.scrollToIndex({
+      index: targetIndex,
+      behavior: 'auto',
+      align: 'start',
+    });
+  } else if (itemRect.right > rootRect.right) {
+    virtuosoHandle.scrollToIndex({
+      index: targetIndex,
+      behavior: 'auto',
+      align: 'end',
+    });
+  } else {
+    // Image is already in view
+  }
+
+  return;
+};
+
+const useScrollableStagingArea = (rootRef: RefObject<HTMLDivElement>) => {
+  const [scroller, scrollerRef] = useState<HTMLElement | null>(null);
+  const [initialize, osInstance] = useOverlayScrollbars({
+    defer: true,
+    events: {
+      initialized(osInstance) {
+        // force overflow styles
+        const { viewport } = osInstance.elements();
+        viewport.style.overflowX = `var(--os-viewport-overflow-x)`;
+        viewport.style.overflowY = `var(--os-viewport-overflow-y)`;
+      },
+    },
+    options: {
+      scrollbars: {
+        visibility: 'auto',
+        autoHide: 'scroll',
+        autoHideDelay: 1300,
+        theme: 'os-theme-dark',
+      },
+      overflow: {
+        y: 'hidden',
+        x: 'scroll',
+      },
+    },
+  });
+
+  useEffect(() => {
+    const { current: root } = rootRef;
+
+    if (scroller && root) {
+      initialize({
+        target: root,
+        elements: {
+          viewport: scroller,
+        },
+      });
+    }
+
+    return () => {
+      osInstance()?.destroy();
+    };
+  }, [scroller, initialize, osInstance, rootRef]);
+
+  return scrollerRef;
+};
 
 export const StagingAreaItemsList = memo(() => {
   const canvasManager = useCanvasManagerSafe();
   const ctx = useCanvasSessionContext();
+  const virtuosoRef = useRef<VirtuosoHandle>(null);
+  const rangeRef = useRef<ListRange>({ startIndex: 0, endIndex: 0 });
+  const rootRef = useRef<HTMLDivElement>(null);
+
   const items = useStore(ctx.$items);
   const selectedItemId = useStore(ctx.$selectedItemId);
 
+  const context = useMemo(() => ({ selectedItemId }), [selectedItemId]);
+  const scrollerRef = useScrollableStagingArea(rootRef);
+
   useEffect(() => {
     if (!canvasManager) {
       return;
@@ -20,21 +151,64 @@ export const StagingAreaItemsList = memo(() => {
     return canvasManager.stagingArea.connectToSession(ctx.$selectedItemId, ctx.$progressData, ctx.$isPending);
   }, [canvasManager, ctx.$progressData, ctx.$selectedItemId, ctx.$isPending]);
 
+  useEffect(() => {
+    return ctx.$selectedItemIndex.listen((index) => {
+      if (!virtuosoRef.current) {
+        return;
+      }
+
+      if (!rootRef.current) {
+        return;
+      }
+
+      if (index === null) {
+        return;
+      }
+
+      scrollIntoView(index, rootRef.current, virtuosoRef.current, rangeRef.current);
+    });
+  }, [ctx.$selectedItemIndex]);
+
+  const onRangeChanged = useCallback((range: ListRange) => {
+    rangeRef.current = range;
+  }, []);
+
   return (
-    <Flex position="relative" maxW="full" w="full" h="72px">
-      <ScrollableContent overflowX="scroll" overflowY="hidden">
-        <Flex gap={2} w="full" h="full" justifyContent="safe center">
-          {items.map((item, i) => (
-            <QueueItemPreviewMini
-              key={`${item.item_id}-mini`}
-              item={item}
-              number={i + 1}
-              isSelected={selectedItemId === item.item_id}
-            />
-          ))}
-        </Flex>
-      </ScrollableContent>
-    </Flex>
+    <Box data-overlayscrollbars-initialize="" ref={rootRef} position="relative" w="full" h="full">
+      <Virtuoso<S['SessionQueueItem'], VirtuosoContext>
+        ref={virtuosoRef}
+        context={context}
+        data={items}
+        horizontalDirection
+        style={virtuosoStyles}
+        itemContent={itemContent}
+        components={components}
+        rangeChanged={onRangeChanged}
+        // Virtuoso expects the ref to be of HTMLElement | null | Window, but overlayscrollbars doesn't allow Window
+        scrollerRef={scrollerRef as VirtuosoProps<S['SessionQueueItem'], VirtuosoContext>['scrollerRef']}
+      />
+    </Box>
   );
 });
 StagingAreaItemsList.displayName = 'StagingAreaItemsList';
+
+const itemContent: ItemContent<S['SessionQueueItem'], VirtuosoContext> = (index, item, { selectedItemId }) => (
+  <QueueItemPreviewMini
+    key={`${item.item_id}-mini`}
+    item={item}
+    index={index}
+    isSelected={selectedItemId === item.item_id}
+  />
+);
+
+const listSx = {
+  '& > * + *': {
+    pl: 2,
+  },
+};
+
+const components: Components<S['SessionQueueItem'], VirtuosoContext> = {
+  List: forwardRef(({ context: _, ...rest }, ref) => {
+    return <Flex ref={ref} sx={listSx} {...rest} />;
+  }),
+};

invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/shared.ts

@@ -13,7 +13,7 @@ export const getProgressMessage = (data?: S['InvocationProgressEvent'] | null) =
 
 export const DROP_SHADOW = 'drop-shadow(0px 0px 4px rgb(0, 0, 0)) drop-shadow(0px 0px 4px rgba(0, 0, 0, 0.3))';
 
-export const getQueueItemElementId = (itemId: number) => `queue-item-status-card-${itemId}`;
+export const getQueueItemElementId = (index: number) => `queue-item-preview-${index}`;
 
 export const getOutputImageName = (item: S['SessionQueueItem']) => {
   const nodeId = Object.entries(item.session.source_prepared_mapping).find(([nodeId]) =>

invokeai/frontend/web/src/features/controlLayers/components/StagingArea/StagingAreaToolbar.tsx

@@ -1,7 +1,6 @@
 import { ButtonGroup, Flex } from '@invoke-ai/ui-library';
 import { useStore } from '@nanostores/react';
 import { useCanvasSessionContext } from 'features/controlLayers/components/SimpleSession/context';
-import { getQueueItemElementId } from 'features/controlLayers/components/SimpleSession/shared';
 import { StagingAreaToolbarAcceptButton } from 'features/controlLayers/components/StagingArea/StagingAreaToolbarAcceptButton';
 import { StagingAreaToolbarDiscardAllButton } from 'features/controlLayers/components/StagingArea/StagingAreaToolbarDiscardAllButton';
 import { StagingAreaToolbarDiscardSelectedButton } from 'features/controlLayers/components/StagingArea/StagingAreaToolbarDiscardSelectedButton';
@@ -12,7 +11,7 @@ import { StagingAreaToolbarPrevButton } from 'features/controlLayers/components/
 import { StagingAreaToolbarSaveSelectedToGalleryButton } from 'features/controlLayers/components/StagingArea/StagingAreaToolbarSaveSelectedToGalleryButton';
 import { StagingAreaToolbarToggleShowResultsButton } from 'features/controlLayers/components/StagingArea/StagingAreaToolbarToggleShowResultsButton';
 import { useCanvasManager } from 'features/controlLayers/contexts/CanvasManagerProviderGate';
-import { memo, useEffect } from 'react';
+import { memo } from 'react';
 import { useHotkeys } from 'react-hotkeys-hook';
 
 import { StagingAreaAutoSwitchButtons } from './StagingAreaAutoSwitchButtons';
@@ -23,16 +22,6 @@ export const StagingAreaToolbar = memo(() => {
 
   const ctx = useCanvasSessionContext();
 
-  useEffect(() => {
-    return ctx.$selectedItemId.listen((id) => {
-      if (id !== null) {
-        document
-          .getElementById(getQueueItemElementId(id))
-          ?.scrollIntoView({ block: 'nearest', inline: 'nearest', behavior: 'auto' });
-      }
-    });
-  }, [ctx.$selectedItemId]);
-
   useHotkeys('meta+left', ctx.selectFirst, { preventDefault: true });
   useHotkeys('meta+right', ctx.selectLast, { preventDefault: true });
 

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

@@ -482,11 +482,6 @@ export const NewGallery = memo(() => {
 
   const context = useMemo<GridContext>(() => ({ imageNames, queryArgs }), [imageNames, queryArgs]);
 
-  // Item content function
-  const itemContent: GridItemContent<string, GridContext> = useCallback((index, imageName) => {
-    return <ImageAtPosition index={index} imageName={imageName} />;
-  }, []);
-
   if (isLoading) {
     return (
       <Flex w="full" h="full" alignItems="center" justifyContent="center" gap={4}>
@@ -553,6 +548,10 @@ const ListComponent: GridComponents<GridContext>['List'] = forwardRef(({ context
 });
 ListComponent.displayName = 'ListComponent';
 
+const itemContent: GridItemContent<string, GridContext> = (index, imageName) => {
+  return <ImageAtPosition index={index} imageName={imageName} />;
+};
+
 const ItemComponent: GridComponents<GridContext>['Item'] = forwardRef(({ context: _, ...rest }, ref) => (
   <GridItem ref={ref} aspectRatio="1/1" {...rest} />
 ));

invokeai/frontend/web/src/features/ui/layouts/StagingArea.tsx

@@ -13,7 +13,7 @@ export const StagingArea = memo(() => {
   }
 
   return (
-    <Flex position="absolute" flexDir="column" bottom={4} gap={2} align="center" justify="center" left={4} right={4}>
+    <Flex position="absolute" flexDir="column" bottom={2} gap={2} align="center" justify="center" left={2} right={2}>
       <StagingAreaItemsList />
       <StagingAreaToolbar />
     </Flex>