fix(ui): change multi image drop to not have selection as payload

This caused a lot of re-rendering whenever the selection changed, which caused a huge performance hit. It also made changing the current image lag a bit. Instead of providing an array of image names as a multi-select dnd payload, there is now no multi-select dnd payload at all - instead, the payload types are used by the `imageDropped` listener to pull the selection out of redux. Now, the only big re-renders are when the selectionCount changes. In the future I'll figure out a good way to do image names as payload without incurring re-renders.
invoke-ai · Jul 5, 2023 · f155887 · f155887
1 parent 1358c5e
commit f155887
Show file tree

Hide file tree

Showing 8 changed files with 100 additions and 62 deletions.
diff --git a/invokeai/frontend/web/src/app/components/ImageDnd/DragPreview.tsx b/invokeai/frontend/web/src/app/components/ImageDnd/DragPreview.tsx
@@ -1,4 +1,8 @@
 import { Box, ChakraProps, Flex, Heading, Image } from '@chakra-ui/react';
+import { createSelector } from '@reduxjs/toolkit';
+import { stateSelector } from 'app/store/store';
+import { useAppSelector } from 'app/store/storeHooks';
+import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
 import { memo } from 'react';
 import { TypesafeDraggableData } from './typesafeDnd';
 
@@ -28,7 +32,24 @@ const STYLES: ChakraProps['sx'] = {
   },
 };
 
+const selector = createSelector(
+  stateSelector,
+  (state) => {
+    const gallerySelectionCount = state.gallery.selection.length;
+    const batchSelectionCount = state.batch.selection.length;
+
+    return {
+      gallerySelectionCount,
+      batchSelectionCount,
+    };
+  },
+  defaultSelectorOptions
+);
+
 const DragPreview = (props: OverlayDragImageProps) => {
+  const { gallerySelectionCount, batchSelectionCount } =
+    useAppSelector(selector);
+
   if (!props.dragData) {
     return;
   }
@@ -57,7 +78,26 @@ const DragPreview = (props: OverlayDragImageProps) => {
     );
   }
 
-  if (props.dragData.payloadType === 'IMAGE_NAMES') {
+  if (props.dragData.payloadType === 'BATCH_SELECTION') {
+    return (
+      <Flex
+        sx={{
+          cursor: 'none',
+          userSelect: 'none',
+          position: 'relative',
+          alignItems: 'center',
+          justifyContent: 'center',
+          flexDir: 'column',
+          ...STYLES,
+        }}
+      >
+        <Heading>{batchSelectionCount}</Heading>
+        <Heading size="sm">Images</Heading>
+      </Flex>
+    );
+  }
+
+  if (props.dragData.payloadType === 'GALLERY_SELECTION') {
     return (
       <Flex
         sx={{
@@ -70,7 +110,7 @@ const DragPreview = (props: OverlayDragImageProps) => {
           ...STYLES,
         }}
       >
-        <Heading>{props.dragData.payload.imageNames.length}</Heading>
+        <Heading>{gallerySelectionCount}</Heading>
         <Heading size="sm">Images</Heading>
       </Flex>
     );

diff --git a/invokeai/frontend/web/src/app/components/ImageDnd/typesafeDnd.tsx b/invokeai/frontend/web/src/app/components/ImageDnd/typesafeDnd.tsx
@@ -77,14 +77,18 @@ export type ImageDraggableData = BaseDragData & {
   payload: { imageDTO: ImageDTO };
 };
 
-export type ImageNamesDraggableData = BaseDragData & {
-  payloadType: 'IMAGE_NAMES';
-  payload: { imageNames: string[] };
+export type GallerySelectionDraggableData = BaseDragData & {
+  payloadType: 'GALLERY_SELECTION';
+};
+
+export type BatchSelectionDraggableData = BaseDragData & {
+  payloadType: 'BATCH_SELECTION';
 };
 
 export type TypesafeDraggableData =
   | ImageDraggableData
-  | ImageNamesDraggableData;
+  | GallerySelectionDraggableData
+  | BatchSelectionDraggableData;
 
 interface UseDroppableTypesafeArguments
   extends Omit<UseDroppableArguments, 'data'> {
@@ -155,11 +159,13 @@ export const isValidDrop = (
     case 'SET_NODES_IMAGE':
       return payloadType === 'IMAGE_DTO';
     case 'SET_MULTI_NODES_IMAGE':
-      return payloadType === 'IMAGE_DTO' || 'IMAGE_NAMES';
+      return payloadType === 'IMAGE_DTO' || 'GALLERY_SELECTION';
     case 'ADD_TO_BATCH':
-      return payloadType === 'IMAGE_DTO' || 'IMAGE_NAMES';
+      return payloadType === 'IMAGE_DTO' || 'GALLERY_SELECTION';
     case 'MOVE_BOARD':
-      return payloadType === 'IMAGE_DTO' || 'IMAGE_NAMES';
+      return (
+        payloadType === 'IMAGE_DTO' || 'GALLERY_SELECTION' || 'BATCH_SELECTION'
+      );
     default:
       return false;
   }

diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/imageDropped.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/imageDropped.ts
@@ -1,24 +1,23 @@
 import { createAction } from '@reduxjs/toolkit';
-import { startAppListening } from '../';
-import { log } from 'app/logging/useLogger';
 import {
   TypesafeDraggableData,
   TypesafeDroppableData,
 } from 'app/components/ImageDnd/typesafeDnd';
-import { imageSelected } from 'features/gallery/store/gallerySlice';
-import { initialImageChanged } from 'features/parameters/store/generationSlice';
+import { log } from 'app/logging/useLogger';
 import {
   imageAddedToBatch,
   imagesAddedToBatch,
 } from 'features/batch/store/batchSlice';
-import { controlNetImageChanged } from 'features/controlNet/store/controlNetSlice';
 import { setInitialCanvasImage } from 'features/canvas/store/canvasSlice';
+import { controlNetImageChanged } from 'features/controlNet/store/controlNetSlice';
+import { imageSelected } from 'features/gallery/store/gallerySlice';
 import {
   fieldValueChanged,
   imageCollectionFieldValueChanged,
 } from 'features/nodes/store/nodesSlice';
-import { boardsApi } from 'services/api/endpoints/boards';
+import { initialImageChanged } from 'features/parameters/store/generationSlice';
 import { boardImagesApi } from 'services/api/endpoints/boardImages';
+import { startAppListening } from '../';
 
 const moduleLog = log.child({ namespace: 'dnd' });
 
@@ -33,6 +32,7 @@ export const addImageDroppedListener = () => {
     effect: (action, { dispatch, getState }) => {
       const { activeData, overData } = action.payload;
       const { actionType } = overData;
+      const state = getState();
 
       // set current image
       if (
@@ -64,9 +64,9 @@ export const addImageDroppedListener = () => {
       // add multiple images to batch
       if (
         actionType === 'ADD_TO_BATCH' &&
-        activeData.payloadType === 'IMAGE_NAMES'
+        activeData.payloadType === 'GALLERY_SELECTION'
       ) {
-        dispatch(imagesAddedToBatch(activeData.payload.imageNames));
+        dispatch(imagesAddedToBatch(state.gallery.selection));
       }
 
       // set control image
@@ -128,14 +128,14 @@ export const addImageDroppedListener = () => {
       // set multiple nodes images (multiple images handler)
       if (
         actionType === 'SET_MULTI_NODES_IMAGE' &&
-        activeData.payloadType === 'IMAGE_NAMES'
+        activeData.payloadType === 'GALLERY_SELECTION'
       ) {
         const { fieldName, nodeId } = overData.context;
         dispatch(
           imageCollectionFieldValueChanged({
             nodeId,
             fieldName,
-            value: activeData.payload.imageNames.map((image_name) => ({
+            value: state.gallery.selection.map((image_name) => ({
               image_name,
             })),
           })

diff --git a/invokeai/frontend/web/src/features/batch/components/BatchImage.tsx b/invokeai/frontend/web/src/features/batch/components/BatchImage.tsx
@@ -19,7 +19,7 @@ const makeSelector = (image_name: string) =>
   createSelector(
     [stateSelector],
     (state) => ({
-      selection: state.batch.selection,
+      selectionCount: state.batch.selection.length,
       isSelected: state.batch.selection.includes(image_name),
     }),
     defaultSelectorOptions
@@ -43,7 +43,7 @@ const BatchImage = (props: BatchImageProps) => {
     [props.imageName]
   );
 
-  const { isSelected, selection } = useAppSelector(selector);
+  const { isSelected, selectionCount } = useAppSelector(selector);
 
   const handleClickRemove = useCallback(() => {
     dispatch(imageRemovedFromBatch(props.imageName));
@@ -63,13 +63,10 @@ const BatchImage = (props: BatchImageProps) => {
   );
 
   const draggableData = useMemo<TypesafeDraggableData | undefined>(() => {
-    if (selection.length > 1) {
+    if (selectionCount > 1) {
       return {
         id: 'batch',
-        payloadType: 'IMAGE_NAMES',
-        payload: {
-          imageNames: selection,
-        },
+        payloadType: 'BATCH_SELECTION',
       };
     }
 
@@ -80,7 +77,7 @@ const BatchImage = (props: BatchImageProps) => {
         payload: { imageDTO },
       };
     }
-  }, [imageDTO, selection]);
+  }, [imageDTO, selectionCount]);
 
   if (isError) {
     return <Icon as={FaExclamationCircle} />;

diff --git a/invokeai/frontend/web/src/features/controlNet/components/ControlNetImagePreview.tsx b/invokeai/frontend/web/src/features/controlNet/components/ControlNetImagePreview.tsx
@@ -1,25 +1,22 @@
-import { memo, useCallback, useMemo, useState } from 'react';
-import { ImageDTO } from 'services/api/types';
-import {
-  ControlNetConfig,
-  controlNetImageChanged,
-  controlNetSelector,
-} from '../store/controlNetSlice';
-import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import { Box, Flex, SystemStyleObject } from '@chakra-ui/react';
-import IAIDndImage from 'common/components/IAIDndImage';
 import { createSelector } from '@reduxjs/toolkit';
-import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
-import { IAILoadingImageFallback } from 'common/components/IAIImageFallback';
-import IAIIconButton from 'common/components/IAIIconButton';
-import { FaUndo } from 'react-icons/fa';
-import { useGetImageDTOQuery } from 'services/api/endpoints/images';
 import { skipToken } from '@reduxjs/toolkit/dist/query';
 import {
   TypesafeDraggableData,
   TypesafeDroppableData,
 } from 'app/components/ImageDnd/typesafeDnd';
+import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
+import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
+import IAIDndImage from 'common/components/IAIDndImage';
+import { IAILoadingImageFallback } from 'common/components/IAIImageFallback';
+import { memo, useCallback, useMemo, useState } from 'react';
+import { useGetImageDTOQuery } from 'services/api/endpoints/images';
 import { PostUploadAction } from 'services/api/thunks/image';
+import {
+  ControlNetConfig,
+  controlNetImageChanged,
+  controlNetSelector,
+} from '../store/controlNetSlice';
 
 const selector = createSelector(
   controlNetSelector,

diff --git a/invokeai/frontend/web/src/features/gallery/components/CurrentImagePreview.tsx b/invokeai/frontend/web/src/features/gallery/components/CurrentImagePreview.tsx
@@ -1,19 +1,19 @@
 import { Box, Flex, Image } from '@chakra-ui/react';
 import { createSelector } from '@reduxjs/toolkit';
-import { useAppSelector } from 'app/store/storeHooks';
-import { isEqual } from 'lodash-es';
-import ImageMetadataViewer from './ImageMetaDataViewer/ImageMetadataViewer';
-import NextPrevImageButtons from './NextPrevImageButtons';
-import { memo, useMemo } from 'react';
-import IAIDndImage from 'common/components/IAIDndImage';
-import { useGetImageDTOQuery } from 'services/api/endpoints/images';
 import { skipToken } from '@reduxjs/toolkit/dist/query';
-import { stateSelector } from 'app/store/store';
-import { selectLastSelectedImage } from 'features/gallery/store/gallerySlice';
 import {
   TypesafeDraggableData,
   TypesafeDroppableData,
 } from 'app/components/ImageDnd/typesafeDnd';
+import { stateSelector } from 'app/store/store';
+import { useAppSelector } from 'app/store/storeHooks';
+import IAIDndImage from 'common/components/IAIDndImage';
+import { selectLastSelectedImage } from 'features/gallery/store/gallerySlice';
+import { isEqual } from 'lodash-es';
+import { memo, useMemo } from 'react';
+import { useGetImageDTOQuery } from 'services/api/endpoints/images';
+import ImageMetadataViewer from './ImageMetaDataViewer/ImageMetadataViewer';
+import NextPrevImageButtons from './NextPrevImageButtons';
 
 export const imagesSelector = createSelector(
   [stateSelector, selectLastSelectedImage],

diff --git a/invokeai/frontend/web/src/features/gallery/components/GalleryImage.tsx b/invokeai/frontend/web/src/features/gallery/components/GalleryImage.tsx
@@ -22,10 +22,10 @@ export const makeSelector = (image_name: string) =>
     [stateSelector],
     ({ gallery }) => {
       const isSelected = gallery.selection.includes(image_name);
-      const selection = gallery.selection;
+      const selectionCount = gallery.selection.length;
       return {
         isSelected,
-        selection,
+        selectionCount,
       };
     },
     defaultSelectorOptions
@@ -44,7 +44,7 @@ const GalleryImage = (props: HoverableImageProps) => {
 
   const localSelector = useMemo(() => makeSelector(image_name), [image_name]);
 
-  const { isSelected, selection } = useAppSelector(localSelector);
+  const { isSelected, selectionCount } = useAppSelector(localSelector);
 
   const dispatch = useAppDispatch();
 
@@ -75,11 +75,10 @@ const GalleryImage = (props: HoverableImageProps) => {
   );
 
   const draggableData = useMemo<TypesafeDraggableData | undefined>(() => {
-    if (selection.length > 1) {
+    if (selectionCount > 1) {
       return {
         id: 'gallery-image',
-        payloadType: 'IMAGE_NAMES',
-        payload: { imageNames: selection },
+        payloadType: 'GALLERY_SELECTION',
       };
     }
 
@@ -90,7 +89,7 @@ const GalleryImage = (props: HoverableImageProps) => {
         payload: { imageDTO },
       };
     }
-  }, [imageDTO, selection]);
+  }, [imageDTO, selectionCount]);
 
   return (
     <Box sx={{ w: 'full', h: 'full', touchAction: 'none' }}>

diff --git a/invokeai/frontend/web/src/features/nodes/components/fields/ImageInputFieldComponent.tsx b/invokeai/frontend/web/src/features/nodes/components/fields/ImageInputFieldComponent.tsx
@@ -7,18 +7,17 @@ import {
 } from 'features/nodes/types/types';
 import { memo, useCallback, useMemo } from 'react';
 
-import { FieldComponentProps } from './types';
-import IAIDndImage from 'common/components/IAIDndImage';
-import { ImageDTO } from 'services/api/types';
 import { Flex } from '@chakra-ui/react';
-import { useGetImageDTOQuery } from 'services/api/endpoints/images';
 import { skipToken } from '@reduxjs/toolkit/dist/query';
 import {
-  NodesImageDropData,
   TypesafeDraggableData,
   TypesafeDroppableData,
 } from 'app/components/ImageDnd/typesafeDnd';
+import IAIDndImage from 'common/components/IAIDndImage';
+import { useGetImageDTOQuery } from 'services/api/endpoints/images';
 import { PostUploadAction } from 'services/api/thunks/image';
+import { ImageDTO } from 'services/api/types';
+import { FieldComponentProps } from './types';
 
 const ImageInputFieldComponent = (
   props: FieldComponentProps<ImageInputFieldValue, ImageInputFieldTemplate>