feat(ui): use custom drag previews for images

invokeai/frontend/web/src/features/dnd/DndDragPreviewMultipleImage.tsx

@@ -0,0 +1,64 @@
+import type { draggable } from '@atlaskit/pragmatic-drag-and-drop/element/adapter';
+import { setCustomNativeDragPreview } from '@atlaskit/pragmatic-drag-and-drop/element/set-custom-native-drag-preview';
+import { Flex, Heading } from '@invoke-ai/ui-library';
+import type { Dnd } from 'features/dnd/dnd';
+import { DND_IMAGE_DRAG_PREVIEW_SIZE, preserveOffsetOnSourceFallbackCentered } from 'features/dnd/dnd';
+import { memo } from 'react';
+import { createPortal } from 'react-dom';
+import { useTranslation } from 'react-i18next';
+import type { ImageDTO } from 'services/api/types';
+import type { Param0 } from 'tsafe';
+
+const DndDragPreviewMultipleImage = memo(({ imageDTOs }: { imageDTOs: ImageDTO[] }) => {
+  const { t } = useTranslation();
+  return (
+    <Flex
+      w={DND_IMAGE_DRAG_PREVIEW_SIZE}
+      h={DND_IMAGE_DRAG_PREVIEW_SIZE}
+      alignItems="center"
+      justifyContent="center"
+      flexDir="column"
+      bg="base.900"
+      borderRadius="base"
+    >
+      <Heading>{imageDTOs.length}</Heading>
+      <Heading size="sm">{t('parameters.images')}</Heading>
+    </Flex>
+  );
+});
+
+DndDragPreviewMultipleImage.displayName = 'DndDragPreviewMultipleImage';
+
+export type DndDragPreviewMultipleImageState = {
+  type: 'multiple-image';
+  container: HTMLElement;
+  imageDTOs: ImageDTO[];
+};
+
+export const createMultipleImageDragPreview = (arg: DndDragPreviewMultipleImageState) =>
+  createPortal(<DndDragPreviewMultipleImage imageDTOs={arg.imageDTOs} />, arg.container);
+
+type SetMultipleDragPreviewArg = {
+  multipleImageDndData: Dnd.types['SourceDataTypeMap']['multipleImage'];
+  setDragPreviewState: (dragPreviewState: DndDragPreviewMultipleImageState | null) => void;
+  onGenerateDragPreviewArgs: Param0<Param0<typeof draggable>['onGenerateDragPreview']>;
+};
+
+export const setMultipleImageDragPreview = ({
+  multipleImageDndData,
+  onGenerateDragPreviewArgs,
+  setDragPreviewState,
+}: SetMultipleDragPreviewArg) => {
+  const { nativeSetDragImage, source, location } = onGenerateDragPreviewArgs;
+  setCustomNativeDragPreview({
+    render({ container }) {
+      setDragPreviewState({ type: 'multiple-image', container, imageDTOs: multipleImageDndData.payload.imageDTOs });
+      return () => setDragPreviewState(null);
+    },
+    nativeSetDragImage,
+    getOffset: preserveOffsetOnSourceFallbackCentered({
+      element: source.element,
+      input: location.current.input,
+    }),
+  });
+};

invokeai/frontend/web/src/features/dnd/DndDragPreviewSingleImage.tsx

@@ -0,0 +1,62 @@
+import type { draggable } from '@atlaskit/pragmatic-drag-and-drop/element/adapter';
+import { setCustomNativeDragPreview } from '@atlaskit/pragmatic-drag-and-drop/element/set-custom-native-drag-preview';
+import { chakra, Flex } from '@invoke-ai/ui-library';
+import type { Dnd } from 'features/dnd/dnd';
+import { DND_IMAGE_DRAG_PREVIEW_SIZE, preserveOffsetOnSourceFallbackCentered } from 'features/dnd/dnd';
+import { memo } from 'react';
+import { createPortal } from 'react-dom';
+import type { ImageDTO } from 'services/api/types';
+import type { Param0 } from 'tsafe';
+
+const ChakraImg = chakra('img');
+
+const DndDragPreviewSingleImage = memo(({ imageDTO }: { imageDTO: ImageDTO }) => {
+  return (
+    <Flex w={DND_IMAGE_DRAG_PREVIEW_SIZE} h={DND_IMAGE_DRAG_PREVIEW_SIZE}>
+      <ChakraImg
+        margin="auto"
+        maxW="full"
+        maxH="full"
+        objectFit="contain"
+        borderRadius="base"
+        src={imageDTO.thumbnail_url}
+      />
+    </Flex>
+  );
+});
+
+DndDragPreviewSingleImage.displayName = 'DndDragPreviewSingleImage';
+
+export type DndDragPreviewSingleImageState = {
+  type: 'single-image';
+  container: HTMLElement;
+  imageDTO: ImageDTO;
+};
+
+export const createSingleImageDragPreview = (arg: DndDragPreviewSingleImageState) =>
+  createPortal(<DndDragPreviewSingleImage imageDTO={arg.imageDTO} />, arg.container);
+
+type SetSingleDragPreviewArg = {
+  singleImageDndData: Dnd.types['SourceDataTypeMap']['singleImage'];
+  setDragPreviewState: (dragPreviewState: DndDragPreviewSingleImageState | null) => void;
+  onGenerateDragPreviewArgs: Param0<Param0<typeof draggable>['onGenerateDragPreview']>;
+};
+
+export const setSingleImageDragPreview = ({
+  singleImageDndData,
+  onGenerateDragPreviewArgs,
+  setDragPreviewState,
+}: SetSingleDragPreviewArg) => {
+  const { nativeSetDragImage, source, location } = onGenerateDragPreviewArgs;
+  setCustomNativeDragPreview({
+    render({ container }) {
+      setDragPreviewState({ type: 'single-image', container, imageDTO: singleImageDndData.payload.imageDTO });
+      return () => setDragPreviewState(null);
+    },
+    nativeSetDragImage,
+    getOffset: preserveOffsetOnSourceFallbackCentered({
+      element: source.element,
+      input: location.current.input,
+    }),
+  });
+};

invokeai/frontend/web/src/features/dnd/DndImage.tsx

@@ -3,6 +3,8 @@ import type { ImageProps, SystemStyleObject } from '@invoke-ai/ui-library';
 import { Image } from '@invoke-ai/ui-library';
 import { useAppStore } from 'app/store/nanostores/store';
 import { Dnd } from 'features/dnd/dnd';
+import type { DndDragPreviewSingleImageState } from 'features/dnd/DndDragPreviewSingleImage';
+import { createSingleImageDragPreview, setSingleImageDragPreview } from 'features/dnd/DndDragPreviewSingleImage';
 import { useImageContextMenu } from 'features/gallery/components/ImageContextMenu/ImageContextMenu';
 import { memo, useEffect, useState } from 'react';
 import type { ImageDTO } from 'services/api/types';
@@ -26,6 +28,7 @@ export const DndImage = memo(({ imageDTO, ...rest }: Props) => {
   const store = useAppStore();
   const [isDragging, setIsDragging] = useState(false);
   const [element, ref] = useState<HTMLImageElement | null>(null);
+  const [dragPreviewState, setDragPreviewState] = useState<DndDragPreviewSingleImageState | null>(null);
 
   useEffect(() => {
     if (!element) {
@@ -40,22 +43,34 @@ export const DndImage = memo(({ imageDTO, ...rest }: Props) => {
       onDrop: () => {
         setIsDragging(false);
       },
+      onGenerateDragPreview: (args) => {
+        if (Dnd.Source.singleImage.typeGuard(args.source.data)) {
+          setSingleImageDragPreview({
+            singleImageDndData: args.source.data,
+            onGenerateDragPreviewArgs: args,
+            setDragPreviewState,
+          });
+        }
+      },
     });
   }, [imageDTO, element, store]);
 
   useImageContextMenu(imageDTO, element);
 
   return (
-    <Image
-      role="button"
-      ref={ref}
-      src={imageDTO.image_url}
-      fallbackSrc={imageDTO.thumbnail_url}
-      w={imageDTO.width}
-      sx={sx}
-      data-is-dragging={isDragging}
-      {...rest}
-    />
+    <>
+      <Image
+        role="button"
+        ref={ref}
+        src={imageDTO.image_url}
+        fallbackSrc={imageDTO.thumbnail_url}
+        w={imageDTO.width}
+        sx={sx}
+        data-is-dragging={isDragging}
+        {...rest}
+      />
+      {dragPreviewState?.type === 'single-image' ? createSingleImageDragPreview(dragPreviewState) : null}
+    </>
   );
 });
 

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

@@ -1,5 +1,8 @@
 /* eslint-disable @typescript-eslint/no-namespace */ // We will use namespaces to organize the Dnd types
 
+import type { Input } from '@atlaskit/pragmatic-drag-and-drop/dist/types/entry-point/types';
+import type { GetOffsetFn } from '@atlaskit/pragmatic-drag-and-drop/dist/types/public-utils/element/custom-native-drag-preview/types';
+import type { SystemStyleObject } from '@invoke-ai/ui-library';
 import { getPrefixedId } from 'features/controlLayers/konva/util';
 import type { CanvasEntityIdentifier } from 'features/controlLayers/store/types';
 import type { BoardId } from 'features/gallery/store/types';
@@ -477,3 +480,35 @@ export const Dnd = {
     },
   },
 };
+
+/**
+ * The size of the image drag preview in theme units.
+ */
+export const DND_IMAGE_DRAG_PREVIEW_SIZE = 32 satisfies SystemStyleObject['w'];
+
+/**
+ * A drag preview offset function that works like the provided `preserveOffsetOnSource`, except when either the X or Y
+ * offset is outside the container, in which case it centers the preview in the container.
+ */
+export function preserveOffsetOnSourceFallbackCentered({
+  element,
+  input,
+}: {
+  element: HTMLElement;
+  input: Input;
+}): GetOffsetFn {
+  return ({ container }) => {
+    const sourceRect = element.getBoundingClientRect();
+    const containerRect = container.getBoundingClientRect();
+
+    let offsetX = input.clientX - sourceRect.x;
+    let offsetY = input.clientY - sourceRect.y;
+
+    if (offsetY > containerRect.height || offsetX > containerRect.width) {
+      offsetX = containerRect.width / 2;
+      offsetY = containerRect.height / 2;
+    }
+
+    return { x: offsetX, y: offsetY };
+  };
+}

invokeai/frontend/web/src/features/gallery/components/ImageGrid/GalleryImage.tsx

@@ -1,15 +1,17 @@
 import { combine } from '@atlaskit/pragmatic-drag-and-drop/combine';
 import { draggable, monitorForElements } from '@atlaskit/pragmatic-drag-and-drop/element/adapter';
-import { preserveOffsetOnSource } from '@atlaskit/pragmatic-drag-and-drop/element/preserve-offset-on-source';
-import { setCustomNativeDragPreview } from '@atlaskit/pragmatic-drag-and-drop/element/set-custom-native-drag-preview';
 import type { SystemStyleObject } from '@invoke-ai/ui-library';
-import { Box, Flex, Heading, Image } from '@invoke-ai/ui-library';
+import { Box, Flex, Image } from '@invoke-ai/ui-library';
 import { createSelector } from '@reduxjs/toolkit';
 import { galleryImageClicked } from 'app/store/middleware/listenerMiddleware/listeners/galleryImageClicked';
 import { useAppStore } from 'app/store/nanostores/store';
 import { useAppSelector } from 'app/store/storeHooks';
 import { useBoolean } from 'common/hooks/useBoolean';
 import { Dnd } from 'features/dnd/dnd';
+import type { DndDragPreviewMultipleImageState } from 'features/dnd/DndDragPreviewMultipleImage';
+import { createMultipleImageDragPreview, setMultipleImageDragPreview } from 'features/dnd/DndDragPreviewMultipleImage';
+import type { DndDragPreviewSingleImageState } from 'features/dnd/DndDragPreviewSingleImage';
+import { createSingleImageDragPreview, setSingleImageDragPreview } from 'features/dnd/DndDragPreviewSingleImage';
 import { useImageContextMenu } from 'features/gallery/components/ImageContextMenu/ImageContextMenu';
 import { GalleryImageHoverIcons } from 'features/gallery/components/ImageGrid/GalleryImageHoverIcons';
 import { getGalleryImageDataTestId } from 'features/gallery/components/ImageGrid/getGalleryImageDataTestId';
@@ -18,8 +20,6 @@ import { $imageViewer } from 'features/gallery/components/ImageViewer/useImageVi
 import { imageToCompareChanged, selectGallerySlice } from 'features/gallery/store/gallerySlice';
 import type { MouseEventHandler } from 'react';
 import { memo, useCallback, useEffect, useId, useMemo, useState } from 'react';
-import ReactDOM from 'react-dom';
-import { useTranslation } from 'react-i18next';
 import type { ImageDTO } from 'services/api/types';
 
 // This class name is used to calculate the number of images that fit in the gallery
@@ -83,16 +83,12 @@ interface Props {
   imageDTO: ImageDTO;
 }
 
-type MultiImageDragPreviewState = {
-  container: HTMLElement;
-  imageDTOs: ImageDTO[];
-  domRect: DOMRect;
-};
-
 export const GalleryImage = memo(({ imageDTO }: Props) => {
   const store = useAppStore();
   const [isDragging, setIsDragging] = useState(false);
-  const [dragPreviewState, setDragPreviewState] = useState<MultiImageDragPreviewState | null>(null);
+  const [dragPreviewState, setDragPreviewState] = useState<
+    DndDragPreviewSingleImageState | DndDragPreviewMultipleImageState | null
+  >(null);
   const [element, ref] = useState<HTMLImageElement | null>(null);
   const dndId = useId();
   const selectIsSelectedForCompare = useMemo(
@@ -144,25 +140,18 @@ export const GalleryImage = memo(({ imageDTO }: Props) => {
             return;
           }
         },
-        // See: https://atlassian.design/components/pragmatic-drag-and-drop/core-package/adapters/element/drag-previews
-        onGenerateDragPreview: ({ nativeSetDragImage, source, location }) => {
-          if (Dnd.Source.multipleImage.typeGuard(source.data)) {
-            const { imageDTOs } = source.data.payload;
-            const domRect = source.element.getBoundingClientRect();
-            setCustomNativeDragPreview({
-              render({ container }) {
-                // Cause a `react` re-render to create your portal synchronously
-                setDragPreviewState({ container, imageDTOs, domRect });
-                // In our cleanup function: cause a `react` re-render to create remove your portal
-                // Note: you can also remove the portal in `onDragStart`,
-                // which is when the cleanup function is called
-                return () => setDragPreviewState(null);
-              },
-              nativeSetDragImage,
-              getOffset: preserveOffsetOnSource({
-                element: source.element,
-                input: location.current.input,
-              }),
+        onGenerateDragPreview: (args) => {
+          if (Dnd.Source.multipleImage.typeGuard(args.source.data)) {
+            setMultipleImageDragPreview({
+              multipleImageDndData: args.source.data,
+              onGenerateDragPreviewArgs: args,
+              setDragPreviewState,
+            });
+          } else if (Dnd.Source.singleImage.typeGuard(args.source.data)) {
+            setSingleImageDragPreview({
+              singleImageDndData: args.source.data,
+              onGenerateDragPreviewArgs: args,
+              setDragPreviewState,
             });
           }
         },
@@ -243,32 +232,10 @@ export const GalleryImage = memo(({ imageDTO }: Props) => {
           <GalleryImageHoverIcons imageDTO={imageDTO} isHovered={isHovered.isTrue} />
         </Flex>
       </Box>
-      {dragPreviewState !== null &&
-        ReactDOM.createPortal(
-          <MultiImagePreview imageDTOs={dragPreviewState.imageDTOs} domRect={dragPreviewState.domRect} />,
-          dragPreviewState.container
-        )}
+      {dragPreviewState?.type === 'multiple-image' ? createMultipleImageDragPreview(dragPreviewState) : null}
+      {dragPreviewState?.type === 'single-image' ? createSingleImageDragPreview(dragPreviewState) : null}
     </>
   );
 });
 
 GalleryImage.displayName = 'GalleryImage';
-
-const MultiImagePreview = memo(({ imageDTOs, domRect }: { imageDTOs: ImageDTO[]; domRect: DOMRect }) => {
-  const { t } = useTranslation();
-  return (
-    <Flex
-      w={domRect.width}
-      h={domRect.height}
-      alignItems="center"
-      justifyContent="center"
-      flexDir="column"
-      bg="base.900"
-    >
-      <Heading>{imageDTOs.length}</Heading>
-      <Heading size="sm">{t('parameters.images')}</Heading>
-    </Flex>
-  );
-});
-
-MultiImagePreview.displayName = 'MultiImagePreview';