Skip to content

Commit ef4d5d7

Browse files
psychedeliciouspsychedelicious
committed
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.
1 parent 6b0dfd8 commit ef4d5d7

File tree

6 files changed

+203
-40
lines changed

6 files changed

+203
-40
lines changed

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

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ const sx = {
2727
alignItems: 'center',
2828
justifyContent: 'center',
2929
flexShrink: 0,
30+
h: 'full',
3031
aspectRatio: '1/1',
3132
borderWidth: 2,
3233
borderRadius: 'base',
@@ -39,11 +40,11 @@ const sx = {
3940

4041
type Props = {
4142
item: S['SessionQueueItem'];
42-
number: number;
43+
index: number;
4344
isSelected: boolean;
4445
};
4546

46-
export const QueueItemPreviewMini = memo(({ item, isSelected, number }: Props) => {
47+
export const QueueItemPreviewMini = memo(({ item, isSelected, index }: Props) => {
4748
const dispatch = useAppDispatch();
4849
const ctx = useCanvasSessionContext();
4950
const { imageLoaded } = useProgressData(ctx.$progressData, item.item_id);
@@ -69,7 +70,7 @@ export const QueueItemPreviewMini = memo(({ item, isSelected, number }: Props) =
6970

7071
return (
7172
<Flex
72-
id={getQueueItemElementId(item.item_id)}
73+
id={getQueueItemElementId(index)}
7374
sx={sx}
7475
data-selected={isSelected}
7576
onClick={onClick}
@@ -78,7 +79,7 @@ export const QueueItemPreviewMini = memo(({ item, isSelected, number }: Props) =
7879
<QueueItemStatusLabel item={item} position="absolute" margin="auto" />
7980
{imageDTO && <DndImage imageDTO={imageDTO} onLoad={onLoad} asThumbnail />}
8081
{!imageLoaded && <QueueItemProgressImage itemId={item.item_id} position="absolute" />}
81-
<QueueItemNumber number={number} position="absolute" top={0} left={1} />
82+
<QueueItemNumber number={index + 1} position="absolute" top={0} left={1} />
8283
<QueueItemCircularProgress itemId={item.item_id} status={item.status} position="absolute" top={1} right={2} />
8384
</Flex>
8485
);

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

Lines changed: 191 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,148 @@
1-
import { Flex } from '@invoke-ai/ui-library';
1+
import { Box, Flex, forwardRef } from '@invoke-ai/ui-library';
22
import { useStore } from '@nanostores/react';
3-
import ScrollableContent from 'common/components/OverlayScrollbars/ScrollableContent';
3+
import { logger } from 'app/logging/logger';
44
import { useCanvasSessionContext } from 'features/controlLayers/components/SimpleSession/context';
55
import { QueueItemPreviewMini } from 'features/controlLayers/components/SimpleSession/QueueItemPreviewMini';
66
import { useCanvasManagerSafe } from 'features/controlLayers/contexts/CanvasManagerProviderGate';
7-
import { memo, useEffect } from 'react';
7+
import { useOverlayScrollbars } from 'overlayscrollbars-react';
8+
import type { CSSProperties, RefObject } from 'react';
9+
import { memo, useCallback, useEffect, useMemo, useRef, useState } from 'react';
10+
import type { Components, ItemContent, ListRange, VirtuosoHandle, VirtuosoProps } from 'react-virtuoso';
11+
import { Virtuoso } from 'react-virtuoso';
12+
import type { S } from 'services/api/types';
13+
14+
import { getQueueItemElementId } from './shared';
15+
16+
const log = logger('system');
17+
18+
const virtuosoStyles = {
19+
width: '100%',
20+
height: '72px',
21+
} satisfies CSSProperties;
22+
23+
type VirtuosoContext = { selectedItemId: number | null };
24+
25+
/**
26+
* Scroll the item at the given index into view if it is not currently visible.
27+
*/
28+
const scrollIntoView = (
29+
targetIndex: number,
30+
rootEl: HTMLDivElement,
31+
virtuosoHandle: VirtuosoHandle,
32+
range: ListRange
33+
) => {
34+
if (range.endIndex === 0) {
35+
// No range is rendered; no need to scroll to anything.
36+
return;
37+
}
38+
39+
const targetItem = rootEl.querySelector(`#${getQueueItemElementId(targetIndex)}`);
40+
41+
if (!targetItem) {
42+
if (targetIndex > range.endIndex) {
43+
virtuosoHandle.scrollToIndex({
44+
index: targetIndex,
45+
behavior: 'auto',
46+
align: 'end',
47+
});
48+
} else if (targetIndex < range.startIndex) {
49+
virtuosoHandle.scrollToIndex({
50+
index: targetIndex,
51+
behavior: 'auto',
52+
align: 'start',
53+
});
54+
} else {
55+
log.debug(
56+
`Unable to find queue item at index ${targetIndex} but it is in the rendered range ${range.startIndex}-${range.endIndex}`
57+
);
58+
}
59+
return;
60+
}
61+
62+
// We found the image in the DOM, but it might be in the overscan range - rendered but not in the visible viewport.
63+
// Check if it is in the viewport and scroll if necessary.
64+
65+
const itemRect = targetItem.getBoundingClientRect();
66+
const rootRect = rootEl.getBoundingClientRect();
67+
68+
if (itemRect.left < rootRect.left) {
69+
virtuosoHandle.scrollToIndex({
70+
index: targetIndex,
71+
behavior: 'auto',
72+
align: 'start',
73+
});
74+
} else if (itemRect.right > rootRect.right) {
75+
virtuosoHandle.scrollToIndex({
76+
index: targetIndex,
77+
behavior: 'auto',
78+
align: 'end',
79+
});
80+
} else {
81+
// Image is already in view
82+
}
83+
84+
return;
85+
};
86+
87+
const useScrollableStagingArea = (rootRef: RefObject<HTMLDivElement>) => {
88+
const [scroller, scrollerRef] = useState<HTMLElement | null>(null);
89+
const [initialize, osInstance] = useOverlayScrollbars({
90+
defer: true,
91+
events: {
92+
initialized(osInstance) {
93+
// force overflow styles
94+
const { viewport } = osInstance.elements();
95+
viewport.style.overflowX = `var(--os-viewport-overflow-x)`;
96+
viewport.style.overflowY = `var(--os-viewport-overflow-y)`;
97+
},
98+
},
99+
options: {
100+
scrollbars: {
101+
visibility: 'auto',
102+
autoHide: 'scroll',
103+
autoHideDelay: 1300,
104+
theme: 'os-theme-dark',
105+
},
106+
overflow: {
107+
y: 'hidden',
108+
x: 'scroll',
109+
},
110+
},
111+
});
112+
113+
useEffect(() => {
114+
const { current: root } = rootRef;
115+
116+
if (scroller && root) {
117+
initialize({
118+
target: root,
119+
elements: {
120+
viewport: scroller,
121+
},
122+
});
123+
}
124+
125+
return () => {
126+
osInstance()?.destroy();
127+
};
128+
}, [scroller, initialize, osInstance, rootRef]);
129+
130+
return scrollerRef;
131+
};
8132

9133
export const StagingAreaItemsList = memo(() => {
10134
const canvasManager = useCanvasManagerSafe();
11135
const ctx = useCanvasSessionContext();
136+
const virtuosoRef = useRef<VirtuosoHandle>(null);
137+
const rangeRef = useRef<ListRange>({ startIndex: 0, endIndex: 0 });
138+
const rootRef = useRef<HTMLDivElement>(null);
139+
12140
const items = useStore(ctx.$items);
13141
const selectedItemId = useStore(ctx.$selectedItemId);
14142

143+
const context = useMemo(() => ({ selectedItemId }), [selectedItemId]);
144+
const scrollerRef = useScrollableStagingArea(rootRef);
145+
15146
useEffect(() => {
16147
if (!canvasManager) {
17148
return;
@@ -20,21 +151,64 @@ export const StagingAreaItemsList = memo(() => {
20151
return canvasManager.stagingArea.connectToSession(ctx.$selectedItemId, ctx.$progressData, ctx.$isPending);
21152
}, [canvasManager, ctx.$progressData, ctx.$selectedItemId, ctx.$isPending]);
22153

154+
useEffect(() => {
155+
return ctx.$selectedItemIndex.listen((index) => {
156+
if (!virtuosoRef.current) {
157+
return;
158+
}
159+
160+
if (!rootRef.current) {
161+
return;
162+
}
163+
164+
if (index === null) {
165+
return;
166+
}
167+
168+
scrollIntoView(index, rootRef.current, virtuosoRef.current, rangeRef.current);
169+
});
170+
}, [ctx.$selectedItemIndex]);
171+
172+
const onRangeChanged = useCallback((range: ListRange) => {
173+
rangeRef.current = range;
174+
}, []);
175+
23176
return (
24-
<Flex position="relative" maxW="full" w="full" h="72px">
25-
<ScrollableContent overflowX="scroll" overflowY="hidden">
26-
<Flex gap={2} w="full" h="full" justifyContent="safe center">
27-
{items.map((item, i) => (
28-
<QueueItemPreviewMini
29-
key={`${item.item_id}-mini`}
30-
item={item}
31-
number={i + 1}
32-
isSelected={selectedItemId === item.item_id}
33-
/>
34-
))}
35-
</Flex>
36-
</ScrollableContent>
37-
</Flex>
177+
<Box data-overlayscrollbars-initialize="" ref={rootRef} position="relative" w="full" h="full">
178+
<Virtuoso<S['SessionQueueItem'], VirtuosoContext>
179+
ref={virtuosoRef}
180+
context={context}
181+
data={items}
182+
horizontalDirection
183+
style={virtuosoStyles}
184+
itemContent={itemContent}
185+
components={components}
186+
rangeChanged={onRangeChanged}
187+
// Virtuoso expects the ref to be of HTMLElement | null | Window, but overlayscrollbars doesn't allow Window
188+
scrollerRef={scrollerRef as VirtuosoProps<S['SessionQueueItem'], VirtuosoContext>['scrollerRef']}
189+
/>
190+
</Box>
38191
);
39192
});
40193
StagingAreaItemsList.displayName = 'StagingAreaItemsList';
194+
195+
const itemContent: ItemContent<S['SessionQueueItem'], VirtuosoContext> = (index, item, { selectedItemId }) => (
196+
<QueueItemPreviewMini
197+
key={`${item.item_id}-mini`}
198+
item={item}
199+
index={index}
200+
isSelected={selectedItemId === item.item_id}
201+
/>
202+
);
203+
204+
const listSx = {
205+
'& > * + *': {
206+
pl: 2,
207+
},
208+
};
209+
210+
const components: Components<S['SessionQueueItem'], VirtuosoContext> = {
211+
List: forwardRef(({ context: _, ...rest }, ref) => {
212+
return <Flex ref={ref} sx={listSx} {...rest} />;
213+
}),
214+
};

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ export const getProgressMessage = (data?: S['InvocationProgressEvent'] | null) =
1313

1414
export const DROP_SHADOW = 'drop-shadow(0px 0px 4px rgb(0, 0, 0)) drop-shadow(0px 0px 4px rgba(0, 0, 0, 0.3))';
1515

16-
export const getQueueItemElementId = (itemId: number) => `queue-item-status-card-${itemId}`;
16+
export const getQueueItemElementId = (index: number) => `queue-item-preview-${index}`;
1717

1818
export const getOutputImageName = (item: S['SessionQueueItem']) => {
1919
const nodeId = Object.entries(item.session.source_prepared_mapping).find(([nodeId]) =>

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

Lines changed: 1 addition & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
import { ButtonGroup, Flex } from '@invoke-ai/ui-library';
22
import { useStore } from '@nanostores/react';
33
import { useCanvasSessionContext } from 'features/controlLayers/components/SimpleSession/context';
4-
import { getQueueItemElementId } from 'features/controlLayers/components/SimpleSession/shared';
54
import { StagingAreaToolbarAcceptButton } from 'features/controlLayers/components/StagingArea/StagingAreaToolbarAcceptButton';
65
import { StagingAreaToolbarDiscardAllButton } from 'features/controlLayers/components/StagingArea/StagingAreaToolbarDiscardAllButton';
76
import { StagingAreaToolbarDiscardSelectedButton } from 'features/controlLayers/components/StagingArea/StagingAreaToolbarDiscardSelectedButton';
@@ -12,7 +11,7 @@ import { StagingAreaToolbarPrevButton } from 'features/controlLayers/components/
1211
import { StagingAreaToolbarSaveSelectedToGalleryButton } from 'features/controlLayers/components/StagingArea/StagingAreaToolbarSaveSelectedToGalleryButton';
1312
import { StagingAreaToolbarToggleShowResultsButton } from 'features/controlLayers/components/StagingArea/StagingAreaToolbarToggleShowResultsButton';
1413
import { useCanvasManager } from 'features/controlLayers/contexts/CanvasManagerProviderGate';
15-
import { memo, useEffect } from 'react';
14+
import { memo } from 'react';
1615
import { useHotkeys } from 'react-hotkeys-hook';
1716

1817
import { StagingAreaAutoSwitchButtons } from './StagingAreaAutoSwitchButtons';
@@ -23,16 +22,6 @@ export const StagingAreaToolbar = memo(() => {
2322

2423
const ctx = useCanvasSessionContext();
2524

26-
useEffect(() => {
27-
return ctx.$selectedItemId.listen((id) => {
28-
if (id !== null) {
29-
document
30-
.getElementById(getQueueItemElementId(id))
31-
?.scrollIntoView({ block: 'nearest', inline: 'nearest', behavior: 'auto' });
32-
}
33-
});
34-
}, [ctx.$selectedItemId]);
35-
3625
useHotkeys('meta+left', ctx.selectFirst, { preventDefault: true });
3726
useHotkeys('meta+right', ctx.selectLast, { preventDefault: true });
3827

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

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -482,11 +482,6 @@ export const NewGallery = memo(() => {
482482

483483
const context = useMemo<GridContext>(() => ({ imageNames, queryArgs }), [imageNames, queryArgs]);
484484

485-
// Item content function
486-
const itemContent: GridItemContent<string, GridContext> = useCallback((index, imageName) => {
487-
return <ImageAtPosition index={index} imageName={imageName} />;
488-
}, []);
489-
490485
if (isLoading) {
491486
return (
492487
<Flex w="full" h="full" alignItems="center" justifyContent="center" gap={4}>
@@ -553,6 +548,10 @@ const ListComponent: GridComponents<GridContext>['List'] = forwardRef(({ context
553548
});
554549
ListComponent.displayName = 'ListComponent';
555550

551+
const itemContent: GridItemContent<string, GridContext> = (index, imageName) => {
552+
return <ImageAtPosition index={index} imageName={imageName} />;
553+
};
554+
556555
const ItemComponent: GridComponents<GridContext>['Item'] = forwardRef(({ context: _, ...rest }, ref) => (
557556
<GridItem ref={ref} aspectRatio="1/1" {...rest} />
558557
));

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ export const StagingArea = memo(() => {
1313
}
1414

1515
return (
16-
<Flex position="absolute" flexDir="column" bottom={4} gap={2} align="center" justify="center" left={4} right={4}>
16+
<Flex position="absolute" flexDir="column" bottom={2} gap={2} align="center" justify="center" left={2} right={2}>
1717
<StagingAreaItemsList />
1818
<StagingAreaToolbar />
1919
</Flex>

0 commit comments

Comments
 (0)