fix(ui): fix gallery display bug, major lag

- Fixed a bug where after you load more, changing boards doesn't work. The offset and limit for the list image query had some wonky logic, now resolved.
- Addressed major lag in gallery when selecting an image.

Both issues were related to the useMultiselect and useGalleryImages hooks, which caused every image in the gallery to re-render on whenever the selection changed. There's no way to memoize away this - we need to know when the selection changes. This is a longstanding issue.

The selection is only used in a callback, though - the onClick handler for an image to select it (or add it to the existing selection). We don't really need the reactivity for a callback, so we don't need to listen for changes to the selection.

The logic to handle multiple selection is moved to a new `galleryImageClicked` listener, which does all the selection right when it is needed.

The result is that gallery images no long need to do heavy re-renders on any selection change.

Besides the multiselect click handler, there was also inefficient use of DND payloads. Previously, the `IMAGE_DTOS` type had a payload of image DTO objects. This was only used to drag gallery selection into a board. There is no need to hold onto image DTOs when we have the selection state already in redux. We were recalculating this payload for every image, on every tick.

This payload is now just the board id (the only piece of information we need for this particular DND event).

- I also removed some unused DND types while making this change.

Author: psychedelicious

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/index.ts

@@ -5,6 +5,7 @@ import type {
   UnknownAction,
 } from '@reduxjs/toolkit';
 import { addListener, createListenerMiddleware } from '@reduxjs/toolkit';
+import { addGalleryImageClickedListener } from 'app/store/middleware/listenerMiddleware/listeners/galleryImageClicked';
 import type { AppDispatch, RootState } from 'app/store/store';
 
 import { addCommitStagingAreaImageListener } from './listeners/addCommitStagingAreaImageListener';
@@ -117,6 +118,9 @@ addImageToDeleteSelectedListener();
 addImagesStarredListener();
 addImagesUnstarredListener();
 
+// Gallery
+addGalleryImageClickedListener();
+
 // User Invoked
 addEnqueueRequestedCanvasListener();
 addEnqueueRequestedNodes();
@@ -135,19 +139,7 @@ addCanvasMergedListener();
 addStagingAreaImageSavedListener();
 addCommitStagingAreaImageListener();
 
-/**
- * Socket.IO Events - these handle SIO events directly and pass on internal application actions.
- * We don't handle SIO events in slices via `extraReducers` because some of these events shouldn't
- * actually be handled at all.
- *
- * For example, we don't want to respond to progress events for canceled sessions. To avoid
- * duplicating the logic to determine if an event should be responded to, we handle all of that
- * "is this session canceled?" logic in these listeners.
- *
- * The `socketGeneratorProgress` listener will then only dispatch the `appSocketGeneratorProgress`
- * action if it should be handled by the rest of the application. It is this `appSocketGeneratorProgress`
- * action that is handled by reducers in slices.
- */
+// Socket.IO
 addGeneratorProgressListener();
 addGraphExecutionStateCompleteListener();
 addInvocationCompleteListener();

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/galleryImageClicked.ts

@@ -0,0 +1,80 @@
+import { createAction } from '@reduxjs/toolkit';
+import { selectListImagesQueryArgs } from 'features/gallery/store/gallerySelectors';
+import { selectionChanged } from 'features/gallery/store/gallerySlice';
+import { imagesApi } from 'services/api/endpoints/images';
+import type { ImageDTO } from 'services/api/types';
+import { imagesSelectors } from 'services/api/util';
+
+import { startAppListening } from '..';
+
+export const galleryImageClicked = createAction<{
+  imageDTO: ImageDTO;
+  shiftKey: boolean;
+  ctrlKey: boolean;
+  metaKey: boolean;
+}>('gallery/imageClicked');
+
+/**
+ * This listener handles the logic for selecting images in the gallery.
+ *
+ * Previously, this logic was in a `useCallback` with the whole gallery selection as a dependency. Every time
+ * the selection changed, the callback got recreated and all images rerendered. This could easily block for
+ * hundreds of ms, more for lower end devices.
+ *
+ * Moving this logic into a listener means we don't need to recalculate anything dynamically and the gallery
+ * is much more responsive.
+ */
+
+export const addGalleryImageClickedListener = () => {
+  startAppListening({
+    actionCreator: galleryImageClicked,
+    effect: async (action, { dispatch, getState }) => {
+      const { imageDTO, shiftKey, ctrlKey, metaKey } = action.payload;
+      const state = getState();
+      const queryArgs = selectListImagesQueryArgs(state);
+      const { data: listImagesData } =
+        imagesApi.endpoints.listImages.select(queryArgs)(state);
+
+      if (!listImagesData) {
+        // Should never happen if we have clicked a gallery image
+        return;
+      }
+
+      const imageDTOs = imagesSelectors.selectAll(listImagesData);
+      const selection = state.gallery.selection;
+
+      if (shiftKey) {
+        const rangeEndImageName = imageDTO.image_name;
+        const lastSelectedImage = selection[selection.length - 1]?.image_name;
+        const lastClickedIndex = imageDTOs.findIndex(
+          (n) => n.image_name === lastSelectedImage
+        );
+        const currentClickedIndex = imageDTOs.findIndex(
+          (n) => n.image_name === rangeEndImageName
+        );
+        if (lastClickedIndex > -1 && currentClickedIndex > -1) {
+          // We have a valid range!
+          const start = Math.min(lastClickedIndex, currentClickedIndex);
+          const end = Math.max(lastClickedIndex, currentClickedIndex);
+          const imagesToSelect = imageDTOs.slice(start, end + 1);
+          dispatch(selectionChanged(selection.concat(imagesToSelect)));
+        }
+      } else if (ctrlKey || metaKey) {
+        if (
+          selection.some((i) => i.image_name === imageDTO.image_name) &&
+          selection.length > 1
+        ) {
+          dispatch(
+            selectionChanged(
+              selection.filter((n) => n.image_name !== imageDTO.image_name)
+            )
+          );
+        } else {
+          dispatch(selectionChanged(selection.concat(imageDTO)));
+        }
+      } else {
+        dispatch(selectionChanged([imageDTO]));
+      }
+    },
+  });
+};

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/imageDeleted.ts

@@ -8,7 +8,7 @@ import {
 import { isControlNetOrT2IAdapter } from 'features/controlAdapters/store/types';
 import { imageDeletionConfirmed } from 'features/deleteImageModal/store/actions';
 import { isModalOpenChanged } from 'features/deleteImageModal/store/slice';
-import { selectListImagesBaseQueryArgs } from 'features/gallery/store/gallerySelectors';
+import { selectListImagesQueryArgs } from 'features/gallery/store/gallerySelectors';
 import { imageSelected } from 'features/gallery/store/gallerySlice';
 import { fieldImageValueChanged } from 'features/nodes/store/nodesSlice';
 import { isImageFieldInputInstance } from 'features/nodes/types/field';
@@ -49,7 +49,7 @@ export const addRequestedSingleImageDeletionListener = () => {
       if (imageDTO && imageDTO?.image_name === lastSelectedImage) {
         const { image_name } = imageDTO;
 
-        const baseQueryArgs = selectListImagesBaseQueryArgs(state);
+        const baseQueryArgs = selectListImagesQueryArgs(state);
         const { data } =
           imagesApi.endpoints.listImages.select(baseQueryArgs)(state);
 
@@ -180,9 +180,9 @@ export const addRequestedMultipleImageDeletionListener = () => {
           imagesApi.endpoints.deleteImages.initiate({ imageDTOs })
         ).unwrap();
         const state = getState();
-        const baseQueryArgs = selectListImagesBaseQueryArgs(state);
+        const queryArgs = selectListImagesQueryArgs(state);
         const { data } =
-          imagesApi.endpoints.listImages.select(baseQueryArgs)(state);
+          imagesApi.endpoints.listImages.select(queryArgs)(state);
 
         const newSelectedImageDTO = data
           ? imagesSelectors.selectAll(data)[0]

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/imageDropped.ts

@@ -12,7 +12,6 @@ import type {
 } from 'features/dnd/types';
 import { imageSelected } from 'features/gallery/store/gallerySlice';
 import { fieldImageValueChanged } from 'features/nodes/store/nodesSlice';
-import { workflowExposedFieldAdded } from 'features/nodes/store/workflowSlice';
 import {
   initialImageChanged,
   selectOptimalDimension,
@@ -35,10 +34,10 @@ export const addImageDroppedListener = () => {
 
       if (activeData.payloadType === 'IMAGE_DTO') {
         log.debug({ activeData, overData }, 'Image dropped');
-      } else if (activeData.payloadType === 'IMAGE_DTOS') {
+      } else if (activeData.payloadType === 'GALLERY_SELECTION') {
         log.debug(
           { activeData, overData },
-          `Images (${activeData.payload.imageDTOs.length}) dropped`
+          `Images (${getState().gallery.selection.length}) dropped`
         );
       } else if (activeData.payloadType === 'NODE_FIELD') {
         log.debug(
@@ -49,19 +48,6 @@ export const addImageDroppedListener = () => {
         log.debug({ activeData, overData }, `Unknown payload dropped`);
       }
 
-      if (
-        overData.actionType === 'ADD_FIELD_TO_LINEAR' &&
-        activeData.payloadType === 'NODE_FIELD'
-      ) {
-        const { nodeId, field } = activeData.payload;
-        dispatch(
-          workflowExposedFieldAdded({
-            nodeId,
-            fieldName: field.name,
-          })
-        );
-      }
-
       /**
        * Image dropped on current image
        */
@@ -207,10 +193,9 @@ export const addImageDroppedListener = () => {
        */
       if (
         overData.actionType === 'ADD_TO_BOARD' &&
-        activeData.payloadType === 'IMAGE_DTOS' &&
-        activeData.payload.imageDTOs
+        activeData.payloadType === 'GALLERY_SELECTION'
       ) {
-        const { imageDTOs } = activeData.payload;
+        const imageDTOs = getState().gallery.selection;
         const { boardId } = overData.context;
         dispatch(
           imagesApi.endpoints.addImagesToBoard.initiate({
@@ -226,10 +211,9 @@ export const addImageDroppedListener = () => {
        */
       if (
         overData.actionType === 'REMOVE_FROM_BOARD' &&
-        activeData.payloadType === 'IMAGE_DTOS' &&
-        activeData.payload.imageDTOs
+        activeData.payloadType === 'GALLERY_SELECTION'
       ) {
-        const { imageDTOs } = activeData.payload;
+        const imageDTOs = getState().gallery.selection;
         dispatch(
           imagesApi.endpoints.removeImagesFromBoard.initiate({
             imageDTOs,

invokeai/frontend/web/src/features/dnd/components/DragPreview.tsx

@@ -1,5 +1,6 @@
 import type { ChakraProps } from '@chakra-ui/react';
 import { Box, Flex, Heading, Image } from '@chakra-ui/react';
+import { useAppSelector } from 'app/store/storeHooks';
 import { InvText } from 'common/components/InvText/wrapper';
 import type { TypesafeDraggableData } from 'features/dnd/types';
 import { memo } from 'react';
@@ -34,6 +35,7 @@ const multiImageStyles: ChakraProps['sx'] = {
 
 const DragPreview = (props: OverlayDragImageProps) => {
   const { t } = useTranslation();
+  const selectionCount = useAppSelector((s) => s.gallery.selection.length);
   if (!props.dragData) {
     return null;
   }
@@ -79,10 +81,10 @@ const DragPreview = (props: OverlayDragImageProps) => {
     );
   }
 
-  if (props.dragData.payloadType === 'IMAGE_DTOS') {
+  if (props.dragData.payloadType === 'GALLERY_SELECTION') {
     return (
       <Flex sx={multiImageStyles}>
-        <Heading>{props.dragData.payload.imageDTOs.length}</Heading>
+        <Heading>{selectionCount}</Heading>
         <Heading size="sm">{t('parameters.images')}</Heading>
       </Flex>
     );

invokeai/frontend/web/src/features/dnd/types/index.ts

@@ -10,6 +10,7 @@ import type {
   useDroppable as useOriginalDroppable,
   UseDroppableArguments,
 } from '@dnd-kit/core';
+import type { BoardId } from 'features/gallery/store/types';
 import type {
   FieldInputInstance,
   FieldInputTemplate,
@@ -51,15 +52,6 @@ export type NodesImageDropData = BaseDropData & {
   };
 };
 
-export type NodesMultiImageDropData = BaseDropData & {
-  actionType: 'SET_MULTI_NODES_IMAGE';
-  context: { nodeId: string; fieldName: string };
-};
-
-export type AddToBatchDropData = BaseDropData & {
-  actionType: 'ADD_TO_BATCH';
-};
-
 export type AddToBoardDropData = BaseDropData & {
   actionType: 'ADD_TO_BOARD';
   context: { boardId: string };
@@ -69,21 +61,14 @@ export type RemoveFromBoardDropData = BaseDropData & {
   actionType: 'REMOVE_FROM_BOARD';
 };
 
-export type AddFieldToLinearViewDropData = BaseDropData & {
-  actionType: 'ADD_FIELD_TO_LINEAR';
-};
-
 export type TypesafeDroppableData =
   | CurrentImageDropData
   | InitialImageDropData
   | ControlAdapterDropData
   | CanvasInitialImageDropData
   | NodesImageDropData
-  | AddToBatchDropData
-  | NodesMultiImageDropData
   | AddToBoardDropData
-  | RemoveFromBoardDropData
-  | AddFieldToLinearViewDropData;
+  | RemoveFromBoardDropData;
 
 type BaseDragData = {
   id: string;
@@ -103,15 +88,15 @@ export type ImageDraggableData = BaseDragData & {
   payload: { imageDTO: ImageDTO };
 };
 
-export type ImageDTOsDraggableData = BaseDragData & {
-  payloadType: 'IMAGE_DTOS';
-  payload: { imageDTOs: ImageDTO[] };
+export type GallerySelectionDraggableData = BaseDragData & {
+  payloadType: 'GALLERY_SELECTION';
+  payload: { boardId: BoardId };
 };
 
 export type TypesafeDraggableData =
   | NodeFieldDraggableData
   | ImageDraggableData
-  | ImageDTOsDraggableData;
+  | GallerySelectionDraggableData;
 
 export interface UseDroppableTypesafeArguments
   extends Omit<UseDroppableArguments, 'data'> {

invokeai/frontend/web/src/features/dnd/util/isValidDrop.ts

@@ -16,8 +16,6 @@ export const isValidDrop = (
   }
 
   switch (actionType) {
-    case 'ADD_FIELD_TO_LINEAR':
-      return payloadType === 'NODE_FIELD';
     case 'SET_CURRENT_IMAGE':
       return payloadType === 'IMAGE_DTO';
     case 'SET_INITIAL_IMAGE':
@@ -28,15 +26,13 @@ export const isValidDrop = (
       return payloadType === 'IMAGE_DTO';
     case 'SET_NODES_IMAGE':
       return payloadType === 'IMAGE_DTO';
-    case 'SET_MULTI_NODES_IMAGE':
-      return payloadType === 'IMAGE_DTO' || 'IMAGE_DTOS';
-    case 'ADD_TO_BATCH':
-      return payloadType === 'IMAGE_DTO' || 'IMAGE_DTOS';
     case 'ADD_TO_BOARD': {
       // If the board is the same, don't allow the drop
 
       // Check the payload types
-      const isPayloadValid = payloadType === 'IMAGE_DTO' || 'IMAGE_DTOS';
+      const isPayloadValid = ['IMAGE_DTO', 'GALLERY_SELECTION'].includes(
+        payloadType
+      );
       if (!isPayloadValid) {
         return false;
       }
@@ -50,12 +46,10 @@ export const isValidDrop = (
         return currentBoard !== destinationBoard;
       }
 
-      if (payloadType === 'IMAGE_DTOS') {
+      if (payloadType === 'GALLERY_SELECTION') {
         // Assume all images are on the same board - this is true for the moment
-        const { imageDTOs } = active.data.current.payload;
-        const currentBoard = imageDTOs[0]?.board_id ?? 'none';
+        const currentBoard = active.data.current.payload.boardId;
         const destinationBoard = overData.context.boardId;
-
         return currentBoard !== destinationBoard;
       }
 
@@ -65,7 +59,9 @@ export const isValidDrop = (
       // If the board is the same, don't allow the drop
 
       // Check the payload types
-      const isPayloadValid = payloadType === 'IMAGE_DTO' || 'IMAGE_DTOS';
+      const isPayloadValid = ['IMAGE_DTO', 'GALLERY_SELECTION'].includes(
+        payloadType
+      );
       if (!isPayloadValid) {
         return false;
       }
@@ -78,11 +74,8 @@ export const isValidDrop = (
         return currentBoard !== 'none';
       }
 
-      if (payloadType === 'IMAGE_DTOS') {
-        // Assume all images are on the same board - this is true for the moment
-        const { imageDTOs } = active.data.current.payload;
-        const currentBoard = imageDTOs[0]?.board_id ?? 'none';
-
+      if (payloadType === 'GALLERY_SELECTION') {
+        const currentBoard = active.data.current.payload.boardId;
         return currentBoard !== 'none';
       }