feat(ui): migrate to pragmatic-drag-and-drop (wip)

invokeai/frontend/web/src/features/controlLayers/components/CanvasDropArea.tsx

@@ -1,34 +1,20 @@
 import { Grid, GridItem } from '@invoke-ai/ui-library';
-import IAIDroppable from 'common/components/IAIDroppable';
-import type {
-  AddControlLayerFromImageDropData,
-  AddGlobalReferenceImageFromImageDropData,
-  AddRasterLayerFromImageDropData,
-  AddRegionalReferenceImageFromImageDropData,
-} from 'features/dnd/types';
+import { DndDropTarget } from 'features/dnd2/DndDropTarget';
+import {
+  addControlLayerFromImageDndTarget,
+  addGlobalReferenceImageFromImageDndTarget,
+  addRasterLayerFromImageDndTarget,
+  addRegionalGuidanceReferenceImageFromImageDndTarget,
+} from 'features/dnd2/types';
 import { useImageViewer } from 'features/gallery/components/ImageViewer/useImageViewer';
 import { memo } from 'react';
 import { useTranslation } from 'react-i18next';
 
-const addRasterLayerFromImageDropData: AddRasterLayerFromImageDropData = {
-  id: 'add-raster-layer-from-image-drop-data',
-  actionType: 'ADD_RASTER_LAYER_FROM_IMAGE',
-};
-
-const addControlLayerFromImageDropData: AddControlLayerFromImageDropData = {
-  id: 'add-control-layer-from-image-drop-data',
-  actionType: 'ADD_CONTROL_LAYER_FROM_IMAGE',
-};
-
-const addRegionalReferenceImageFromImageDropData: AddRegionalReferenceImageFromImageDropData = {
-  id: 'add-control-layer-from-image-drop-data',
-  actionType: 'ADD_REGIONAL_REFERENCE_IMAGE_FROM_IMAGE',
-};
-
-const addGlobalReferenceImageFromImageDropData: AddGlobalReferenceImageFromImageDropData = {
-  id: 'add-control-layer-from-image-drop-data',
-  actionType: 'ADD_GLOBAL_REFERENCE_IMAGE_FROM_IMAGE',
-};
+const addRasterLayerFromImageDndTargetData = addRasterLayerFromImageDndTarget.getData({});
+const addControlLayerFromImageDndTargetData = addControlLayerFromImageDndTarget.getData({});
+const addRegionalGuidanceReferenceImageFromImageDndTargetData =
+  addRegionalGuidanceReferenceImageFromImageDndTarget.getData({});
+const addGlobalReferenceImageFromImageDndTargetData = addGlobalReferenceImageFromImageDndTarget.getData({});
 
 export const CanvasDropArea = memo(() => {
   const { t } = useTranslation();
@@ -51,28 +37,28 @@ export const CanvasDropArea = memo(() => {
         pointerEvents="none"
       >
         <GridItem position="relative">
-          <IAIDroppable
-            dropLabel={t('controlLayers.canvasContextMenu.newRasterLayer')}
-            data={addRasterLayerFromImageDropData}
+          <DndDropTarget
+            label={t('controlLayers.canvasContextMenu.newRasterLayer')}
+            targetData={addRasterLayerFromImageDndTargetData}
           />
         </GridItem>
         <GridItem position="relative">
-          <IAIDroppable
-            dropLabel={t('controlLayers.canvasContextMenu.newControlLayer')}
-            data={addControlLayerFromImageDropData}
+          <DndDropTarget
+            label={t('controlLayers.canvasContextMenu.newControlLayer')}
+            targetData={addControlLayerFromImageDndTargetData}
           />
         </GridItem>
 
         <GridItem position="relative">
-          <IAIDroppable
-            dropLabel={t('controlLayers.canvasContextMenu.newRegionalReferenceImage')}
-            data={addRegionalReferenceImageFromImageDropData}
+          <DndDropTarget
+            label={t('controlLayers.canvasContextMenu.newRegionalReferenceImage')}
+            targetData={addRegionalGuidanceReferenceImageFromImageDndTargetData}
           />
         </GridItem>
         <GridItem position="relative">
-          <IAIDroppable
-            dropLabel={t('controlLayers.canvasContextMenu.newGlobalReferenceImage')}
-            data={addGlobalReferenceImageFromImageDropData}
+          <DndDropTarget
+            label={t('controlLayers.canvasContextMenu.newGlobalReferenceImage')}
+            targetData={addGlobalReferenceImageFromImageDndTargetData}
           />
         </GridItem>
       </Grid>

invokeai/frontend/web/src/features/controlLayers/components/CanvasRightPanel.tsx

@@ -1,16 +1,18 @@
-import { useDndContext } from '@dnd-kit/core';
+import { combine } from '@atlaskit/pragmatic-drag-and-drop/combine';
+import { dropTargetForElements, monitorForElements } from '@atlaskit/pragmatic-drag-and-drop/element/adapter';
 import { Box, Button, Spacer, Tab, TabList, TabPanel, TabPanels, Tabs } from '@invoke-ai/ui-library';
-import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
-import IAIDropOverlay from 'common/components/IAIDropOverlay';
+import { useAppDispatch, useAppSelector, useAppStore } from 'app/store/storeHooks';
 import { CanvasLayersPanelContent } from 'features/controlLayers/components/CanvasLayersPanelContent';
 import { CanvasManagerProviderGate } from 'features/controlLayers/contexts/CanvasManagerProviderGate';
 import { selectEntityCountActive } from 'features/controlLayers/store/selectors';
+import { DndDropOverlay } from 'features/dnd2/DndDropOverlay';
+import type { DndState } from 'features/dnd2/types';
 import GalleryPanelContent from 'features/gallery/components/GalleryPanelContent';
 import { useImageViewer } from 'features/gallery/components/ImageViewer/useImageViewer';
 import { useRegisteredHotkeys } from 'features/system/components/HotkeysModal/useHotkeyData';
 import { selectActiveTabCanvasRightPanel } from 'features/ui/store/uiSelectors';
 import { activeTabCanvasRightPanelChanged } from 'features/ui/store/uiSlice';
-import { memo, useCallback, useMemo, useRef, useState } from 'react';
+import { memo, useCallback, useEffect, useMemo, useRef, useState } from 'react';
 import { useTranslation } from 'react-i18next';
 
 export const CanvasRightPanel = memo(() => {
@@ -79,37 +81,12 @@ CanvasRightPanel.displayName = 'CanvasRightPanel';
 
 const PanelTabs = memo(() => {
   const { t } = useTranslation();
-  const activeTab = useAppSelector(selectActiveTabCanvasRightPanel);
+  const store = useAppStore();
   const activeEntityCount = useAppSelector(selectEntityCountActive);
-  const tabTimeout = useRef<number | null>(null);
-  const dndCtx = useDndContext();
-  const dispatch = useAppDispatch();
-  const [mouseOverTab, setMouseOverTab] = useState<'layers' | 'gallery' | null>(null);
-
-  const onOnMouseOverLayersTab = useCallback(() => {
-    setMouseOverTab('layers');
-    tabTimeout.current = window.setTimeout(() => {
-      if (dndCtx.active) {
-        dispatch(activeTabCanvasRightPanelChanged('layers'));
-      }
-    }, 300);
-  }, [dndCtx.active, dispatch]);
-
-  const onOnMouseOverGalleryTab = useCallback(() => {
-    setMouseOverTab('gallery');
-    tabTimeout.current = window.setTimeout(() => {
-      if (dndCtx.active) {
-        dispatch(activeTabCanvasRightPanelChanged('gallery'));
-      }
-    }, 300);
-  }, [dndCtx.active, dispatch]);
-
-  const onMouseOut = useCallback(() => {
-    setMouseOverTab(null);
-    if (tabTimeout.current) {
-      clearTimeout(tabTimeout.current);
-    }
-  }, []);
+  const [layersTabDndState, setLayersTabDndState] = useState<DndState>('idle');
+  const [galleryTabDndState, setGalleryTabDndState] = useState<DndState>('idle');
+  const layersTabRef = useRef<HTMLDivElement>(null);
+  const galleryTabRef = useRef<HTMLDivElement>(null);
 
   const layersTabLabel = useMemo(() => {
     if (activeEntityCount === 0) {
@@ -118,23 +95,131 @@ const PanelTabs = memo(() => {
     return `${t('controlLayers.layer_other')} (${activeEntityCount})`;
   }, [activeEntityCount, t]);
 
+  /**
+   * Handle dnd events for the tabs. When a tab is hovered for a certain amount of time, switch to that tab.
+   */
+  useEffect(() => {
+    if (!layersTabRef.current || !galleryTabRef.current) {
+      return;
+    }
+
+    let tabTimeout: number | null = null;
+
+    const layersTabCleanup = combine(
+      dropTargetForElements({
+        element: layersTabRef.current,
+        onDragEnter: () => {
+          // If we are already on the layers tab, do nothing
+          if (selectActiveTabCanvasRightPanel(store.getState()) === 'layers') {
+            return;
+          }
+
+          // Else set the state to active and switch to the layers tab after a timeout
+          setLayersTabDndState('active');
+          tabTimeout = window.setTimeout(() => {
+            tabTimeout = null;
+            store.dispatch(activeTabCanvasRightPanelChanged('layers'));
+            // When we switch tabs, the other tab should be pending
+            setLayersTabDndState('idle');
+            setGalleryTabDndState('pending');
+          }, 300);
+        },
+        onDragLeave: () => {
+          // Set the state to idle or pending depending on the current tab
+          if (selectActiveTabCanvasRightPanel(store.getState()) === 'layers') {
+            setLayersTabDndState('idle');
+          } else {
+            setLayersTabDndState('pending');
+          }
+          // Abort the tab switch if it hasn't happened yet
+          if (tabTimeout !== null) {
+            clearTimeout(tabTimeout);
+          }
+        },
+      }),
+      monitorForElements({
+        // Only monitor if we are not already on the layers tab
+        canMonitor: () => selectActiveTabCanvasRightPanel(store.getState()) !== 'layers',
+        onDragStart: () => {
+          // Set the state to pending when a drag starts
+          setLayersTabDndState('pending');
+        },
+      })
+    );
+
+    const galleryTabCleanup = combine(
+      dropTargetForElements({
+        element: galleryTabRef.current,
+        onDragEnter: () => {
+          // If we are already on the gallery tab, do nothing
+          if (selectActiveTabCanvasRightPanel(store.getState()) === 'gallery') {
+            return;
+          }
+
+          // Else set the state to active and switch to the gallery tab after a timeout
+          setGalleryTabDndState('active');
+          tabTimeout = window.setTimeout(() => {
+            tabTimeout = null;
+            store.dispatch(activeTabCanvasRightPanelChanged('gallery'));
+            // When we switch tabs, the other tab should be pending
+            setGalleryTabDndState('idle');
+            setLayersTabDndState('pending');
+          }, 300);
+        },
+        onDragLeave: () => {
+          // Set the state to idle or pending depending on the current tab
+          if (selectActiveTabCanvasRightPanel(store.getState()) === 'gallery') {
+            setGalleryTabDndState('idle');
+          } else {
+            setGalleryTabDndState('pending');
+          }
+          // Abort the tab switch if it hasn't happened yet
+          if (tabTimeout !== null) {
+            clearTimeout(tabTimeout);
+          }
+        },
+      }),
+      monitorForElements({
+        // Only monitor if we are not already on the gallery tab
+        canMonitor: () => selectActiveTabCanvasRightPanel(store.getState()) !== 'gallery',
+        onDragStart: () => {
+          // Set the state to pending when a drag starts
+          setGalleryTabDndState('pending');
+        },
+      })
+    );
+
+    const sharedCleanup = monitorForElements({
+      onDrop: () => {
+        // Reset the dnd state when a drop happens
+        setGalleryTabDndState('idle');
+        setLayersTabDndState('idle');
+      },
+    });
+
+    return () => {
+      layersTabCleanup();
+      galleryTabCleanup();
+      sharedCleanup();
+      if (tabTimeout !== null) {
+        clearTimeout(tabTimeout);
+      }
+    };
+  }, [store]);
+
   return (
     <>
-      <Tab position="relative" onMouseOver={onOnMouseOverLayersTab} onMouseOut={onMouseOut} w={32}>
+      <Tab ref={layersTabRef} position="relative" w={32}>
         <Box as="span" w="full">
           {layersTabLabel}
         </Box>
-        {dndCtx.active && activeTab !== 'layers' && (
-          <IAIDropOverlay isOver={mouseOverTab === 'layers'} withBackdrop={false} />
-        )}
+        <DndDropOverlay dndState={layersTabDndState} withBackdrop={false} />
       </Tab>
-      <Tab position="relative" onMouseOver={onOnMouseOverGalleryTab} onMouseOut={onMouseOut} w={32}>
+      <Tab ref={galleryTabRef} position="relative" w={32}>
         <Box as="span" w="full">
           {t('gallery.gallery')}
         </Box>
-        {dndCtx.active && activeTab !== 'gallery' && (
-          <IAIDropOverlay isOver={mouseOverTab === 'gallery'} withBackdrop={false} />
-        )}
+        <DndDropOverlay dndState={galleryTabDndState} withBackdrop={false} />
       </Tab>
     </>
   );

invokeai/frontend/web/src/features/controlLayers/components/ControlLayer/ControlLayer.tsx

@@ -1,5 +1,4 @@
 import { Spacer } from '@invoke-ai/ui-library';
-import IAIDroppable from 'common/components/IAIDroppable';
 import { CanvasEntityContainer } from 'features/controlLayers/components/common/CanvasEntityContainer';
 import { CanvasEntityHeader } from 'features/controlLayers/components/common/CanvasEntityHeader';
 import { CanvasEntityHeaderCommonActions } from 'features/controlLayers/components/common/CanvasEntityHeaderCommonActions';
@@ -11,7 +10,8 @@ import { ControlLayerSettings } from 'features/controlLayers/components/ControlL
 import { ControlLayerAdapterGate } from 'features/controlLayers/contexts/EntityAdapterContext';
 import { EntityIdentifierContext } from 'features/controlLayers/contexts/EntityIdentifierContext';
 import type { CanvasEntityIdentifier } from 'features/controlLayers/store/types';
-import type { ReplaceLayerImageDropData } from 'features/dnd/types';
+import { DndDropTarget } from 'features/dnd2/DndDropTarget';
+import { replaceLayerWithImageDndTarget, type ReplaceLayerWithImageDndTargetData } from 'features/dnd2/types';
 import { memo, useMemo } from 'react';
 import { useTranslation } from 'react-i18next';
 
@@ -25,10 +25,11 @@ export const ControlLayer = memo(({ id }: Props) => {
     () => ({ id, type: 'control_layer' }),
     [id]
   );
-  const dropData = useMemo<ReplaceLayerImageDropData>(
-    () => ({ id, actionType: 'REPLACE_LAYER_WITH_IMAGE', context: { entityIdentifier } }),
-    [id, entityIdentifier]
+  const targetData = useMemo<ReplaceLayerWithImageDndTargetData>(
+    () => replaceLayerWithImageDndTarget.getData({ entityIdentifier }),
+    [entityIdentifier]
   );
+
   return (
     <EntityIdentifierContext.Provider value={entityIdentifier}>
       <ControlLayerAdapterGate>
@@ -43,7 +44,7 @@ export const ControlLayer = memo(({ id }: Props) => {
           <CanvasEntitySettingsWrapper>
             <ControlLayerSettings />
           </CanvasEntitySettingsWrapper>
-          <IAIDroppable data={dropData} dropLabel={t('controlLayers.replaceLayer')} />
+          <DndDropTarget targetData={targetData} label={t('controlLayers.replaceLayer')} />
         </CanvasEntityContainer>
       </ControlLayerAdapterGate>
     </EntityIdentifierContext.Provider>

invokeai/frontend/web/src/features/controlLayers/components/RasterLayer/RasterLayer.tsx

@@ -1,5 +1,4 @@
 import { Spacer } from '@invoke-ai/ui-library';
-import IAIDroppable from 'common/components/IAIDroppable';
 import { CanvasEntityContainer } from 'features/controlLayers/components/common/CanvasEntityContainer';
 import { CanvasEntityHeader } from 'features/controlLayers/components/common/CanvasEntityHeader';
 import { CanvasEntityHeaderCommonActions } from 'features/controlLayers/components/common/CanvasEntityHeaderCommonActions';
@@ -8,7 +7,9 @@ import { CanvasEntityEditableTitle } from 'features/controlLayers/components/com
 import { RasterLayerAdapterGate } from 'features/controlLayers/contexts/EntityAdapterContext';
 import { EntityIdentifierContext } from 'features/controlLayers/contexts/EntityIdentifierContext';
 import type { CanvasEntityIdentifier } from 'features/controlLayers/store/types';
-import type { ReplaceLayerImageDropData } from 'features/dnd/types';
+import { DndDropTarget } from 'features/dnd2/DndDropTarget';
+import type { ReplaceLayerWithImageDndTargetData } from 'features/dnd2/types';
+import { replaceLayerWithImageDndTarget } from 'features/dnd2/types';
 import { memo, useMemo } from 'react';
 import { useTranslation } from 'react-i18next';
 
@@ -19,9 +20,9 @@ type Props = {
 export const RasterLayer = memo(({ id }: Props) => {
   const { t } = useTranslation();
   const entityIdentifier = useMemo<CanvasEntityIdentifier<'raster_layer'>>(() => ({ id, type: 'raster_layer' }), [id]);
-  const dropData = useMemo<ReplaceLayerImageDropData>(
-    () => ({ id, actionType: 'REPLACE_LAYER_WITH_IMAGE', context: { entityIdentifier } }),
-    [id, entityIdentifier]
+  const targetData = useMemo<ReplaceLayerWithImageDndTargetData>(
+    () => replaceLayerWithImageDndTarget.getData({ entityIdentifier }),
+    [entityIdentifier]
   );
 
   return (
@@ -34,7 +35,7 @@ export const RasterLayer = memo(({ id }: Props) => {
             <Spacer />
             <CanvasEntityHeaderCommonActions />
           </CanvasEntityHeader>
-          <IAIDroppable data={dropData} dropLabel={t('controlLayers.replaceLayer')} />
+          <DndDropTarget targetData={targetData} label={t('controlLayers.replaceLayer')} />
         </CanvasEntityContainer>
       </RasterLayerAdapterGate>
     </EntityIdentifierContext.Provider>

invokeai/frontend/web/src/features/dnd2/DndDropOverlay.tsx

@@ -0,0 +1,68 @@
+import { Flex, Text } from '@invoke-ai/ui-library';
+import type { DndState } from 'features/dnd2/types';
+import { memo } from 'react';
+
+type Props = {
+  dndState: DndState;
+  label?: string;
+  withBackdrop?: boolean;
+};
+
+export const DndDropOverlay = memo((props: Props) => {
+  const { dndState, label, withBackdrop = true } = props;
+
+  if (dndState === 'idle') {
+    return null;
+  }
+
+  return (
+    <Flex position="absolute" top={0} right={0} bottom={0} left={0}>
+      <Flex
+        position="absolute"
+        top={0}
+        right={0}
+        bottom={0}
+        left={0}
+        bg={withBackdrop ? 'base.900' : 'transparent'}
+        opacity={0.7}
+        borderRadius="base"
+        alignItems="center"
+        justifyContent="center"
+        transitionProperty="common"
+        transitionDuration="0.1s"
+      />
+
+      <Flex
+        position="absolute"
+        top={0.5}
+        right={0.5}
+        bottom={0.5}
+        left={0.5}
+        opacity={1}
+        borderWidth={1.5}
+        borderColor={dndState === 'active' ? 'invokeYellow.300' : 'base.500'}
+        borderRadius="base"
+        borderStyle="dashed"
+        transitionProperty="common"
+        transitionDuration="0.1s"
+        alignItems="center"
+        justifyContent="center"
+      >
+        {label && (
+          <Text
+            fontSize="lg"
+            fontWeight="semibold"
+            color={dndState === 'active' ? 'invokeYellow.300' : 'base.500'}
+            transitionProperty="common"
+            transitionDuration="0.1s"
+            textAlign="center"
+          >
+            {label}
+          </Text>
+        )}
+      </Flex>
+    </Flex>
+  );
+});
+
+DndDropOverlay.displayName = 'DndDropOverlay';

invokeai/frontend/web/src/features/dnd2/DndDropTarget.tsx

@@ -0,0 +1,94 @@
+import { combine } from '@atlaskit/pragmatic-drag-and-drop/combine';
+import { dropTargetForElements, monitorForElements } from '@atlaskit/pragmatic-drag-and-drop/element/adapter';
+import { Box } from '@invoke-ai/ui-library';
+import { dndDropped } from 'app/store/middleware/listenerMiddleware/listeners/dnd';
+import { useAppDispatch } from 'app/store/storeHooks';
+import { DndDropOverlay } from 'features/dnd2/DndDropOverlay';
+import type { DndState, DndTargetData } from 'features/dnd2/types';
+import { isDndSourceData, isValidDrop } from 'features/dnd2/types';
+import { memo, useEffect, useRef, useState } from 'react';
+
+type Props = {
+  label?: string;
+  disabled?: boolean;
+  targetData: DndTargetData;
+};
+
+export const DndDropTarget = memo((props: Props) => {
+  const { label, targetData, disabled } = props;
+  const [dndState, setDndState] = useState<DndState>('idle');
+  const ref = useRef<HTMLDivElement>(null);
+  const dispatch = useAppDispatch();
+
+  useEffect(() => {
+    if (!ref.current) {
+      return;
+    }
+
+    return combine(
+      dropTargetForElements({
+        element: ref.current,
+        canDrop: (args) => {
+          if (disabled) {
+            return false;
+          }
+          const sourceData = args.source.data;
+          if (!isDndSourceData(sourceData)) {
+            return false;
+          }
+          return isValidDrop(sourceData, targetData);
+        },
+        onDragEnter: () => {
+          setDndState('active');
+        },
+        onDragLeave: () => {
+          setDndState('pending');
+        },
+        getData: () => targetData,
+        onDrop: (args) => {
+          const sourceData = args.source.data;
+          if (!isDndSourceData(sourceData)) {
+            return;
+          }
+          dispatch(dndDropped({ sourceData, targetData }));
+        },
+      }),
+      monitorForElements({
+        canMonitor: (args) => {
+          if (disabled) {
+            return false;
+          }
+          const sourceData = args.source.data;
+          if (!isDndSourceData(sourceData)) {
+            return false;
+          }
+          return isValidDrop(sourceData, targetData);
+        },
+        onDragStart: () => {
+          setDndState('pending');
+        },
+        onDrop: () => {
+          setDndState('idle');
+        },
+      })
+    );
+  }, [targetData, disabled, dispatch]);
+
+  return (
+    <Box
+      ref={ref}
+      position="absolute"
+      top={0}
+      right={0}
+      bottom={0}
+      left={0}
+      w="full"
+      h="full"
+      pointerEvents={dndState === 'idle' ? 'none' : 'auto'}
+    >
+      <DndDropOverlay dndState={dndState} label={label} />
+    </Box>
+  );
+});
+
+DndDropTarget.displayName = 'DndDropTarget';

invokeai/frontend/web/src/features/dnd2/types.ts

@@ -0,0 +1,297 @@
+import type { CanvasEntityIdentifier } from 'features/controlLayers/store/types';
+import type { BoardId } from 'features/gallery/store/types';
+import type { ImageDTO } from 'services/api/types';
+
+export type DndData = Record<string | symbol, unknown>;
+const _buildDataTypeGuard =
+  <T extends DndData>(key: symbol) =>
+  (data: DndData): data is T => {
+    return Boolean(data[key]);
+  };
+const _buildDataGetter =
+  <T extends DndData>(key: symbol) =>
+  (data: Omit<T, typeof key>): T => {
+    return {
+      [key]: true,
+      ...data,
+    } as T;
+  };
+const buildDndSourceApi = <T extends DndData>(key: symbol) =>
+  ({ key, typeGuard: _buildDataTypeGuard<T>(key), getData: _buildDataGetter<T>(key) }) as const;
+
+//#region DndSourceData
+const _SingleImageDndSourceDataKey = Symbol('SingleImageDndSourceData');
+export type SingleImageDndSourceData = {
+  [_SingleImageDndSourceDataKey]: true;
+  imageDTO: ImageDTO;
+};
+export const singleImageDndSource = buildDndSourceApi<SingleImageDndSourceData>(_SingleImageDndSourceDataKey);
+
+const _MultipleImageDndSourceDataKey = Symbol('MultipleImageDndSourceData');
+export type MultipleImageDndSourceData = {
+  [_MultipleImageDndSourceDataKey]: true;
+  imageDTOs: ImageDTO[];
+  boardId: BoardId;
+};
+export const multipleImageDndSource = buildDndSourceApi<MultipleImageDndSourceData>(_MultipleImageDndSourceDataKey);
+
+/**
+ * A union of all possible DndSourceData types.
+ */
+const sourceApis = [singleImageDndSource, multipleImageDndSource] as const;
+export type DndSourceData = SingleImageDndSourceData | MultipleImageDndSourceData;
+export const isDndSourceData = (data: DndData): data is DndSourceData => {
+  for (const sourceApi of sourceApis) {
+    if (sourceApi.typeGuard(data)) {
+      return true;
+    }
+  }
+  return false;
+};
+
+//#endregion
+
+//#region DndTargetData
+const buildDndTargetApi = <T extends DndData>(
+  key: symbol,
+  validateDrop: (sourceData: DndSourceData, targetData: T) => boolean
+) => ({ key, typeGuard: _buildDataTypeGuard<T>(key), getData: _buildDataGetter<T>(key), validateDrop }) as const;
+
+const _SetGlobalReferenceImageDndTargetDataKey = Symbol('SetGlobalReferenceImageDndTargetData');
+export type SetGlobalReferenceImageDndTargetData = {
+  [_SetGlobalReferenceImageDndTargetDataKey]: true;
+  globalReferenceImageId: string;
+};
+export const setGlobalReferenceImageDndTarget = buildDndTargetApi<SetGlobalReferenceImageDndTargetData>(
+  _SetGlobalReferenceImageDndTargetDataKey,
+  singleImageDndSource.typeGuard
+);
+
+const _SetRegionalGuidanceReferenceImageDndTargetDataKey = Symbol('SetRegionalGuidanceReferenceImageDndTargetData');
+export type SetRegionalGuidanceReferenceImageDndTargetData = {
+  [_SetRegionalGuidanceReferenceImageDndTargetDataKey]: true;
+  regionalGuidanceId: string;
+  referenceImageId: string;
+};
+export const setRegionalGuidanceReferenceImageDndTarget =
+  buildDndTargetApi<SetRegionalGuidanceReferenceImageDndTargetData>(
+    _SetRegionalGuidanceReferenceImageDndTargetDataKey,
+    singleImageDndSource.typeGuard
+  );
+
+const _AddRasterLayerFromImageDndTargetDataKey = Symbol('AddRasterLayerFromImageDndTargetData');
+export type AddRasterLayerFromImageDndTargetData = {
+  [_AddRasterLayerFromImageDndTargetDataKey]: true;
+};
+export const addRasterLayerFromImageDndTarget = buildDndTargetApi<AddRasterLayerFromImageDndTargetData>(
+  _AddRasterLayerFromImageDndTargetDataKey,
+  singleImageDndSource.typeGuard
+);
+
+const _AddControlLayerFromImageDndTargetDataKey = Symbol('AddControlLayerFromImageDndTargetData');
+export type AddControlLayerFromImageDndTargetData = {
+  [_AddControlLayerFromImageDndTargetDataKey]: true;
+};
+export const addControlLayerFromImageDndTarget = buildDndTargetApi<AddControlLayerFromImageDndTargetData>(
+  _AddControlLayerFromImageDndTargetDataKey,
+  singleImageDndSource.typeGuard
+);
+
+const _AddInpaintMaskFromImageDndTargetDataKey = Symbol('AddInpaintMaskFromImageDndTargetData');
+export type AddInpaintMaskFromImageDndTargetData = {
+  [_AddInpaintMaskFromImageDndTargetDataKey]: true;
+};
+export const addInpaintMaskFromImageDndTarget = buildDndTargetApi<AddInpaintMaskFromImageDndTargetData>(
+  _AddInpaintMaskFromImageDndTargetDataKey,
+  singleImageDndSource.typeGuard
+);
+
+const _AddRegionalGuidanceFromImageDndTargetDataKey = Symbol('AddRegionalGuidanceFromImageDndTargetData');
+export type AddRegionalGuidanceFromImageDndTargetData = {
+  [_AddRegionalGuidanceFromImageDndTargetDataKey]: true;
+};
+export const addRegionalGuidanceFromImageDndTarget = buildDndTargetApi<AddRegionalGuidanceFromImageDndTargetData>(
+  _AddRegionalGuidanceFromImageDndTargetDataKey,
+  singleImageDndSource.typeGuard
+);
+
+const _AddRegionalGuidanceReferenceImageFromImageDndTargetDataKey = Symbol(
+  'AddRegionalGuidanceReferenceImageFromImageDndTargetData'
+);
+export type AddRegionalGuidanceReferenceImageFromImageDndTargetData = {
+  [_AddRegionalGuidanceReferenceImageFromImageDndTargetDataKey]: true;
+};
+export const addRegionalGuidanceReferenceImageFromImageDndTarget =
+  buildDndTargetApi<AddRegionalGuidanceReferenceImageFromImageDndTargetData>(
+    _AddRegionalGuidanceReferenceImageFromImageDndTargetDataKey,
+    singleImageDndSource.typeGuard
+  );
+
+const _AddGlobalReferenceImageFromImageDndTargetDataKey = Symbol('AddGlobalReferenceImageFromImageDndTargetData');
+export type AddGlobalReferenceImageFromImageDndTargetData = {
+  [_AddGlobalReferenceImageFromImageDndTargetDataKey]: true;
+};
+export const addGlobalReferenceImageFromImageDndTarget =
+  buildDndTargetApi<AddGlobalReferenceImageFromImageDndTargetData>(
+    _AddGlobalReferenceImageFromImageDndTargetDataKey,
+    singleImageDndSource.typeGuard
+  );
+
+const _ReplaceLayerWithImageDndTargetDataKey = Symbol('ReplaceLayerWithImageDndTargetData');
+export type ReplaceLayerWithImageDndTargetData = {
+  [_ReplaceLayerWithImageDndTargetDataKey]: true;
+  entityIdentifier: CanvasEntityIdentifier<'control_layer' | 'raster_layer' | 'inpaint_mask' | 'regional_guidance'>;
+};
+export const replaceLayerWithImageDndTarget = buildDndTargetApi<ReplaceLayerWithImageDndTargetData>(
+  _ReplaceLayerWithImageDndTargetDataKey,
+  singleImageDndSource.typeGuard
+);
+
+const _SetUpscaleInitialImageFromImageDndTargetDataKey = Symbol('SetUpscaleInitialImageFromImageDndTargetData');
+export type SetUpscaleInitialImageFromImageDndTargetData = {
+  [_SetUpscaleInitialImageFromImageDndTargetDataKey]: true;
+};
+export const setUpscaleInitialImageFromImageDndTarget = buildDndTargetApi<SetUpscaleInitialImageFromImageDndTargetData>(
+  _SetUpscaleInitialImageFromImageDndTargetDataKey,
+  singleImageDndSource.typeGuard
+);
+
+const _SetNodeImageFieldDndTargetDataKey = Symbol('SetNodeImageFieldDndTargetData');
+export type SetNodeImageFieldDndTargetData = {
+  [_SetNodeImageFieldDndTargetDataKey]: true;
+  nodeId: string;
+  fieldName: string;
+};
+export const setNodeImageFieldDndTarget = buildDndTargetApi<SetNodeImageFieldDndTargetData>(
+  _SetNodeImageFieldDndTargetDataKey,
+  singleImageDndSource.typeGuard
+);
+
+const _SelectForCompareDndTargetDataKey = Symbol('SelectForCompareDndTargetData');
+export type SelectForCompareDndTargetData = {
+  [_SelectForCompareDndTargetDataKey]: true;
+  firstImageName?: string | null;
+  secondImageName?: string | null;
+};
+export const selectForCompareDndTarget = buildDndTargetApi<SelectForCompareDndTargetData>(
+  _SelectForCompareDndTargetDataKey,
+  singleImageDndSource.typeGuard
+);
+
+const _AddToBoardDndTargetDataKey = Symbol('AddToBoardDndTargetData');
+export type AddToBoardDndTargetData = {
+  [_AddToBoardDndTargetDataKey]: true;
+  boardId: string;
+};
+export const addToBoardDndTarget = buildDndTargetApi<AddToBoardDndTargetData>(
+  _AddToBoardDndTargetDataKey,
+  (sourceData, targetData) => {
+    if (singleImageDndSource.typeGuard(sourceData)) {
+      const { imageDTO } = sourceData;
+      const currentBoard = imageDTO.board_id ?? 'none';
+      const destinationBoard = targetData.boardId;
+      return currentBoard !== destinationBoard;
+    }
+
+    if (multipleImageDndSource.typeGuard(sourceData)) {
+      const currentBoard = sourceData.boardId;
+      const destinationBoard = targetData.boardId;
+      return currentBoard !== destinationBoard;
+    }
+
+    return false;
+  }
+);
+
+const _RemoveFromBoardDndTargetDataKey = Symbol('RemoveFromBoardDndTargetData');
+export type RemoveFromBoardDndTargetData = {
+  [_RemoveFromBoardDndTargetDataKey]: true;
+};
+export const removeFromBoardDndTarget = buildDndTargetApi<RemoveFromBoardDndTargetData>(
+  _RemoveFromBoardDndTargetDataKey,
+  (sourceData) => {
+    if (singleImageDndSource.typeGuard(sourceData)) {
+      const currentBoard = sourceData.imageDTO.board_id ?? 'none';
+      return currentBoard !== 'none';
+    }
+
+    if (multipleImageDndSource.typeGuard(sourceData)) {
+      const currentBoard = sourceData.boardId;
+      return currentBoard !== 'none';
+    }
+
+    return false;
+  }
+);
+
+const targetApis = [
+  setGlobalReferenceImageDndTarget,
+  setRegionalGuidanceReferenceImageDndTarget,
+  // Add layer from image
+  addRasterLayerFromImageDndTarget,
+  addControlLayerFromImageDndTarget,
+  addGlobalReferenceImageFromImageDndTarget,
+  addRegionalGuidanceReferenceImageFromImageDndTarget,
+  //
+  addRegionalGuidanceFromImageDndTarget,
+  addInpaintMaskFromImageDndTarget,
+  //
+  replaceLayerWithImageDndTarget,
+  setUpscaleInitialImageFromImageDndTarget,
+  setNodeImageFieldDndTarget,
+  selectForCompareDndTarget,
+  // Board ops
+  addToBoardDndTarget,
+  removeFromBoardDndTarget,
+] as const;
+
+/**
+ * A union of all possible DndTargetData types.
+ */
+export type DndTargetData =
+  | SetGlobalReferenceImageDndTargetData
+  | SetRegionalGuidanceReferenceImageDndTargetData
+  | AddRasterLayerFromImageDndTargetData
+  | AddControlLayerFromImageDndTargetData
+  | AddInpaintMaskFromImageDndTargetData
+  | AddRegionalGuidanceFromImageDndTargetData
+  | AddRegionalGuidanceReferenceImageFromImageDndTargetData
+  | AddGlobalReferenceImageFromImageDndTargetData
+  | ReplaceLayerWithImageDndTargetData
+  | SetUpscaleInitialImageFromImageDndTargetData
+  | SetNodeImageFieldDndTargetData
+  | AddToBoardDndTargetData
+  | RemoveFromBoardDndTargetData
+  | SelectForCompareDndTargetData;
+
+export const isDndTargetData = (data: DndData): data is DndTargetData => {
+  for (const targetApi of targetApis) {
+    if (targetApi.typeGuard(data)) {
+      return true;
+    }
+  }
+  return false;
+};
+//#endregion
+
+/**
+ * Validates whether a drop is valid.
+ * @param sourceData The data being dragged.
+ * @param targetData The data of the target being dragged onto.
+ * @returns Whether the drop is valid.
+ */
+export const isValidDrop = (sourceData: DndSourceData, targetData: DndTargetData): boolean => {
+  for (const targetApi of targetApis) {
+    if (targetApi.typeGuard(targetData)) {
+      /**
+       * TS cannot narrow the type of the targetApi and will error in the validator call.
+       * We've just checked that targetData is of the right type, though, so this cast to `any` is safe.
+       */
+      /* eslint-disable-next-line @typescript-eslint/no-explicit-any */
+      return targetApi.validateDrop(sourceData, targetData as any);
+    }
+  }
+  return false;
+};
+
+export type DndState = 'idle' | 'pending' | 'active';

invokeai/frontend/web/src/features/gallery/components/Boards/BoardsList/GalleryBoard.tsx

@@ -2,8 +2,9 @@ import type { SystemStyleObject } from '@invoke-ai/ui-library';
 import { Box, Flex, Icon, Image, Text, Tooltip } from '@invoke-ai/ui-library';
 import { skipToken } from '@reduxjs/toolkit/query';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
-import IAIDroppable from 'common/components/IAIDroppable';
-import type { AddToBoardDropData } from 'features/dnd/types';
+import { DndDropTarget } from 'features/dnd2/DndDropTarget';
+import type { AddToBoardDndTargetData } from 'features/dnd2/types';
+import { addToBoardDndTarget } from 'features/dnd2/types';
 import { AutoAddBadge } from 'features/gallery/components/Boards/AutoAddBadge';
 import BoardContextMenu from 'features/gallery/components/Boards/BoardContextMenu';
 import { BoardEditableTitle } from 'features/gallery/components/Boards/BoardsList/BoardEditableTitle';
@@ -45,12 +46,8 @@ const GalleryBoard = ({ board, isSelected }: GalleryBoardProps) => {
     }
   }, [selectedBoardId, board.board_id, autoAssignBoardOnClick, autoAddBoardId, dispatch]);
 
-  const droppableData: AddToBoardDropData = useMemo(
-    () => ({
-      id: board.board_id,
-      actionType: 'ADD_TO_BOARD',
-      context: { boardId: board.board_id },
-    }),
+  const targetData: AddToBoardDndTargetData = useMemo(
+    () => addToBoardDndTarget.getData({ boardId: board.board_id }),
     [board.board_id]
   );
 
@@ -85,7 +82,7 @@ const GalleryBoard = ({ board, isSelected }: GalleryBoardProps) => {
           </Tooltip>
         )}
       </BoardContextMenu>
-      <IAIDroppable data={droppableData} dropLabel={t('gallery.move')} />
+      <DndDropTarget targetData={targetData} label={t('gallery.move')} />
     </Box>
   );
 };

invokeai/frontend/web/src/features/gallery/components/Boards/BoardsList/NoBoardBoard.tsx

@@ -1,8 +1,9 @@
 import type { SystemStyleObject } from '@invoke-ai/ui-library';
 import { Box, Flex, Icon, Text, Tooltip } from '@invoke-ai/ui-library';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
-import IAIDroppable from 'common/components/IAIDroppable';
-import type { RemoveFromBoardDropData } from 'features/dnd/types';
+import { DndDropTarget } from 'features/dnd2/DndDropTarget';
+import type { RemoveFromBoardDndTargetData } from 'features/dnd2/types';
+import { removeFromBoardDndTarget } from 'features/dnd2/types';
 import { AutoAddBadge } from 'features/gallery/components/Boards/AutoAddBadge';
 import { BoardTooltip } from 'features/gallery/components/Boards/BoardsList/BoardTooltip';
 import NoBoardBoardContextMenu from 'features/gallery/components/Boards/NoBoardBoardContextMenu';
@@ -43,13 +44,7 @@ const NoBoardBoard = memo(({ isSelected }: Props) => {
     }
   }, [dispatch, autoAssignBoardOnClick]);
 
-  const droppableData: RemoveFromBoardDropData = useMemo(
-    () => ({
-      id: 'no_board',
-      actionType: 'REMOVE_FROM_BOARD',
-    }),
-    []
-  );
+  const targetData: RemoveFromBoardDndTargetData = useMemo(() => removeFromBoardDndTarget.getData({}), []);
 
   const { t } = useTranslation();
 
@@ -102,7 +97,7 @@ const NoBoardBoard = memo(({ isSelected }: Props) => {
           </Tooltip>
         )}
       </NoBoardBoardContextMenu>
-      <IAIDroppable data={droppableData} dropLabel={t('gallery.move')} />
+      <DndDropTarget targetData={targetData} label={t('gallery.move')} />
     </Box>
   );
 });

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

@@ -7,7 +7,6 @@ import MultipleSelectionMenuItems from 'features/gallery/components/ImageContext
 import SingleSelectionMenuItems from 'features/gallery/components/ImageContextMenu/SingleSelectionMenuItems';
 import { selectSelectionCount } from 'features/gallery/store/gallerySelectors';
 import { map } from 'nanostores';
-import type { RefObject } from 'react';
 import { memo, useCallback, useEffect, useRef } from 'react';
 import type { ImageDTO } from 'services/api/types';
 
@@ -60,12 +59,13 @@ const getImageDTOFromMap = (target: Node): ImageDTO | undefined => {
  * @param imageDTO The image DTO to register the context menu for.
  * @param targetRef The ref of the target element that should trigger the context menu.
  */
-export const useImageContextMenu = (imageDTO: ImageDTO | undefined, targetRef: RefObject<HTMLDivElement>) => {
+export const useImageContextMenu = (imageDTO: ImageDTO | undefined, targetRef: HTMLDivElement | null
+) => {
   useEffect(() => {
-    if (!targetRef.current || !imageDTO) {
+    if (!targetRef || !imageDTO) {
       return;
     }
-    const el = targetRef.current;
+    const el = targetRef;
     elToImageMap.set(el, imageDTO);
     return () => {
       elToImageMap.delete(el);

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

@@ -1,40 +1,83 @@
+import { combine } from '@atlaskit/pragmatic-drag-and-drop/combine';
+import { draggable, monitorForElements } from '@atlaskit/pragmatic-drag-and-drop/element/adapter';
 import type { SystemStyleObject } from '@invoke-ai/ui-library';
-import { Box, Flex, Text, useShiftModifier } from '@invoke-ai/ui-library';
+import { Box, Flex, Image, Skeleton, Text, useShiftModifier } from '@invoke-ai/ui-library';
 import { useStore } from '@nanostores/react';
 import { createSelector } from '@reduxjs/toolkit';
+import { galleryImageClicked } from 'app/store/middleware/listenerMiddleware/listeners/galleryImageClicked';
 import { $customStarUI } from 'app/store/nanostores/customStarUI';
+import { useAppStore } from 'app/store/nanostores/store';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
-import IAIDndImage from 'common/components/IAIDndImage';
 import IAIDndImageIcon from 'common/components/IAIDndImageIcon';
 import IAIFillSkeleton from 'common/components/IAIFillSkeleton';
+import { useBoolean } from 'common/hooks/useBoolean';
 import { imagesToDeleteSelected } from 'features/deleteImageModal/store/slice';
-import type { GallerySelectionDraggableData, ImageDraggableData, TypesafeDraggableData } from 'features/dnd/types';
+import { multipleImageDndSource, singleImageDndSource } from 'features/dnd2/types';
+import { useImageContextMenu } from 'features/gallery/components/ImageContextMenu/ImageContextMenu';
 import { getGalleryImageDataTestId } from 'features/gallery/components/ImageGrid/getGalleryImageDataTestId';
 import { useImageViewer } from 'features/gallery/components/ImageViewer/useImageViewer';
-import { useMultiselect } from 'features/gallery/hooks/useMultiselect';
 import { useScrollIntoView } from 'features/gallery/hooks/useScrollIntoView';
-import { selectSelectedBoardId } from 'features/gallery/store/gallerySelectors';
 import { imageToCompareChanged, selectGallerySlice } from 'features/gallery/store/gallerySlice';
-import type { MouseEvent } from 'react';
-import { memo, useCallback, useMemo, useState } from 'react';
+import type { MouseEvent, MouseEventHandler } from 'react';
+import { memo, useCallback, useEffect, useMemo, useState } from 'react';
 import { useTranslation } from 'react-i18next';
 import { PiArrowsOutBold, PiStarBold, PiStarFill, PiTrashSimpleFill } from 'react-icons/pi';
 import { useStarImagesMutation, useUnstarImagesMutation } from 'services/api/endpoints/images';
 import type { ImageDTO } from 'services/api/types';
 
 // This class name is used to calculate the number of images that fit in the gallery
-export const GALLERY_IMAGE_CLASS_NAME = 'gallery-image';
+export const GALLERY_IMAGE_CONTAINER_CLASS_NAME = 'gallery-image-container';
 
-const imageSx: SystemStyleObject = { w: 'full', h: 'full' };
-const boxSx: SystemStyleObject = {
+const galleryImageContainerSX = {
   containerType: 'inline-size',
-};
-
-const badgeSx: SystemStyleObject = {
-  '@container (max-width: 80px)': {
-    '&': { display: 'none' },
+  w: 'full',
+  h: 'full',
+  '.gallery-image-size-badge': {
+    '@container (max-width: 80px)': {
+      '&': { display: 'none' },
+    },
   },
-};
+  '.gallery-image': {
+    touchAction: 'none',
+    userSelect: 'none',
+    webkitUserSelect: 'none',
+    position: 'relative',
+    justifyContent: 'center',
+    alignItems: 'center',
+    aspectRatio: '1/1',
+    '::before': {
+      content: '""',
+      display: 'inline-block',
+      position: 'absolute',
+      top: 0,
+      left: 0,
+      right: 0,
+      bottom: 0,
+      pointerEvents: 'none',
+      borderRadius: 'base',
+    },
+    '&[data-selected=true]::before': {
+      boxShadow:
+        'inset 0px 0px 0px 3px var(--invoke-colors-invokeBlue-500), inset 0px 0px 0px 4px var(--invoke-colors-invokeBlue-800)',
+    },
+    '&[data-selected-for-compare=true]::before': {
+      boxShadow:
+        'inset 0px 0px 0px 3px var(--invoke-colors-invokeGreen-300), inset 0px 0px 0px 4px var(--invoke-colors-invokeGreen-800)',
+    },
+    '&:hover::before': {
+      boxShadow:
+        'inset 0px 0px 0px 2px var(--invoke-colors-invokeBlue-300), inset 0px 0px 0px 3px var(--invoke-colors-invokeBlue-800)',
+    },
+    '&:hover[data-selected=true]::before': {
+      boxShadow:
+        'inset 0px 0px 0px 3px var(--invoke-colors-invokeBlue-400), inset 0px 0px 0px 4px var(--invoke-colors-invokeBlue-800)',
+    },
+    '&:hover[data-selected-for-compare=true]::before': {
+      boxShadow:
+        'inset 0px 0px 0px 3px var(--invoke-colors-invokeGreen-200), inset 0px 0px 0px 4px var(--invoke-colors-invokeGreen-800)',
+    },
+  },
+} satisfies SystemStyleObject;
 
 interface HoverableImageProps {
   imageDTO: ImageDTO;
@@ -57,87 +100,125 @@ export const GalleryImage = memo(({ index, imageDTO }: HoverableImageProps) => {
 GalleryImage.displayName = 'GalleryImage';
 
 const GalleryImageContent = memo(({ index, imageDTO }: HoverableImageProps) => {
-  const dispatch = useAppDispatch();
-  const selectedBoardId = useAppSelector(selectSelectedBoardId);
+  const store = useAppStore();
+  const [isDragging, setIsDragging] = useState(false);
+  const [element, ref] = useState<HTMLImageElement | null>(null);
+  const imageViewer = useImageViewer();
   const selectIsSelectedForCompare = useMemo(
     () => createSelector(selectGallerySlice, (gallery) => gallery.imageToCompare?.image_name === imageDTO.image_name),
     [imageDTO.image_name]
   );
   const isSelectedForCompare = useAppSelector(selectIsSelectedForCompare);
-  const { handleClick, isSelected, areMultiplesSelected } = useMultiselect(imageDTO);
+  const selectIsSelected = useMemo(
+    () =>
+      createSelector(selectGallerySlice, (gallery) =>
+        gallery.selection.some((i) => i.image_name === imageDTO.image_name)
+      ),
+    [imageDTO.image_name]
+  );
+  const isSelected = useAppSelector(selectIsSelected);
 
-  const imageContainerRef = useScrollIntoView(isSelected, index, areMultiplesSelected);
+  useScrollIntoView(element, isSelected, index);
 
-  const draggableData = useMemo<TypesafeDraggableData | undefined>(() => {
-    if (areMultiplesSelected) {
-      const data: GallerySelectionDraggableData = {
-        id: 'gallery-image',
-        payloadType: 'GALLERY_SELECTION',
-        payload: { boardId: selectedBoardId },
-      };
-      return data;
+  useEffect(() => {
+    if (!element) {
+      return;
     }
+    return combine(
+      draggable({
+        element,
+        getInitialData: () => {
+          const { gallery } = store.getState();
+          // When we have multiple images selected, and the dragged image is part of the selection, initiate a
+          // multi-image drag.
+          if (gallery.selection.length > 1 && gallery.selection.includes(imageDTO)) {
+            return multipleImageDndSource.getData({ imageDTOs: gallery.selection, boardId: gallery.selectedBoardId });
+          }
 
-    if (imageDTO) {
-      const data: ImageDraggableData = {
-        id: 'gallery-image',
-        payloadType: 'IMAGE_DTO',
-        payload: { imageDTO },
-      };
-      return data;
-    }
-  }, [imageDTO, selectedBoardId, areMultiplesSelected]);
+          // Otherwise, initiate a single-image drag
+          return singleImageDndSource.getData({ imageDTO });
+        },
+        // This is a "local" drag start event, meaning that it is only called when this specific image is dragged.
+        onDragStart: (args) => {
+          // When we start dragging a single image, set the dragging state to true. This is only called when this
+          // specific image is dragged.
+          if (singleImageDndSource.typeGuard(args.source.data)) {
+            setIsDragging(true);
+            return;
+          }
+        },
+      }),
+      monitorForElements({
+        // This is a "global" drag start event, meaning that it is called for all drag events.
+        onDragStart: (args) => {
+          // When we start dragging multiple images, set the dragging state to true if the dragged image is part of the
+          // selection. This is called for all drag events.
+          if (multipleImageDndSource.typeGuard(args.source.data) && args.source.data.imageDTOs.includes(imageDTO)) {
+            setIsDragging(true);
+          }
+        },
+        onDrop: () => {
+          // Always set the dragging state to false when a drop event occurs.
+          setIsDragging(false);
+        },
+      })
+    );
+  }, [imageDTO, element, store]);
 
-  const [isHovered, setIsHovered] = useState(false);
+  const isHovered = useBoolean(false);
 
-  const handleMouseOver = useCallback(() => {
-    setIsHovered(true);
-  }, []);
+  const onClick = useCallback<MouseEventHandler<HTMLDivElement>>(
+    (e) => {
+      store.dispatch(
+        galleryImageClicked({
+          imageDTO,
+          shiftKey: e.shiftKey,
+          ctrlKey: e.ctrlKey,
+          metaKey: e.metaKey,
+          altKey: e.altKey,
+        })
+      );
+    },
+    [imageDTO, store]
+  );
 
-  const imageViewer = useImageViewer();
-  const onDoubleClick = useCallback(() => {
+  const onDoubleClick = useCallback<MouseEventHandler<HTMLDivElement>>(() => {
     imageViewer.open();
-    dispatch(imageToCompareChanged(null));
-  }, [dispatch, imageViewer]);
-
-  const handleMouseOut = useCallback(() => {
-    setIsHovered(false);
-  }, []);
+    store.dispatch(imageToCompareChanged(null));
+  }, [imageViewer, store]);
 
   const dataTestId = useMemo(() => getGalleryImageDataTestId(imageDTO.image_name), [imageDTO.image_name]);
 
-  if (!imageDTO) {
-    return <IAIFillSkeleton />;
-  }
+  useImageContextMenu(imageDTO, element);
 
   return (
-    <Box w="full" h="full" className={GALLERY_IMAGE_CLASS_NAME} data-testid={dataTestId} sx={boxSx}>
+    <Box
+      className={GALLERY_IMAGE_CONTAINER_CLASS_NAME}
+      data-testid={dataTestId}
+      sx={galleryImageContainerSX}
+      opacity={isDragging ? 0.3 : 1}
+    >
       <Flex
-        ref={imageContainerRef}
-        userSelect="none"
-        position="relative"
-        justifyContent="center"
-        alignItems="center"
-        aspectRatio="1/1"
-        onMouseOver={handleMouseOver}
-        onMouseOut={handleMouseOut}
+        role="button"
+        className="gallery-image"
+        onMouseOver={isHovered.setTrue}
+        onMouseOut={isHovered.setFalse}
+        onClick={onClick}
+        onDoubleClick={onDoubleClick}
+        data-selected={isSelected}
+        data-selected-for-compare={isSelectedForCompare}
       >
-        <IAIDndImage
-          onClick={handleClick}
-          onDoubleClick={onDoubleClick}
-          imageDTO={imageDTO}
-          draggableData={draggableData}
-          isSelected={isSelected}
-          isSelectedForCompare={isSelectedForCompare}
-          minSize={0}
-          imageSx={imageSx}
-          isDropDisabled={true}
-          isUploadDisabled={true}
-          thumbnail={true}
-          withHoverOverlay
-        >
-          <HoverIcons imageDTO={imageDTO} isHovered={isHovered} />
-        </IAIDndImage>
+        <Image
+          ref={ref}
+          src={imageDTO.thumbnail_url}
+          fallback={<SizedSkeleton width={imageDTO.width} height={imageDTO.height} />}
+          w={imageDTO.width}
+          objectFit="contain"
+          maxW="full"
+          maxH="full"
+          borderRadius="base"
+        />
+        <HoverIcons imageDTO={imageDTO} isHovered={isHovered.isTrue} />
       </Flex>
     </Box>
   );
@@ -220,21 +301,17 @@ const StarIcon = memo(({ imageDTO }: { imageDTO: ImageDTO }) => {
   const [unstarImages] = useUnstarImagesMutation();
 
   const toggleStarredState = useCallback(() => {
-    if (imageDTO) {
-      if (imageDTO.starred) {
-        unstarImages({ imageDTOs: [imageDTO] });
-      }
-      if (!imageDTO.starred) {
-        starImages({ imageDTOs: [imageDTO] });
-      }
+    if (imageDTO.starred) {
+      unstarImages({ imageDTOs: [imageDTO] });
+    } else {
+      starImages({ imageDTOs: [imageDTO] });
     }
   }, [starImages, unstarImages, imageDTO]);
 
   const starIcon = useMemo(() => {
     if (imageDTO.starred) {
       return customStarUi ? customStarUi.on.icon : <PiStarFill />;
-    }
-    if (!imageDTO.starred) {
+    } else {
       return customStarUi ? customStarUi.off.icon : <PiStarBold />;
     }
   }, [imageDTO.starred, customStarUi]);
@@ -242,11 +319,9 @@ const StarIcon = memo(({ imageDTO }: { imageDTO: ImageDTO }) => {
   const starTooltip = useMemo(() => {
     if (imageDTO.starred) {
       return customStarUi ? customStarUi.off.text : 'Unstar';
-    }
-    if (!imageDTO.starred) {
+    } else {
       return customStarUi ? customStarUi.on.text : 'Star';
     }
-    return '';
   }, [imageDTO.starred, customStarUi]);
 
   return (
@@ -266,6 +341,7 @@ StarIcon.displayName = 'StarIcon';
 const SizeBadge = memo(({ imageDTO }: { imageDTO: ImageDTO }) => {
   return (
     <Text
+      className="gallery-image-size-badge"
       position="absolute"
       background="base.900"
       color="base.50"
@@ -277,10 +353,15 @@ const SizeBadge = memo(({ imageDTO }: { imageDTO: ImageDTO }) => {
       px={2}
       lineHeight={1.25}
       borderTopEndRadius="base"
-      sx={badgeSx}
       pointerEvents="none"
     >{`${imageDTO.width}x${imageDTO.height}`}</Text>
   );
 });
 
 SizeBadge.displayName = 'SizeBadge';
+
+const SizedSkeleton = memo(({ width, height }: { width: number; height: number }) => {
+  return <Skeleton w={`${width}px`} h="auto" objectFit="contain" aspectRatio={`${width}/${height}`} />;
+});
+
+SizedSkeleton.displayName = 'SizedSkeleton';

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

@@ -14,7 +14,7 @@ import { PiImageBold, PiWarningCircleBold } from 'react-icons/pi';
 import { useListImagesQuery } from 'services/api/endpoints/images';
 
 import { GALLERY_GRID_CLASS_NAME } from './constants';
-import { GALLERY_IMAGE_CLASS_NAME, GalleryImage } from './GalleryImage';
+import { GALLERY_IMAGE_CONTAINER_CLASS_NAME, GalleryImage } from './GalleryImage';
 
 const GalleryImageGrid = () => {
   useGalleryHotkeys();
@@ -79,7 +79,7 @@ const Content = () => {
       // Managing refs for dynamically rendered components is a bit tedious:
       // - https://react.dev/learn/manipulating-the-dom-with-refs#how-to-manage-a-list-of-refs-using-a-ref-callback
       // As a easy workaround, we can just grab the first gallery image element directly.
-      const imageEl = document.querySelector(`.${GALLERY_IMAGE_CLASS_NAME}`);
+      const imageEl = document.querySelector(`.${GALLERY_IMAGE_CONTAINER_CLASS_NAME}`);
       if (!imageEl) {
         // No images in gallery?
         return;

invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageComparisonDroppable.tsx

@@ -1,7 +1,8 @@
 import { Flex } from '@invoke-ai/ui-library';
-import { useAppSelector } from 'app/store/storeHooks';
-import IAIDroppable from 'common/components/IAIDroppable';
-import type { SelectForCompareDropData } from 'features/dnd/types';
+import { useAppStore } from 'app/store/storeHooks';
+import { DndDropTarget } from 'features/dnd2/DndDropTarget';
+import type { SelectForCompareDndTargetData } from 'features/dnd2/types';
+import { selectForCompareDndTarget } from 'features/dnd2/types';
 import { memo, useMemo } from 'react';
 import { useTranslation } from 'react-i18next';
 
@@ -9,22 +10,18 @@ import { selectComparisonImages } from './common';
 
 export const ImageComparisonDroppable = memo(() => {
   const { t } = useTranslation();
-  const { firstImage, secondImage } = useAppSelector(selectComparisonImages);
-  const selectForCompareDropData = useMemo<SelectForCompareDropData>(
-    () => ({
-      id: 'image-comparison',
-      actionType: 'SELECT_FOR_COMPARE',
-      context: {
-        firstImageName: firstImage?.image_name,
-        secondImageName: secondImage?.image_name,
-      },
-    }),
-    [firstImage?.image_name, secondImage?.image_name]
-  );
+  const store = useAppStore();
+  const targetData = useMemo<SelectForCompareDndTargetData>(() => {
+    const { firstImage, secondImage } = selectComparisonImages(store.getState());
+    return selectForCompareDndTarget.getData({
+      firstImageName: firstImage?.image_name,
+      secondImageName: secondImage?.image_name,
+    });
+  }, [store]);
 
   return (
     <Flex position="absolute" top={0} right={0} bottom={0} left={0} gap={2} pointerEvents="none">
-      <IAIDroppable data={selectForCompareDropData} dropLabel={t('gallery.selectForCompare')} />
+      <DndDropTarget targetData={targetData} label={t('gallery.selectForCompare')} />
     </Flex>
   );
 });

invokeai/frontend/web/src/features/gallery/hooks/useGalleryNavigation.ts

@@ -2,7 +2,7 @@ import { useAltModifier } from '@invoke-ai/ui-library';
 import { createSelector } from '@reduxjs/toolkit';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import { GALLERY_GRID_CLASS_NAME } from 'features/gallery/components/ImageGrid/constants';
-import { GALLERY_IMAGE_CLASS_NAME } from 'features/gallery/components/ImageGrid/GalleryImage';
+import { GALLERY_IMAGE_CONTAINER_CLASS_NAME } from 'features/gallery/components/ImageGrid/GalleryImage';
 import { getGalleryImageDataTestId } from 'features/gallery/components/ImageGrid/getGalleryImageDataTestId';
 import { virtuosoGridRefs } from 'features/gallery/components/ImageGrid/types';
 import { useGalleryImages } from 'features/gallery/hooks/useGalleryImages';
@@ -29,7 +29,7 @@ import type { ImageDTO } from 'services/api/types';
  * Gets the number of images per row in the gallery by grabbing their DOM elements.
  */
 const getImagesPerRow = (): number => {
-  const imageEl = document.querySelector(`.${GALLERY_IMAGE_CLASS_NAME}`);
+  const imageEl = document.querySelector(`.${GALLERY_IMAGE_CONTAINER_CLASS_NAME}`);
   const gridEl = document.querySelector(`.${GALLERY_GRID_CLASS_NAME}`);
 
   if (!imageEl || !gridEl) {

invokeai/frontend/web/src/features/gallery/hooks/useMultiselect.ts

@@ -8,24 +8,21 @@ import type { MouseEvent } from 'react';
 import { useCallback, useMemo } from 'react';
 import type { ImageDTO } from 'services/api/types';
 
-export const useMultiselect = (imageDTO?: ImageDTO) => {
+export const useMultiselect = (imageDTO: ImageDTO) => {
   const dispatch = useAppDispatch();
   const areMultiplesSelected = useAppSelector(selectHasMultipleImagesSelected);
   const selectIsSelected = useMemo(
     () =>
       createSelector(selectGallerySlice, (gallery) =>
-        gallery.selection.some((i) => i.image_name === imageDTO?.image_name)
+        gallery.selection.some((i) => i.image_name === imageDTO.image_name)
       ),
-    [imageDTO?.image_name]
+    [imageDTO.image_name]
   );
   const isSelected = useAppSelector(selectIsSelected);
   const isMultiSelectEnabled = useFeatureStatus('multiselect');
 
-  const handleClick = useCallback(
+  const onClick = useCallback(
     (e: MouseEvent<HTMLDivElement>) => {
-      if (!imageDTO) {
-        return;
-      }
       if (!isMultiSelectEnabled) {
         dispatch(selectionChanged([imageDTO]));
         return;
@@ -47,6 +44,6 @@ export const useMultiselect = (imageDTO?: ImageDTO) => {
   return {
     areMultiplesSelected,
     isSelected,
-    handleClick,
+    onClick,
   };
 };

invokeai/frontend/web/src/features/gallery/hooks/useScrollIntoView.ts

@@ -1,7 +1,9 @@
+import { useAppSelector } from 'app/store/storeHooks';
 import { virtuosoGridRefs } from 'features/gallery/components/ImageGrid/types';
+import { selectHasMultipleImagesSelected } from 'features/gallery/store/gallerySelectors';
 import { getIsVisible } from 'features/gallery/util/getIsVisible';
 import { getScrollToIndexAlign } from 'features/gallery/util/getScrollToIndexAlign';
-import { useEffect, useRef } from 'react';
+import { useEffect } from 'react';
 
 /**
  * Scrolls an image into view when it is selected. This is necessary because
@@ -11,13 +13,13 @@ import { useEffect, useRef } from 'react';
  * Also handles when an image is selected programmatically - for example, when
  * auto-switching the new gallery images.
  *
+ * @param imageContainerRef The ref to the image container.
  * @param isSelected Whether the image is selected.
  * @param index The index of the image in the gallery.
- * @param selectionCount The number of images selected.
  * @returns
  */
-export const useScrollIntoView = (isSelected: boolean, index: number, areMultiplesSelected: boolean) => {
-  const imageContainerRef = useRef<HTMLDivElement>(null);
+export const useScrollIntoView = (imageContainerRef: HTMLElement | null, isSelected: boolean, index: number) => {
+  const areMultiplesSelected = useAppSelector(selectHasMultipleImagesSelected);
 
   useEffect(() => {
     if (!isSelected || areMultiplesSelected) {
@@ -33,7 +35,7 @@ export const useScrollIntoView = (isSelected: boolean, index: number, areMultipl
       return;
     }
 
-    const itemRect = imageContainerRef.current?.getBoundingClientRect();
+    const itemRect = imageContainerRef?.getBoundingClientRect();
     const rootRect = root.getBoundingClientRect();
 
     if (!itemRect || !getIsVisible(itemRect, rootRect)) {
@@ -42,7 +44,5 @@ export const useScrollIntoView = (isSelected: boolean, index: number, areMultipl
         align: getScrollToIndexAlign(index, range),
       });
     }
-  }, [isSelected, index, areMultiplesSelected]);
-
-  return imageContainerRef;
+  }, [isSelected, index, areMultiplesSelected, imageContainerRef]);
 };